关于 nemo-mbridge-perf-parallelism-strategies
为 Megatron Bridge 开发者提供模型并行策略选择指南,按模型规模、硬件拓扑和序列长度给出起始配置,并说明 MoE 的 GPU 数量估算方法。指南强调这些配置是启发式起点,需通过性能分析验证。
主要功能
- 01按模型规模选择并行维度
- 02按硬件拓扑配置 TP 与 PP
- 03按序列长度选择 SP 与 CP
- 04MoE 最小 GPU 数量计算
- 05显存估算与配置验证要点
适用场景
- 01为新模型确定并行配置起点
- 02估算 MoE 训练所需 GPU 数量
- 03检查跨节点 TP 与网络拓扑取舍
- 04验证组合并行配置能否初始化
安装
复制下面的命令,在终端中运行。
npx skills add https:/ / github.com/ promptingcompany/ nv-skills --skill nemo-mbridge-perf-parallelism-strategies安装后在 Claude Code 中输入/nemo-mbridge-perf-parallelism-strategies确认已加载。
包含的文件
除 SKILL.md 外,安装时一起下载的参考资料和脚本。
BENCHMARK.mdcard.yamlskill-card.mdskill.oms.sigevals/evals.json详细介绍
并行策略选择技能
有关各种并行类型的稳定背景信息,请参阅:
- @docs/parallelisms.md
- @skills/nemo-mbridge-perf-parallelism-strategies/card.yaml
按模型规模选择
稠密模型
| 模型规模 | GPU 数量 | 推荐起点 |
|---|---|---|
| < 1B | 1-8 | 仅 DP |
| 1-10B | 8-16 | TP=2-4 + DP |
| 10-70B | 16-64 | TP=4-8 + PP=2-4 + DP |
| 70-175B | 64-256 | TP=8 + PP=4-8 + DP |
| 175-500B | 256-1024 | TP=8 + PP=8-16 + CP=2 + DP |
MoE 模型
MoE 并行与稠密模型不同。由于每个 token 只激活一部分参数,TP 通常可以保持在 1 或 2——已激活的参数分片可以放在单个 GPU 上。EP 是主要的扩展维度,PP 则负责跨节点分布层。
| 模型(总参数 / 激活参数) | TP | PP | EP | 备注 |
|---|---|---|---|---|
| OLMoE 7B / 1B | 1 | 1 | 8 | 仅 EP,可容纳于单节点 |
| Moonlight 16B / 3B | 2 | 1 | 8 | 对共享层使用较小的 TP |
| DeepSeek-V2 236B / 21B | 1 | 4 | 32 | 完全不使用 TP |
| GLM-4.5 Air 106B / 12B | 1 | 4 | 8 | 完全不使用 TP |
| Qwen3 30B-A3B | 4 | 2 | 4 | |
| GLM-4.5 355B / 32B | 2 | 8 | 16 | |
| Qwen3 235B-A22B | 4 | 16 | 8 | 预训练使用 CP=2 |
| DeepSeek-V3 671B / 37B | 2 | 16 | 64 | TP=2,而非 8 |
| Kimi-K2 1T | 2 | 16 | 32 |
主要规律:
- TP 的配置取决于激活参数,而非总参数。一个激活参数为 37B 的 671B MoE 模型,所需 TP 远低于 70B 稠密模型。
- EP 随专家数量扩展。常见配置:EP = num_experts 或 num_experts / experts_per_gpu。
- PP 负责模型深度。大型 MoE 模型会在节点间使用 PP=8-16。
- ETP(专家张量并行)很少使用。Llama 4 是一个例外(ETP=4)。
这些只是起点,而非硬性规则。务必对首次迭代进行性能分析,以验证内存和通信情况。
按硬件拓扑选择
使用 NVLink 的单节点:
cfg.model.tensor_model_parallel_size = 8
使用 InfiniBand 的多节点:
cfg.model.tensor_model_parallel_size = 8
cfg.model.pipeline_model_parallel_size = N
网络受限(以太网):
cfg.model.tensor_model_parallel_size = 4
cfg.model.pipeline_model_parallel_size = M
稳定规则是:将 TP 限制在单个 NVLink 域内。跨节点扩展时使用 PP 或 DP。跨节点使用 TP 几乎总会导致性能下降。
按序列长度选择
| 序列长度 | 建议 |
|---|---|
| < 2K | 标准 TP + PP + DP |
| 2K-8K | 添加 SP(sequence_parallel=True) |
| 8K-32K | 添加 CP=2 |
| 32K+ | 添加 CP=4-8;对于较大的 CP,可考虑 a2a+p2p |
组合启用并行
3D 并行(TP + PP + DP):
cfg.model.tensor_model_parallel_size = 4
cfg.model.pipeline_model_parallel_size = 4
cfg.model.sequence_parallel = True
4D 并行(TP + PP + CP + DP):
cfg.model.tensor_model_parallel_size = 8
cfg.model.pipeline_model_parallel_size = 8
cfg.model.context_parallel_size = 2
cfg.model.sequence_parallel = True
使用 EP + PP 的 MoE(例如,在 128 个 GPU 上运行 DeepSeek-V2 236B):
cfg.model.tensor_model_parallel_size = 1
cfg.model.pipeline_model_parallel_size = 4
cfg.model.expert_model_parallel_size = 32
cfg.model.sequence_parallel = False
使用小规模 TP + PP + EP 的 MoE(例如,在 256 个 GPU 上运行 DeepSeek-V3 671B):
cfg.model.tensor_model_parallel_size = 2
cfg.model.pipeline_model_parallel_size = 16
cfg.model.expert_model_parallel_size = 64
cfg.model.sequence_parallel = True
DP 大小始终为隐式值:
data_parallel_size = world_size / (TP * PP * CP) # dense path
expert_data_parallel_size = world_size / (PP * EP * ETP) # MoE path
最少 GPU 数量
运行某项配置所需的最少 GPU 数量(即 DP=1、EDP=1)并非所有并行维度的乘积。稠密路径使用 TP*CP 网格,MoE 路径使用 EP*ETP 网格;在每个 PP 阶段,这两个网格共用同一组 GPU——它们相互重叠,而不是相乘。因此:
min_gpus = PP * max(TP * CP, EP * ETP)
常见的简化算法(错误): PP * TP * CP * EP * ETP。这会导致 GPU 分配过多,而且常见于许多 README 和 slurm 资源配置表中。不要沿用这种算法。
注意力与 MoE 并行的解耦(稠密路径和专家路径使用不同的网格形状,但共用同一组 PP 阶段 GPU)详见 Pangu Ultra MoE(arXiv:2504.14960)。
示例
| 配置 | 错误(PP·TP·CP·EP·ETP) | 正确(PP·max(TP·CP, EP·ETP)) |
|---|---|---|
| PP=1, TP=2, CP=1, EP=8, ETP=1 | 16 | 8(1 个节点) |
| PP=1, TP=4, CP=1, EP=8, ETP=1 | 32 | 8(max(4, 8)) |
| PP=1, TP=2, CP=2, EP=8, ETP=1 | 32 | 8(max(4, 8)) |
| PP=1, TP=2, CP=4, EP=8, ETP=1 | 64 | 8(max(8, 8)) |
| PP=2, TP=2, CP=1, EP=8, ETP=1 | 32 | 16(2 · max(2, 8)) |
| PP=1, TP=2, CP=1, EP=4, ETP=2 | 16 | 8(max(2, 8)) |
扩展到最低配置以上
增加 GPU 会扩展 DP 和/或 EDP(world_size 必须同时满足这两个公式)。在 min_gpus 下,较大网格一侧的 DP(或 EDP)为 1,较小网格一侧则吸收余量。
示例 — TP=2、CP=1、EP=8、ETP=1、PP=1:
- 8 个 GPU(
min_gpus):稠密模型DP = 8/2 = 4,MoEEDP = 8/8 = 1 - 16 个 GPU:稠密模型
DP = 8,MoEEDP = 2→ 全局批次大小扩大 2 倍 - 32 个 GPU:稠密模型
DP = 16,MoEEDP = 4→ 全局批次大小扩大 4 倍
配置 slurm 脚本时,根据 min_gpus 计算 --nodes(若要通过 DP/EDP 提高吞吐量,则使用其倍数)。
回答 MoE 配置问题时,请包含以下检查清单:
- 使用请求的值计算
min_gpus = PP * max(TP * CP, EP * ETP) - 明确指出错误的
PP * TP * CP * EP * ETP全量乘积不可用 - 给出两个 DP 公式:稠密模型
world_size / (TP * PP * CP)和 MoEworld_size / (PP * EP * ETP) - 提及 TP 拓扑、SP、CP 可整除性,以及长序列的 CP 建议
内存估算
不使用并行(70B 模型,FP16):
parameters: 140 GB
gradients: 140 GB
optimizer states: 280 GB (Adam)
activations: 48 GB (batch=1, seq=4K)
total: 608 GB
使用 TP=4、PP=4、DP=4(64 个 GPU):
parameters: 8.75 GB per GPU
gradients: 8.75 GB per GPU
optimizer states: 17.50 GB per GPU
activations: 3.00 GB per GPU
total: ~38 GB per GPU
代码锚点
在模型提供程序中设置并行维度:
model_config = GPTModelProvider(
tensor_model_parallel_size=2,
# ... other model parameters
)
DP 大小计算:
data_parallel_size = world_size / (tensor_model_parallel_size × pipeline_model_parallel_size × context_parallel_size)
Bridge 初始化会将并行设置接入进程组:
parallel_state.initialize_model_parallel(
tensor_model_parallel_size=model_config.tensor_model_parallel_size,
pipeline_model_parallel_size=model_config.pipeline_model_parallel_size,
...
context_parallel_size=model_config.context_parallel_size,
hierarchical_context_parallel_sizes=model_config.hierarchical_context_parallel_sizes,
expert_model_parallel_size=model_config.expert_model_parallel_size,
...
)
注意事项
-
跨节点使用 TP 会严重降低吞吐量。始终将 TP 限制在单个 NVLink 域内。
-
不交错使用 PP 会产生较大的流水线气泡。尽可能使用
virtual_pipeline_model_parallel_size。 -
SP 要求
tensor_model_parallel_size > 1。单独启用 SP 而不使用 TP 会导致配置错误。 -
CP 要求
seq_length % (2 * context_parallel_size) == 0。 -
EP 仅适用于 MoE 模型。在稠密模型上设置
expert_model_parallel_size不会生效,或会报错。 -
上方的模型大小与并行度对照表只是初始启发式建议。务必对首次迭代进行性能分析,检查内存和通信情况。
-
CUDA_DEVICE_MAX_CONNECTIONS等相关环境变量会与重叠设置相互影响。请参阅 @skills/nemo-mbridge-perf-tp-dp-comm-overlap/SKILL.md。 -
MoE 配置所需的最少 GPU 数量为
PP * max(TP*CP, EP*ETP),而非所有维度的乘积。每个 PP 阶段中,稠密模型的TP*CP网格和 MoE 的EP*ETP网格共用同一批 GPU。请参阅上方的“最少 GPU 数量”部分。
验证
使用可用的最小配置并覆盖并行设置,快速检查组合并行是否能正确初始化:
CUDA_VISIBLE_DEVICES=0,1,2,3 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 \
scheduler.lr_warmup_iters=0 \
validation.eval_iters=0 validation.eval_interval=0 \
checkpoint.save_interval=0 \
logger.log_interval=1
成功标准:
- 退出代码为 0
- 第 3 次迭代时损失值有限(例如
lm loss: 1.003808E+01) - 日志显示 TP=2、PP=2、DP=1 的布局,共有 4 个 rank