混合专家包

Megatron-Core 提供丰富的并行映射，结合了专家并行与张量并行、数据并行、序列并行和流水线并行。这使 Mixtral 8X7B bf16 训练达到 468 TFLOPS（截至 MCore v0.9）。

并行性

• 专家并行

  • MoE 模型的特定并行方法，其中专家被分区到不同的工作器上，每个工作器处理不同的训练样本批次，每个工作器为每个 MoE 层处理一个或多个专家。

• 3D 并行：数据并行、张量并行、流水线并行

  • 注意：当将 MoE 与专家并行和张量并行结合使用时，必须启用序列并行。

• 上下文并行:

  • 拆分序列维度以支持长上下文训练。

• 更丰富的并行映射：EP 可以与 DP/TP/PP/CP 结合使用，以处理更大的 MoE 变体。

• 完全分布式优化器支持。

路由器和负载均衡

• 路由器类型

  • Top-K MLP 路由器

• 负载均衡算法

  • Sinkhorn (S-BASE)

  • 辅助损失 / 负载均衡损失

性能优化

• 当本地专家数量 > 1 时的 GroupedGEMM

  • 支持的数据类型：bf16

  • 针对更大 MoE 模型的性能改进

• 为 MoE 启用 --tp-comm-overlap

• FP8 训练支持

令牌分发机制

• 无丢弃 / 无令牌丢弃

• 令牌丢弃，带或不带填充到容量

易用性

• Mixtral 模型的检查点转换器，详情请参阅示例。

• 分布式检查点

• 逐层日志记录

• Upcycling 支持

• 细粒度 upcycling

即将推出的特性

• 用于大规模 MoE 训练的新并行性

• GroupedGEMM 的 FP8 支持

• 令牌置换 / 反置换融合

• TopK 路由器融合

• MoE 层频率

MoE 相关参数

用法

快速入门

要训练具有 8 个专家和辅助损失的 top-2 MoE 模型，请包含以下参数

            

            
--num-experts 8
--expert-model-parallel-size 8
--moe-grouped-gemm
--moe-router-load-balancing-type aux_loss # options: aux_loss, sinkhorn, none. Default is aux_loss.
--moe-router-topk 2
--moe-aux-loss-coeff 1e-2
--use-distributed-optimizer
--moe-token-dispatcher-type alltoall
        

要启用令牌丢弃机制（例如 GShard 和 SwitchTransformer），请包含以下参数

            

            
--moe-expert-capacity-factor 1.0
--moe-pad-expert-input-to-capacity # Optional
        

下图说明了 MCore 中不同的丢弃策略

1. 默认的无丢弃策略不会丢弃或填充任何令牌。

2. 通过设置 --moe-expert-capacity-factor，超过专家容量的令牌将根据其选定的概率被丢弃。当 EP > 1 时，丢弃操作在 EP rank 之间的令牌交换操作之前执行。容量的公式为 capacity = num_tokens_per_rank * topk * capacity_factor / num_experts。

3. 通过设置 --moe-pad-expert-input-to-capacity，令牌少于容量的专家将被填充到容量。

微调 Mixtral 模型

Megatron-Core 完全支持 Mixtral MoE 模型，我们为 Mixtral 模型提供了从 huggingface 格式到 MCore 格式的检查点转换器。

分布式检查点

MCore v0.7 引入了完全并行和异步的保存功能到分布式检查点，这解决了传统检查点保存方法效率低下的问题。它还解决了传统格式中不同并行映射的检查点之间的不兼容问题。借助新的分布式检查点解决方案，MCore 可以通过保存和加载统一格式的检查点来实现灵活的并行配置。与原生 PyTorch 解决方案相比，MCore 在检查点开销方面最多可减少 50 倍。

从 MCore v0.8 开始，MoE 支持分布式检查点，这意味着用户可以使用任意并行组合进行保存和加载，并且目前可用，包括专家并行。

1. 版本 0.8 支持使用 TPxCPxEPxPP 重分片加载权重和分布式优化器状态以及 SequentialMLP。

2. 版本 0.8.0 支持 GroupedMLP 权重重分片，版本 0.10.0 支持优化器状态重分片。加载和保存时在 GroupedMLP/SequentialMLP 之间切换部分受支持。

3. TEGroupedMLP 完全支持分布式检查点，并且在版本 0.9.0 中与 SequentialMLP 完全可互换。

4. 由于优化器类型不同，优化器状态重分片无法跨越 EP=1 和 EP>1。

用法

• --ckpt-format torch_dist 主要参数，它将尝试使用分布式检查点进行保存和加载。

• --auto-detect-ckpt-format 有了这个参数，它可以加载分布式检查点和旧式检查点。

