[BUG] MOE + external-cuda-graph-score-attn + not playing nicely with overlap-grad-reduce
- 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
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