训练 MOE 模型

我们已将 Megatron-Infinigence 训练工具组合打包为容器镜像，用户可直接拉取镜像使用，无需手动配置环境。

镜像名称： infini-ai/megatron-infinigence:v1-ngc25.04-20250725

查看训练框架目录结构

shell

root@c3c6f53a836b:~# tree -L 2 /root/megatron-infinigence/
/root/megatron-infinigence/
├── Megatron-LM -> Megatron-LMs/Megatron-LM_core_r0.8.0
├── Megatron-LMs
│   └── Megatron-LM_core_r0.8.0
├── README.md
├── dist
│   └── megatron_infini-0.0.1-cp312-cp312-linux_x86_64.whl
├── docs
│   ├── MoEtutorial.md
│   └── profile.md
├── megatron_infini
│   ├── __init__.py
│   ├── apply_patcher.cpython-312-x86_64-linux-gnu.so
│   ├── common
│   ├── core
│   ├── examples
│   ├── hardware
│   ├── hetero
│   ├── models
│   ├── multimodal
│   ├── pretrain_deepseek.py
│   ├── pretrain_llama.py
│   ├── pretrain_llava.py
│   ├── pretrain_qwen_moe.py
│   ├── schedulers
│   └── utils
├── megatron_infini.egg-info
│   ├── PKG-INFO
│   ├── SOURCES.txt
│   ├── dependency_links.txt
│   └── top_level.txt
├── requirements.txt
└── tools
    ├── ckpt_convert
    ├── loss_align
    ├── plot_scripts
    └── preprocess_data

训练参数

本节详细介绍 Megatron-Infinigence MoE 训练框架的各种参数配置，包括模型配置参数和高级优化特性。这些参数帮助用户根据具体需求和硬件环境优化训练性能。

模型配置参数

以下是针对不同模型架构的核心配置参数，包括 MLA、YARN RoPE 和 MOE 等关键组件的参数设置。

MLA 参数

针对 Multi-head Latent Attention (MLA) 架构的专用配置参数：

--q-lora-rank: MLA 的 Q latent 维度。针对 16B 的 DeepSeek V2 Lite 不使用
--kv-lora-rank: MLA 的 K & V latent 维度
--qk-rope-head-dim: MLA 的 Q 和 K 的 RoPE 部分的维度
--qk-nope-head-dim: MLA 的 Q 和 K 的没有 RoPE 部分的维度
--v-head-dim: MLA 的 V 的 head 维度
--mla-replicate-l1: MLA 的 replicate L1 优化。针对 TP 切分生效。当开启时，Q/K/V 的 down_proj 不切分，只切分 up_proj，减少一次通信；关闭时，Q/K/V 的 down_proj 和 up_proj 都切分。建议开启，提升 1-3% 性能
--mla-fuse-qkv: MLA 的计算中，Q/K/V 的 down_proj 计算合并。对性能几乎无影响，但可以支持 LayerNorm 和 down_proj 的 fuse，重算以节省显存。只支持有 Q-LoRA 的场景。不建议开启，除非显存严重不足，接受降低 3-7% 性能节省 3% 显存

YARN RoPE 参数

针对 DeepSeek 模型中 YARN RoPE 的配置参数：

--rotary-scaling-factor: DeepSeek 中 YARN RoPE 的缩放因子
--rotary-mscale: DeepSeek 中 YARN RoPE 的 mscale 参数
--rotary-mscale-all-dim: DeepSeek 中 YARN RoPE 的全维度 mscale 参数
--original-max-position-embeddings: DeepSeek 中 YARN RoPE 的原始最大位置嵌入
--rotary-beta-fast: DeepSeek 中 YARN RoPE 的快速 beta 参数
--rotary-beta-slow: DeepSeek 中 YARN RoPE 的慢速 beta 参数

MOE 参数

针对 Mixture of Experts (MOE) 架构的专用配置参数：

--first-k-dense-replace: 针对 DeepSeek 模型，前 K 层使用 Dense 而非 MOE。仅在 DeepSeek 模型中生效
--moe-router-topk-scaling-factor: Router score 的 scaling 因子。DeepSeek V2 中使用，V3 中不使用。仅在 DeepSeek 模型中生效
--moe-router-deepseekv3: 使用 DeepSeek V3 的 MOE router。支持 loss-free 的 router load balance 策略；尚不支持 device-limited 的 router 策略
--moe-router-score-function: 可选 router score function，支持 softmax 和 sigmoid。DeepSeek V3 中使用 sigmoid。仅当 moe-router-deepseekv3 开启时生效
--moe-router-enable-expert-bias: MLA 的 MOE router 的 enable expert bias。仅当 moe-router-deepseekv3 开启时生效
--moe-router-bias-update-rate: Loss-free 的 router load balance 策略中 update rate。仅当 moe-router-deepseekv3 开启时生效
--enable-shared-expert-gate: 开启 shared expert gate。仅在 Qwen 模型中生效
--moe-token-dispatcher-type-patch: 优化后的 dispatcher。配合 DualpipeV 使用时，当使用 overlap_allgather_fused 和 overlap_deepep 时，支持通信和计算重叠。仅在 DeepSeek 和 Qwen 模型中生效

高级优化特性

