训练框架账：Megatron-LM 与 DeepSpeed 的并行抽象

PyTorch 给了原语（DDP/FSDP），大规模训练框架把这些原语编排成体系——Megatron-LM 强在张量/流水并行，DeepSpeed 强在 ZeRO 与内核优化。本文对比两套框架的并行抽象与选型逻辑，衔接 分布式并行全景 与 PyTorch 分布式引擎。

1. 框架的四层抽象

层	职责	代表实现
调度层	数据加载、步进、checkpoint、日志	Trainer / 脚本
并行层	DP/TP/PP/SP/ZeRO 的编排	Megatron、DeepSpeed
通信层	All-Reduce/All-Gather/All-to-All 封装	NCCL 之上封装
内核层	算子融合、混合精度、显存优化	FlashAttention、ZeRO-Offload

判断：框架之争主要是"并行层 + 内核层"的抽象之争——底层都是 NCCL + 同一批算子。

2. Megatron-LM：张量与流水并行为主

能力	做法	适用
张量并行（TP）	单层参数按列/行切到多卡	单卡放不下单层
流水并行（PP）	按层切分，微批次流水	层数深、模型大
序列并行（SP）	注意力部分按序列切	长序列
虚拟流水（VPP）	细粒度切层缓解气泡	PP 气泡优化

关键：TP 通信密集（每层都要 All-Reduce），必须在同一节点/NVLink 域内；PP 通信稀疏，可跨节点（见 拓扑账）。

3. DeepSpeed：ZeRO 与内核优化为主

能力	做法	适用
ZeRO-1/2/3	状态分片，显存换通信	显存瓶颈（见 ZeRO 账）
流水并行	层切分 + 负载均衡	大模型
Offload	优化器/权重卸载 CPU/NVMe	显存极限场景
内核优化	算子融合、FlashAttention 接入	吞吐提升

关键：DeepSpeed 的差异化在 ZeRO 与 Offload——与 Megatron 的 TP/PP 是互补关系而非竞争。

4. 选型矩阵

场景	推荐	理由
≤32 卡、显存紧张	DeepSpeed ZeRO-2/3	部署简单，显存收益直接
70B+ 单节点放不下	Megatron TP+PP	TP 跨节点受限，先 TP 后 PP
千卡级超大模型	Megatron+DeepSpeed 组合	TP×PP×DP×ZeRO 四维
长序列训练	Megatron SP + FlashAttention	序列切分 + 注意力优化

判断：不是二选一——生产大模型通常 Megatron 管 TP/PP、DeepSpeed 管 ZeRO，框架组合是常态（例如 Megatron-DeepSpeed 融合实现）。

5. 可复算 Python：TP/PP 组合的卡数约束

def tp_pp_layout(gpus, tp, pp):
    """TP/PP 组合合法性：TP 必须整除卡数，PP 分组跨节点"""
    if gpus % tp != 0:
        return "TP 不整除卡数，非法"
    n_nodes = gpus // tp // pp
    return f"{gpus}卡 → TP{tp}×PP{pp}×DP{n_nodes}，每 TP 组 {tp} 卡"

tp_pp_layout(64, 8, 4)
# 64卡 → TP8×PP4×DP2：TP8 需同节点（8卡/节点），PP4 跨节点流水

6. 交付清单

1. ≤32 卡先上 DeepSpeed ZeRO，别急着上 TP/PP；
2. 70B+ 用 Megatron TP+PP：TP 组钉同一节点，PP 跨节点；
3. 千卡级用 TP×PP×DP×ZeRO 组合，每维单独验效率；
4. 长序列补 Megatron SP + FlashAttention；
5. 每次改并行布局跑 验收基线，看吞吐与通信占比。

---

本文为本站整理思考后撰写，用于知识库学习交流；公式与算例为本站自行推导整理，具体版本与实测数据以现场为准。