SequentialMLP、GroupedMLP 和 TEGroupedMLP 之间的检查点兼容性

            

            
    ┌───────────────┐          ┌───────────────┐          ┌───────────────┐
    │   GroupedMLP  │          │ SequentialMLP │          │ TEGroupedMLP  │
    │               │          │               │          │               │
    │               │          │               │          │               │
    │ ┌───────────┐ │          │ ┌───────────┐ │          │ ┌───────────┐ │
    │ │legacy ckpt│ │          │ │legacy ckpt│ │          │ │legacy ckpt│ │
    │ └─────┬─────┘ │          │ └─────┬─────┘ │          │ └─────┬─────┘ │
    │       ▼       │          │       ▼       │          │       ▼       │
    │  ┌─────────┐  │          │  ┌─────────┐  │          │  ┌─────────┐  │
    │  │dist ckpt│  │          │  │dist ckpt│  │          │  │dist ckpt│  │
┌──►│  │ weight  │  │◄────────►│  │ weight  │  │◄────────►│  │ weight  │  │◄──┐
│   │  └─────────┘  │          │  └─────────┘  │          │  └─────────┘  │   │
└───┼───────────────┼──────────┼───────────────┼──────────┼───────────────┼───┘
    │┌─────────────┐│          │┌─────────────┐│          │┌─────────────┐│
    ││  dist ckpt  ││          ││  dist ckpt  ││          ││  dist ckpt  ││
    ││optim states ││          ││optim states ││◄────────►││optim states ││
    │└─────────────┘│          │└─────────────┘│          │└─────────────┘│
    └───────────────┘          └───────────────┘          └───────────────┘
        

分布式检查点的最佳实践

1. 将旧式检查点转换为分布式检查点。为了实现这一点，我们可以同时添加 --ckpt-format torch_dist --auto-detect-ckpt-format，然后它将加载旧式检查点，并在训练进度尝试保存检查点时将其另存为分布式检查点格式。

2. 将旧式 GroupedMLP 的检查点转换为 TEGroupedMLP。这仅支持权重部分。为了实现这一点，我们可以使用上述方法将旧式检查点转换为旧式 GroupedMLP 的分布式检查点。更新库并使用 TEGroupedMLP 后，我们可以通过添加参数 --no-load-optim 直接加载先前保存的检查点。

共享专家

MCore v0.9 引入了共享专家特性。我们可以通过设置合适的 --moe-shared-expert-intermediate-size 来启用此特性。

共享专家的并行模式遵循密集部分的设置，即注意力模块。共享专家不进行分布式处理，但在 EP rank 中复制。

我们还有一个实验性特性，试图重叠共享专家和分发器中的通信和计算。我们可以设置 --moe-shared-expert-overlap 并使用 alltoall 分发器来启用它。重叠依赖于环境设置 CUDA_DEVICE_MAX_CONNECTIONS=1。共享专家中的 AllGather 和 ReduceScatter 通信与分发器中的 permute/unpermute 重叠。共享专家中的 MLP 计算部分与分发器中的 AlltoAll 通信重叠。前向和后向传递都可以重叠。但是要在后向传递中获得重叠，PyTorch 版本应 >= 2.2.0。

Upcycling

使用 --moe-use-upcycling 启用 upcycling，它从 --load 目录加载密集模型，在运行时将其转换为 MoE 模型，并开始训练。转换后的模型在训练开始前保存到 --save 路径。Upcycling 构建在分布式检查点之上，支持与现有密集检查点不同的并行模式，例如 upcycling 期间的任意专家并行。

我们目前仅支持默认的 upcycling 策略，该策略将现有的 MLP 复制到多个专家，每个专家都从 MLP 的副本开始。未来，我们将支持更多最先进的 upcycling 策略，例如来自我们最近的研究工作的细粒度 upcycling。

注意：MoE 模型结构通过脚本参数定义。所有 MoE 相关参数（例如 --num-experts）都可以自定义；但是，其他模型结构参数必须与密集模型的参数一致。

MoE 训练示例

            

            
#!/bin/bash

# Runs Mixtral 8x7B model on 32 H100/A100 GPUs
# The Dropless MoE suffers from an imbalanced token distribution at the early stage of training (the first few hundred iterations), which may lead to poor performance and out-of-memory (OOM) issues.
# To check the performance of a Dropless MoE model, we should run the model for at least 500 iterations or resume from trained checkpoints.

export CUDA_DEVICE_MAX_CONNECTIONS=1

GPUS_PER_NODE=8
# Change for multinode config
MASTER_ADDR=${MASTER_ADDR:-"localhost"}
MASTER_PORT=${MASTER_PORT:-"6000"}
NNODES=${NNODES:-"1"}
NODE_RANK=${RANK:-"0"}
WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))