以下是 Megatron-Infinigence 框架提供的高级优化特性，这些特性可以显著提升 MoE 模型的训练性能和内存效率。用户可以根据硬件配置和性能需求选择合适的优化组合。

DualpipeV 流水线调度

Dualpipev 是基于 Deepseek 开源的 dualpipe 变体的一种流水线调度方式。相较于 1f1b 的流水方式，该方式可以重叠两个 micro batch 之间的前向反向计算与 MLP 层前后 dispatch 和 combine 的通信。

使用方式：

在 megatron-infinigence 中打开如下开关：

bash

--scheduler dualpipev

计算与通信重叠优化

在统一框架中，如果想要重叠前后计算与 MLP 层前后的 allgather 通信，需要使用特殊的 dispatch 组件。

使用方式：

该开关需要配合 dualpipev 一起使用，在 megatron-infinigence 中打开如下开关：

bash

--scheduler dualpipev
--moe-token-dispatcher-type-patch overlap_allgather_fused

MLP 部分重算优化

在统一框架中，为了节省显存的前提下并且保持原本的计算性能，我们针对部分显存消耗较高的组件进行了重算。针对 MLP 部分的前后两次 fc1 和 fc2，我们只重算第一个 fc1 而不对第二个 fc2 进行重算，这个操作我们称为对 MLP 层的 recompute。

使用方式：

在 megatron-infinigence 中打开如下开关：

bash

--mlp-recompute

TEgroupgemm 加速支持

在统一框架中，目前集成了 TEgroupgemm 组件，该组件使用 transformer-engine 库中提供的高效 groupgemm 算子，将 MLP 层多个 expert 矩阵乘计算融合为一个 groupgemm 算子来提高矩阵乘算子的计算性能。

使用方式：

在 megatron-infinigence 中打开如下开关：

bash

--moe-grouped-gemm-te

Flash Attention 3 支持

目前 flashattn 发布了 flashattn3 的算子专门针对 NVIDIA Hopper 系列的 GPU 提高 flash-attn 的性能。

使用方式：

注意如果使用 transformer-engine 内置使用的是 te 自带的 flash-attn，因此需要切换到 local 的 transformer-impl 才能调用到 flash-attn3 算子。

在 megatron-infinigence 中打开如下开关：

bash

export PYTHONPATH=${PROJECT_PATH}:${PROJECT_PATH}/Megatron-LM:$PYTHONPATH:/workspace/flash-attention
--transformer-impl local

训练脚本

本节提供完整的训练脚本示例，帮助用户快速启动 DeepSeekV3-671B 和 Qwen3-235B 模型的训练。这些脚本包含了所有必要的配置参数和优化设置。

DeepSeekV3-671B 训练脚本

以下是 deepseekv3.sh 的完整脚本内容：

bash

#!/bin/bash

## 硬件相关环境设置
# CUDA
export CUDA_DEVICE_MAX_CONNECTIONS=32
export TORCH_NCCL_AVOID_RECORD_STREAMS=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export NVTE_BWD_LAYERNORM_SM_MARGIN=16

## 软件相关环境变量(python / pytorch / megatron)
export OMP_NUM_THREADS=10
export PYTORCH_ENABLE_SAME_RAND_A100=1
export MHA_BWD_NO_ATOMIC_F64=1
export MAX_JOBS=20
export PROJECT_PATH=/root/megatron-infinigence
export PYTHONPATH=${PROJECT_PATH}:${PROJECT_PATH}/Megatron-LM:$PYTHONPATH:/workspace/flash-attention
export ARCH="NVIDIA_H100"

## 软件相关环境变量(python / pytorch / megatron)
##############

# Distributed training variables
export NNODES=${WORLD_SIZE}
GPUS_PER_NODE=8
GPU_NUM=$((${GPUS_PER_NODE}*${NNODES}))
WORLD_SIZE=$((${GPUS_PER_NODE}*${NNODES}))
NODE_RANK=$NODE_RANK
MASTER_ADDR=${MASTER_ADDR:-localhost}
MASTER_PORT=${MASTER_PORT:-7000}
SEED=${SEED:-1234}
GROUP_GEMM=${GROUP_GEMM:-0}
GROUP_GEMM_TE=${GROUP_GEMM_TE:-0}
TD=${TD:-allgather_fused}
RECOMPUTE=${RECOMPUTE:-0}
CPU_ADAM=${CPU_ADAM:-0}
SCHED=${SCHED:-1f1b}

## ORIG:
NUM_LAYERS=${NUM_LAYERS:-61}
HIDDEN_SIZE=7168
NUM_ATTN_HEADS=128
FFN_HIDDEN_SIZE=18432
MOE_FFN_HIDDEN_SIZE=2048
MAX_POSITION_EMBEDDINGS=163840
EXTRA_VOCAB_SIZE=467
RMS_NORM_EPS=1e-6
MAX_SEQ_LEN=4096
MAX_PAD_LEN=4096
MICRO_BATCH_SIZE=1
GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-4096}

# mla
QK_NOPE_HEAD_DIM=128
QK_ROPE_HEAD_DIM=64
V_HEAD_DIM=128
ROPE_THETA=10000
SCALE_FACTOR=40
ORIG_MAX_POSITION_EMBEDDINGS=4096
Q_LORA_RANK=1536
KV_LORA_RANK=512

