transformers

415e6d2f - Enable FSDP2 + expert parallelism via a 2-D (fsdp, tp) device mesh (#48516)

Enable FSDP2 + expert parallelism via a 2-D (fsdp, tp) device mesh (#48516) * Fix NaN gradients in expert-parallel training: mask uninitialized grouped_mm rows Under EP, sentinel token-expert slots sit beyond offsets[-1] and torch._grouped_mm leaves those output rows (fwd output and bwd d_input) uninitialized. The forward relied on a single post-mask plus a single pre-mask, letting NaN/Inf from uninitialized memory transit the activation and down-projection backward. The gate product's backward (act_fn(gate) * up) turns 0 x Inf into NaN (torch.autograd anomaly mode names this exact Mul), and it escapes into finite gradients: full fine-tuning of any EP-sharded MoE produced nan grad_norm on the second step (the first step survives only because freshly-allocated CUDA memory happens to be zeroed) and the loss collapsed to 0. Mask the sentinel-tail rows after each grouped GEMM instead. Full fine-tuning of OLMoE-1B-7B under ep=4 now matches the single-GPU loss trajectory. * Fix wrong gradients for all non-expert parameters in expert-parallel training Under EP the router hook zeroes the routing scores of non-local experts, so in backward each rank's score gradient covers only the slots of its local experts, and nothing sums the per-rank partial gradients: the gate weights and, through the gate's input, every parameter upstream of each MoE block receive gradients missing the contributions that flow through remote experts. The existing _AllReduceBackward on the experts' hidden input covers the dispatch branch, and the top_k_weights branch is explicitly skipped when is_expert_parallel -- but under EP it is exactly as partial as under TP-MoE. Measured against a single-GPU reference (OLMoE-1B-7B, one batch, fp32 so rounding noise vanishes): before the fix, 3/179 parameters agree (relative max-abs errors 0.3-2.5 on attention, norms, embeddings and router gates, 10-100x above the run-to-run noise floor; only the last layer's experts and the final norm -- the parameters backward reaches before crossing an expert block -- are correct). After the fix: 179/179 agree, max relative error 2.7e-5. Fix: allreduce-sum the score gradient in the EP router hook, before the non-local mask (each slot has exactly one owning rank, so the sum is exact). * Trim comments * Gate the router-score backward all-reduce on grad mode; drop the post-mask superseded by the per-mm masks * Fix expert parallelism through Trainer (#48208) * Enable FSDP2 + expert parallelism via a 2-D (fsdp, tp) device mesh DistributedConfig(tp_size=E, fsdp_size=D, enable_expert_parallel=True) builds a 2-D mesh: experts are sharded across tp, everything else is fully sharded across fsdp. The Trainer mirrors both dimensions into accelerate's ParallelismConfig, averages the expert gradients over fsdp (FSDP2 only reduces what it shards), computes the gradient norm across parameters on different meshes, and gathers the DTensor state dict on save. * Gather the sharded state dict through the model's own save helpers * Document combining expert parallelism with FSDP2 * Address review: drop the dead expert-gradient sync, torch-native mixed-mesh norm, single collective save FSDP2 composes over the tp-sharded experts and shards them across fsdp as well, so nothing is replicated over fsdp and the Trainer-side gradient averaging never ran; remove it and describe the actual layout. The gradient norm is now one get_total_norm per mesh, each reduced over its own mesh. save_model runs save_pretrained on every rank so its gather is collective and only rank 0 writes; the FSDP branch of the gather (full state dict on rank 0 only) now also covers the 2-D mesh. ParallelismConfig keeps a user-supplied config and only claims what the model was loaded with; pipeline parallelism is rejected together with tp/fsdp; optimizer checkpoints are refused for models sharded at load time since they cannot be resumed. * Give each device mesh its own optimizer param group instead of stepping per parameter Fused/foreach AdamW cannot span parameters on different meshes, but it can run per mesh: 189 ms/step per-parameter vs 41 ms fused per mesh group on Qwen3-30B-A3B at tp=4 x fsdp=2. * Test the 2-D mesh with a batch split across the fsdp ranks and compare gradient norms * Refresh the throughput table * Cosmetic: docstring layout, comment placement, doc link * Update docs/source/en/expert_parallelism.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/expert_parallelism.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/expert_parallelism.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/expert_parallelism.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Drop the comment on the parallel sizes * Trim the parallelism config comment --------- Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>