Implement SimKO to add entropy in TopK token sampling during RL #13
Conversation
…into add-grpo # Conflicts: # .github/CODEOWNERS # .github/workflows/integration_test_8gpu_h100.yaml # .github/workflows/integration_test_8gpu_models.yaml # .github/workflows/integration_test_8gpu_torchft.yaml # torchtitan/components/checkpoint.py # torchtitan/experiments/__init__.py # torchtitan/models/attention.py # torchtitan/models/deepseek_v3/infra/parallelize.py # torchtitan/models/llama3/__init__.py # torchtitan/models/llama3/model/args.py # torchtitan/models/llama3/model/model.py # torchtitan/tools/logging.py # torchtitan/train.py
add grpo :)
add inference logp IS
feat: qwen3-next
feat: seed-oss support
…ing; add compute_token_entropy and apply_simko_adjustment functions in GRPO step.
|
|
||
| return loss | ||
|
|
||
| def compute_token_entropy(pred, mask): |
There was a problem hiding this comment.
we have a vocab parallel entropy fn in utils, that should be used instead
There was a problem hiding this comment.
sure will do in the following commit.
| def apply_simko_adjustment( | ||
| ratio, pred, labels, reward, mask, | ||
| alpha=0.01, K=3, lambda_top1=1.1, tau_percentile=80 | ||
| ): |
There was a problem hiding this comment.
this fn does not work with tensor parallel
| Returns: | ||
| adjusted_ratio: Ratio with SimKO adjustments | ||
| """ | ||
| # 1. Identify forking tokens (high-entropy) |
There was a problem hiding this comment.
does this need backpropagated through, or can you wrap it with no_grad?
There was a problem hiding this comment.
No this won't need a backprop and not seen anywhere in paper to mention so, yeah no grad is way to go
|
|
||
| # 2. Get top-K tokens and compute their ratios | ||
| _, topk_indices = torch.topk(pred, k=K, dim=-1) | ||
| new_log_probs_full = torch.nn.functional.log_softmax(pred, dim=-1) |
There was a problem hiding this comment.
sure, do you think like doing log softmax one time would do? If we are doing a no grad on this, it could be reused for old poilicy top k probs as well
There was a problem hiding this comment.
well this needs grads, since the grads are flowing through this softmax, we need to combine the way we're currently getting logprobs with the top-k logprob method, as the logprobs we need for the GRPO/GSPO loss are not guaranteed to be within the top-k. the quick and dirty way may be to wrap it in a checkpoint and chunk it, that way we "save" the memory by having it re-execute on backward in manageable chunks, even if it's not optimal, it should still just be a tiny portion of the compute so we can add to the backlog as i'm mostly worried about memory 😎
jquesnelle
force-pushed
the
dev-updated-again
branch
from
921ffb2 to
94d6abd
Compare
This Impl is directly inspired from this paper - https://arxiv.org/pdf/2510.14807v1
I'm still actively testing out and will share results on it, the topK sampling metrics during training are yet to get logged in W&B