已合并
[mindspore][master]add grok1 sh #3383
AtomGit-Bot创建于 2025年9月24日
[mindspore][master]add grok1 sh #3383
已合并
从refs/pull/3383/head合入到master
共 1 个文件变更+147-0
| @@ -0,0 +1,147 @@ | |||
| 1 | +#!/bin/bash | ||
| 2 | +# To check the performance of a Dropless MoE model, we should run the model for at least 500 iterations or resume from trained checkpoints. | ||
| 3 | +export CUDA_DEVICE_MAX_CONNECTIONS=1 | ||
| 4 | +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True | ||
| 5 | + | ||
| 6 | +NPUS_PER_NODE=8 | ||
| 7 | +MASTER_ADDR=localhost | ||
| 8 | +MASTER_PORT=6000 | ||
| 9 | +NNODES=4 | ||
| 10 | +NODE_RANK=0 | ||
| 11 | +WORLD_SIZE=$(($NPUS_PER_NODE*$NNODES)) | ||
| 12 | + | ||
| 13 | +CKPT_SAVE_DIR="your model save ckpt path" | ||
| 14 | +DATA_PATH="your data path" | ||
| 15 | +TOKENIZER_PATH="your tokenizer path" | ||
| 16 | +CKPT_LOAD_DIR="your model ckpt path" | ||
| 17 | + | ||
| 18 | +TP=8 | ||
| 19 | +PP=2 | ||
| 20 | +VPP=1 | ||
| 21 | +EP=2 | ||
| 22 | +CP=1 | ||
| 23 | +CP_TYPE=ulysses_cp_algo | ||
| 24 | +NUM_LAYERS=12 | ||
| 25 | +SEQ_LEN=8192 | ||
| 26 | +MBS=1 | ||
| 27 | +GBS=128 | ||
| 28 | + | ||
| 29 | +DISTRIBUTED_ARGS=" | ||
| 30 | + --master_addr $MASTER_ADDR \ | ||
| 31 | + --node_rank $NODE_RANK \ | ||
| 32 | + --worker_num $WORLD_SIZE \ | ||
| 33 | + --local_worker_num $NPUS_PER_NODE \ | ||
| 34 | + --master_port $MASTER_PORT \ | ||
| 35 | + --log_dir=msrun_log \ | ||
| 36 | + --join=False \ | ||
| 37 | + --cluster_time_out=300 \ | ||
| 38 | + --bind_core=True \ | ||
| 39 | +" | ||
| 40 | + | ||
| 41 | + | ||
| 42 | +MOE_ARGS=" | ||
| 43 | + --moe-permutation-async-comm \ | ||
| 44 | + --moe-token-dispatcher-type alltoall_seq \ | ||
| 45 | + --moe-grouped-gemm \ | ||
| 46 | + --num-experts 8 \ | ||
| 47 | + --moe-router-load-balancing-type aux_loss \ | ||
| 48 | + --moe-router-topk 2 \ | ||
| 49 | + --moe-aux-loss-coeff 1e-2 \ | ||
| 50 | + --embedding-multiplier-scale 78.38367176906169 \ | ||
| 51 | + --output-multiplier-scale 0.5773502691896257 \ | ||
| 52 | + --input-jitter \ | ||
| 53 | + --moe-layer-freq -1 \ | ||
| 54 | + --first-k-dense-replace -1 \ | ||
| 55 | +" | ||
| 56 | + | ||
| 57 | +GPT_ARGS=" | ||
| 58 | + --spec mindspeed_llm.tasks.models.spec.grok_spec layer_spec \ | ||
| 59 | + --use-mcore-models \ | ||
| 60 | + --tensor-model-parallel-size ${TP} \ | ||
| 61 | + --pipeline-model-parallel-size ${PP} \ | ||
| 62 | + --num-layers-per-virtual-pipeline-stage ${VPP} \ | ||
| 63 | + --expert-model-parallel-size ${EP} \ | ||
| 64 | + --sequence-parallel \ | ||
| 65 | + --context-parallel-size ${CP} \ | ||
| 66 | + --context-parallel-algo ${CP_TYPE} \ | ||
| 67 | + --num-layers ${NUM_LAYERS} \ | ||
| 68 | + --hidden-size 6144 \ | ||
| 69 | + --ffn-hidden-size 32768 \ | ||
| 70 | + --num-attention-heads 48 \ | ||
| 71 | + --group-query-attention \ | ||
| 72 | + --num-query-groups 8 \ | ||
| 73 | + --post-norm \ | ||
| 74 | + --tokenizer-type PretrainedFromHF \ | ||
| 75 | + --tokenizer-name-or-path ${TOKENIZER_PATH} \ | ||
| 76 | + --seq-length ${SEQ_LEN} \ | ||
| 77 | + --max-position-embeddings 8192 \ | ||
| 78 | + --use-distributed-optimizer \ | ||
| 79 | + --gemm-gradient-accumulation-fusion \ | ||
| 80 | + --reuse-fp32-param \ | ||
| 81 | + --use-flash-attn \ | ||
| 82 | + --make-vocab-size-divisible-by 1 \ | ||
| 83 | + --untie-embeddings-and-output-weights \ | ||
| 84 | + --disable-bias-linear \ | ||
| 85 | + --attention-dropout 0.0 \ | ||
| 86 | + --init-method-std 0.02 \ | ||
| 87 | + --hidden-dropout 0.0 \ | ||
| 88 | + --position-embedding-type rope \ | ||
| 89 | + --normalization RMSNorm \ | ||
| 90 | + --use-fused-rotary-pos-emb \ | ||
| 91 | + --use-rotary-position-embeddings \ | ||
| 92 | + --use-fused-swiglu \ | ||
| 93 | + --use-fused-rmsnorm \ | ||
| 94 | + --swiglu \ | ||
| 95 | + --no-masked-softmax-fusion \ | ||
| 96 | + --attention-softmax-in-fp32 \ | ||
| 97 | + --vocab-size 131072 \ | ||
| 98 | + --rotary-base 10000 \ | ||
| 99 | +" | ||
| 100 | + | ||
| 101 | +CKPT_ARGS=" | ||
| 102 | + --no-load-optim \ | ||
| 103 | + --no-load-rng \ | ||
| 104 | + --no-save-optim \ | ||
| 105 | + --no-save-rng \ | ||
| 106 | + --seed 1234 \ | ||
| 107 | + --save $CKPT_SAVE_DIR \ | ||
| 108 | + --load $CKPT_LOAD_DIR \ | ||
| 109 | +" | ||
| 110 | + | ||
| 111 | +TRAIN_ARGS=" | ||
| 112 | + --micro-batch-size ${MBS} \ | ||
| 113 | + --global-batch-size ${GBS} \ | ||
| 114 | + --lr 1.0e-5 \ | ||
| 115 | + --train-iters 2000 \ | ||
| 116 | + --lr-decay-iters 1280 \ | ||
| 117 | + --lr-decay-style cosine \ | ||
| 118 | + --min-lr 1.0e-6 \ | ||
| 119 | + --weight-decay 0.1 \ | ||
| 120 | + --lr-warmup-iters 2 \ | ||
| 121 | + --clip-grad 1.0 \ | ||
| 122 | + --bf16 | ||
| 123 | +" | ||
| 124 | + | ||
| 125 | +DATA_ARGS=" | ||
| 126 | + --no-shared-storage \ | ||
| 127 | + --data-path $DATA_PATH \ | ||
| 128 | + --split 949,50,1 \ | ||
| 129 | +" | ||
| 130 | + | ||
| 131 | +OUTPUT_ARGS=" | ||
| 132 | + --log-interval 1 \ | ||
| 133 | + --save-interval 2000 \ | ||
| 134 | + --eval-interval 2000 \ | ||
| 135 | + --eval-iters 0 \ | ||
| 136 | +" | ||
| 137 | + | ||
| 138 | +msrun $DISTRIBUTED_ARGS pretrain_gpt.py \ | ||
| 139 | + $GPT_ARGS \ | ||
| 140 | + $DATA_ARGS \ | ||
| 141 | + $OUTPUT_ARGS \ | ||
| 142 | + $CKPT_ARGS \ | ||
| 143 | + $TRAIN_ARGS \ | ||
| 144 | + $MOE_ARGS \ | ||
| 145 | + --distributed-backend nccl \ | ||
| 146 | + --ai-framework mindspore \ | ||
| 147 | + | tee logs/pretrain_grok1_mcore_40b.log | ||