CHECKPOINT_PATH=$1
TOKENIZER_MODEL=$2
DATA_PATH=$3

DISTRIBUTED_ARGS=(
    --nproc_per_node $GPUS_PER_NODE
    --nnodes $NNODES
    --node_rank $NODE_RANK
    --master_addr $MASTER_ADDR
    --master_port $MASTER_PORT
)

MODEL_ARGS=(
    --disable-bias-linear
    --seq-length 4096
    --max-position-embeddings 32768
    --num-layers 32
    --hidden-size 4096
    --ffn-hidden-size 14336
    --num-attention-heads 32
    --init-method-std 0.01
    --attention-dropout 0.0
    --hidden-dropout 0.0
    --normalization RMSNorm
    --position-embedding-type rope
    --swiglu
    --untie-embeddings-and-output-weights
    --group-query-attention
    --num-query-groups 8
    --no-masked-softmax-fusion
    --no-position-embedding
)

MOE_ARGS=(
    --num-experts 8
    --expert-model-parallel-size 8
    --moe-router-load-balancing-type aux_loss # options: aux_loss, sinkhorn, None. Default is aux_loss.
    --moe-router-topk 2
    --moe-aux-loss-coeff 1e-2
    --moe-grouped-gemm
)

DATA_ARGS=(
    --tokenizer-type Llama2Tokenizer
    --tokenizer-model ${TOKENIZER_MODEL}
    --data-path $DATA_PATH
    --split 99990,8,2
)

TRAINING_ARGS=(
    --micro-batch-size 1
    --global-batch-size 128
    --lr 1e-4
    --train-iters 500000
    --lr-decay-iters 320000
    --lr-decay-style cosine
    --min-lr 1.0e-5
    --weight-decay 0.1
    --lr-warmup-iters 500
    --clip-grad 1.0
    --bf16
    --overlap-grad-reduce
    --overlap-param-gather
)

MODEL_PARALLEL_ARGS=(
    --tensor-model-parallel-size 1
    --pipeline-model-parallel-size 4
    --num-layers-per-virtual-pipeline-stage 8
    --sequence-parallel
    --use-distributed-optimizer
)

LOGGING_ARGS=(
    --log-interval 1 \
    --save-interval 10000 \
    --eval-interval 1000 \
    --eval-iters 10 \
    --save $CHECKPOINT_PATH \
    --load $CHECKPOINT_PATH \
    --tensorboard-dir "${CHECKPOINT_PATH}/tensorboard" \
    --no-load-optim \
    --no-load-rng
)

if [ -n "${WANDB_API_KEY}" ]; then
    LOGGING_ARGS+=(
        --wandb-project ${WANDB_PROJECT:-"Mixtral-Finetuning"}
        --wandb-exp-name ${WANDB_NAME:-"Mixtral_8x7B"}
    )
fi

torchrun ${DISTRIBUTED_ARGS[@]} pretrain_gpt.py \
    ${MODEL_ARGS[@]} \
    ${MOE_ARGS[@]} \
    ${DATA_ARGS[@]} \
    ${TRAINING_ARGS[@]} \
    ${MODEL_PARALLEL_ARGS[@]} \
    ${LOGGING_ARGS[@]}
        

并行映射的调优指南

要找到一个良好的并行映射，以帮助您实现新模型的高吞吐量，以下是一些可能有所帮助的通用规则。以下是每种并行策略在不同方面的属性概述。

对于特定模型，最佳并行映射因模型架构、训练序列长度和硬件平台而异。以下是一些获得更好性能的通用规则

1. 尽可能保持模型并行规模较小。

   • 对于大型语言模型，通常需要模型并行以防止 OOM，但这会带来通信开销并损害性能。

   • 使用分布式优化器时，主权重和优化器状态将在所有 DP rank 中分片，并产生轻微的通信开销。因此，当训练期间有大量空闲 GPU 内存时，请尝试减小模型并行规模并增加数据并行规模。

2. 确保 EPxTP 通信在 NVLink 域内。

   • EP 和 TP 的通信应尽可能保持在 NVLink 域内，因为两者都是通信密集型的。

   • 如果模型太大并且需要跨多个节点进行扩展，请在 TP 和 EP 之前考虑 PP。有关详细信息，请参阅第 3 项。

3. 使用流水线并行进一步扩展模型。

   • 当 PP_size >= 2 时，通过设置 num_layers_per_virtual_pipeline_stage 启用虚拟流水线并行 (VPP) 以减少 pp 气泡。

   • VPP_size 调优：vpp_size 的合法值是 num_layers/pp_size 的所有公约数，例如，num_layers=24，pp_size=4，那么我们可以从 {1, 2, 3, 6} 中选择 vpp_size。vpp_size 越大，流水线气泡越少，而每个 PP 阶段之间的 P2P 通信数量越大。经验上，中间值通常给出最佳折衷方案。VPP_size=num_layers / PP_size / num_layers_per_virtual_pipeline_stage

