Skip to content

Cap GRPO sampled-logprob context - #28

Merged
bingran-you merged 1 commit into
mainfrom
bry/cap-grpo-logprob-context
Jul 14, 2026
Merged

Cap GRPO sampled-logprob context#28
bingran-you merged 1 commit into
mainfrom
bry/cap-grpo-logprob-context

Conversation

@bingran-you

@bingran-you bingran-you commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

What changed

  • Adds --max-logprob-context-tokens, defaulting to 24,576, to the OpenCode model bridge.
  • Applies the stricter cap only when the request explicitly asks for sampled-token logprobs.
  • Keeps the ordinary 49,152-token context for baseline, SFT, gate, and final evaluation.
  • Reuses the existing tool-output-only context fitting, so system/user instructions remain intact.

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, and grad_norm=0 through six completed steps, so the stage can restart from the saved SFT checkpoint without losing a policy update.

Validation

  • 224 package contract tests pass.
  • 26 focused bridge/CLI tests pass.
  • Ruff check/format, Python compilation, and git diff --check pass.
  • Live replay of the known 59k-token follow-up with logprobs=true fitted to exactly 20,480 prompt tokens, returned from the real Qwen3.5-9B TRL server, and produced a sidecar with the same 20,480 exact prompt IDs.
  • Ordinary evaluation requests still use the separate 49,152-token cap.

@bingran-you
bingran-you merged commit 04485dc into main Jul 14, 2026
1 check passed
@bingran-you
bingran-you deleted the bry/cap-grpo-logprob-context branch July 14, 2026 21:10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant