[BUG] Load Balancing loss discrepancy with/without CUDA Graphs
- 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
Assessment
This issue has not been assessed yet.