Repository navigation
fix(models): pass contiguous router-replay experts to compiled layers - #3947
Merged
Merged
Conversation
Standalone torch script for the torch 2.13.0 bug where a compiled function returns an alias of a misaligned input at storage offset 0 once automatic dynamic shapes recompile it. Prints BUG REPRODUCED on torch 2.13.0 and is clean on torch >= 2.14 (pytorch/pytorch#191002). Needs one CUDA device. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
With router replay, each MoE model handed its compiled decoder layer the strided view `routed_experts[:, :, layer_idx]`. When top_k is not a multiple of 4, that int32 view starts at a non-16-byte-aligned offset for some layers. Inductor clones misaligned inputs, and once automatic dynamic shapes kick in, AOTAutograd regenerates the returned alias of the input at storage offset 0, so those layers silently replay layer 0's experts. On DeepSeek V4 (top_k 6) this made compiled trainer logprobs diverge from eager from the second step on (per-token |dlogprob| 0.13-0.31 vs a 0.025 floor). A contiguous slice is a fresh, aligned allocation and removes it. glm4_moe already did this. Applied to every model that passed the view. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
The reproducer was only needed to document the torch bug behind the contiguous routed-experts fix. Keep it reachable in history instead of on main. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
garrett361
marked this pull request as ready for review
October 8, 2026 19:03
garrett361
changed the title
mikasenghaas
approved these changes
Oct 8, 2026
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.
On torch 2.13.0 with compile and router replay, some MoE layers silently replay layer 0's routing instead of their own from the second training step on. This PR calls
.contiguous()on each layer's routed-experts slice, which removes the bug.Each decoder layer
lreceivesrouted_experts[:, :, l], a view into one shared int32 tensor with axes (batch, sequence, layer, expert slot). Whentop_kis not a multiple of 4, that view starts at a byte offset that is not 16-byte aligned for some layers (odd layers on DeepSeek V4,top_k = 6).Inductor copies a misaligned input into a fresh aligned buffer, which starts at offset 0. Once automatic dynamic shapes recompile the layer, AOTAutograd mistakenly applies that offset 0 to the original shared tensor when it rebuilds the layer's routed-experts output. Offset 0 of the shared tensor holds layer 0's experts. A contiguous slice is its own aligned tensor, so Inductor never makes the copy.
glm4_moealready called.contiguous().top_k = 8(Qwen3-30B-A3B, Trinity-Mini, GLM-4.5-Air) were unaffected.DeepSeek V4 Flash RL probe (compile on, router replay on, lr 0), mismatch KL mean, lower is better:
With compile off, the same probe spans 0.0008-0.0032 over steps 1-10.
Verification
A standalone torch reproducer was added in 66641b0 and removed again at the end of this branch, so main never carries it. Run it on one CUDA device from that commit with
TORCHINDUCTOR_FORCE_DISABLE_CACHES=1 uv run python tools/repro_misaligned_replay_alias.py. In earlier runs it printedBUG REPRODUCEDon torch 2.13.0 and was clean on the 2.14.0.dev20260809 and 2.16.0.dev20261008 nightlies.🤖 Generated with Claude Code