NVIDIA / NVIDIA/Megatron-LM

[QUESTION] Slow pretraining performance for MoE model on HPC cluster with A100 GPUs

Open
#1,841 2 comments 0 reactions 0 assignees View on GitHub
community-request module: moe question
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Is this training speed normal? Am I making some big mistakes?**
I am experiencing slow pretraining performance while training a Mixture-of-Experts (MoE) LLM using Megatron-LM on a HPC cluster with NVIDIA A100 GPUs (64GB VRAM). I am using 4 nodes × 4 GPUs for testing, but I plan to scale up to 256 nodes. I also tried with 64 nodes, but the performance did not improve and slightly decreased.

After 1100 iteration the logs report:
iteration 1100/1562 | consumed samples: 140800 | elapsed time per iteration (ms): 28837.4 | throughput per GPU (TFLOP/s/GPU): 19.8 | learning rate: 6.918239E-06 | global batch size: 128 | lm loss: 5.367071E+00 | loss scale: 1.0 | grad norm: 3.296 | num zeros: 82575440.0 | params norm: 720.186 | number of skipped iterations: 0 | number of nan iterations: 0

I am concerned about the efficiency and potential bottlenecks, I'm planning to scale up to 256 nodes.
I considered a peak capacity of 312 TFLOP/s/GPU, the effective utilization (MFU) reached 6.3%.
Environment:
- GPUs: NVIDIA A100 64GB
- Singularity image: nvcr.io/nvidia/pytorch:25.06-py3
- Megatron Core: 0.15.0rc5
- HPC Cluster: multi-node (4 nodes × 4 GPUs in test)

```
MODEL_ARCHITECTURE="\
--num-layers 24 \
--hidden-size 2880 \
--ffn-hidden-size 8640 \
--num-attention-heads 64 \
--num-query-groups 8 \
--group-query-attention \
--seq-length 4096 \
--max-position-embeddings 32768 \
--disable-bias-linear \
--swiglu \
--attention-dropout 0.0 \
--hidden-dropout 0.0 \
--qk-head-dim 128 \
--v-head-dim 128"

MOE_ARGS="\
--moe-router-topk 8 \
--num-experts 64 \
--moe-ffn-hidden-size 1080 \
--moe-grouped-gemm \
--moe-layer-freq 1"

DISTRIBUTED_TRAINING="\
--sequence-parallel \
--tensor-model-parallel-size 1 \
--pipeline-model-parallel-size 2 \
--context-parallel-size 1 \
--num-layers-per-virtual-pipeline-stage 12 \
--expert-model-parallel-size 8 \
--expert-tensor-parallel-size 1"

BATCH_SIZE=" \
--micro-batch-size 1 \
--global-batch-size 128"

TRAINING_SCHEDULE="\
--train-samples 200000 \
--lr-warmup-samples 407040 \
--lr-decay-samples 166809600"

LEARNING_RATE="\
--lr 2.0e-5 \
--min-lr 2.0e-6 \
--lr-decay-style cosine"

OPTIMIZER="\
--weight-decay 0.1 \
--adam-beta1 0.9 \
--adam-beta2 0.95 \
--clip-grad 1.0 \
--loss-scale 1.0"

INITIALIZATION="\
--init-method-std 0.006 \
--seed ${RANDOM}"

DATA="\
--train-data-path ${TRAIN_DATA} \
--tokenizer-type HuggingFaceTokenizer \
--tokenizer-model /tokenizer/mistral-nemo"

CHECKPOINTING="\
--save-interval 200 \
--save ${CHECKPOINT_DIR}"

EVALUATION="\
--eval-interval ${EVAL_INTERVAL} \
--eval-iters -1"

LOGGING="\
--log-interval 1 \
--log-params-norm \
--log-num-zeros-in-grad \
--log-throughput \
--log-validation-ppl-to-tensorboard \
--tensorboard-dir ${TENSORBOARD_DIR}"

MEMORY_OPTIMIZATION="\
--recompute-activations \
--use-flash-attn"
```

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.