Repository navigation
revert(trainer): compile around the checkpoint again for MoE blocks with an fp32 router - #4037
Merged
Merged
Conversation
…ith an fp32 router This reverts commit f3827c4 (#3925). #3925 wrapped the compiled block in an eager checkpoint, AC(compile(block)), for checkpointed blocks holding a nested FSDP unit (the fp32 router). With that layout the checkpoint recompute re-enters Dynamo on the autograd thread, and MoE RL runs with context parallelism, activation checkpointing and compile crash in an early backward with "aten.empty_like.default / aten.empty.memory_format encountered during backward but not found in storage". Reverting #3925 removes the crash on the same batches. The root cause is not yet understood, so restore the previous compile(AC(block)) layout and AC -> Compile -> FSDP order until it is. This gives back #3925's speedup: MoE blocks with an fp32 router and activation checkpointing run eager again. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
garrett361
marked this pull request as ready for review
October 12, 2026 13:53
hallerite
approved these changes
Oct 12, 2026
This was referenced Oct 22, 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.
Reverts #3925. With an fp32 router and activation checkpointing, #3925 changed checkpointed MoE blocks from
compile(AC(block))toAC(compile(block)), which makes the checkpoint recompute re-enter Dynamo on the autograd thread. Since then, MoE RL runs with context parallelism, activation checkpointing and compile crash in an early backward withaten.empty_like.defaultoraten.empty.memory_format encountered during backward but not found in storage. That happens with both ring and ulysses CP, which shows that reverting #3925 removes the crash on the same batches).The runs in #3925 (no CP) never hit this, and the root cause is not yet understood. This PR restores the previous layout and the
AC -> Compile -> FSDPorder until it is. The cost is #3925's speedup: MoE blocks with an fp32 router and AC run eager again (GLM-4.5-Air about 5.2 to 6.2 s/step, Qwen3-30B-A3B 4.5 to 6.0 s/step in #3925's benchmarks). A follow-up will root-cause the crash and decide whether to re-land #3925 together with #3951.This conflicts with #4029, which edits the unwrap block removed here.
🤖 Generated with Claude Code