并行策略及优化指南

并行策略及优化指南#

LoongForge 基于 Megatron-LM 构建，完全兼容所有现有的 Megatron-LM 优化策略。 在此基础上，我们增加了多项增强。本文档介绍基本的并行策略及其优化的启用方式。 这些策略可以按需组合，以在数百至数千 GPU上高效训练十亿至万亿参数的模型。

1. 并行策略#

策略

并行维度

主要应用场景

数据并行（DP）

批次维度

标准训练；默认启用

张量并行（TP）

单层算子 / 权重

大隐藏层尺寸或显存受限场景

流水线并行（PP）

模型深度

超深模型（层数多）

上下文并行（CP）

序列长度

长序列训练（8K+）

专家并行（EP）

MoE 专家

混合专家模型

1.1 数据并行（DP）#

  • 并行对象：不同的微批次样本

  • 核心思想：每个 rank 保留完整的模型副本；梯度同步

在 DP 中，每个 GPU 处理批次的一个子集。根据配置，模型相关状态可以完全复制或沿 DP 维度分片以节省显存。

标准 DP（无分片）#

torchrun --nproc_per_node=8 pretrain_gpt.py \
    --data-parallel-sharding-strategy no_shard

  • 每个 GPU 存储完整的参数、梯度和优化器状态

  • 每个 GPU 仅处理批次的一部分

  • 梯度通过 All-Reduce 同步

分片数据并行#

使用 --data-parallel-sharding-strategy 沿 DP 维度分片部分状态，降低每张 GPU 的显存占用。

--data-parallel-sharding-strategy {no_shard | optim | optim_grads | optim_grads_params}

策略

每个 DP rank 上保留的内容

no_shard（默认）

参数 + 梯度 + 优化器状态

optim

参数 + 梯度 + 分片的优化器状态

optim_grads

参数 + 分片的梯度 + 分片的优化器状态

optim_grads_params

分片的参数 + 分片的梯度 + 分片的优化器状态

1.2 张量并行（TP）#

  • 并行对象：单层内的矩阵运算

  • 核心思想：将大矩阵沿某一维度切分到多个 GPU

--tensor-model-parallel-size 4   # 4 路张量并行
--sequence-parallel              # 推荐开启

--sequence-parallel 在 LayerNorm 和 Dropout 中分片序列维度以减少激活显存；通常与 TP 配合使用。

1.3 流水线并行（PP）#

  • 并行对象：模型深度（层维度）

  • 核心思想：不同 GPU 拥有不同的阶段；以微批次流水线方式执行

--pipeline-model-parallel-size 8              # 8 个阶段
--num-layers-per-virtual-pipeline-stage 4     # 虚拟阶段用于负载均衡

1.4 上下文并行（CP）#

  • 并行对象：序列长度（token 维度）

  • 核心思想：将长序列切分到多个 GPU

--context-parallel-size 2   # 2 路上下文并行
--cp-comm-type p2p          # 通信方式

1.5 专家并行（EP）#

  • 并行对象：MoE 层中的专家

  • 核心思想：不同 GPU 持有不同专家；token 通过 All-to-All 分发

--expert-model-parallel-size 8   # 8 路专家并行

2. 性能优化#

2.1 通信优化#

  1. 梯度归约重叠

    --overlap-grad-reduce

    将梯度的 All-Reduce 与反向计算重叠。

  2. 参数收集重叠

    --overlap-param-gather

    将参数的 All-Gather 与前向计算重叠。

  3. TP 通信重叠

    --tp-comm-overlap

    将张量并行通信与计算重叠。

  4. EP 通信重叠

    --overlap-moe-expert-parallel-comm

    将 MoE All-to-All 与计算重叠。

  5. DeepEP 优化 DeepEP 是来自 DeepSeek 的高性能 MoE token 分发/合并库。它大幅降低跨节点 All-to-All的调度和同步开销，是 DeepSeek-V3 及类似模型的推荐配置。

    --moe-token-dispatcher-type=flex   # {allgather | alltoall | flex}
--moe-enable-deepep
--moe-deepep-num-sms N             # DeepEP 可使用的 SM 数量

    DeepEP 仅与 flex 分发器配合使用。

    • allgather（默认）：通过 All-Gather 收集 token

    • alltoall：在专家之间直接交换 token

    • flex：允许使用高性能后端如 DeepEP