4. 如果可能，专家层优先选择 EP 而不是 TP

   • TP 比 EP 节省更多内存，但 EP 可以实现比 TP 更好的 GEMM 效率和更少的通信开销。

   • 如果 EP 大小增加到专家数量，则省略专家计算的本地令牌置换/反置换。

   • 简化 MoE 层的计算图，更便于执行潜在的通信-计算重叠。

   • 在实践中，对于 8x7B，EP8TP1 比 EP4TP2 更好。

5. 为长上下文训练启用上下文并行。

   • CP 的效率很大程度上取决于其通信是否可以与计算重叠。

   • 经验上，当序列长度 >= 8K 时使用 CP。

MoE 并行折叠

MoE 并行折叠将 MoE 相关并行组与密集组分离。

1. 传统的 MoE 并行组通过使用具有默认顺序 tp-cp-ep-dp-pp 的 5 维并行组生成器与密集组缠绕在一起。MoE 中的 EP 组是注意力机制中 DP 的子组。

2. 使用 MoE 并行折叠，我们对注意力机制使用具有 tp-cp-dp-pp 的并行组生成器，对 MoE 使用另一个具有 tp-ep-dp-pp 的并行组生成器。MoE 中的 EPxTP 组是注意力机制中 DPxCPxTP 的子组。

通过设置 --expert-tensor-parallel-size，我们可以设置 MoE 特定的 TP 大小。

MoE 并行折叠的优势

1. CP 和 EP 组默认折叠在一起，这样

   1. 它减少了同时启用 CP 和 EP 所需的最少 GPU 数量。例如，传统的 (CP=8, EP=8) 方式至少需要 64 个 GPU，现在只需要 8 个 GPU。

   2. CP 和 EP 通信都可以放在 NVLink 域中。

2. 我们可以为注意力机制和 MoE 部分设置不同的 TP 大小。

   1. 对于 MoE，EP 通常比 TP 更有效。但在传统方式中，仅使用 EP 对于大多数模型来说可能会导致 OOM。

   2. 使用 MoE 并行折叠，我们可以为注意力机制部分启用 TP，并为 MoE 模型设置 TP=1，这通常可以获得更好的 MFU。

端到端训练实践

使用最新的 NVIDIA PyTorch 或 NeMo Docker 镜像

• NGC PyTorch 镜像

• NGC NeMo 镜像

令牌分发器选择

• 令牌分发器将令牌发送到指定的专家，涉及张量重排和通信。

• 分发器 allgather 是默认选项。当仅使用张量并行或 Top-k 值非常大时，它可以实现更好的性能和效率。

• 如果应用专家并行，建议使用分发器 alltoall。

• 分发器 alltoall_seq 是 alltoall 的原始实现，保留它是为了潜在的兼容性风险。

启用通信重叠

• 使用分布式优化器启用 --overlap-param-gather 和 --overlap-grad-reduce。

• 当 TP>1 时启用 --tp-comm-overlap。

• 当 PP > 1 时，通过设置 num_layers_per_virtual_pipeline_stage 启用 p2p 通信重叠。

当 num_local_experts>1 时，使用 --moe-grouped-gemm 启用 GroupedGEMM

• GroupedGEMM 比每个专家的普通顺序 GEMM 具有更高的效率。

• 建议使用 TE 版本的 Grouped GEMM（通过升级到 MCore v0.8 和 TE v1.9），它支持梯度累积融合和 FP8 训练。

从头开始训练时，由令牌分布不均引起的 OOM 当路由器训练不足时，MoE 会遇到严重的负载不均衡问题，导致模型容易耗尽内存 (OOM)，这通常发生在从头开始训练的前 100~300 步中。因此，在前 200 步中有两种推荐的方法来避免 OOM 问题，在令牌分布更稳定后可以移除这些方法

1. 增加 expert-tensor-parallel-size 并减少 expert-model-parallel-size 以在 MoELayer 中用 TP 替换 EP，这可以防止 EP rank 之间的负载不均衡。由于当前的 ETP 实现有一些内存开销，您可以进一步添加 --moe-layer-recompute 仅为 MoE Layer 启用激活重计算。

2. 通过添加 --moe-token-capacity-factor 1.0 将容量因子设置为相对较小的数字，例如 1.0。

参考最佳并行映射

以下是 MCore v0.8 中 Mixtral 8x7B 和 8x22B 模型的参考并行映射

详细基准信息：服务器

• 8xH100 80GB HBM3

• NVLink 第四代

• InfiniBand 8x400 Gbit/s

Docker 镜像

• PyTorch 24.09 with TransformerEngine v1.11