fix(r3): align routed-expert metadata and padding

[dyurk-lila]

Summary

Fix routed-expert replay (R3) correctness for RL training, and introduce a shared token-metadata layout that later routed-expert and sampler-support work builds on. This fixes a RouterReplay lifecycle bug that has been latent in the R3 implementation since it was added; recent padding work only made it easier to trigger.

This PR is the first in a small stack of routed-expert / sampler-support improvements; it is the correctness base the rest build on.

Results

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. Enabling routed-expert replay both improves mean training reward and sharply reduces the gap between rollout (inference) logprobs and the trainer's recomputed logprobs: routes actually taken during sampling are replayed at train time instead of being re-derived, so training and inference stay aligned. Curves are lightly EMA-smoothed over the raw per-step values.

Mean training reward	Rollout vs. training logprob gap

With replay on, reward pulls ahead over the back half of training while the mean absolute rollout-vs-train logprob difference drops to roughly a third of the baseline (from ~0.023 to ~0.008 by end of training).

Where the bugs arose

1. Forward-only R3 leaked backward replay state into the next schedule

The R3 path calls setup_per_microbatch_replay_forward() from both forward-only logprob passes and training passes. Megatron's set_target_indices() appends every target tensor to replay_backward_list, which activation recomputation later consumes in FIFO order. A forward-only schedule has no backward phase, so its FIFO entries were never consumed, and clear_router_replay() was not called at schedule boundaries.

The next training backward could therefore consume routes from an earlier forward-only microbatch. This is data-dependent: equal aligned token counts let the wrong routes pass silently, while different counts can leave fewer routing-map assignments than Megatron's dropless all-to-all expects and fail with Split sizes doesn't match total dim 0 size.

Fix: scope global RouterReplay state to one Megatron pipeline schedule. Clear once before the schedule and again in finally after success or failure. Do not clear between training microbatches, whose backward FIFO intentionally spans the schedule.

2. vLLM routes describe a captured prefix, not every training token

vLLM records routes for tokens actually executed by the rollout model. The final trajectory can be longer: the last sampled token has no later decode forward, a synthetic EOS may be appended, and multi-turn generation can append observations. Intermediate observations are captured when the next turn replays the prefix; a terminal uncaptured suffix is not.

The old tensorization gave Megatron no way to distinguish an uncaptured row from a real route. The PPO loss mask does not solve this: the router and expert-bias accounting do not read it, and loss-masked observations that condition later actions still need their captured routes replayed.

Fix: keep rollout_expert_indices ragged and use its length as the captured-prefix length. Derive router_padding_mask after left padding, marking only left/alignment padding and the uncaptured suffix. Carry that mask through TrainingInput, replay experiences, microbatch padding, and the Megatron model call.

Synthetic padding rows use distinct dummy experts [0, ..., topk - 1]. The mask excludes them from expert-bias accounting, while distinct indices preserve Megatron's dropless tokens * topk dispatcher invariant without a dispatcher patch.

3. RL packing needs the same row layout for routes and masks

With microbatch size greater than one, Megatron's RL packing produces layouts such as [seq0, pad0, seq1, pad1], where each input row is one sequence and each row receives TP/CP alignment padding.

Fix: build one TokenMetadataLayout per microbatch and apply it to both routes and router_padding_mask. The layout owns sequence lengths, row/segment alignment, and CP front/back placement. Generic construction, alignment, next-token shifting, and packed-output restoration live in skyrl/utils/token_metadata.py; RouterReplay-specific installation stays in replay_utils.py.

4. Expert-bias padding mask mis-broadcasts

TopKRouter._apply_expert_bias combines a [tokens, experts] routing map with a one-dimensional [tokens] mask, which does not broadcast over experts. Freezing router parameters does not disable the separate expert-bias token counters, so dummy and uncaptured rows would corrupt bias updates when expert bias is enabled.

Fix: pass Megatron's padding_mask through the model and apply one narrow compatibility shim that reshapes [tokens] to [tokens, 1] before calling the original method. GPTModel already scatters this mask beside its embedding on the first PP stage; HybridModel scatters only the embedding, so the mask is scattered explicitly. Intermediate PP stages are scattered to match sequence-parallel hidden states.

5. Dynamic sampling could detach routes from their trajectory

Replacement copied a hard-coded subset of GeneratorOutput fields, while filtering rebuilt a different hard-coded subset, so rollout_expert_indices could remain from a rejected sample or disappear.

Fix: slice every per-trajectory generator field generically for replacement and filtering, keeping route metadata, vision fields, trajectory IDs, and future per-row fields owned by the selected sample.

Deliberately not included

• No rollout_expert_num_captured_tokens field; len(rollout_expert_indices[i]) is authoritative.
• No sorting or deduplication of captured real-token routes. Only synthetic padding rows are constructed with deterministic distinct experts.
• No change to synthetic-EOS loss semantics.

Testing

• pytest tests/utils/test_token_metadata.py
• pytest tests/backends/skyrl_train/utils/test_replay_utils.py tests/backends/skyrl_train/test_train_batch.py tests/backends/skyrl_train/test_token_based_batching_utils.py
• pytest tests/train/dataset/test_preprocess.py tests/train/generators/test_skyrl_gym_generator.py tests/train/test_trainer_utils.py

Regressions cover captured prefixes, loss-masked observations, microbatch > 1 row packing, CP/SP layout transforms, packed next-token restoration, mocked GPTModel/HybridModel mask handling, expert-bias accounting, dynamic replacement/filter ownership, dummy batch padding, and route preprocessing. The GPU/Megatron path is exercised by tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_router_replay.py.

[gemini-code-assist]

Warning

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