# moe
NUM_EXPERTS=${NUM_EXPERTS:-256}
ROUTER_TOPK=8
NUM_SHARED_EXPERTS=1
FIRST_K_DENSE_REPLACE=${FIRST_K_DENSE_REPLACE:-3}

## trains
TRAIN_TOKENS=1000000000
WARMUP_TOKENS=10000
TRAIN_ITERS=$(( ${TRAIN_TOKENS} / ${GLOBAL_BATCH_SIZE} / ${MAX_SEQ_LEN} ))
LR_WARMUP_ITERS=$(( ${WARMUP_TOKENS}  / ${GLOBAL_BATCH_SIZE} / ${MAX_SEQ_LEN} ))
LR_DECAY_ITERS=$(( ${TRAIN_TOKENS} /  ${GLOBAL_BATCH_SIZE} / ${MAX_SEQ_LEN} ))
LR=5e-6
MIN_LR=1e-6

# Paths
SRC_PATH=$PROJECT_PATH/megatron_infini/pretrain_deepseek.py
DATA_PATH=${DATA_PATH:-/workspace/datasets/DeepSeek-V3/deepseekv3_text_document}
TOKENIZER_PATH=/workspace/DeepSeek-V3

PARALLEL_PERFORMANCE_ARGS=" \
    --bf16 \
    --use-distributed-optimizer \
    --sequence-parallel \
    --transformer-impl transformer_engine \
    --module-impl {\"core_attn\":\"local\"} \
    --use-flash-attn \
    --tensor-model-parallel-size ${TP} \
    --tensor-model-parallel-for-expert ${TPE} \
    --expert-model-parallel-size ${EP} \
    --pipeline-model-parallel-size ${PP} \
    --moe-token-dispatcher-type-patch ${TD} \
    --mla-replicate-l1 \
    --no-bias-swiglu-fusion \
    "

