NVIDIA / NVIDIA/Megatron-LM

[BUG] Load Balancing loss discrepancy with/without CUDA Graphs

Open
#1,494 6 comments 1 reaction 0 assignees View on GitHub
bug community-request waiting-on-customer
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Describe the bug**
I am seeing discrepancy in reduction values with and without CUDA graphs turned on. Can you also please look into why that discrepancy happens?

Changing [this](https://github.com/NVIDIA/Megatron-LM/blob/core_r0.11.0/megatron/core/transformer/moe/moe_utils.py#L579) to the following (to add print statements):
```
def reduce_aux_losses_tracker_across_ranks():
"""Collect and reduce the auxiliary losses across ranks."""
tracker = parallel_state.get_moe_layer_wise_logging_tracker()
aux_losses = {k: v['values'].float() for k, v in tracker.items()}
if torch.distributed.get_rank() == 0:
print(f"first {aux_losses=}")
for name in tracker:
values = tracker[name]["values"]
# Collect aux losses across PP.
torch.distributed.all_reduce(
values, group=parallel_state.get_pipeline_model_parallel_group()
)
if torch.distributed.get_rank() == 0:
print(f"{values=}")
# Reduce aux losses across ranks.
if tracker[name].get('reduce_group') is not None:
torch.distributed.all_reduce(values, group=tracker[name].get('reduce_group'))
if torch.distributed.get_rank() == 0:
print(f"reduced vals {values=}")
if tracker[name].get('avg_group') is not None:
torch.distributed.all_reduce(
values, group=tracker[name]['avg_group'], op=torch.distributed.ReduceOp.AVG
)
```

I see the following results with and without cuda graphs enabled:
with CUDA graphs:
```
first aux_losses={'z_loss': tensor([0.0000, 1.0703, 0.0000, 0.9629, 0.8457, 0.6562, 0.9121, 1.1836],
device='cuda:0'), 'load_balancing_loss': tensor([0.0000, 1.0056, 0.0000, 1.0083, 1.0372, 1.1118, 1.0406, 1.0901],
device='cuda:0')}
values=tensor([0.0000, 1.0703, 0.0000, 0.9629, 0.8457, 0.6562, 0.9121, 1.1836],
device='cuda:0')
values=tensor([0.0000, 1.0056, 0.0000, 1.0083, 1.0372, 1.1118, 1.0406, 1.0901],
device='cuda:0')
```

without CUDA graphs:
```
first aux_losses={'z_loss': tensor([0.0000, 1.0703, 0.0000, 0.9629, 0.8457, 0.6592, 0.9102, 1.1797],
device='cuda:0'), 'load_balancing_loss': tensor([0.0000, 1.0056, 0.0000, 1.0083, 1.0370, 1.1117, 1.0406, 1.0903],
device='cuda:0')}
values=tensor([0.0000, 1.0703, 0.0000, 0.9629, 0.8457, 0.6592, 0.9102, 1.1797],
device='cuda:0')
values=tensor([0.0000, 1.0056, 0.0000, 1.0083, 1.0370, 1.1117, 1.0406, 1.0903],
device='cuda:0')
reduced vals values=tensor([0.0000, 2.0175, 0.0000, 2.0102, 2.0742, 2.2586, 2.1155, 2.1638],
device='cuda:0')
```

As can be seen, there is a doubling of the final reduced values btw the 2 implementations. This doubling is with TP=2. I assume that the value will be 'n' times depending on the TP value.

I enabled CUDA graphs using the following flag:
```
use_te_rng_tracker: True
enable_cuda_graph: True
```

**To Reproduce**
I am using the `from megatron.core.models.gpt import GPTModel` with the following moe configs:

```
num_layers: 8
moe_layer_freq: [0, 1, 0, 1, 1, 1, 1, 1]
moe_use_legacy_grouped_gemm: False
use_te_rng_tracker: True
enable_cuda_graph: True
```
and the following 5-d parallel configs:
```
context_parallel_size: 1
tensor_model_parallel_size: 2
pipeline_model_parallel_size: 1
num_moe_experts: 2
expert_model_parallel_size: 2
```

I hope that helps with the repro. Happy to provide any additional specific attribute values.

For the non graph version, I unset the `use_te_rng_tracker` and `enable_cuda_graph` to False.

**Expected behavior**
The load balancing loss should be the same btw the 2 settings.

**Environment (please complete the following information):**
Version info:
Megatron-LM.git@aa6207e
NeMo.git@7192a2c
TransformerEngine.git@8eb1712

**Additional context**
- What I have found out uptil now is that this is because when CUDA graphs is enabled, the code never comes within the `if` condition [here](https://github.com/NVIDIA/Megatron-LM/blob/core_r0.11.0/megatron/core/transformer/moe/moe_utils.py#L590), which makes me believe that there is some discrepancy in the groups [here](https://github.com/NVIDIA/Megatron-LM/blob/core_r0.11.0/megatron/core/transformer/moe/router.py#L250)

- To provide more context, the numbers for the vals are same for the first batch, but then start differing from the 2nd batch onwards. I presume this is because in the 1st batch, CUDA graphs arent used.

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.