[QUESTION]Memory increase issue introduced by Megatron-FSDP during training.
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
I used the same Qwen3 script to train with Megatron-FSDP on the core_v0.14.0 and core_v0.16.0 branches of Megatron-LM, and found that the core_v0.16.0 branch consumes one-third more memory than core_v0.14.0. Can you provide any explanation for this?
qwen3脚本
torchrun \
--nproc_per_node 8 \
--nnodes 1 \
--node_rank 0 \
--master_addr localhost \
--master_port 29600 pretrain_gpt.py \
--data-path /data/qwen-data/openweb-qwen_text_document \
--tokenizer-model /data/models/Qwen2-57B-A14B \
--num-layers 48 \
--hidden-size 2048 \
--num-attention-heads 32 \
--ffn-hidden-size 6144 \
--num-query-groups 4 \
--group-query-attention \
--attention-softmax-in-fp32 \
--distributed-backend nccl \
--attention-dropout 0 \
--hidden-dropout 0 \
--use-rotary-position-embeddings \
--rotary-base 1000000 \
--num-experts 128 \
--moe-router-topk 8 \
--kv-channels 128 \
--moe-token-dispatcher-type alltoall \
--moe-ffn-hidden-size 768 \
--moe-router-load-balancing-type aux_loss \
--moe-aux-loss-coeff 0.001 \
--untie-embeddings-and-output-weights \
--swiglu \
--normalization RMSNorm \
--qk-layernorm \
--disable-bias-linear \
--lr 5e-6 \
--lr-decay-style cosine \
--min-lr 5e-7 \
--adam-beta1 0.9 \
--adam-beta2 0.95 \
--adam-eps 1e-5 \
--weight-decay 0.1 \
--clip-grad 1 \
--lr-warmup-iters 0 \
--tokenizer-type HuggingFaceTokenizer \
--split 949,50,1 \
--data-cache-path ./cache \
--log-throughput \
--log-interval 1 \
--save-interval 10000 \
--eval-interval 1000 \
--no-load-optim \
--no-load-rng \
--make-vocab-size-divisible-by 100 \
--position-embedding-type rope \
--no-position-embedding \
--use-distributed-optimizer \
--use-flash-attn \
--norm-epsilon 1e-6 \
--no-save-optim \
--lr-decay-iters 1000 \
--moe-grouped-gemm \
--bf16 \
--train-iters 10 \
--eval-iters 1 \
--tensor-model-parallel-size 1 \
--pipeline-model-parallel-size 1 \
--context-parallel-size 1 \
--expert-model-parallel-size 1 \
--seq-length 4096 \
--max-position-embeddings 40960 \
--micro-batch-size 1 \
--global-batch-size 32 \
--transformer-impl transformer_engine \
--moe-router-pre-softmax \
--no-gradient-accumulation-fusion \
--use-megatron-fsdp \
--data-parallel-sharding-strategy optim_grads \
--ckpt-format fsdp_dtensor \
--use-distributed-optimizer \
--num-layers 6 \
--manual-gc \
--manual-gc-interval 300
Contributor guide
Assessment
This issue has not been assessed yet.