NVIDIA / NVIDIA/Megatron-LM

[BUG] MOE + external-cuda-graph-score-attn + not playing nicely with overlap-grad-reduce

Open
#1,667 0 comments 0 reactions 0 assignees View on GitHub
bug community-request module: moe
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 3h
Merged PRs (30d)
272

Description

**Describe the bug**
A clear and concise description of what the bug is.

I want to try to reduce the CPU overhead of my MOE. However, I want it to be dropless so I rely on capturing just the attn scope.

I only recompute the MOE/MLP layers, purposefully leaving out the core-attn layer.

I have the following flags enabled
```bash
--cuda-graph-scope=attn
--external-cuda-graph
--te-rng-tracker
```

I also have the following flags enabled, but I don't think are relevant here. I purposefully avoid any recomputation in core_attn. Does cudagraph does not work with ANY selective activation here? If so, that should be documented:
FYI @ko3n1g

```bash
--recompute-activations
--recompute-granularity selective #moe mlp #full
--recompute-modules mlp moe
```

I finally also have
```bash
--overlap-grad-reduce
--use-distributed-optimizer
```

```bash
--tensor-model-parallel-size 1
--pipeline-parallel-size 8
```
is also enabled if that has any effect here.

**To Reproduce**
Steps to reproduce the behavior. The easier it is to reproduce the faster it will get maintainer attention.

**Expected behavior**
A clear and concise description of what you expected to happen.
CUDAGRAPH works with overlapping grad-reduce or at least gives a better error.

**Stack trace/logs**
If applicable, add the stack trace or logs from the time of the error.

```
[rank239]: pretrain(
[rank239]: File "/mnt/weka/home/aaron.gokaslan/git/Megatron-LM-main/megatron/training/training.py", line 864, in pretrain
[rank239]: iteration, num_floating_point_operations_so_far = train(
[rank239]: ^^^^^^
[rank239]: File "/mnt/weka/home/aaron.gokaslan/git/Megatron-LM-main/megatron/training/training.py", line 2278, in train
[rank239]: ) = train_step(
[rank239]: ^^^^^^^^^^^
[rank239]: File "/mnt/weka/home/aaron.gokaslan/git/Megatron-LM-main/megatron/training/training.py", line 1395, in train_step
[rank239]: losses_reduced = forward_backward_func(
[rank239]: ^^^^^^^^^^^^^^^^^^^^^^
[rank239]: File "/mnt/weka/home/aaron.gokaslan/git/Megatron-LM-main/megatron/core/pipeline_parallel/schedules.py", line 1982, in forward_backward_pipelining_without_interleaving
[rank239]: input_tensor_grad = backward_step(
[rank239]: ^^^^^^^^^^^^^^
[rank239]: File "/mnt/weka/home/aaron.gokaslan/git/Megatron-LM-main/megatron/core/pipeline_parallel/schedules.py", line 409, in backward_step
[rank239]: custom_backward(output_tensor[0], output_tensor_grad[0])
[rank239]: File "/mnt/weka/home/aaron.gokaslan/git/Megatron-LM-main/megatron/core/pipeline_parallel/schedules.py", line 160, in custom_backward
[rank239]: Variable._execution_engine.run_backward(
[rank239]: File "/mnt/weka/home/aaron.gokaslan/git/Megatron-LM-main/megatron/core/distributed/distributed_data_parallel.py", line 501, in hook
[rank239]: param.grad is not None
[rank239]: AssertionError: param.grad being None is not safe when overlap_grad_reduce is True
```

**Environment (please complete the following information):**
- Megatron-LM commit ID - dc65034354a21447ce5ea99cb0e1f37fd2a2e864
- PyTorch version - NGC Docker container 26.05
- CUDA version - Same as above
- NCCL version - Same above

**Proposed fix**
If you have a proposal for how to fix the issue state it here or link to a PR.

**Additional context**
Add any other context about the problem here.

Contributor guide

Open the contributing guide

Research direction

Reproduce the configuration using the listed CUDA-graph, selective-recompute, and overlap-grad-reduce flags. Start at megatron/core/distributed/distributed_data_parallel.py:501 and trace the backward path through megatron/core/pipeline_parallel/schedules.py; done means the configuration no longer asserts on a missing gradient or reports a clearer incompatibility.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.