fix: stride the pad-token replay fill across experts

[EazyReal]

Summary

Routing replay reuses stored top-k selections during training. When replay data is padded to the training token layout, fully padded rows are represented as all -1; those sentinels must be replaced with in-range indices before downstream gather and dispatch.

The existing repair writes the same arange(topk) % n_cols row into every fully padded row. For MoE routing, n_cols is the number of experts, so every pad token goes to experts 0..topk-1. Loss masking prevents those rows from affecting the objective, but it does not skip the forward expert dispatch: with topk <= experts-per-rank, rank 0 receives topk x pad_count synthetic assignments (and larger top-k values concentrate them on the first EP ranks).

This surfaced as a trainer OOM on a 128x H100 DeepSeek-V4-Flash run at 262,144-token context. Modeling the constant fill predicted a hot-rank receive count of 719,029 tokens versus 718,848 observed at the OOM site (0.025% error).

Fix

Number only the fully invalid rows and fill them by pad-row ordinal:

fill[pad_row, col] = (pad_row * topk + col) % n_cols

This makes the synthetic assignments independent of where padding appears in the batch. That distinction matters for BSHD batches, where per-sample padding can recur at regular token positions and an absolute token-position stride can alias back to the same expert block. The arange is also sized to the invalid rows rather than the full [n_tokens, topk] tensor, and the replay tensor is cloned before assignment.

Fully padded rows are loss-masked, so any in-range selection is semantically valid. Selections remain distinct within each row whenever topk <= n_cols. Partially invalid rows remain unchanged for the indexer path, where -1 has separate meaning; for fully padded indexer rows, n_cols denotes KV positions and the same fill remains in range.

Verification

• The regression test uses noncontiguous pad rows at the same per-sample position. It fails both the constant fill and an absolute token-position stride, and pins the pad-row-ordinal result.
• The existing partial-padding test verifies that rows containing valid indexer selections are not rewritten.
• python -m pytest -q tests/fast/utils/test_replay_base.py tests/fast/backends/test_fsdp_replay_routers.py: 13 passed on CPU-only torch.
• ruff, isort, and black pass on both changed files.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.