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

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:

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
• 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

Start with the pretrain_gpt.py entry point and reproduce the reported configuration using TP size 4, context parallel size 2, FP8, and TransformerEngine 1.13.0. Compare the layernorm weight backward gradient across single-device, TP-only, CP-only, and TP+CP runs; done means the combined configuration matches the single-device oracle without changing main_grad behavior.