Star212
322
代码
代码
Issues68
Pull Requests92
项目讨论
Wiki
分析
Star212
322
已合并
update qwen2moe update for faster #2334
jzh6229创建于 2025年3月4日
update qwen2moe update for faster #2334
已合并
jzh6229创建于 2025年3月4日
从refs/pull/2334/head合入到master
共 4 个文件变更+314-5
Aexamples/mcore/qwen2_moe/pretrain_qwen2_57b_a14b_4k_A3_ptd.sh+149-0
@@ -0,0 +1,149 @@
1+#!/bin/bash
2+export HCCL_CONNECT_TIMEOUT=1800
3+export CUDA_DEVICE_MAX_CONNECTIONS=1
4+export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
5+ 
6+NPUS_PER_NODE=16
7+MASTER_ADDR=localhost #主节点IP
8+MASTER_PORT=6000
9+NNODES=4
10+NODE_RANK=0
11+WORLD_SIZE=$(($NPUS_PER_NODE*$NNODES))
12+ 
13+# please fill these path configurations
14+CKPT_SAVE_DIR="your model save ckpt path"
15+DATA_PATH="your data path"
16+TOKENIZER_PATH="your tokenizer path"
17+CKPT_LOAD_DIR="your model ckpt path"
18+ 
19+TP=1
20+PP=4
21+EP=16
22+CP=2
23+MBS=1
24+GBS=768
25+CP_TYPE='megatron_cp_algo'
26+SEQ_LENGTH=4096
27+TRAIN_ITERS=2000
28+ROUTER_BALANCING_TYPE='aux_loss'
29+ 
30+DISTRIBUTED_ARGS="
31+ --nproc_per_node $NPUS_PER_NODE \
32+ --nnodes $NNODES \
33+ --node_rank $NODE_RANK \
34+ --master_addr $MASTER_ADDR \
35+ --master_port $MASTER_PORT
36+"
37+ 
38+MOE_ARGS="
39+ --num-experts 64 \
40+ --moe-router-topk 8 \
41+ --n-shared-experts 8 \
42+ --shared-expert-gate \
43+ --moe-intermediate-size 2560 \
44+ --moe-router-load-balancing-type ${ROUTER_BALANCING_TYPE} \
45+ --moe-grouped-gemm \
46+ --moe-token-dispatcher-type alltoall \
47+ --moe-aux-loss-coeff 0.001 \
48+ --moe-permutation-async-comm \
49+ --moe-alltoall-overlap-comm \
50+ --use-fused-moe-token-permute-and-unpermute \
51+"
52+ 
53+OPTIMIZE_ARGS="
54+ --use-flash-attn \
55+ --use-fused-rotary-pos-emb \
56+ --use-rotary-position-embeddings \
57+ --use-fused-swiglu \
58+ --use-fused-rmsnorm \
59+ --no-masked-softmax-fusion \
60+ --use-distributed-optimizer \
61+ --overlap-grad-reduce \
62+ --overlap-param-gather
63+"
64+ 
65+TRAIN_ARGS="
66+ --micro-batch-size ${MBS} \
67+ --global-batch-size ${GBS} \
68+ --lr 1.25e-6 \
69+ --lr-decay-style cosine \
70+ --min-lr 1.25e-7 \
71+ --weight-decay 1e-1 \
72+ --lr-warmup-fraction 0.01 \
73+ --attention-dropout 0.0 \
74+ --init-method-std 0.01 \
75+ --hidden-dropout 0.0 \
76+ --clip-grad 1.0 \
77+ --adam-beta1 0.9 \
78+ --adam-beta2 0.95 \
79+ --initial-loss-scale 4096 \
80+ --seed 42 \
81+ --bf16 \
82+ --train-iters ${TRAIN_ITERS} \
83+ --seq-length ${SEQ_LENGTH} \
84+ --no-shared-storage
85+"
86+ 
87+MODEL_PARALLEL_ARGS="
88+ --sequence-parallel \
89+ --tensor-model-parallel-size ${TP} \
90+ --pipeline-model-parallel-size ${PP} \
91+ --expert-model-parallel-size ${EP} \
92+ --context-parallel-size ${CP} \
93+ --context-parallel-algo ${CP_TYPE}
94+"
95+ 
96+GPT_ARGS="
97+ --use-cp-send-recv-overlap \
98+ --use-fused-ring-attention-update \
99+ --no-gradient-accumulation-fusion \
100+ --gemm-gradient-accumulation-fusion \
101+ --reuse-fp32-param \
102+ --use-mcore-models \
103+ --tokenizer-name-or-path ${TOKENIZER_PATH} \
104+ --max-position-embeddings ${SEQ_LENGTH} \
105+ --noop-layers 4,5,30,31 \
106+ --num-layers 32 \
107+ --hidden-size 3584 \
108+ --ffn-hidden-size 18944 \
109+ --num-attention-heads 28 \
110+ --tokenizer-type PretrainedFromHF \
111+ --make-vocab-size-divisible-by 1 \
112+ --padded-vocab-size 151936 \
113+ --rotary-base 1000000 \
114+ --untie-embeddings-and-output-weights \
115+ --disable-bias-linear \
116+ --position-embedding-type rope \
117+ --normalization RMSNorm \
118+ --swiglu \
119+ --attention-softmax-in-fp32 \
120+ --add-qkv-bias \
121+ --group-query-attention \
122+ --num-query-groups 4
123+"
124+ 
125+DATA_ARGS="
126+ --data-path $DATA_PATH \
127+ --split 100,0,0
128+"
129+ 
130+OUTPUT_ARGS="
131+ --log-interval 1 \
132+ --save-interval ${TRAIN_ITERS} \
133+ --eval-interval ${TRAIN_ITERS} \
134+ --eval-iters 0 \
135+ --no-load-optim \
136+ --no-load-rng
137+"
138+ 
139+torchrun $DISTRIBUTED_ARGS pretrain_gpt.py \
140+ $GPT_ARGS \
141+ $DATA_ARGS \
142+ $MOE_ARGS \
143+ $OUTPUT_ARGS \
144+ $OPTIMIZE_ARGS \
145+ $TRAIN_ARGS \
146+ $MODEL_PARALLEL_ARGS \
147+ --load $CKPT_LOAD_DIR \
148+ --distributed-backend nccl \
149+ | tee logs/train_mcore_qwen2_57b_a14b_4k.log
Mmindspeed_llm/training/arguments.py+2-0
@@ -1096,6 +1096,8 @@ def _validate_moe_args(args):
1096 raise AssertionError('`--moe-zero-memory` only supports `--moe-alltoall-overlap-comm` for now.')1096 raise AssertionError('`--moe-zero-memory` only supports `--moe-alltoall-overlap-comm` for now.')
1097 if args.moe_zero_memory != "disable" and args.recompute_method is not None:1097 if args.moe_zero_memory != "disable" and args.recompute_method is not None:
1098 raise AssertionError('`--moe-zero-memory` does not support full recomputation for now.')1098 raise AssertionError('`--moe-zero-memory` does not support full recomputation for now.')
1099+ if args.shared_expert_gate and args.gradient_accumulation_fusion:
1100+ raise AssertionError('args.shared_expert_gate does not support gradient_accumulation_fusion.')
1099 1101 
1100 1102 
1101def _validate_mla(args):1103def _validate_mla(args):
Atests/poc/qwen2_moe/pretrain_qwen2_57b_a14b_4k_64die_A3.sh+152-0
@@ -0,0 +1,152 @@
1+#!/bin/bash
2+export HCCL_CONNECT_TIMEOUT=1800
3+export CUDA_DEVICE_MAX_CONNECTIONS=1
4+export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
5+ 
6+NPUS_PER_NODE=16
7+MASTER_ADDR=localhost #主节点IP
8+MASTER_PORT=6000
9+NNODES=4
10+NODE_RANK=0
11+WORLD_SIZE=$(($NPUS_PER_NODE*$NNODES))
12+ 
13+# please fill these path configurations
14+CKPT_SAVE_DIR="your model save ckpt path"
15+DATA_PATH="your data path"
16+TOKENIZER_PATH="your tokenizer path"
17+CKPT_LOAD_DIR="your model ckpt path"
18+ 
19+TP=1
20+PP=4
21+EP=16
22+CP=1
23+MBS=1
24+GBS=768
25+CP_TYPE='megatron_cp_algo'
26+SEQ_LENGTH=4096
27+TRAIN_ITERS=2000
28+ROUTER_BALANCING_TYPE='aux_loss'
29+ 
30+DISTRIBUTED_ARGS="
31+ --nproc_per_node $NPUS_PER_NODE \
32+ --nnodes $NNODES \
33+ --node_rank $NODE_RANK \
34+ --master_addr $MASTER_ADDR \
35+ --master_port $MASTER_PORT
36+"
37+ 
38+MOE_ARGS="
39+ --num-experts 64 \
40+ --moe-router-topk 8 \
41+ --n-shared-experts 8 \
42+ --shared-expert-gate \
43+ --moe-intermediate-size 2560 \
44+ --moe-router-load-balancing-type ${ROUTER_BALANCING_TYPE} \
45+ --moe-grouped-gemm \
46+ --moe-token-dispatcher-type alltoall \
47+ --moe-aux-loss-coeff 0.001 \
48+ --moe-permutation-async-comm \
49+ --moe-alltoall-overlap-comm \
50+ --use-fused-moe-token-permute-and-unpermute \
51+"
52+ 
53+OPTIMIZE_ARGS="
54+ --use-flash-attn \
55+ --use-fused-rotary-pos-emb \
56+ --use-rotary-position-embeddings \
57+ --use-fused-swiglu \
58+ --use-fused-rmsnorm \
59+ --no-masked-softmax-fusion \
60+ --use-distributed-optimizer \
61+ --overlap-grad-reduce \
62+ --overlap-param-gather
63+"
64+ 
65+TRAIN_ARGS="
66+ --micro-batch-size ${MBS} \
67+ --global-batch-size ${GBS} \
68+ --lr 1.25e-6 \
69+ --lr-decay-style cosine \
70+ --min-lr 1.25e-7 \
71+ --weight-decay 1e-1 \
72+ --lr-warmup-fraction 0.01 \
73+ --attention-dropout 0.0 \
74+ --init-method-std 0.01 \
75+ --hidden-dropout 0.0 \
76+ --clip-grad 1.0 \
77+ --adam-beta1 0.9 \
78+ --adam-beta2 0.95 \
79+ --initial-loss-scale 4096 \
80+ --seed 42 \
81+ --bf16 \
82+ --train-iters ${TRAIN_ITERS} \
83+ --seq-length ${SEQ_LENGTH} \
84+ --no-shared-storage
85+"
86+ 
87+MODEL_PARALLEL_ARGS="
88+ --sequence-parallel \
89+ --tensor-model-parallel-size ${TP} \
90+ --pipeline-model-parallel-size ${PP} \
91+ --expert-model-parallel-size ${EP} \
92+ --context-parallel-size ${CP} \
93+ --context-parallel-algo ${CP_TYPE}
94+"
95+ 
96+GPT_ARGS="
97+ --fix-router \
98+ --load $CKPT_LOAD_DIR \
99+ --no-gradient-accumulation-fusion \
100+ --gemm-gradient-accumulation-fusion \
101+ --reuse-fp32-param \
102+ --recompute-granularity full \
103+ --recompute-method block \
104+ --recompute-num-layers 2 \
105+ --use-mcore-models \
106+ --tokenizer-name-or-path ${TOKENIZER_PATH} \
107+ --max-position-embeddings ${SEQ_LENGTH} \
108+ --noop-layers 4,5,30,31 \
109+ --num-layers 32 \
110+ --hidden-size 3584 \
111+ --ffn-hidden-size 18944 \
112+ --num-attention-heads 28 \
113+ --tokenizer-type PretrainedFromHF \
114+ --make-vocab-size-divisible-by 1 \
115+ --padded-vocab-size 151936 \
116+ --rotary-base 1000000 \
117+ --untie-embeddings-and-output-weights \
118+ --disable-bias-linear \
119+ --position-embedding-type rope \
120+ --normalization RMSNorm \
121+ --swiglu \
122+ --attention-softmax-in-fp32 \
123+ --add-qkv-bias \
124+ --no-gradient-accumulation-fusion \
125+ --group-query-attention \
126+ --num-query-groups 4
127+"
128+ 
129+DATA_ARGS="
130+ --data-path $DATA_PATH \
131+ --split 100,0,0
132+"
133+ 
134+OUTPUT_ARGS="
135+ --log-interval 1 \
136+ --save-interval ${TRAIN_ITERS} \
137+ --eval-interval ${TRAIN_ITERS} \
138+ --eval-iters 0 \
139+ --no-load-optim \
140+ --no-load-rng
141+"
142+ 
143+torchrun $DISTRIBUTED_ARGS pretrain_gpt.py \
144+ $GPT_ARGS \
145+ $DATA_ARGS \
146+ $MOE_ARGS \
147+ $OUTPUT_ARGS \
148+ $OPTIMIZE_ARGS \
149+ $TRAIN_ARGS \
150+ $MODEL_PARALLEL_ARGS \
151+ --distributed-backend nccl \
152+ | tee logs/train_mcore_qwen2_57b_a14b_4k.log
Mtests/st/shell_scripts/qwen2_moe_tp1_pp2_ep2_cp2_32k.sh+11-5
@@ -25,7 +25,7 @@ CP=2
25SEQ_LENGTH=3276825SEQ_LENGTH=32768
26TRAIN_ITERS=1526TRAIN_ITERS=15
27CP_TYPE='ulysses_cp_algo'27CP_TYPE='ulysses_cp_algo'
28-ROUTER_BALANCING_TYPE='softmax_topk'28+ROUTER_BALANCING_TYPE='pai_megatron_aux_loss'
29 29 
30DISTRIBUTED_ARGS="30DISTRIBUTED_ARGS="
31 --nproc_per_node $NPUS_PER_NODE \31 --nproc_per_node $NPUS_PER_NODE \
@@ -44,19 +44,23 @@ MOE_ARGS="
44 --moe-intermediate-size 320 \44 --moe-intermediate-size 320 \
45 --moe-grouped-gemm \45 --moe-grouped-gemm \
46 --moe-permutation-async-comm \46 --moe-permutation-async-comm \
47- --moe-token-dispatcher-type allgather \47+ --moe-alltoall-overlap-comm \
48- --moe-aux-loss-coeff 0.00148+ --use-fused-moe-token-permute-and-unpermute \
49+ --moe-token-dispatcher-type alltoall \
50+ --moe-aux-loss-coeff 0.001 \
51+ --reuse-fp32-param \
49"52"
50 53 
51OPTIMIZE_ARGS="54OPTIMIZE_ARGS="
52- --use-mc2 \
53 --use-flash-attn \55 --use-flash-attn \
54 --use-fused-rotary-pos-emb \56 --use-fused-rotary-pos-emb \
55 --use-rotary-position-embeddings \57 --use-rotary-position-embeddings \
56 --use-fused-swiglu \58 --use-fused-swiglu \
57 --use-fused-rmsnorm \59 --use-fused-rmsnorm \
58 --no-masked-softmax-fusion \60 --no-masked-softmax-fusion \
59- --use-distributed-optimizer61+ --use-distributed-optimizer \
62+ --overlap-grad-reduce \
63+ --overlap-param-gather
60"64"
61 65 
62TRAIN_ARGS="66TRAIN_ARGS="
@@ -83,6 +87,7 @@ TRAIN_ARGS="
83"87"
84 88 
85MODEL_PARALLEL_ARGS="89MODEL_PARALLEL_ARGS="
90+ --sequence-parallel \
86 --tensor-model-parallel-size ${TP} \91 --tensor-model-parallel-size ${TP} \
87 --pipeline-model-parallel-size ${PP} \92 --pipeline-model-parallel-size ${PP} \
88 --expert-model-parallel-size ${EP} \93 --expert-model-parallel-size ${EP} \
@@ -102,6 +107,7 @@ GPT_ARGS="
102 --make-vocab-size-divisible-by 1 \107 --make-vocab-size-divisible-by 1 \
103 --padded-vocab-size 151936 \108 --padded-vocab-size 151936 \
104 --rotary-base 1000000 \109 --rotary-base 1000000 \
110+ --attention-softmax-in-fp32 \
105 --untie-embeddings-and-output-weights \111 --untie-embeddings-and-output-weights \
106 --disable-bias-linear \112 --disable-bias-linear \
107 --position-embedding-type rope \113 --position-embedding-type rope \