并行与分布式推理：Megatron-LM / GPipe / ZeRO

当模型参数超过单卡显存（70B 模型仅 FP16 权重就约 140GB，A100 只有 80GB），答案只有一条：把模型拆开，让多卡一起扛。拆的方向不同，就有了本文的三篇论文——Megatron-LM（层内拆：张量并行）、GPipe（层间拆：流水并行）、DeepSpeed ZeRO（状态拆：参数/梯度/优化器分区）。它们本是训练技术，如今却是大模型推理的标配。

读之前建议先过一遍 GPU 与硬件（NVLink、显存层级）和 大模型推理优化。

并行策略总览

策略	拆什么	解决的问题	适用
数据并行（DP）	数据拆给多卡，每卡一份完整模型	训练吞吐（并行度）	训练为主；推理不用（每卡重复存权重浪费）
张量并行（TP）	层内把矩阵按行/列切分给多卡	单层太大，放不进单卡	推理 + 训练（vLLM -tp 参数）
流水并行（PP）	层间把层分给多卡，串成流水线	整体放不下，切层分摊	推理 + 训练（vLLM -pp 参数）
ZeRO	训练状态（优化器/梯度/参数）分区	训练显存冗余；推理侧 = 参数分区/Offload	训练为主；推理用 ZeRO-Inference 变体

并行度公式

一张卡能装下的模型大小 = 单卡显存 × 卡数 ÷ 并行度开销。常见部署组合：DP × PP × TP（三维并行）。Megatron 论文把它叫「1D/2D/3D 并行」——只做 TP 是 1D，TP+PP 是 2D，TP+PP+DP 是 3D。

一、Megatron-LM：层内张量并行（Shoeybi et al., 2019）

一句话贡献

用几条通信原语实现层内张量并行：把 Transformer 层的矩阵乘按行/列切到多卡，每层只多两次 all-reduce，即可训练/推理单层放不进单卡的大模型；8.3B 参数在 512 卡上跑到 15.1 PFLOPS、76% 扩展效率。

背景与动机

2019 年的并行训练靠数据并行，但模型大到一定程度，一层的权重 + 激活都放不下一张卡——数据并行无解（每卡都要完整模型）。Megatron 的选择是「层内切」：既然矩阵乘 Y = X·W，那把 W 沿列切成两块、沿行切成两块，让不同卡各算一半，最后拼起来——不需要新编译器，只在 PyTorch 里插入几条通信语句。

方法核心

以 Transformer 的 MLP 和注意力为例：

张量并行：把权重矩阵 W 切成两块，X 不变

Y = X · W        ┌──────────┬──────────┐
                  │  W1(列切) │  W2(列切) │     → Y = [Y1, Y2]
                  └──────────┴──────────┘
 卡0: Y1 = X·W1   卡1: Y2 = X·W2   （f 切法：计算后通信拼接）

Z = W · X        ┌──────────┐
                  │  W1(行切) │           → Z = Z1 + Z2（先各自算，再 all-reduce 求和）
                  ├──────────┤
                  │  W2(行切) │
                  └──────────┘
 卡0: Z1 = W1·X   卡1: Z2 = W2·X  （g 切法：先通信再计算）

• f / g 两种切法交替：列切（f）输出拼接、行切（g）输出求和，保证每个 Transformer 层只引入两次 all-reduce 通信，且通信可与计算重叠。
• 1D → 2D → 3D：张量并行（层内）与流水并行（层间）、数据并行（数据）正交组合；论文用 8 卡 TP × 64 卡 DP 等配置训练。
• 配套工程：fused 算子（bias+gelu 融合）、scaled masked softmax、通信与计算 overlap。

关键结果

指标	数字
规模	8.3B 参数 GPT-2 类模型，512 卡训练
吞吐	15.1 PFLOPS（全应用），单卡基线 39 TFLOPS（约峰值 30%）
扩展效率	76%（相对单卡强基线，512 卡）
效果	WikiText-103 困惑度 10.8、LAMBADA 66.5%——当时 SOTA

局限

• 通信开销随 TP 卡数增长：每层两次 all-reduce，TP 太大时通信反噬——所以 TP 一般限制在单节点内（NVLink 带宽高），跨节点用 PP。
• 需要每层能被切成整数块（head 数、hidden 维度对齐）。

