feat(sample-support): score policy logprobs over the recorded support end to end - #2084
feat(sample-support): score policy logprobs over the recorded support end to end#2084dyurk-lila wants to merge 19 commits into
Conversation
Extract the single-request HTTP generation path out of RemoteInferenceClient into a standalone RemoteGenerateClient plus a RemoteGenerateResult dataclass, so routed-expert results can be obtained without constructing the full inference/control-plane client. RemoteInferenceClient now owns an internal RemoteGenerateClient and delegates session management, _post, and _generate_single to it. Endpoint routing, retry/backoff, cache_salt handling, serialization, and lifecycle behavior are unchanged.
For multi-turn rollouts, build routed-expert (R3) trace data incrementally as the conversation grows instead of re-gathering the whole conversation's routes on every turn. - Add `routed_experts_prompt_starts` to `InferenceEngineInput` and thread a per-request `routed_experts_prompt_start` through `RemoteInferenceClient` and `RemoteInferenceGenerator.generate()`, forwarded as a sampling param so the engine only returns routes for the newly generated suffix. - Introduce `TokenMetadataTrace` (token-aligned array accumulator) and `RoutedExpertTrace`, which records each generation's routes and finalizes a full per-token routed-expert array with loss-mask-aware terminal padding. - Track a `RoutedExpertTrace` on `AgentLoopState` and record routes per turn, replacing the previous whole-conversation re-gather in `SkyRLGymGenerator.agent_loop`. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support-conditioned scoring for bounded sampler replay, allowing the renormalization of rollout logprobs over the sampler's recorded bounded support. It adds a new PackedTensor utility to efficiently manage ragged token-aligned batch fields, integrates sample support capture into the vLLM server and remote client, and implements the necessary validation and configuration options. Additionally, it enforces compatibility constraints between rollout router replay (R3), sample support, step-wise training, and virtual pipeline parallelism. The review feedback suggests improving the readability and clarity of the documentation regarding these compatibility constraints.
| ``` | ||
|
|
||
| To enable rollout router replay, set `generator.inference_engine.enable_return_routed_experts=True`, `trainer.policy.megatron_config.moe_enable_routing_replay=True`, and use the `mp` distributed_executor_backend for vLLM. Note that | ||
| R3 requires per-token routes that line up with the trained tokens, so `SkyRLGymGenerator` refuses it together with `generator.step_wise_trajectories`, `generator.use_conversation_multi_turn=False`, a custom `generator.chat_template`, and `generator.vision_language_generator` (each breaks that alignment; see the [step-wise training](../tutorials/step-wise-training) page for the step-wise reason). Routes are captured for train batches only. Note that |
There was a problem hiding this comment.
This sentence is quite long and lists several distinct points. For better readability, consider breaking it down into smaller sentences and using a list. For example:
R3 requires per-token routes that line up with the trained tokens. This alignment is broken by certain generator configurations. Therefore, `SkyRLGymGenerator` refuses to enable R3 when used with any of the following:
- `generator.step_wise_trajectories` (see the [step-wise training](../tutorials/step-wise-training) page for details)
- `generator.use_conversation_multi_turn=False`
- a custom `generator.chat_template`
- `generator.vision_language_generator`| | `trajectory_ids` | `List[TrajectoryID]` | Associates each step-sample with its parent trajectory. All steps of the same trajectory share the same `TrajectoryID`. | | ||
| | `rollout_logprobs` | `List[List[float]]` | Per-token logprobs from the inference engine, aligned with `response_ids`. Required for TIS. | | ||
| | `rollout_sample_support` | `List[List[List[int]]]` | Per generated token, the bounded top-k vocab IDs the sampler drew from, right-padded with `SAMPLE_SUPPORT_PADDING` (`-1`). One dense `[tokens, top_k]` block per row, aligned with `response_ids`. Every position must be present: an observation token or a synthetic (loop-appended) EOS carries an all-padding row rather than being absent, so the block stays rectangular and aligned. Enabled with `generator.inference_engine.enable_return_sample_support_set`, and requested per request (train batches only — eval's greedy `top_k=-1` params cannot satisfy the capture contract). | | ||
| | `rollout_expert_indices` | `Optional[List[RoutedExpertIndices]]` | Per token, the `[layers, topk]` MoE routes vLLM recorded, as one `[tokens, layers, topk]` integer array per trajectory (rollout router replay, R3). **Refused with `step_wise_trajectories=True`**: each step-wise row's prompt is the whole history so far while routes are recorded for that step's generated tokens only, so a step's routes would replay onto the first N prompt tokens of its row with no length mismatch to assert on, silently training against routing that does not match the rollout. | |
There was a problem hiding this comment.
The explanation for why R3 is refused with step_wise_trajectories=True is a bit dense. For improved clarity, you could rephrase it. For example:
Refused with
step_wise_trajectories=True: This is because in step-wise training, each row's prompt consists of the entire history, but routes are only recorded for the tokens generated in that step. This would cause a mismatch where the step's routes are replayed onto the initial prompt tokens of the training row. Without a length mismatch to trigger an error, this would lead to silently training against incorrect routing information.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.
Reviewed by Cursor Bugbot for commit 6be01c8. Configure here.
| "sample-support replay permits at most one loss-bearing token without recorded support per " | ||
| f"trajectory (the appended EOS), got counts {per_trajectory_count[offenders].tolist()} for " | ||
| f"trajectories {offenders.tolist()}" | ||
| ) |
There was a problem hiding this comment.
CP skips multi-EOS rejection
Medium Severity
The synthetic-EOS guard counts unsupported loss-bearing tokens only on each context-parallel rank's local shard. With context_parallel_size > 1, two such tokens on different ranks never exceed the per-rank limit of one, so the batch is not rejected and both are scored as full-vocabulary fallbacks after the CP allgather.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit 6be01c8. Configure here.
| if len(rollout_sample_support) != num_samples: | ||
| raise ValueError("rollout_sample_support must contain support for every trajectory") | ||
|
|
||
| sample_support_tensor = build_sample_support(rollout_sample_support, response_lens) |
There was a problem hiding this comment.
Missing sampled-in-support check
Medium Severity
Preprocessing packs rollout_sample_support without verifying that each loss-bearing response token appears in its support row, and the scorer never checks either. Rows that only repair this at capture can still reach training; the scorer then treats valid_mask as true and emits a score that is not a log-probability under the recorded support.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit 6be01c8. Configure here.
2 participants
Problem
PRs 30–31 record and transport the sampler’s bounded support, but Megatron still computes policy and reference log-probabilities over the full vocabulary. This PR closes that loop by renormalizing each sampled token over the support recorded during rollout.
The feature is enabled explicitly with
trainer.algorithm.enable_sample_support_replay. In this PR the scorer is wired into the Megatron policy, policy-recompute, and reference forwards; FSDP support follows in PR 33.Support-conditioned scoring
For each token, the scorer gathers logits only for recorded support IDs and computes:
Trailing
-1members are ignored. A row with recorded support must contain the sampled token, an invariant established at capture and checked during preprocessing.Under tensor parallelism, each rank scores the support members and sampled token that fall inside its vocabulary shard. The normalization uses one MAX reduction followed by one combined SUM reduction for the denominator and sampled score. The straight-through collective construction preserves gradients without introducing a third reduction.
Fused LM-head path
When the Megatron forward returns hidden states instead of vocabulary logits, the scorer projects only the selected
(token position, vocabulary row)pairs._ChunkedCandidateProjectionbounds the temporary activation bylogprobs_chunk_sizeand recomputes selected chunks during backward, avoiding a retained[all candidate pairs, hidden size]tensor.The same scoring API therefore accepts either sharded vocabulary logits or hidden states plus the local LM-head weight.
Alignment and synthetic EOS
Support row IDs are aligned through the same
TokenMetadataLayoutused by routed-expert replay. The implementation covers ordinary left-padded batches, controller-packed batches with one trajectory per segment, and context-parallel local segments. Controller-packed rows containing multiple subsequences are rejected because they do not preserve the one-support-segment-per-trajectory contract.The generator may append an EOS after vLLM returns, so that loss-bearing token has no captured support. The scorer permits at most one such row per trajectory and computes its ordinary full-vocabulary log-probability. Capacity is fixed at one candidate slot per trajectory, including packed context-parallel layouts; unused slots are masked after scoring. Any other loss-bearing row without support is rejected instead of being trained at a default score.
Configuration and diagnostics
Enabling replay requires support capture and a compatible Megatron strategy. Missing support produces an error that identifies both the capture configuration and the custom-generator forwarding contract. Feature-disabled behavior retains the existing full-vocabulary scorer.
Testing
Note
High Risk
Changes how policy and reference log-probabilities are computed (including TP collectives and fused LM-head backward), which directly affects training gradients. Misalignment or missing support would silently or loudly corrupt the loss.
Overview
When
trainer.algorithm.enable_sample_support_replayis on, Megatron policy, recompute, and reference forwards renormalize each sampled token over the bounded top-k support captured at rollout instead of the full vocabulary.Scoring gathers only recorded support IDs (padding
-1ignored), thenlog p(sampled | support)via a TP MAX plus one combined SUM for the denominator and sampled score. The fused LM-head path projects selected(position, vocab row)pairs in chunks so it never materializes full logits or a retained[pairs, hidden]activation.Support rows are aligned with the same token layout as router replay. Multi-subsequence packed rows are rejected. At most one loss-bearing token per trajectory may lack support (appended EOS) and is scored over the full vocab; any other missing support fails loudly.
Requires capture (
enable_return_sample_support_set) and Megatron. Disabled runs keep the existing full-vocab scorer. The stacked diff also packs per-token side channels and wires capture through vLLM generate.Reviewed by Cursor Bugbot for commit 6be01c8. Bugbot is set up for automated code reviews on this repo. Configure here.