[BUG] reduce_aux_losses_tracker_accross_ranks hangs if first pipeline stage has no moe layers
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 3h
- Merged PRs (30d)
- 272
Description
**Describe the bug**
https://github.com/NVIDIA/Megatron-LM/blob/8a5521ac4226fbefeeb2a102ebecac32a01d4852/megatron/core/transformer/moe/moe_utils.py#L586-L588
`reduce_aux_losses_tracker_across_ranks` do all_reduce accross ` _PIPELINE_MODEL_PARALLEL_GROUP`. If some pipeline stage has no moe layers, all_reduce will hangs.
**To Reproduce**
modeling with:
```shell
--tensor-model-parallel-size 1
--pipeline-model-parallel-size 8
--expert-model-parallel-size 1
--expert-tensor-parallel-size 1
--num-layers 16
--moe-layer-freq "([0]*3+[1]*13)"
```
will hangs because fisrt pp stage has no tracker info while other stage has tracker info like `{'load_balancing_loss': {'values': tensor([...])}}`
**Expected behavior**
First pp stage should have zero padding values.
**Stack trace/logs**
N/A
**Environment (please complete the following information):**
- Megatron-LM commit ID: 8a5521ac4226fbefeeb2a102ebecac32a01d4852
- PyTorch version: 2.5.1
- CUDA version: 12.4
- NCCL version: 2.21.5
**Proposed fix**
N/A
**Additional context**
N/A
Contributor guide
Research direction
Read reduce_aux_losses_tracker_across_ranks in megatron/core/transformer/moe/moe_utils.py around lines 586-588, then reproduce the hang with the provided tensor, pipeline, expert, layer, and moe-layer-freq settings. Done means the first pipeline stage supplies zero padding values and the cross-rank reduction completes without hanging.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 55/100