NVIDIA / NVIDIA/Megatron-LM

[BUG] Expert Parallelism deadlock with Gemma4 (EP=4, DP≥16)

Open
#5,749 10 comments 0 reactions 1 assignee Claimed by @YangFei1990 View on GitHub
bug community-request waiting-on-maintainers
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

## Description

When fine-tuning `Gemma4-26B-A4B-it` model with Expert Parallelism (EP=4) on multi-node clusters (64–128 GPUs), all ranks deadlock after 92–852 training steps on a large-scale dataset with multi-image samples. The hang occurs in the NCCL allreduce/reduce-scatter kernel during gradient synchronization, as the EP alltoall collective for MoE token dispatch conflicts with DP gradient reduce-scatter on the same NCCL stream.

The deadlock seems to be caused by the alltoall token dispatcher. An identical configuration with TP=4/EP=1 (eliminating the alltoall collective) completes 1000 steps without hanging on the exact same setup, but suffers from slow step times (7.54s/step v/s 3.44s/step) so EP≥4 is desirable.

## Root Cause Analysis

Progressive patching across 4 runs identified the failure site:

| Run | Patches applied | Hung at step | Blocking site |
| -- | -- | -- | -- |
| v0 | `logical_and_across_model_parallel_group` bypassed | 534 | first CPU sync after stuck GPU |
| v1 | * `reduce_max_stat` bypassed | 720 | next CPU sync |
| v2 | * `check_for_nan_in_grad` disabled | 852 | `finalize_model_grads` (grad allreduce) |
| v3 | all Python sync points removed | 584 | GPU completely stuck in prior allreduce |

The non-deterministic step (534/584/719/852) confirms a timing-dependent NCCL deadlock, not data or iteration-count dependent. py-spy captures show all 128 ranks blocked symmetrically in `logical_and_across_model_parallel_group` (the first Python-level CPU sync that discovers the stuck GPU).

Another experiment on the same cluster, dataset, world_size (128 GPUs), and dependency stack but with `--expert_model_parallel_size 1 --tensor_model_parallel_size 4` (no EP alltoall) completes 1000/1000 steps without hanging. This confirms the source of deadlock to be alltoall collective used for EP.

### Mechanism

The `alltoall` token dispatcher dispatches tokens to EP-remote experts via `torch.distributed.all_to_all`. With EP=4 and `use_distributed_optimizer=True`, the DP gradient reduce-scatter uses `model_parallel_group` communicators that partially overlap with EP communicators. Over hundreds of steps, GPU scheduling jitter causes ranks to enter their respective reduce-scatter collectives out of sync. When one rank starts a new bucket's collective while another is still in alltoall on the same NCCL stream → deadlock.

## Reproduction

### Environment

* megatron-core: 0.17.1 (also tested 0.18.1 — same deadlock, hangs earlier)
* Training framework: [ms-swift]() 4.3.1 + [mcore-bridge]() 1.5.2 (thin integration layer that calls megatron-core's distributed training APIs)
* PyTorch: 2.8.0a0 (NGC 25.05-py3)
* NCCL: 2.24 (NGC 25.05 bundled); also tested 2.21.5 (NGC 24.10) — same result
* Transformer Engine: 2.3.0 (NGC 25.05 bundled)
* GPU: H200 141GB SXM × 128 (16 nodes × 8 GPUs, EFA 3.2 Tbps)
* Also reproduced on: H200 × 64 (8 nodes × 8 GPUs)

### Model

* Gemma4-26B-A4B-it: 30 layers, 128 experts, 25.8B parameters, hybrid sliding-window + global attention

### Parallelism Configuration

```
TP=1, EP=4, PP=1, DP=32
micro_batch_size=1, global_batch_size=128 (grad_accum=4)
use_distributed_optimizer=True
moe_token_dispatcher_type=alltoall
```

### Minimum Command to Reproduce

```bash
NPROC_PER_NODE=8 NNODES=16 MASTER_ADDR= \
megatron sft \
--model \
--dataset \
--tensor_model_parallel_size 1 \
--expert_model_parallel_size 4 \
--pipeline_model_parallel_size 1 \
--micro_batch_size 1 \
--global_batch_size 128 \
--attention_backend auto \
--recompute_granularity full \
--recompute_method uniform \
--recompute_num_layers 1 \
--finetune true \
--max_length 20480 \
--train_iters 1000
```

`megatron sft` is the [ms-swift]() CLI for Megatron-backend SFT. The key megatron-core settings are `expert_model_parallel_size=4`, `use_distributed_optimizer=True` (default), and `moe_token_dispatcher_type=alltoall` (default). Any training harness that exercises these codepaths with DP≥16 should reproduce. Training starts normally (3.44 s/step with fused attention) but hangs after 92–852 steps depending on version and timing.

### What Works

* `--expert_model_parallel_size 1 --tensor_model_parallel_size 4 --sequence_parallel true` → **completes 1000 steps** (7.54 s/step, no alltoall) but suffers from slow step times (v/s 3.44s/step for EP=4/TP=1).

### What Doesn't Work

* EP=4 with `alltoall` on 128 GPUs (16 nodes): hangs at step 92–852 (5 separate runs)
* EP=4 with `alltoall` on 64 GPUs (8 nodes): hangs at step 92–139 (2 separate runs)
* megatron-core 0.17.1: hangs at step 852 (3/3 runs, same step each time with same version)
* megatron-core 0.18.1: hangs at step 257 (bucket draining changes timing, triggers race earlier)
* `--overlap_grad_reduce false`: still hangs at step 852
* `--moe_router_num_groups 4 --moe_router_group_topk 1` (balanced routing): hangs at step 103 (worse)

### What Was Ruled Out

| Hypothesis | How tested | Result |
| -- | -- | -- |
| NCCL version | Tested 2.21.5 and 2.24 | Both deadlock |
| Data/seed dependent | 3 different seeds | Same outcome |
| Scale-specific (128 GPU only) | Tested at 64 GPUs (8 nodes) | Also deadlocks |
| `overlap_grad_reduce` timing | `--overlap_grad_reduce false` | Still hangs |
| Token routing asymmetry | Balanced routing (`moe_router_num_groups=4`) | Hung faster |
| Training framework Python sync points | Patched out all 5 blocking `.item()` syncs in training loop | Shifts step, doesn't fix |
| megatron-core bucket management | 0.18.1 `previous_grad_reduce_bucket_group` draining | Hung earlier |
| Distributed optimizer | Can't disable (OOM: 206 GiB Adam states) | Mandatory for 26B |
| EP group size reduction | EP=1 OOMs, EP=2 marginal OOM | Can't reduce |

## Related Issues

* NVIDIA/Megatron-LM#5100 PP deadlock with `variable_seq_lengths=True` and `batch_p2p_comm=True`. Same symptom class (collective ordering conflict causing NCCL deadlock at scale) but different codepath (P2P shape exchange v/s EP alltoall + DP reduce-scatter).

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.