☰
Megatron-Bridge 并行化策略完整指南:TP、PP、EP、CP 混合并行如何组合最优
2026/10/23 4:13:23 网站建设 项目流程
  • 人工智能
  • 大模型
  • 预训练
  • 微调
  • LoRA
  • 分布式训练
  • 多模态

【免费下载链接】Megatron-Bridge

Training library for Megatron-based models with bidirectional Hugging Face conversion capability

项目地址:https://gitcode.com/gh_mirrors/me/Megatron-Bridge
点击查看免费下载

Megatron-Bridge 是一款面向 Megatron 架构的大模型训练库,支持 Hugging Face 双向转换,其核心能力之一就是把**张量并行(TP)、流水线并行(PP)、专家并行(EP)、上下文并行(CP)与数据并行(DP)**自由混合组合,帮助新手也能在有限 GPU 上高效训练大模型。本文用通俗的语言讲清每种并行的原理、配置参数,以及混合并行的选型技巧。

Megatron-Bridge 在 NeMo 框架中的位置:覆盖数据准备、预训练/SFT/LoRA、后训练对齐与推理全流程

一、为什么大模型训练需要混合并行

一张 GPU 显存再大,也装不下几十 B 甚至上千 B 参数模型的全部参数 + 梯度 + 优化器状态 + 激活值。以 70B 模型 FP16 训练为例,无并行时总显存需求高达608 GB;而 TP=4、PP=4、DP=4(64 张 GPU)组合后,单卡只需约38 GB。

并行策略大体分两类:

类别策略一句话原理
数据并行DP / DDP / 分布式优化器每张卡存完整模型副本,数据切开,梯度同步
模型并行TP / PP / EP把模型本身切开摊到多张卡上
激活切分SP / CP沿序列维度切分激活,省长序列显存

数据并行(DDP):多卡各存一份模型,训练前用 all-reduce 汇总梯度,是 Megatron-Bridge 的默认部署方式

Megatron-Bridge 中 DDP 是默认开启的:总卡数对应 DP 组大小,一旦启用模型并行,DP 组大小自动缩小。再配合分布式优化器(use_distributed_optimizer=True),可把优化器状态分片到各 DP 卡上,大幅降低显存。

二、四种核心模型并行策略逐个拆解

1. 张量并行 TP:把一层的权重切开

TP 将单个 Transformer 层内的权重矩阵横向切分到多张 GPU 上,同时降低参数显存和激活显存。

TP=2 时,每层的权重矩阵被切成两份分别放在两张卡上

启用方式:在模型 provider 中设置tensor_model_parallel_size,源码入口见 src/megatron/bridge/models/gpt_provider.py:

