feat(megatron): add bounded sampler-support replay

[dyurk-lila]

Note on the diff: This PR is part of a routed-expert-replay / sampler-support series and builds on the PRs below. GitHub can't show the intermediate branches here, so the diff is cumulative on top of main — the changes new to this PR sit on top of:

• fix(r3): align routed-expert metadata and padding #1908 — fix(r3) align routed-expert metadata and padding
• perf(r3): accelerate routed-expert transport with packed arrays #1909 — perf(r3) packed-array transport
• perf(r3): expand replay routes to local layers only #1910 — perf(r3) expand routes to local layers only
• refactor(inference): extract reusable RemoteInferenceGenerator #1911 — refactor extract RemoteInferenceGenerator
• refactor(generators): assemble routed-expert traces incrementally #1912 — refactor incremental routed-expert traces

Reviewing in PR order (lowest number first) shows each incremental change cleanly.

This adds the core of bounded sampler-support replay: the data plumbing, inference capture contract, and Megatron dense-path trainer support. It builds on the shared incremental token-metadata trace.

Problem

With bounded sampling, rollout probabilities are normalized over the candidates retained by vLLM's processed sampler rather than the full vocabulary. Megatron currently recomputes full-vocabulary log-probabilities, so the trainer cannot reproduce the rollout distribution for importance ratios or KL terms.

vLLM capture contract

When capture is enabled, the request asks vLLM for flat_logprobs with logprobs=top_k. Each token contributes one fixed row:

[sampled token, top candidate 0, ..., top candidate top_k-1]

The first column supplies the existing sampled-token logprob. The remaining columns become replay support IDs in one NumPy reshape; candidates whose processed logprob is -inf are stored as -1. This supports top-k, top-p, and min-p filtering without per-token searches, deduplication, or a second dict-format path.

Configuration rejects combinations whose rollout distribution cannot be replayed exactly: capture must be enabled, temperature > 0, top_k > 1, repetition penalty must remain 1.0, and arbitrary sampler additional_kwargs are unsupported. Replay is Megatron-only in this change; a follow-up adds the FSDP path.

Generator and incremental metadata contract

The RemoteInferenceGenerator owns the single-request HTTP boundary. This extends that typed interface:

• RemoteInferenceGenerator.generate(return_sample_support=True) selects /skyrl/v1/generate and adds the boolean capture field;
• RemoteGenerateResult.sample_support carries the decoded support rows alongside routed experts and token logprobs;
• RemoteInferenceClient batches those results into InferenceEngineOutput.rollout_sample_support.

Multi-turn accumulation follows the shared TokenMetadataTrace, the same backbone used by incremental routed-expert capture. Each generation contributes one contiguous int32 NumPy chunk. Synthetic EOS and observation positions contribute fixed-width -1 rows. The trace is finalized once per trajectory and converted to the list response only at the generator-output boundary, keeping routed-expert and sample-support alignment on the same lifecycle.

Metadata ownership

Support rows follow response tokens through stepwise merging, dynamic replacement/filtering, truncation, replay buffering, and microbatch padding. Preprocessing validates one fixed support width, trailing -1 padding, int32 vocab IDs, and presence of every loss-bearing sampled token.

An EOS may be appended after vLLM returns, so vLLM cannot have captured support for that token. An empty loss-bearing row is accepted only for that synthetic EOS; every other missing loss-bearing support row is rejected. At most one such EOS is allowed per trajectory, matching the generator contract.

Megatron replay

Megatron renormalizes sampled-token scores over the recorded support with fixed-shape gathers. TP uses one MAX reduction and one combined SUM reduction for the denominator and sampled score. The fused LM-head path projects only selected candidate pairs and chunks by logprobs_chunk_size, bounding the temporary [candidate pairs, hidden] allocation.

Synthetic-EOS full-vocabulary fallback also uses fixed capacity: one device-side candidate slot per trajectory, including packed CP-local segments. Unused slots are masked after computation, avoiding host synchronization, boolean compaction, ragged TP shapes, and per-microbatch CPU-to-GPU length transfer.

For ordinary unpacked RL, logits and targets use direct next-token slicing. For controller-packed RL microbatches, sample support and routed-expert replay share the TokenMetadataLayout, built once from the attention mask, applying the same per-segment padding, next-token shift, CP front/back sharding, and final scatter into canonical batch positions. Controller-packed multi-subsequence/SFT metadata is rejected rather than adding a second packing model.

Results

Bounded sampler-support replay lets top-$p$ / top-$k$ sampling be used during rollouts without breaking trainer/inference alignment. These curves come from an RL run training GLM-4.7-Flash (an MoE model) on the easy split of NVIDIA's open Nemotron-Terminal-Synthetic-Tasks dataset — a sandboxed terminal/coding agent environment. We compare routed-expert replay alone against routed-expert replay plus support-set replay with top-$k$=20, top-$p$=0.95 sampling. Training reward tracks closely between the two arms, while the top-$p$ arm retains substantially more policy entropy instead of collapsing. Curves are lightly EMA-smoothed over the raw per-step values; both arms are clipped to the step range they share (11–151).

Mean training reward	Policy entropy

Reward is nearly indistinguishable between the arms, but by the end of training the top-$p$ arm holds entropy around ~0.14 while the R3-only arm collapses to ~0.05.

That retained entropy shows up as sampling diversity at evaluation. The two arms are essentially tied on Pass@1, but the higher-entropy top-$p$ arm produces more diverse completions and pulls ahead at higher $k$:

Arm	Pass@1	Pass@4	Pass@8
R3 only	0.514	0.605	0.631
R3 + top-$p$	0.512	0.620	0.641

So preserving support during rollouts trades a negligible Pass@1 difference for a consistent Pass@4 / Pass@8 gain — the extra entropy widens coverage exactly where multi-sample metrics benefit.

Testing

• pytest tests/backends/skyrl_train/utils/test_sample_support_replay.py tests/backends/skyrl_train/inference_servers/test_vllm_sample_support.py tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py tests/backends/skyrl_train/inference_servers/test_remote_inference_client.py tests/backends/skyrl_train/test_token_based_batching_utils.py tests/backends/skyrl_train/test_train_batch.py tests/train/dataset/test_sample_support_preprocess.py tests/train/generators/test_generator_output_utils.py tests/train/generators/test_skyrl_gym_generator.py tests/train/test_config.py tests/train/test_trainer.py — all passing (synthetic-EOS value/gradient parity, no-EOS rows, packed CP segments, fused chunk bounds, and preprocessing invariants).
• Ruff and Black clean on all changed files.

[gemini-code-assist]

Warning

Gemini encountered an error creating the review. You can try again by commenting /gemini review.