对今天的启示

• 推理侧的第一选择：vLLM、TensorRT-LLM、DeepSpeed 的 tensor-parallel-size 就是它；70B 模型用 2-4 张 A100 靠 TP 就能跑起来，见 vLLM 实战。
• 判断依据：单层放不下 → 必须 TP；层多整体放不下 → TP+PP。

二、GPipe：层间流水并行（Huang et al., 2018）

一句话贡献

把模型按层切成 K 段、每段放一块 GPU，用微批（micro-batch）把流水线灌满——通过「切微批」让并行度不再等于模型批大小，6B 参数 128 层 Transformer 接近线性加速。

背景与动机

按层切分（把 1-32 层放卡0、33-64 层放卡1…）看似简单，但朴素的「整批一次过流水线」会形成严重的 bubble（气泡）：每块卡大部分时间在等上游算完，利用率惨不忍睹。GPipe 的关键改进：把 batch 切成多个微批（micro-batch）流水执行，让所有卡同时处于忙碌状态。

方法核心

K=4 块卡，M=8 个微批，流水线执行（以第 0 层为起点示意）：

时间 →
卡0: [mb1][mb2][mb3][mb4][mb5][mb6][mb7][mb8]
卡1:      [mb1][mb2][mb3][mb4][mb5][mb6][mb7][mb8]
卡2:           [mb1][mb2][mb3][mb4][mb5][mb6][mb7][mb8]
卡3:                [mb1][mb2][mb3][mb4][mb5][mb6][mb7][mb8]
                     └──────── 尾部气泡 ────────┘

微批越多，灌得越满，空闲（气泡）占比越小

• 气泡开销 = (K-1)/(M+K-1)：卡数 K 越少、微批 M 越多，气泡占比越低（M→∞ 时趋近 0）。
• 切微批让「并行度」与「batch size」解耦：以前要并行度 8 就得 batch=8，现在 batch=1 切成 8 个微批也能跑 8 卡。
• 激活重计算（re-materialization）降低训练显存峰值。

关键结果

指标	数字
模型	6B 参数、128 层 Transformer（多语言翻译）；557M AmoebaNet（ImageNet 84.4% top-1）
加速	切微批后接近线性的加速比（气泡被灌满）
适用范围	任何「层序列」模型都适用，与架构解耦

局限

• 气泡无法完全消除（首尾灌不满）。
• 训练侧激活重计算拖慢一点速度；推理侧 PP 增加首 token 延迟（请求要「流过」整条流水线）。
• 各层计算量不均时（embedding、输出层 vs 中间层）负载不平衡。

对今天的启示

• 推理侧 PP 与 TP 组合使用：TP 解决「一层放不下」，PP 解决「层太多」；vLLM 已支持 PP。
• 部署时注意：PP 带来的首 token 延迟增加、以及跨节点通信瓶颈，小规模场景优先 TP。
• 气泡分析（(K-1)/(M+K-1)）是评估任何流水线方案的第一张计算纸。

三、DeepSpeed ZeRO：状态分区（Rajbhandari et al., 2019）

一句话贡献

把数据并行里「每卡都存一份」的冗余去掉：优化器状态、梯度、参数三个阶段逐层分区（Zero Redundancy），让显存随卡数线性扩展——400 卡训练 100B+ 参数达到 15 PFLOPS、超线性加速。

背景与动机

数据并行看似简单，但每张卡都存完整模型 + 完整优化器状态——对 7.5B 模型，纯数据并行每卡要 ~120GB，其中模型权重只占 30GB，冗余占了大头（优化器状态 75GB、梯度 30GB）。ZeRO 的洞察：这些冗余不是必须的，可以像模型并行一样分区存储，但保留数据并行的通信模式（通信量低、粒度大）。

方法核心：三阶段

Stage 1：优化器状态分区    → 每卡只存 1/N 优化器状态（显存 ÷N）
Stage 2：+ 梯度分区        → 每卡只存 1/N 梯度（显存 ÷N）
Stage 3：+ 参数分区        → 每卡只存 1/N 参数（显存 ÷N）
         （计算某层时才 all-gather 该层参数，用完即丢）

Stage 3 让「每卡显存 ÷ N」：100 张卡 = 100 倍可训练规模
推理侧对应物：ZeRO-Inference / 参数 offload 到 CPU

