[BUG] Inconsistent LayerNorm Parameter Gradient with TP+CP+FP8
- 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
Assessment
This issue has not been assessed yet.