NVIDIA / NVIDIA/Megatron-LM

[BUG] Wrong attention gradient in Transformer Engine

Open
#1,503 4 comments 0 reactions 1 assignee Claimed by @asolergi-nv View on GitHub
bug community-request module: transformer engine waiting-on-customer
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Describe the bug**

The attention gradient dQ calculation in TE is wrong.

When training GPT using Megatron with `--transformer-impl transformer_engine` on same parameters and data, the single card result of dQ and tensor-parallel result of dQ do not align. The relative error of dQ is as large as 0.5 for a 24 layer GPT model, and can reach 0.9 for a 128 layer GPT model.

Then when checking with the correct `--transformer-impl local` baseline implementation, both TE single card and TE tensor-parallel's dQ result has a relative error 1.1 compared with the baseline `local`'s dQ result. They are all trained with the same parameters and data.

The calling path for Transformer Engine implementation is `DotProductAttention.forward` -> `FusedAttnFunc.apply` -> `FusedAttnFunc.backward` -> `fused_attn_bwd`.

**To Reproduce**

Run the following script for TE single card training. The TP version just adds torchrun-related distributed arguments and sets `--tensor-model-parallel-size 2`. The `local` baseline just sets `--transformer-impl local`.

```bash
ARGS="
--tensor-model-parallel-size 1 \
--pipeline-model-parallel-size 1 \
--context-parallel-size 1 \
--num-layers 24 \
--hidden-size 1024 \
--num-attention-heads 16 \
--seq-length 1024 \
--max-position-embeddings 1024 \
--attention-dropout 0.0 \
--hidden-dropout 0.0 \
--attention-softmax-in-fp32 \
--bf16 \
--clip-grad 1.0 \
--micro-batch-size 4 \
--global-batch-size 16 \
--lr 0.00015 \
--min-lr 1.0e-5 \
--train-iters 1 \
--lr-warmup-fraction 0.01 \
--lr-decay-iters 320000 \
--lr-decay-style cosine \
--weight-decay 1e-2 \
--use-mcore-models \
--no-gradient-accumulation-fusion \
--transformer-impl transformer_engine \
--data-path /workspace/dataset/wikitext_text_document \
--vocab-file /workspace/dataset/gpt2-vocab.json \
--merge-file /workspace/dataset/gpt2-merges.txt \
--split 949,50,1 \
--log-interval 100 \
--save-interval 10000 \
--eval-interval 1000 \
--eval-iters 0 \
--save /workspace/checkpoints \
--load /workspace/checkpoints
"

torchrun \
/home/ubuntu/repos/Megatron-LM/pretrain_gpt.py \
$ARGS

```

**Expected behavior**

The gradient of attention (dQ) in the TE implementation should be the same as single card when TP is on, and should be the same with the `local` implementation.

When it comes to bf16, the relative error of the correctly calculated tensor should be at the magnitude of 1e-2, but the current implementation produces wrong results whose relative error is more than 1 when comparing with correct results.

**Stack trace/logs**

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

**Environment (please complete the following information):**
- Megatron-LM commit ID: [`11996c9f`](https://github.com/NVIDIA/Megatron-LM/commit/11996c9f)
- PyTorch version: 2.5.1+cu124
- CUDA version: 12.4
- NCCL version: 2.21.5
- TransformerEngine: both 1.13.0 and 2.1.0 from pip install has this problem (`pip install "transformer_engine[pytorch]"`)
- cuDNN: 9.8.0.87-1

**Proposed fix**
See the function call path above.

**Additional context**
Please fix it soon. This is serious bug and has been for a long time.

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.