2.2 流水线负载均衡#

流水线负载均衡是一种针对 PP / VPP 的高级划分机制，允许用户通过显式布局字符串精确指定每层如何映射到流水线阶段。 它解决以下问题：

  • 模型结构不均匀（如解码器层不可整除、MTP/loss 层）

  • 默认切分导致的阶段间负载不均衡

  • 流水线气泡大、GPU 利用率低

使用 --pipeline-model-parallel-layout 为每个阶段指定层类型和数量。

--pipeline-model-parallel-size 16
--pipeline-model-parallel-layout "Et*3|(tt|)*29,m|L"

  • 布局按 | 分隔每个阶段

  • *N 表示重复一个块 N 次

  • 支持的符号

    • E：Embedding

    • t：Transformer 解码器

    • m：MTP 层

    • L：Loss 计算

2.3 算子融合#

MoE 排列融合#

--moe-permute-fusion

融合 token 重排算子以减少显存访问。

3. 显存优化#

3.1 重计算（Activation Checkpointing）#

在 GPU 显存紧张时，用额外的反向计算换取更低的激活显存。

通过三个正交旋钮控制：

维度

标志

选项

方法

--recompute-method

uniform / block

粒度

--recompute-granularity

full / selective

层数

--recompute-num-layers

正整数

方法#

--recompute-method uniform   # 将模型均分为等大的重计算单元
--recompute-method block     # 仅重计算选定的 Transformer 层

粒度#

--recompute-granularity full       # 重计算整个 Transformer 层
--recompute-granularity selective  # 仅重计算列出的子模块

层数#

--recompute-num-layers N

  • uniform：每个重计算单元的层数

  • block：每个 rank / PP 阶段上重计算的层数

选择性子模块（仅与 selective 配合使用）#

--recompute-modules core_attn moe_act mlp

支持的模块： core_attn、mlp、moe、moe_act、shared_experts、routed_experts、 layernorm、mla_up_proj、 a2a_overlap_attn、a2a_overlap_post_attn、a2a_overlap_mlp（后三个需要 EP A2A 重叠）

3.2 激活卸载#

在前向过程中将选定的激活张量卸载到 CPU，在反向过程中按需取回，以降低峰值 GPU 显存。

四个控制标志：

维度

标志

选项

启用

--fine-grained-activation-offloading

on / off

模块

--offload-modules

模块列表

张量

--offload-tensors

张量标签列表

最小大小

--min-offloaded-tensor-size

字节数（int）

 
--fine-grained-activation-offloading
--offload-modules expert_fc1 core_attn
--offload-tensors dispatched_input
--min-offloaded-tensor-size 1048576

支持的模块： attn_norm、core_attn、attn_proj、mlp_norm、expert_fc1、moe_act

支持的张量标签：

  • dispatched_input（MoE token 分发输出）

  • pre_mlp_layernorm_output（MLP 前的 LayerNorm 输出）

3.3 优化器状态 CPU 卸载#

将优化器状态（如 Adam 动量/方差）从 GPU 移至 CPU 内存，大幅减少 GPU 显存，代价是额外的 CPU 与 GPU 间数据传输。 可与重计算、激活卸载和通信重叠配合使用。

--optimizer-cpu-offload
--optimizer-offload-fraction 1.0   # (0, 1]；1.0 = 全部卸载

  • fraction < 1.0 允许在显存节省与开销之间进行权衡

3.4 融合线性交叉熵#

将输出层的线性投影（hidden @ weight.T）与交叉熵损失融合为单一操作，并结合沿词表维度的分块计算，消除完整 logits 张量带来的峰值显存尖峰。对于典型配置（num_tokens=16384, vocab_size=129280），可节省高达 ~40 GB 的 logits 相关显存。

框架根据 GPU 架构自动选择实现方式：

  • 非 Blackwell GPU：带缓冲区复用和在线 Softmax 的纯 PyTorch 实现——在超过原生 Torch 实现性能的同时显著降低峰值显存

--cross-entropy-loss-fusion \
--cross-entropy-fusion-impl linear

词表分块大小可通过环境变量调整（默认 3072）：

export LCE_GENERIC_FWD_VOCAB_SPLIT_SIZE=3072
export LCE_GENERIC_BWD_VOCAB_SPLIT_SIZE=3072

详见融合线性交叉熵。