# [BUG] Pipeline Parallelism deadlock with `variable_seq_lengths=True` and `batch_p2p_comm=True` (PP>=4)
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
## Description
When training large MoE models (e.g., GLM-5.1 with DSA attention, 78 layers, 256 experts) using Pipeline Parallelism (PP=4) with `variable_seq_lengths=True` (required for sequence packing/padding-free mode), all ranks deadlock during the first training step.
The hang occurs in `_communicate_shapes()` in `p2p_communication.py`, where tensor shapes are exchanged between PP stages before actual data transfer.
## Root Cause Analysis
The P2P operation ordering in `_communicate_shapes()` and `_batched_p2p_ops()` uses a fixed order:
```python
# _communicate_shapes (line 235-255) and _batched_p2p_ops (line 27-47)
ops = [send_prev, recv_prev, send_next, recv_next]
```
However, the correct implementation `_p2p_ops()` (line 79+) uses **rank-dependent ordering** with even/odd alternation:
```python
# _p2p_ops - correct implementation
if group.rank() % 2 == 0:
# send_next, recv_prev, send_prev, recv_next
else:
# recv_prev, send_next, recv_next, send_prev
```
Additionally, `_p2p_ops()` uses separate process groups (`even_send_odd_recv_group` vs `even_recv_odd_send_group`) when `group.size() == 2` to allow overlapping independent communications.
The fixed ordering in `_communicate_shapes()` and `_batched_p2p_ops()` causes send/recv pairs between adjacent PP stages to mismatch, leading to a circular-wait deadlock — especially with PP>=4 where the 1F1B schedule has complex communication patterns during steady-state to cooldown transitions.
## py-spy Analysis
Dumping stacks across multiple nodes during the hang reveals different ranks stuck at different communication phases:
**Most ranks** — stuck in shape exchange:
```
irecv (distributed_c10d.py:2536)
batch_isend_irecv (distributed_c10d.py:2854)
_communicate_shapes (p2p_communication.py:257)
_communicate (p2p_communication.py:323)
send_forward_recv_backward (p2p_communication.py:545)
forward_backward_pipelining_without_interleaving (schedules.py:2324)
```
**Some ranks** — stuck in data sync (past shape exchange):
```
synchronize (cuda/__init__.py:1108)
_communicate (p2p_communication.py:419)
send_forward_recv_backward (p2p_communication.py:545)
```
**Other ranks** — already in cooldown phase:
```
synchronize (cuda/__init__.py:1108)
_communicate (p2p_communication.py:419)
send_backward_recv_forward (p2p_communication.py:578)
forward_backward_pipelining_without_interleaving (schedules.py:2510)
```
This confirms that different PP stages are stuck in different communication phases — some still exchanging shapes while others have moved past to data transfer or cooldown, creating a deadlock.
## Reproduction
### Environment
* megatron-core: 0.18.0 (main branch, commit 3f6a2ed84)
* mcore-bridge: 1.5.0.dev0
* ms-swift: 4.3.0.dev0
* PyTorch: 2.10.0
* NCCL: 2.27.5
* GPU: H800 80GB × 192 (24 nodes × 8 GPUs)
### Model
* GLM-5.1 (ZhipuAI): 78 layers, 256 experts, \~700B total parameters, DSA attention
### Parallelism Configuration
```
TP=4, EP=8, PP=4, SP=True
micro_batch_size=1, global_batch_size=48
```
### Minimum Command to Reproduce
```bash
megatron sft \
--model \
--dataset \
--tensor_model_parallel_size 4 \
--expert_model_parallel_size 8 \
--pipeline_model_parallel_size 4 \
--decoder_last_pipeline_num_layers 18 \
--moe_grouped_gemm true \
--sequence_parallel true \
--packing true \
--micro_batch_size 1 \
--global_batch_size 48 \
--recompute_granularity full \
--recompute_method uniform \
--recompute_num_layers 1 \
--finetune true \
--max_length 8192
```
### What works
* `--packing false` (disables `variable_seq_lengths` path, avoids `_communicate_shapes`) → **training succeeds**
* PP=1 → no PP communication needed → works (but OOM on 80GB cards for this model)
### What doesn't work
* `--packing true` with PP>=4 → **deadlock at first training step**
* Tested with megatron-core 0.17.0 and 0.18.0 (main branch) — both deadlock
* Tested with dev branch — also deadlocks (same `_communicate_shapes` code)
## Related Issues
* NVIDIA/Megatron-LM#1450 — Same P2P ordering bug reported for PP=2 + VPP=2 with `variable_seq_lengths=True`. PR [#1451]() proposed a fix (reorder to `send_prev, recv_next, send_next, recv_prev`) but was **never merged**.
* [#1485]() — Follow-up PR, also closed without merge.
## Proposed Fix
Align `_communicate_shapes()` and `_batched_p2p_ops()` with the same even/odd rank-dependent ordering already used in `_p2p_ops()`. Alternatively, make `_communicate_shapes()` use `_p2p_ops()` internally instead of `batch_isend_irecv`.
A simpler workaround: allow users to set `batch_p2p_comm=False` so that all P2P communication (both shape exchange and data transfer) uses `_p2p_ops()`. Currently `_communicate_shapes()` ignores the `batch_p2p_comm` flag and always uses `batch_isend_irecv`.
Contributor guide
Assessment
This issue has not been assessed yet.