model_config = GPTModelProvider( tensor_model_parallel_size=4, # 在 4 张 GPU 间切分张量 )

⚠️ 经验法则:TP 尽量限制在单机 NVLink 域内(TP ≤ 8),跨节点做 TP 几乎总是性能损失。

2. 流水线并行 PP:把连续层段分配到不同 GPU

PP 将模型按层切成若干段,每段放在一组 GPU 上,像流水线一样依次处理微批。

PP 将连续层段分配到多组 GPU,微批在流水线中依次推进

启用方式:设置pipeline_model_parallel_size。普通 PP 会产生"流水线气泡"(空闲等待),可通过virtual_pipeline_model_parallel_size启用交错流水线调度显著缩小气泡:

model_config = GPTModelProvider( pipeline_model_parallel_size=4, virtual_pipeline_model_parallel_size=2, # 每级流水线再分 2 个层块 )

PP 可以跨节点工作,是跨机扩展的主力策略,但需要仔细调节 batch size。

3. 专家并行 EP:MoE 模型的扩展主维度

MoE(混合专家)模型中每个 token 只激活部分专家,**专家并行(EP)**把这些专家分布到不同 GPU 上——注意 EP 只作用于专家层,不影响其余层的并行映射。

EP=4 时,8 个专家被均匀分到 4 张 GPU,每张卡承载 2 个专家

启用方式:num_moe_experts必须能被expert_model_parallel_size整除:

model_config = GPTModelProvider( num_moe_experts=8, expert_model_parallel_size=4, # 8 个专家分到 4 卡 expert_tensor_parallel_size=2, # 可选:专家内部再做 TP )

MoE 的关键认知:TP 按"激活参数"而非总参数来选。DeepSeek-V3(总参 671B / 激活 37B)只需 TP=2,而不是像 70B 稠密模型那样上 TP=8;EP 才是扩展主维度(EP=64),PP 负责跨节点分深度(PP=16)。

4. 上下文并行 CP:长序列训练的关键

CP 沿序列维度把所有层的激活切分到多张 GPU——与只切部分层激活的序列并行(SP)不同,CP 是应对长上下文(8K~32K+)的核心手段。前向时每卡只存自己那段的 KV,反向时通过环形拓扑通信重组。

与 TP 配套的序列并行(SP):把注意力中的 LayerNorm/Dropout 等沿序列维切分,需tensor_model_parallel_size > 1

启用方式:

model_config = GPTModelProvider( context_parallel_size=2, # 序列切到 2 张 GPU sequence_parallel=True, # 2K-8K 序列建议开启 SP )

⚠️ 两个易踩的坑:SP 必须搭配 TP>1;CP 要求seq_length % (2 × context_parallel_size) == 0。

三、混合并行:多策略组合与 GPU 数计算

Megatron-Bridge 允许把以上策略任意混合。典型 4D/5D 组合配置(完整示例见 docs/parallelisms.md):

model_config = GPTModelProvider( # 模型并行 tensor_model_parallel_size=2, pipeline_model_parallel_size=4, virtual_pipeline_model_parallel_size=2, # 激活切分 sequence_parallel=True, context_parallel_size=2, # MoE 专家并行 num_moe_experts=8, expert_model_parallel_size=4, )

DP 组大小自动推导,无需手动配置:

  • 稠密模型:DP = world_size / (TP × PP × CP),例如 32 卡、TP=2、PP=4、CP=2 → DP=2
  • MoE 模型:EDP = world_size / (PP × EP × ETP)

💡常见误区:MoE 配置的最小 GPU 数不是各维度连乘,而是min_gpus = PP × max(TP×CP, EP×ETP)——稠密 TP×CP 网格与 MoE EP×ETP 网格在每个 PP stage 内共享同一批 GPU。例如 PP=1、TP=2、EP=8 只需8 卡而非 16 卡。并行状态在 src/megatron/bridge/training/initialize.py 中统一初始化,整体配置容器为 src/megatron/bridge/training/config.py 中的ConfigContainer。

四、最优并行组合怎么选:一张表搞定

按模型规模选(稠密模型)

模型规模GPU 数推荐起步配置
< 1B1-8仅 DP
1-10B8-16TP=2-4 + DP
10-70B16-64TP=4-8 + PP=2-4 + DP
70-175B64-256TP=8 + PP=4-8 + DP
175-500B256-1024TP=8 + PP=8-16 + CP=2 + DP

按模型规模选(MoE 模型,总参/激活参)

规模典型布局
< 20B仅 EP(TP=1, PP=1)
20-100BTP=1-2 + PP=2-4 + EP=8-16
100-500BTP=2-4 + PP=8-16 + EP=8-32
500B+TP=2 + PP=16 + EP=32-64

真实案例:DeepSeek-V3 671B → TP=2 / PP=16 / EP=64;Kimi-K2 1T → TP=2 / PP=16 / EP=32。

按硬件拓扑与序列长度微调

  • 单机 NVLink:最大化节点内 TP(至多 TP=8)
  • 多机 InfiniBand:TP 留在节点内,用 PP 跨节点扩展
  • 以太网等受限网络:最小化 TP,跨节点优先 PP
  • 序列 <2K:标准 TP+PP+DP;2K-8K:加 SP;8K-32K:加 CP=2;32K+:CP=4-8

更多选型细节与验证命令见 skills/nemo-mbridge-perf-parallelism-strategies/SKILL.md。

五、避坑清单:新手最容易犯的错误

  1. TP 跨节点——吞吐暴跌,务必限制在单个 NVLink 域内
  2. PP 不开交错——气泡严重,尽量设置virtual_pipeline_model_parallel_size
  3. 只开 SP 不开 TP——sequence_parallel=True但tensor_model_parallel_size=1是配置错误
  4. EP 用于稠密模型——expert_model_parallel_size对非 MoE 模型无效或报错
  5. 序列长度不整除——CP 要求seq_length % (2 × CP) == 0
  6. 盲目照搬表格——选型表是起点,务必先跑第一步 profiling 验证显存与通信

进阶提示:当通信成为瓶颈时,可进一步开启通信重叠(如tp_comm_overlap、overlap_moe_expert_parallel_comm)把通信隐藏在计算之下,指南见 docs/training/communication-overlap.md。

六、快速上手:4 卡验证你的并行布局

配好并行后,用最小 recipe 做一次冒烟测试即可确认布局正确:

uv run python -m torch.distributed.run --nproc_per_node=4 \ scripts/training/run_recipe.py \ --recipe llama32_1b_pretrain_config \ model.tensor_model_parallel_size=2 \ model.pipeline_model_parallel_size=2 \ model.sequence_parallel=True \ train.train_iters=3 train.global_batch_size=8 train.micro_batch_size=1

成功标准:退出码 0、第 3 步 loss 为有限值、日志显示TP=2 PP=2 DP=1布局。更多并行化实践(MoE 优化、长上下文、通信重叠)可参考 docs/training/ 下的各篇专题文档。

结语

记住三句话就能用好 Megatron-Bridge 的混合并行:TP 看激活参数、留在节点内;PP 管跨节点深度、记得开交错;EP 是 MoE 扩展主维度,CP 专治长序列。DP 会按公式自动补齐,先用推荐表起步,再用 profiling 微调,就能为你的模型找到显存与吞吐的最优平衡点。

  • 人工智能
  • 大模型
  • 预训练
  • 微调
  • LoRA
  • 分布式训练
  • 多模态

【免费下载链接】Megatron-Bridge

Training library for Megatron-based models with bidirectional Hugging Face conversion capability

项目地址:https://gitcode.com/gh_mirrors/me/Megatron-Bridge
点击查看免费下载
上一篇:MuPDF JavaScript 绑定中的 Page 对象:渲染、文本提取、链接管理与条码识别的完整 API 指南
下一篇:Dolphin 映像转换:RVZ + Zstd 帮你省一半空间

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询