DeepSpeed MoE + ZeRO-3 Hang. Missing Leaf Module Marking

DeepSpeed ZeRO-3 with Mixture-of-Experts (MoE) models hangs during training because MoE blocks need to be marked as leaf modules for ZeRO-3 parameter gathering. Without this marking, ZeRO-3 attempts to gather parameters recursively through the MoE router, causing deadlocks. Denpex detects the MoE+ZeRO-3 hang pattern and recommends the leaf module fix.

DeepSpeed MoE + ZeRO-3 Hang. Missing Leaf Module Marking is a Distributed Training failure seen during ML training runs. DeepSpeed ZeRO-3 with Mixture-of-Experts (MoE) models hangs during training because MoE blocks need to be marked as leaf modules for ZeRO-3 parameter gathering. Without this marking, ZeRO-3 attempts to gather parameters recursively through the MoE router, causing deadlocks. Denpex detects the MoE+ZeRO-3 hang pattern and recommends the leaf module fix. Common tags: Deepspeed, Moe, Zero3, Leaf Module.

ZeRO-3 parameter gathering traverses the model hierarchy recursively, entering MoE router logic during the gather phase. MoE routers have conditional parameter access (only active experts per token), which creates deadlocks when ZeRO-3 tries to gather all expert parameters simultaneously. Without leaf module marking, ZeRO-3 doesn't know to treat MoE blocks as atomic units that should not be recursively decomposed. The hang occurs because rank 0 is waiting for rank 1's expert parameters while rank 1 is stuck in the MoE router. Taken together, these mechanisms explain why the failure is reproducible, why it tends to surface on specific workloads or scales, and why generic mitigation attempts often fall short without addressing the underlying cause.

• Training with DeepSpeed ZeRO-3 and MoE models hangs at the first forward pass
• NCCL timeout occurs during parameter gathering for MoE layers
• The hang is specific to MoE models. Non-MoE models with the same ZeRO-3 config work fine

• DeepSpeed ZeRO-3 with Mixture-of-Experts models (Mixtral, Switch Transformer, etc.)
• DeepSpeed with expert-parallel training
• Models using Megablocks or ScatterMoE with DeepSpeed ZeRO-3
• DeepSpeed versions before v0.14.0

Use this checklist to test the hypothesis against a small reproduction. No single line proves the root cause, so preserve the preceding events and compare one variable at a time.

• ✓Reproduce the failure from a clean checkpoint/seed: the symptom must appear without warm-up state from a previous run.
• ✓Verified signal present: Training hangs with NCCL timeout during the first forward pass through MoE layers
• ✓Verified signal present: nvidia-smi shows 0% GPU utilization on all ranks during the hang
• ✓Verified signal present: DeepSpeed logs show parameter gathering initiated but never completing for MoE parameters
• ✓Verified signal present: No error output. The job just hangs indefinitely
• ✓A targeted fix from the "How to fix it" section eliminates or substantially reduces the symptom within one validation pass.

Select the traceback or open the failed terminal, then run Denpex locally to see the initiating rank, collateral failures, exact fix, and verification command without uploading the log.

DeepSpeed changes when parameters, gradients and optimizer state are created, partitioned, gathered and offloaded. The hub separates ZeRO, memory, checkpoint and pipeline failures by lifecycle phase.

• ZeRO-3 parameter gathering traverses the model hierarchy recursively, entering MoE router logic during the gather phase
• MoE routers have conditional parameter access (only active experts per token), which creates deadlocks when ZeRO-3 tries to gather all expert parameters simultaneously
• Without leaf module marking, ZeRO-3 doesn't know to treat MoE blocks as atomic units that should not be recursively decomposed
• The hang occurs because rank 0 is waiting for rank 1's expert parameters while rank 1 is stuck in the MoE router