Repository navigation
perf(trainer): compile MoE blocks with fullgraph by keeping EP experts in the block's FSDP unit - #3752
Closed
MarioSieg wants to merge 4 commits into
Closed
perf(trainer): compile MoE blocks with fullgraph by keeping EP experts in the block's FSDP unit#3752MarioSieg wants to merge 4 commits into
MarioSieg wants to merge 4 commits into
Conversation
FA4's flash_attn_varlen_func is a Python autograd.Function that calls torch._guards.active_fake_mode(), which Dynamo refuses to trace. Under activation checkpointing each block is one checkpointed region, so that graph break drops the whole block to eager for every FA4 model, dense or MoE. Wrap FA4's varlen forward/backward in prime_rl_attn::flash_attn_4_varlen (with fake impls and autograd) and add the namespace to the selective checkpointing save list so attention is not recomputed in backward.
…_funcs Keep _compute_attention identical to main and register the custom op as the FA4 entry of _funcs, with the same call signature and (out, lse) return value as flash_attn_varlen_func. Code that reads or swaps self.func / _funcs keeps working, and CP (ring, ulysses) still replaces _compute_attention as before.
… compile with fullgraph Under expert parallelism the experts were their own FSDP unit nested inside each transformer block. FSDP2 wraps its hooks in torch._dynamo.disable, so dynamo graph-broke at the experts and the MoE dispatch/combine glue and the AC policy ran eagerly, which the Qwen3-235B profile showed as ~2.7 s/step of unfused elementwise and index kernels. Shard the expert parameters over the EP-complement mesh from the block's own shard_placement_fn via ShardPlacementResult instead. FSDP2 groups the block's parameters by process group, so the experts still all-gather and reduce-scatter over dp_shard_mod_ep, and the block-level gradient divide factor is the dense parameters' default divisor. fullgraph=true now needs moe_router_dtype=bfloat16 on custom MoE models: the fp32 router is still its own FSDP unit, and setup_fsdp fails early with that message.
…ck compiles with fullgraph The fused forward was a custom op wrapped in a Python autograd.Function that saved x.to(bfloat16).contiguous() for backward. When the block compiles with fullgraph under activation checkpointing, dynamo rejects that saved tensor because it aliases the autograd.Function's input inside the checkpoint subgraph. Make the backward a custom op too and register autograd on prime_rl::mega_moe_forward, saving the op's own inputs. Both ops take the caller's dtypes, run the kernels in bf16, and return gradients in the inputs' dtypes, so the expert dW dtype is unchanged.
Member
Author
|
Closing: splitting into independent PRs against main instead of stacking on the Mega MoE branch. |
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.
Stacked on #3651 (Mega MoE). Includes the two FA4 custom-op commits from #3676, which fullgraph needs.
Why
The Qwen3-235B SFT profile against torchtitan put prime-rl's vanilla MoE dispatch at 7.0 s/step vs torchtitan's 5.0 s. ~2.7 s of that was unfused eager MoE dispatch/combine glue plus a per-op Python AC policy: the block could not compile with
fullgraph=Truebecause the EP experts were a nested FSDP unit inside each block, and FSDP2 wraps its hooks intorch._dynamo.disable. torchtitan avoids this by sharding the experts over the EP-complement mesh from the block's ownfully_shardcall.What
shard_placement_fnreturns aShardPlacementResultwith thedp_shard_mod_epmesh for expert parameters. FSDP2 (torch 2.13) groups a unit's parameters by process group, so the experts still all-gather/reduce-scatter over the EP-complement mesh. The block-level gradient divide factor is the dense parameters' default divisor, so dense numerics are unchanged. Prefetch lists drop the experts module.prime_rl::mega_moe_backwardop) instead of a Pythonautograd.Function, which dynamo rejected inside the checkpoint region (saved input alias, then a buffer-registry write; the buffer is now registered at construction).fullgraph=truenow requiresmoe_router_dtype="bfloat16"on custom MoE models: the fp32 router is still its own FSDP unit.setup_fsdpraises early with that message; the fp32-router default is unchanged.Benchmark
Qwen3-235B-A22B-Thinking-2507 SFT, 4 nodes x 8 B300, FSDP over 32 GPUs, ep=8, seq 16384 (one packed sample per GPU per step), balanced routing, AdamW fp32 states, full AC unless noted, 20 steps, steady state = steps 5..20. Same MFU formula in every row (true head_dim, 280.8 GFLOP/token). New runs use
fullgraph=true,moe_router_dtype=bfloat16.Per-step losses match the pre-change runs to 3 decimals over the first steps (same data order). A 1-node Qwen3-30B-A3B fullgraph smoke run (job 1453) and
tests/unit/train/models/test_moe.pypass. Full tables:outputs/bench_20261002_q235_4n_s16k_fullgraph.mdon the bench machine.Open
perf.py/sft/train.pyis not part of this PR.