NVIDIA / NVIDIA/Megatron-LM

[BUG] reduce_aux_losses_tracker_accross_ranks hangs if first pipeline stage has no moe layers

Open
#1,444 6 comments 0 reactions 0 assignees View on GitHub
bug community-request module: moe waiting-on-customer
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.