feat(fsdp): score policy logprobs over the recorded support under Ulysses SP - #2085
feat(fsdp): score policy logprobs over the recorded support under Ulysses SP#2085dyurk-lila wants to merge 21 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 a new feature for sample-support capture in the inference engine, allowing for the renormalization of rollout logprobs over a bounded sampler support set. The changes include updates to the inference server's communication protocol (using packed side-channel arrays), new utilities for managing packed tensors and metadata traces, and updates to the training pipeline to support this new metadata. I have provided feedback on improving the robustness of the network response parsing and adding necessary null checks for configuration parameters.
| try: | ||
| raw = await resp.read() | ||
| body = load_packed_body(raw) if packed_side_channels else orjson.loads(raw) | ||
| except orjson.JSONDecodeError as exc: |
There was a problem hiding this comment.
If load_packed_body(raw) raises a ValueError (which can happen if the response is truncated or corrupted), the exception will propagate immediately and bypass the retry loop. Catching both orjson.JSONDecodeError and ValueError ensures that transient network/transfer issues causing malformed payloads are retried gracefully.
| try: | |
| raw = await resp.read() | |
| body = load_packed_body(raw) if packed_side_channels else orjson.loads(raw) | |
| except orjson.JSONDecodeError as exc: | |
| try: | |
| raw = await resp.read() | |
| body = load_packed_body(raw) if packed_side_channels else orjson.loads(raw) | |
| except (orjson.JSONDecodeError, ValueError) as exc: |
| # Eval requests opt out of capture and do not use these constraints. | ||
| if self.generator.inference_engine.enable_return_sample_support_set: | ||
| sampling_params = self.generator.sampling_params | ||
| if sampling_params.temperature <= 0: |
There was a problem hiding this comment.
If sampling_params.temperature is None, the comparison sampling_params.temperature <= 0 will raise a TypeError. Checking for None explicitly prevents this and ensures a clear ValueError is raised instead.
| if sampling_params.temperature <= 0: | |
| if sampling_params.temperature is None or sampling_params.temperature <= 0: |
| def load_packed_body(raw: bytes, *, fields: tuple[str, ...] = PACKED_SIDE_CHANNEL_FIELDS) -> dict[str, Any]: | ||
| """Parse a response after replacing registered base64 blobs with views. | ||
|
|
||
| Null fields pass through. An envelope layout the scan cannot splice raises | ||
| instead of falling back to materializing the base64 as a Python string. | ||
| """ |
There was a problem hiding this comment.
The load_packed_body function relies on exact byte-level matching of the JSON payload (assuming compact formatting from orjson). If any intermediate proxy, load balancer, or API gateway reformats or pretty-prints the JSON response in transit, the byte-scan will fail to match, causing _restore_packed_data to raise a ValueError and crash the client. Consider documenting this strict compact-JSON requirement or implementing a more robust fallback parser if intermediate proxies are expected.
2 participants
Problem
PR 32 applies support-conditioned replay to Megatron. An FSDP policy or reference model must compute log-probabilities over the same recorded distribution; otherwise policy ratios and KL terms still compare incompatible normalizations.
This PR makes the replay scorer backend-neutral and wires it through the Hugging Face/FSDP policy, policy-recompute, and reference forwards, including padding removal and Ulysses sequence parallelism.
Shared scorer boundary
The common
score_aligned_sample_supporthelper now accepts already-aligned model outputs, sampled IDs, support row IDs, loss masks, and optional trajectory IDs. Both backends use it for support normalization and the synthetic-EOS full-vocabulary fallback.Megatron continues to derive trajectory structure from its
TokenMetadataLayout. FSDP supplies explicit trajectory IDs after its own padding-removal and sequence-parallel transforms. This keeps backend-specific layout mechanics outside the numerical scorer.FSDP alignment
The FSDP wrapper builds three token-aligned metadata channels in canonical left-padded coordinates:
-1where no row exists;Those channels follow the model input through:
nnz_indicesused for token IDs;Support row IDs are already placed at the logit positions that predict each response token, so they are not rolled. The loss mask is token-positioned and is rolled by one alongside the sampled targets. Keeping this distinction explicit avoids an off-by-one that can otherwise pass shape checks.
Ulysses metadata support
ulysses_pad_and_slice_inputsnow accepts[batch, sequence, ...]metadata rather than assuming a two-dimensional token-ID tensor. Padding is applied along the sequence dimension and accepts a field-specific value:0for ordinary IDs and masks,-1for missing support rows and padding trajectory IDs.This lets all token-aligned channels take exactly the same sequence shard without constructing a dense support tensor on every rank.
Policy and reference behavior
megatronandfsdpstrategies.Testing
Note
High Risk
Changes how policy and reference log-probabilities are computed when replay is on, which directly affects PPO ratios, KL, and gradients. Alignment bugs under packing or Ulysses would silently train against the wrong distribution.
Overview
Makes FSDP/HF scoring match Megatron’s support-conditioned replay: when
enable_sample_support_replayis on, policy training, policy recompute, and reference forwards score each token over the recorded sampler top-k (plus a full-vocab fallback for the synthetic EOS) instead of the full vocabulary.The shared
score_aligned_sample_supporthelper now takes already-aligned logits, sampled IDs, support row IDs, loss masks, and optional trajectory IDs. FSDP builds those channels in left-padded coordinates, then follows the same unpad / next-token / Ulysses shard path as the tokens. Support row IDs stay at logit positions (not rolled); the loss mask is rolled with the targets.ulysses_pad_and_slice_inputspads along the sequence dim for higher-rank metadata and accepts per-channel sentinels (-1for missing support rows). Reference forwards use the same support and policy temperature so KL compares matching distributions. Disabled replay keeps the previous FSDP scoring path.Reviewed by Cursor Bugbot for commit 95be5ba. Bugbot is set up for automated code reviews on this repo. Configure here.