NVIDIA / NVIDIA/Megatron-LM

[BUG] Possible Deadlock Isue on A2A_Overlap overlap-moe-expert-parallel-comm on large runs with uneven stages in pipeline layout

Open
#1,810 8 comments 0 reactions 1 assignee Claimed by @Wohox View on GitHub
bug community-request module: moe
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 3h
Merged PRs (30d)
272

Description

Deadlock when `--overlap-moe-expert-parallel-comm` is enabled (single shared CUDA event causes NCCL order mismatch)

## Summary
Enabling `--overlap-moe-expert-parallel-comm` can deadlock at scale. The overlap scheduler currently uses **one shared CUDA event** as a “baton” for *both* compute and comm streams across all nodes in a model chunk. This lets **multiple comm ops become eligible at once** on each rank, so different ranks may launch collectives in **different orders** (e.g., `[C1, C2]` vs `[C2, C1]`). That violates NCCL’s “same collectives, same order per communicator” invariant and leads to a hang.

The deadlock appears to happen at the beginning of the backward pass.

## Affected code (current behavior)
- `megatron/core/pipeline_parallel/utils.py`
`stream_acquire_context(stream, event)` uses **the same event** for both `wait` and `record`.
- `megatron/core/models/common/model_chunk_schedule_plan.py`
A single `torch.cuda.Event()` is constructed and passed to **all** nodes (compute and comm). Both comm nodes (`moe_dispatch`, `moe_combine`) run on the comm stream and **both** gate on the same event.

This makes it possible for the run loop to enqueue **two comm sequences** like:
```
comm_stream: wait(E) ; NCCL #1 ; record(E)
comm_stream: wait(E) ; NCCL #2 ; record(E)
```
Once compute records `E`, *both* waits are satisfied and the two NCCL ops can start in **rank-local enqueue order**, which can differ across ranks → deadlock.

## Root cause (why it deadlocks)
- **NCCL invariant:** all ranks must call the **same collectives in the same order** for a given communicator.
- With a **single shared event**:
1. Compute records `E`.
2. Multiple comm nodes (e.g., `dispatch(k)` and `combine(k-1)` or `dispatch(k)` and `dispatch(k+1)`) were already queued behind `wait(E)`.
3. Tiny host/GPU timing differences cause **different enqueue orders** across ranks.
4. Each rank launches a different “first” collective → **order mismatch** → hang.

## Proposed fix
Replace the single baton with **two CUDA events** per model chunk and make the contract explicit:

- **Two events per chunk**
- `comp_event`: recorded by compute nodes.
- `comm_event`: recorded by comm nodes.

- **Alternating dependency**
- **Compute nodes:** `wait(comm_event)` → compute → `record(comp_event)`
- **Comm nodes:** `wait(comp_event)` → NCCL all-to-all → `record(comm_event)`

- **Seed once**
On init, record `comm_event` on the compute stream once so the first compute step isn’t blocked.

- **APIs**
- Change `stream_acquire_context` to `stream_acquire_context(stream, wait_event, record_event)`.
- `ScheduleNode(..., event_pair=(wait_event, record_event))` (remove single-event legacy).

This preserves compute/comm overlap but prevents “two comm ops eligible at once” races.

**To Reproduce**
Steps to reproduce the behavior. The easier it is to reproduce the faster it will get maintainer attention.

Train a large MoE (like the DeepSeek 671B using the default configs from Megatron Model Zoo on a large number of GPUs (128 nodes per the configuration). It will reliably deadlock when using the latest main when pretraining from scratch as the computation between experts is more unbalanced near the beginning of training. I am trying to craft a small repro that occurs on a small number of GPUs, but I have been unsuccessful so far. The issue occurs on H100s and H200s.

**Expected behavior**
A clear and concise description of what you expected to happen.

**Stack trace/logs**
If applicable, add the stack trace or logs from the time of the error.

**Environment (please complete the following information):**
- Megatron-LM commit ID - 15a0d4738f062b93b0f698dfc0a3b1bb63f26878
- PyTorch version - Nemo 25.07 container
- CUDA version - 12.9
- NCCL version - version associated with Nemo 25.07 container

**Proposed fix**
If you have a proposal for how to fix the issue state it here or link to a PR.

**Additional context**
Add any other context about the problem here.

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.