• 保留数据并行的通信模式：不是模型并行那套每层 all-reduce，而是「需要时 all-gather、算完即可丢」，通信量低、实现简单。
• ZeRO-Offload：把参数/优化器状态进一步放到 CPU/内存，单卡也能训超大模型（推理侧即「模型放 CPU、算在 GPU」的 offload）。

关键结果

指标	数字
规模	400 卡训练 100B+ 参数，吞吐 15 PFLOPS
加速	超线性扩展（显存解放带来的批大小收益反超通信开销）
对比	相对 SOTA：模型规模提升 8x、性能提升 10x
易用性	13B 参数无需模型并行即可训练（只靠 ZeRO）

局限

• Stage 3 的按需 all-gather 在低带宽环境下通信占比高（超线性收益依赖高速互联）。
• 训练为主；推理侧直接套 ZeRO 意义有限（推理不存梯度/优化器状态），要改用 ZeRO-Inference（权重分区 + offload）。

对今天的启示

• 推理侧：ZeRO-Inference 允许把权重分区到多卡或 offload 到 CPU 内存，让 70B 级模型在 2 卡甚至单卡 + CPU 内存上跑起来——牺牲速度换「装得下」。
• 判断：要吞吐 → TP/PP（权重常驻显存、低延迟）；要「装得下」→ ZeRO 分区/offload（慢一点，但能跑）。

四、对比表

方法	通信开销	显存效果	实现复杂度	适用场景
张量并行（Megatron）	每层 2 次 all-reduce，随 TP 卡数增长	单层权重/激活分散到多卡	中（需算子级切分支持）	推理首选：单层放不下；vLLM -tp
流水并行（GPipe）	仅阶段间传激活，通信低	每卡只存部分层	中（需层切分 + 微批）	训练 + 大模型推理（与 TP 叠加）
ZeRO（DeepSpeed）	Stage 3 按需 all-gather，高	参数/梯度/优化器随卡数线性分摊	高	训练为主；推理用 offload 变体

五、推理侧并行实践

大模型推理的常用组合是 TP + PP（vLLM 直接支持）：

• 单节点、层能放下：纯 TP（如 2-4 卡 NVLink 内），通信快、延迟低。
• 模型太大 / 多节点：TP（节点内）× PP（节点间），节点间只传激活，规避跨节点通信瓶颈。
• 极度缺显存：ZeRO-Inference / offload，让权重躺在 CPU 内存，GPU 按需取用——吞吐低但「能跑」。
• 部署细节与参数（--tensor-parallel-size、--pipeline-parallel-size）见 vLLM 实战案例。

并行度不是越大越好

TP 卡数翻倍 → 通信量翻倍；PP 层数翻倍 → 气泡与首 token 延迟翻倍。并行度只加到「模型放得下 + 延迟可接受」为止，多出来的卡开多个副本（DP）接流量，往往吞吐更高——这正是 部署架构模式 里的水平扩展思路。

六、共同局限

1. 通信瓶颈：跨节点带宽（如 25/50/100 Gbps 网卡）远低于 NVLink（600GB/s 级），TP 跨节点会明显降速——TP 留在节点内是铁律。
2. 负载不均：PP 各层计算量不均、TP 切分不整除都会造成「木桶效应」，切分要按计算量/显存精细化。
3. 复杂度：三维并行的调度、容错、显存规划是自研引擎的高门槛——这也是为什么优先用 vLLM、TensorRT-LLM、DeepSpeed 这类成熟框架，见框架对比。

延伸阅读

• 大模型推理优化 —— KV Cache、调度与并行在 LLM 推理中的完整拼图
• vLLM 实战案例 —— -tp/-pp 参数的真实配置与选型
• PagedAttention：vLLM 系统论文 —— 显存管理如何与并行叠加
• GPU 与硬件选型 —— NVLink、显存带宽如何决定并行方案
• 推理服务系统：Clipper / Orca / Nexus 等 —— 分布式服务的调度视角
• 部署架构模式 —— 单卡放得下时为何优先水平扩展

参考资料

• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism（arXiv 1909.08053）
• GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism（arXiv 1811.06965）
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models（arXiv 1910.02054）
• DeepSpeed-Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale（arXiv 2207.00032）