[QUESTION] Slow pretraining performance for MoE model on HPC cluster with A100 GPUs
- 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
Assessment
This issue has not been assessed yet.