NVIDIA / NVIDIA/Megatron-LM

[BUG] Inconsistent LayerNorm Parameter Gradient with TP+CP+FP8

Open
#1,528 1 comment 0 reactions 0 assignees View on GitHub
bug module: transformer engine
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Describe the bug**

When using FP8 training with TP and CP, the gradient of layernorm weight is TP size times larger than the correct gradient. The gradient here is the gradient calculated in the backward propagation, not the `main_grad`.

When using TP alone or using CP alone, the gradient is consistent with the single-device oracle.

**To Reproduce**

Use TransformerEngine 1.13.0 (installed from pip), because 2.1.0 and 2.0.0 has bug on FP8 (see https://github.com/NVIDIA/TransformerEngine/issues/1585).

Turn on both TP and CP.

Training script:

```bash
ARCH_ARGS="
--num-layers 4 \
--hidden-size 128 \
--ffn-hidden-size 128 \
--num-attention-heads 4 \
--seq-length 16 \
--max-position-embeddings 16 \
--position-embedding-type rope \
--rotary-base 500000 \
--rotary-percent 1.0 \
--use-rope-scaling \
--rope-scaling-factor 32 \
--apply-query-key-layer-scaling \
--attention-dropout 0.0 \
--hidden-dropout 0.0 \
--normalization RMSNorm \
--make-vocab-size-divisible-by 128 \
--disable-bias-linear \
--no-bias-swiglu-fusion \
--no-gradient-accumulation-fusion \
"

PRECISION_ARGS="
--attention-softmax-in-fp32 \
--bf16 \
--fp8-format hybrid \
"

TRAINING_ARGS="
--clip-grad 1.0 \
--init-method-std 0.02 \
--micro-batch-size 2 \
--global-batch-size 8 \
--lr 0.00015 \
--min-lr 1.0e-5 \
--train-iters 1 \
--lr-decay-iters 320000 \
--lr-decay-style cosine \
--weight-decay 1e-2 \
--use-mcore-models \
--no-gradient-accumulation-fusion \
--transformer-impl transformer_engine \
"

IO_ARGS="
--data-path /workspace/dataset/wikitext/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/
"

DISTRIBUTED_ARGS="
--nproc_per_node $nproc_per_node \
--nnodes 1 \
--node_rank 0 \
--master_addr localhost \
--master_port 6000 \
"

PARALLEL_ARGS="
--tensor-model-parallel-size 4 \
--pipeline-model-parallel-size 1 \
--context-parallel-size 2 \
"

torchrun $DISTRIBUTED_ARGS \
/path/to/Megatron-LM/pretrain_gpt.py \
$PARALLEL_ARGS $ARCH_ARGS $PRECISION_ARGS $TRAINING_ARGS $IO_ARGS
```

**Expected behavior**

The gradient of layernorm weight when both CP and TP are on should be the same as single-device or the ones only with TP or the ones only with CP.

**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: [`1e1675d`](https://github.com/sunnyszy/Megatron-LM/commit/1e1675d81ac392af11b8afdf32aa9dbdef026fdb)
- PyTorch version: 2.5.1+cu124
- CUDA version: 12.4
- NCCL version: 2.21.5
- TransformerEngine: 1.13.0 from pypi

**Proposed fix**
Move the division of TP size before updating parameters with main grads to the backward calculation kernel.

**Additional context**
N/A

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.