Cap GRPO sampled-logprob context - #28
Merged
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changed
--max-logprob-context-tokens, defaulting to24,576, to the OpenCode model bridge.49,152-token context for baseline, SFT, gate, and final evaluation.Why
After raising the aggregate rollout limit, a valid ~37k-token data-agent trajectory reached TRL but OOMed during trainer-side policy-logprob recomputation. PyTorch attempted a 34.33 GiB logits allocation on top of 53.8 GiB already used.
Evaluation can serve the full context on the inference H100, but GRPO must also recompute logits on the trainer H100. A 24,576-token sampled-logprob context bounds that trainer path while preserving the full evaluation context and exact sidecar token IDs.
The failed GRPO run had zero reward variance,
loss=0, andgrad_norm=0through six completed steps, so the stage can restart from the saved SFT checkpoint without losing a policy update.Validation
224package contract tests pass.26focused bridge/CLI tests pass.git diff --checkpass.logprobs=truefitted to exactly20,480prompt tokens, returned from the real Qwen3.5-9B TRL server, and produced a sidecar with the same20,480exact prompt IDs.