if [ "$GROUP_GEMM" -eq 1 ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --moe-grouped-gemm \
    "
elif [ "$GROUP_GEMM_TE" -eq 1 ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --moe-grouped-gemm-te \
    "
fi

if [ -n "$FORCE_DROP_AND_PADDING" ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --force-drop-and-padding \
    "
fi

if [ -n "$VPP" ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --num-layers-per-virtual-pipeline-stage ${VPP} \
    "
fi

if [ "$CUDA_DEVICE_MAX_CONNECTIONS" -gt 1 ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --no-async-tensor-model-parallel-allreduce \
    "
fi

if [ "$SCHED" != "1f1b" ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --scheduler ${SCHED} \
    "
fi

if [ -n "$HPP" ]; then
    if [ "$SCHED" == "1f1b" ]; then
        PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
            --hetero-pipeline-stages $HPP \
        "
    else
        PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
            --virtual-hetero-pipeline-stages $HPP \
        "
    fi
fi

if [ "$RECOMPUTE" == "1" ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --recompute-method uniform \
        --recompute-num-layers 1 \
        --recompute-granularity full \
    "
elif [ "$MLP_RECOMPUTE" == "1" ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --mlp-recompute \
    "
fi

if [ "$CPU_ADAM" == "1" ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --optimizer hybridadam \
        --optimizer-offload-policy static \
        --optimizer-offload-fraction 1 \
        --optimizer-enable-pin \
    "
fi

if [ "$DP_OVERLAP" == "1" ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --overlap-grad-reduce \
    "
fi

  
MOE_ARGS=" \
    --num-experts ${NUM_EXPERTS} \
    --moe-router-topk ${ROUTER_TOPK} \
    --moe-ffn-hidden-size ${MOE_FFN_HIDDEN_SIZE} \
    --enable-shared-expert \
    --num-shared-experts ${NUM_SHARED_EXPERTS} \
    --first-k-dense-replace ${FIRST_K_DENSE_REPLACE} \
    --moe-router-deepseekv3 \
    --moe-router-score-function sigmoid \
    --moe-router-enable-expert-bias True \
    --moe-router-topk-scaling-factor 2.5 \
    --moe-router-bias-update-rate 1e-3 \
    --moe-aux-loss-coeff 1e-4 \
    "

MLA_ARGS=" \
    --q-lora-rank ${Q_LORA_RANK} \
    --kv-lora-rank ${KV_LORA_RANK} \
    --qk-nope-head-dim ${QK_NOPE_HEAD_DIM} \
    --qk-rope-head-dim ${QK_ROPE_HEAD_DIM} \
    --v-head-dim ${V_HEAD_DIM} \
    --kv-channels ${V_HEAD_DIM} \
    --qk-layernorm \
    "

OTHER_NETWORK_ARGS=" \
    --use-mcore-models \
    --disable-bias-linear \
    --patch-tokenizer-type DeepSeekTokenizer \
    --tokenizer-model ${TOKENIZER_PATH} \
    --extra-vocab-size ${EXTRA_VOCAB_SIZE} \
    --max-padding-length ${MAX_PAD_LEN} \
    --swiglu \
    --normalization RMSNorm \
    --norm-epsilon ${RMS_NORM_EPS} \
    --use-rotary-position-embeddings \
    --no-rope-fusion \
    --position-embedding-type rope \
    --untie-embeddings-and-output-weights \
    --rotary-base ${ROPE_THETA} \
    --rotary-scaling-factor ${SCALE_FACTOR} \
    --rotary-seq-len-interpolation-factor 1 \
    --rotary-mscale 1.0 \
    --rotary-mscale-all-dim 1.0 \
    --original-max-position-embeddings ${ORIG_MAX_POSITION_EMBEDDINGS} \
    --rotary-beta-fast 32 \
    --rotary-beta-slow 1 \
    "

NETWORK_SIZE_ARGS="  \
    --num-layers ${NUM_LAYERS} \
    --hidden-size ${HIDDEN_SIZE} \
    --num-attention-heads ${NUM_ATTN_HEADS} \
    --ffn-hidden-size ${FFN_HIDDEN_SIZE} \
    --max-position-embeddings ${MAX_POSITION_EMBEDDINGS} \
    --seq-length ${MAX_SEQ_LEN} \
    "

TRAINING_ARGS=" \
    --micro-batch-size ${MICRO_BATCH_SIZE} \
    --global-batch-size ${GLOBAL_BATCH_SIZE} \
    --train-iters ${TRAIN_ITERS} \
    --eval-interval 10000 \
    --eval-iters 0 \
    "

LEARNING_ARGS=" \
    --lr ${LR} \
    --min-lr ${MIN_LR} \
    --lr-decay-style cosine \
    --lr-decay-iters ${LR_DECAY_ITERS} \
    --lr-warmup-iters ${LR_WARMUP_ITERS} \
    --attention-dropout 0.0 \
    --hidden-dropout 0.0 \
    --weight-decay 0.1 \
    --adam-beta1 0.9 \
    --adam-beta2 0.95 \
    --clip-grad 1.0 \
    --init-method-std 0.008 \
    --seed ${SEED} \
    "

LOAD_SAVE_ARGS=" \
    --no-load-optim \
    --no-load-rng \
    --num-workers 8 \
    --no-save-optim \
    "

if [ -n "$LOAD_PATH" ]; then
    LOAD_SAVE_ARGS="$LOAD_SAVE_ARGS \
        --load ${LOAD_PATH} \
    "
fi

if [ -n "$SAVE_PATH" ]; then
    LOAD_SAVE_ARGS="$LOAD_SAVE_ARGS \
        --save ${SAVE_PATH}
    "
fi

LOGGING_ARGS=" \
    --log-interval 1 \
    --log-throughput \
    --save-interval 500 \
    --timing-log-level 2 \
    "

INTERVAL_ARGS=" \
    --save-interval 10000 \
    --eval-interval 1000 \
    --eval-iters -1 \
    "

DATASET_ARGS=" \
    --num-workers 8 \
    --data-path ${DATA_PATH} \
    --split 99,1,0 \
    --dataset LLama-Pretrain-Idxmap \
    "

LAUNCHER=" \
    torchrun \
    --nproc_per_node ${GPUS_PER_NODE} \
    --nnodes ${NNODES} \
    --node_rank ${NODE_RANK} \
    --master_addr ${MASTER_ADDR} \
    --master_port ${MASTER_PORT} \
    "

RUN_CMD="${LAUNCHER} ${SRC_PATH} \
    ${PARALLEL_PERFORMANCE_ARGS} \
    ${MIXED_PRECISION_ARGS} \
    ${MOE_ARGS} \
    ${MLA_ARGS} \
    ${OTHER_NETWORK_ARGS} \
    ${NETWORK_SIZE_ARGS} \
    ${TRAINING_ARGS} \
    ${LEARNING_ARGS} \
    ${LOAD_SAVE_ARGS} \
    ${LOGGING_ARGS} \
    ${INTERVAL_ARGS} \
    ${DATASET_ARGS} \
    "

echo ${RUN_CMD}
${RUN_CMD}

Qwen3-235B 训练脚本

以下是 qwen3-235b.sh 的完整脚本内容：

bash

#!/bin/bash

## 硬件相关环境设置
# CUDA
export CUDA_DEVICE_MAX_CONNECTIONS=32
export TORCH_NCCL_AVOID_RECORD_STREAMS=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export NVTE_BWD_LAYERNORM_SM_MARGIN=16

## 软件相关环境变量(python / pytorch / megatron)
export OMP_NUM_THREADS=10
export PYTORCH_ENABLE_SAME_RAND_A100=1
export MHA_BWD_NO_ATOMIC_F64=1
export MAX_JOBS=20
export PROJECT_PATH=/root/megatron-infinigence
export PYTHONPATH=${PROJECT_PATH}:${PROJECT_PATH}/Megatron-LM:$PYTHONPATH:/workspace/flash-attention
export ARCH="NVIDIA_H100"

## 软件相关环境变量(python / pytorch / megatron)
##############

# Distributed training variables
export NNODES=${WORLD_SIZE}
GPUS_PER_NODE=8
GPU_NUM=$((${GPUS_PER_NODE}*${NNODES}))
WORLD_SIZE=$((${GPUS_PER_NODE}*${NNODES}))
NODE_RANK=$NODE_RANK
MASTER_ADDR=${MASTER_ADDR:-localhost}
MASTER_PORT=${MASTER_PORT:-7000}
SEED=${SEED:-1234}
GROUP_GEMM=${GROUP_GEMM:-0}
GROUP_GEMM_TE=${GROUP_GEMM_TE:-0}
TD=${TD:-allgather_fused}
RECOMPUTE=${RECOMPUTE:-0}
CPU_ADAM=${CPU_ADAM:-0}
SCHED=${SCHED:-1f1b}

## ORIG:
NUM_LAYERS=${NUM_LAYERS:-94}
HIDDEN_SIZE=4096
FFN_HIDDEN_SIZE=12288
EXTRA_VOCAB_SIZE=421
RMS_NORM_EPS=1e-5
MAX_SEQ_LEN=4096
MAX_PAD_LEN=4096

# gqa
NUM_ATTN_HEADS=64
NUM_KEY_VALUE_HEADS=4

# rope
MAX_POSITION_EMBEDDINGS=131072
ROPE_THETA=1000000
ROPE_PERCENT=1.0

# moe
MOE_FFN_HIDDEN_SIZE=1536
NUM_EXPERTS=${NUM_EXPERTS:-128}
ROUTER_TOPK=8

# bz
MICRO_BATCH_SIZE=${MICRO_BATCH_SIZE:-1}
GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-4096}

## trains
TRAIN_TOKENS=1000000000
WARMUP_TOKENS=10000
TRAIN_ITERS=$(( ${TRAIN_TOKENS} / ${GLOBAL_BATCH_SIZE} / ${MAX_SEQ_LEN} ))
LR_WARMUP_ITERS=$(( ${WARMUP_TOKENS}  / ${GLOBAL_BATCH_SIZE} / ${MAX_SEQ_LEN} ))
LR_DECAY_ITERS=$(( ${TRAIN_TOKENS} /  ${GLOBAL_BATCH_SIZE} / ${MAX_SEQ_LEN} ))
LR=5e-6
MIN_LR=1e-6

# Paths
SRC_PATH=$PROJECT_PATH/megatron_infini/pretrain_qwen_moe.py
DATA_PATH=${DATA_PATH:-/workspace/datasets/Qwen3-30B-A3B/qwen3_text_document}
TOKENIZER_PATH=/workspace/Qwen3-30B-A3B

PARALLEL_PERFORMANCE_ARGS="
    --bf16 \
    --use-distributed-optimizer \
    --sequence-parallel \
    --transformer-impl transformer_engine \
    --use-flash-attn \
    --tensor-model-parallel-size ${TP} \
    --tensor-model-parallel-for-expert ${TPE} \
    --expert-model-parallel-size ${EP} \
    --pipeline-model-parallel-size ${PP} \
    --moe-token-dispatcher-type-patch ${TD} \
    --use-local-multi-tensor-scale \
"

if [ "$GROUP_GEMM" -eq 1 ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --moe-grouped-gemm \
    "
elif [ "$GROUP_GEMM_TE" -eq 1 ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --moe-grouped-gemm-te \
    "
fi

if [ -n "$FORCE_DROP_AND_PADDING" ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --force-drop-and-padding \
    "
fi

if [ -n "$VPP" ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --num-layers-per-virtual-pipeline-stage ${VPP} \
    "
fi

if [ "$CUDA_DEVICE_MAX_CONNECTIONS" -gt 1 ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --no-async-tensor-model-parallel-allreduce \
    "
fi

if [ "$SCHED" != "1f1b" ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --scheduler ${SCHED} \
    "
fi

if [ -n "$HPP" ]; then
    if [ "$SCHED" == "1f1b" ]; then
        PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
            --hetero-pipeline-stages $HPP \
        "
    else
        PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
            --virtual-hetero-pipeline-stages $HPP \
        "
    fi
fi

if [ "$RECOMPUTE" == "1" ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --recompute-method uniform \
        --recompute-num-layers 1 \
        --recompute-granularity full \
    "
elif [ "$MLP_RECOMPUTE" == "1" ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --mlp-recompute \
    "
fi

if [ "$CPU_ADAM" == "1" ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --optimizer hybridadam \
        --optimizer-offload-policy static \
        --optimizer-offload-fraction 1 \
        --optimizer-enable-pin \
    "
fi

if [ "$DP_OVERLAP" == "1" ]; then
    PARALLEL_PERFORMANCE_ARGS="$PARALLEL_PERFORMANCE_ARGS \
        --overlap-grad-reduce \
    "
fi

MOE_ARGS=" \
    --num-experts ${NUM_EXPERTS} \
    --moe-router-topk ${ROUTER_TOPK} \
    --moe-ffn-hidden-size ${MOE_FFN_HIDDEN_SIZE} \
    --moe-router-topk-scaling-factor 2.5 \
    --moe-aux-loss-coeff 1e-4 \
    --moe-router-pre-softmax \
    --cross-entropy-loss-fusion \
"

GQA_ARGS=" \
    --group-query-attention \
    --num-query-groups ${NUM_KEY_VALUE_HEADS} \
    --num-attention-heads ${NUM_ATTN_HEADS} \
    --add-qkv-bias \
"

OTHER_NETWORK_ARGS=" \
    --use-mcore-models \
    --disable-bias-linear \
    --patch-tokenizer-type Qwen3Tokenizer \
    --tokenizer-model ${TOKENIZER_PATH} \
    --extra-vocab-size ${EXTRA_VOCAB_SIZE} \
    --swiglu \
    --normalization RMSNorm \
    --norm-epsilon ${RMS_NORM_EPS} \
    --use-rotary-position-embeddings \
    --position-embedding-type rope \
    --rotary-percent ${ROPE_PERCENT} \
    --rotary-base ${ROPE_THETA} \
    --rotary-emb-use-cache \
    --max-position-embeddings ${MAX_POSITION_EMBEDDINGS} \
    --untie-embeddings-and-output-weights \
"

NETWORK_SIZE_ARGS=" \
    --num-layers ${NUM_LAYERS} \
    --hidden-size ${HIDDEN_SIZE} \
    --seq-length ${MAX_SEQ_LEN} \
"

TRAINING_ARGS=" \
    --micro-batch-size ${MICRO_BATCH_SIZE} \
    --global-batch-size ${GLOBAL_BATCH_SIZE} \
    --train-iters ${TRAIN_ITERS} \
    --eval-interval 10000 \
    --eval-iters 0 \
"

LEARNING_ARGS=" \
    --lr ${LR} \
    --min-lr ${MIN_LR} \
    --lr-decay-style cosine \
    --lr-decay-iters ${LR_DECAY_ITERS} \
    --lr-warmup-iters ${LR_WARMUP_ITERS} \
    --attention-dropout 0.0 \
    --hidden-dropout 0.0 \
    --weight-decay 0.1 \
    --adam-beta1 0.9 \
    --adam-beta2 0.95 \
    --clip-grad 1.0 \
    --init-method-std 0.008 \
    --seed ${SEED} \
"

LOAD_SAVE_ARGS=" \
    --no-load-optim \
    --no-load-rng \
    --no-save-optim \
    --no-save-rng \
"

if [ -n "$LOAD_PATH" ]; then
    LOAD_SAVE_ARGS="$LOAD_SAVE_ARGS \
        --load ${LOAD_PATH}
    "
fi

if [ -n "$SAVE_PATH" ]; then
    LOAD_SAVE_ARGS="$LOAD_SAVE_ARGS \
        --save ${SAVE_PATH}
    "
fi

LOGGING_ARGS=" \
    --log-interval 1 \
    --log-throughput \
    --save-interval 500 \
    --timing-log-level 2 \
    "

INTERVAL_ARGS=" \
    --save-interval 10000 \
    --eval-interval 1000 \
    --eval-iters -1 \
    "

DATASET_ARGS=" \
    --num-workers 8 \
    --data-path ${DATA_PATH} \
    --split 99,1,0 \
    --dataset LLama-Pretrain-Idxmap \
    "

LAUNCHER=" \
    torchrun \
    --nproc_per_node ${GPUS_PER_NODE} \
    --nnodes ${NNODES} \
    --node_rank ${NODE_RANK} \
    --master_addr ${MASTER_ADDR} \
    --master_port ${MASTER_PORT} \
    "

RUN_CMD="${LAUNCHER} ${SRC_PATH} \
    ${PARALLEL_PERFORMANCE_ARGS} \
    ${MIXED_PRECISION_ARGS} \
    ${MOE_ARGS} \
    ${GQA_ARGS} \
    ${OTHER_NETWORK_ARGS} \
    ${NETWORK_SIZE_ARGS} \
    ${TRAINING_ARGS} \
    ${LEARNING_ARGS} \
    ${LOAD_SAVE_ARGS} \
    ${LOGGING_ARGS} \
    ${INTERVAL_ARGS} \
    ${DATASET_ARGS} \
    "

echo ${RUN_CMD}
${RUN_CMD}

准备数据

AIStudio 任务在多个可用区可用。您的模型及所有训练数据必须上传至训练任务所在可用区的共享高性能存储目录中，并完成相应的预处理。

注意

为避免存储访问问题，请确保您的数据与训练任务位于同一可用区。不同可用区之间的存储资源相互隔离，无法跨区直接使用。请按照下文指导进行配置。

请自行下载模型 weight 以及 tokenizer，以 DeepseekV3 为例。

shell

git clone https://modelscope.cn/models/deepseek-ai/DeepSeek-V3

注意

由于模型与数据集下载受多种因素影响。为保证高效利用资源，建议提前下载模型与数据集。您可以直接利用 AICoder 下载公有模型与数据集，或通过 AICoder 上传自有数据至选定可用区的共享存储中。

预处理训练数据

下载好模型后，需要预先处理数据集为 Megatron 框架支持的格式。建议您通过 AIStudio 数据处理任务自行完成数据转换。

数据处理示例

Megatron-infinigence 容器镜像中预置了数据集处理实例（位置 tools/preprocee_data/preprocess_data.sh），以下通过使用该工具简述数据预处理要求。

建议您提前通过 AICoder 下载数据集至共享高性能存储。

shell

wget https://atp-modelzoo.oss-cn-hangzhou.aliyuncs.com/release/datasets/WuDaoCorpus2.0_base_sample.tgz
tar zxvf WuDaoCorpus2.0_base_sample.tgz

示例脚本在镜像中路径为 megatron-infinigence/tools/preprocess_data/preprocess_data.sh/preprocess_data.sh。修改如下环境变量：

shell

WUDAO_DATA_PATH=WUDAO_DATA_PATH # download dataset path
WORKSHOP_PATH=WORKSHOP_PATH # megatron-infinigence file path
OUTPUT_FILE=OUTPUT_FILE # ouput megatron-infinigence dataset file path 
TOKENIZER_PATH=TOKENIZER_PATH # tokenzier file path
PATCH_TOKENIZER_TYPE=PATCH_TOKENIZER_TYPE # tokenzier type
EXTRA_VOCAB_SIZE=EXTRA_VOCAB_SIZE # extra vocab size

执行预处理脚本：
shell
```
bash megatron-infinigence/tools/preprocess_data/preprocess_data.sh
```
运行成功后得到处理好的数据集。.bin文件储存 Token 数据，.idx 文件存储文档元信息（包括index头、版本号、数据类型、数据大小等信息）。
- OUTPUT_FILE_text_document.bin
- OUTPUT_FILE_text_document.idx
注意
测试时可以直接使用镜像中已经处理好的数据集：
- /workspace/datasets/DeepSeek-V3
- /workspace/datasets/Qwen3-30B-A3B

发起数据处理任务

在真实场景中，我们建议使用 AIStudio 任务功能的数据处理任务完成预处理工作。该类型任务使用租户自购的包年包月资源池中空闲的 CPU 资源，并且优先保障与 GPU 相关任务不受影响。

在智算云平台任务列表页面点击 创建任务，进入创建任务界面，可创建单机或分布式任务。
创建任务
在任务页面，选择任务类型为数据处理任务。根据页面提示，完成其他所有步骤中的配置。
注意
如果 worker 规格列表中没有可选规格，您可以联系商务或售后服务。
在启动命令中，填写您的数据转换命令。您可以使用自己的数据转换工具，也可以参考镜像中 megatron-infinigence/tools/preprocess_data/ 目录下的脚本进行修改和使用。

转换模型格式

在训练前，需要先确定训练的并行配置，然后进行权重转换处理。您可以使用 megatron-infinigence 框架中的权重转换工具（megatron-infinigence/tools/ckpt_convert/）。

警告

请确认好之后需要跑的并行配置，不同并行配置之间的 ckpt 不能通用。在使用前，请详细阅读 megatron-infinigence/tools/ckpt_convert/README.md。

shell

SCRIPT_PATH=/磁盘路径/megatron-infinigence/tools/ckpt_convert/src
CKPT_PATH_HF=/磁盘路径/ckpt/DeepSeek-V3
CKPT_PATH_MF=/磁盘路径/ckpt/DeepSeek-V3_middle_file
CKPT_PATH_MG=/磁盘路径/ckpt/DeepSeek-V3_mg_tp1_pp2
TP_SIZE=1
PP_SIZE=2
# # download
# modelscope download --model 'deepseek-ai/DeepSeek-V3' --local_dir "$CKPT_PATH_HF"

rm -rf $CKPT_PATH_MF
python $SCRIPT_PATH/convert_hf_to_middle_file.py \    
        --load-path $CKPT_PATH_HF \    
        --save-path $CKPT_PATH_MF \    
        --model 'DeepSeek-V3' \    
        --use-gpu-num 0 \    
        --process-num 16rm -rf $CKPT_PATH_MG
        
python $SCRIPT_PATH/convert_middle_file_to_mg.py \    
        --load-path $CKPT_PATH_MF \    
        --save-path $CKPT_PATH_MG \    
        --model 'DeepSeek-V3' \    
        --tp-size $TP_SIZE \    
        --tpe-size 1 \    
        --ep-size 1 \    
        --pp-size $PP_SIZE \    
        --use-gpu-num 0 \    
        --process-num 16rm -rf $CKPT_PATH_MF

在实际使用时，可以在 AIStudio 平台使用以下方式完成权重转换：

使用 AIStudio 任务功能的数据处理任务完成预处理工作。该类型任务使用租户自购的包年包月资源池中空闲的 CPU 资源，并且优先保障与 GPU 相关任务不受影响。
使用 AIStudio 的开发机功能，在 GPU 环境中进行权重转换。

运行任务与示例

预置测试脚本

我们在镜像预置了测试脚本，您可以直接使用这些脚本运行任务进行测试。

注意

请注意修改其中的相关并行参数并导入对应的 tokenzier，数据集以及转换好的 ckpt。

root@c3c6f53a836b:/# tree -L 3 /root/megatron-infinigence/megatron_infini/examples/
/root/megatron-infinigence/megatron_infini/examples/
└── moe
    └── nvidia
        ├── deepseekv3_671b.sh
        ├── qwen3_235b.sh
        ├── test_deepseek.sh
        ├── test_deepseek_node16.sh
        ├── test_deepseek_node16_baseline.sh
        ├── test_deepseek_node32.sh
        ├── test_deepseek_node32_baseline.sh
        ├── test_deepseek_node8.sh
        ├── test_deepseek_node8_baseline.sh
        ├── test_qwen3.sh
        ├── test_qwen3_node16.sh
        ├── test_qwen3_node16_baseline.sh
        ├── test_qwen3_node32.sh
        ├── test_qwen3_node32_baseline.sh
        ├── test_qwen3_node8.sh
        └── test_qwen3_node8_baseline.sh

发起训练任务

在智算云平台任务列表页面点击 创建任务，进入创建任务界面，可创建单机或分布式任务。
创建任务
在任务页面，选择任务类型为训练任务。配置任务 Worker 规格、Worker 数量，设置「分布式框架」为 Pytorch DDP。根据页面提示，完成其他所有步骤中的配置。
选择镜像为 infini-ai/megatron-infinigence:v1-ngc25.04-20250725。

在启动命令中，填写测试脚本：

shell

bash megatron-infinigence/megatron_infini/examples/moe/nvidia/test_deepseek.sh

性能测试

单机脚本验证已经内置在该路径下，由于单机显存限制，对模型进行削减：

/root/megatron-infinigence/megatron_infini/examples/moe/nvidia/test_deepseek.sh
/root/megatron-infinigence/megatron_infini/examples/moe/nvidia/test_qwen3.sh

目前提供8，16，32机的性能复现脚本如下：

# 运行分布式代码前，需要将存储在镜像内部的 dataset 复制至共享磁盘
if [[ $RANK == 0 ]]; then
  cp -r /workspace/datasets /共享高性能存储路径/
fi

export DATA_PATH=/共享磁盘/datasets/DeepSeek-V3/deepseekv3_text_document
/root/megatron-infinigence/megatron_infini/examples/moe/nvidia/test_deepseek_node8.sh
/root/megatron-infinigence/megatron_infini/examples/moe/nvidia/test_deepseek_node8_baseline.sh
/root/megatron-infinigence/megatron_infini/examples/moe/nvidia/test_deepseek_node16.sh
/root/megatron-infinigence/megatron_infini/examples/moe/nvidia/test_deepseek_node16_baseline.sh
/root/megatron-infinigence/megatron_infini/examples/moe/nvidia/test_deepseek_node32.sh
/root/megatron-infinigence/megatron_infini/examples/moe/nvidia/test_deepseek_node32_baseline.sh

export DATA_PATH=/共享磁盘/datasets/Qwen3-30B-A3B/qwen3_text_document
/root/megatron-infinigence/megatron_infini/examples/moe/nvidia/test_qwen3_node8.sh
/root/megatron-infinigence/megatron_infini/examples/moe/nvidia/test_qwen3_node8_baseline.sh
/root/megatron-infinigence/megatron_infini/examples/moe/nvidia/test_qwen3_node16.sh
/root/megatron-infinigence/megatron_infini/examples/moe/nvidia/test_qwen3_node16_baseline.sh
/root/megatron-infinigence/megatron_infini/examples/moe/nvidia/test_qwen3_node32.sh
/root/megatron-infinigence/megatron_infini/examples/moe/nvidia/test_qwen3_node32_baseline.sh

以下是不同调度器配置下的性能对比数据：

Deepseek-v3 Tp 1 tpe1 expert 256 7 层 8机 H100 (性能提升 45%)

scheduler	ep	pp	hpp	cpuadam	attn	recompute	groupgemm	Tflops / tgs
1f1b	16	2	4 3	NA	fa3	无	te	241.3 / 7183
dualpipev	16	2	2 2 1 2	NA	fa3	无	te	350 / 10469.135

qwen3 tp1 tpe1 expert 128 30层 8机 H100 (性能提升 41%)

scheduler	ep	pp	hpp	cpuadam	attn	recompute	groupgemm	Tflops / tgs
1f1b	8	4	2 3 3 3 2 3 3 3	NA	fla3	无	te	226.1 / 7392.604
1f1b	16	4	7 8 8 7	NA	fla3	无	te	210.8 / 5224.775
dualpipev overlap	8	4	2 3 3 3 2 3 3 3	NA	fla3	无	te	313.0 / 10233.981
dualpipev overlap	16	4	3 4 4 4 3 4 4 4	NA	fla3	无	te	304.2 / 7535.244

比较 dualpipev 和 1f1b 的 loss 在单机脚本上 Loss 曲线对比如下：

警告

测试过程中需要加载一致的 checkpoint。如果使用初始化 weight，不同的流水方式如 1f1b 和 dualpipev 的 loss 曲线会由于初始化 weight 不一致导致无法对齐 loss。

训练 MOE 模型 ​

训练参数 ​

模型配置参数 ​

MLA 参数 ​

YARN RoPE 参数 ​

MOE 参数 ​

高级优化特性 ​

DualpipeV 流水线调度 ​

计算与通信重叠优化 ​

MLP 部分重算优化 ​

TEgroupgemm 加速支持 ​

Flash Attention 3 支持 ​

训练脚本 ​

DeepSeekV3-671B 训练脚本 ​

Qwen3-235B 训练脚本 ​

准备数据 ​

预处理训练数据 ​

数据处理示例 ​

发起数据处理任务 ​

转换模型格式 ​

运行任务与示例 ​

预置测试脚本 ​

发起训练任务 ​

性能测试 ​

训练 MOE 模型

训练参数

模型配置参数

MLA 参数

YARN RoPE 参数

MOE 参数

高级优化特性

DualpipeV 流水线调度

计算与通信重叠优化

MLP 部分重算优化

TEgroupgemm 加速支持

Flash Attention 3 支持

训练脚本

DeepSeekV3-671B 训练脚本

Qwen3-235B 训练脚本

准备数据

预处理训练数据

数据处理示例

发起数据处理任务

转换模型格式

运行任务与示例

预置测试脚本

发起训练任务

性能测试