Repository navigation
Conversation
Packs were padded with token 0. Identical pad tokens get identical hidden states, so the MoE router sends every pad token of a pack to the same experts, and the expert-parallel rank that hosts them receives the whole pad run. On a large-context MoE SFT run, a pack that was mostly padding pushed that rank to its memory limit, where it stalled until the NCCL watchdog killed the job. Repeating the pack's own tokens routes padding like real text. Pads stay loss-masked and follow every real token, so the loss is unchanged. Multimodal packs keep token 0, so no placeholder id escapes its item. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
garrett361
added a commit
that referenced
this pull request
Oct 20, 2026
…oid MoE hot spots (#4184) This PR gives RL the same `pad_to_seq_len` flag that SFT supports (#4160). The flag pads every packed row to `seq_len`, so torch.compile and DeepGEMM stop recompiling on new lengths. Full-length padding makes an existing MoE problem worse (identical pad tokens all route to the same experts), so this PR also changes RL and SFT to pad with the row's own tokens instead of a constant id. - RL used to pad with token 1, and router replay padded with zero bytes, which sent every pad slot to expert 0. This PR pads text rows with their own tokens and replay with their real routes. - Multimodal rows keep a constant pad id, because a repeated image placeholder would have no image feature to receive. - This PR supersedes #3932 / #4149 (SFT pad content, credit Mika). Please keep the `Co-authored-by: Mika Senghaas` trailer when squashing. Related: #3350. We ran FP8 e2e checks on one H200 node per run, each with an empty DeepGEMM cache (harness in a31524e): | Run | Dynamo recompiles per rank, total (after step 1) | DeepGEMM compiles, total (after step 1) | | --- | --- | --- | | SFT GLM-4.5-Air 4L `--data.pad-to-seq-len` | 1 (0) | 22 (0) | | SFT GLM-4.5-Air 4L | 3 (0) | 34 (7) | | RL Qwen3-30B-A3B + replay `--orchestrator.pad-to-seq-len` | 0 (0) | 17 (0) | | RL Qwen3-30B-A3B + replay | 1 (0) | 54 (28) | The flag only makes the row length static: the number of micro batches per step, `max_seqlen`, and multimodal shapes still vary. The throughput and MFU metrics count padding, so they are not comparable with and without the flag. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Co-authored-by: Mika Senghaas <mika@primeintellect.ai>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fresh rebase of #3932 (by @mikasenghaas, closed as superseded by #3948, but #3948 only changed the pad length, not the pad content).
SFT packs are padded with token 0. Identical pad tokens get identical hidden states, so the MoE router sends every pad token to the same experts and the EP rank hosting them takes the whole pad run. This change fills the padding with the pack's own tokens (cycled), so pads route like real text. Pads stay loss-masked, so the loss is unchanged. Multimodal packs keep token 0, because placeholder ids must not appear outside their item.
On main, padding is now only up to the next multiple of
cp, so the pad runs are much shorter than when #3932 was opened. The change is still cheap and removes the routing hot spot.Original commits are cherry-picked with authorship preserved. The only conflict was in tests: the new test now passes
pad_to_multiple_of.tests/unit/train/sft/test_sft_dataset.py: 74 passed.Co-authored-by: Mika Senghaas mika@primeintellect.ai
Note
Low Risk
Training-data padding only; loss masking is preserved and multimodal packing keeps the old pad-id behavior.
Overview
SFT pack padding no longer fills alignment gaps with repeated token
0for text-only packs.CatDataset._finalize_packnow extendsinput_idsby cycling the pack’s real tokens (itertools.cycle/islice), so MoE routers see varied hidden states instead of routing every pad to the same experts (which could OOM an expert-parallel rank). Loss behavior is unchanged: padded positions still getloss_mask=Falseandtarget_ids=0.Multimodal packs are unchanged and still pad with
0, because repeating placeholder token ids would break vision feature scattering.A unit test asserts the cycled pad pattern, masks, and
position_idsfor a short text pack padded topad_to_multiple_of=8.Reviewed by Cursor Bugbot for commit c0095db. Bugbot is set up for automated code reviews on this repo. Configure here.