百亿参数模型的并行训练：节点内张量并行、节点间数据并行

现在训练一个 1000 亿参数的 Transformer 模型已经算不上什么新鲜事。GPT-3 有 1750 亿参数，Llama 2 最大版本达 700 亿，许多团队现在随口就把"100B 作为基准"挂在嘴边。但第一次真正上手训练的团队，往往会在一个意想不到的地方撞墙：不是算力，是内存。

几乎所有人都会觉得模型放不进 GPU 内存就买更多 GPU，训练慢就换更快的卡。如果你这样想买再多 H100 也没用，因为要不清楚如何把模型参数、激活值（activations）和优化器状态（optimizer state）的内存占用分散到这些卡上，结果就是网络把算力拖死，以峰值吞吐量的零头跑着任务。100B 参数这个量级模型的核心问题是：哪些组件可以跨设备做切片（shard），哪些可以复制（replicate），哪些因为切片的通信开销比省下来的计算还贵，必须老老实实留在一张卡上。

Meta 在训练其最大规模的开源模型时，投入了大量工程资源处理这个问题。他们的技术博客，谈"扩展（scaling）"时用的语言几乎清一色是内存效率，而不是绝对吞吐量。并行策略的选择——数据并行（data parallelism）、张量并行（tensor parallelism）、流水线并行（pipeline parallelism）还是组合——本质上是在决定对模型的哪个维度做切分，而每种选择都会产生一套通信模式，这套模式如果不契合硬件，就会把整个集群卡死。

内存问题：为什么 100B 参数放不进一张 GPU

从算术说起，一个 100B 参数的 Transformer 模型，仅权重在 float32 下就占约 400 GB 显存。H100 只有 80 GB。换成 bfloat16（训练中的常用格式）也需要 200 GB，显然是不够的。而且这还只是权重，训练过程中还有三类额外的内存消耗：

1. 优化器状态——使用 Adam 时，每个参数需要维护一阶矩和二阶矩两个值，内存占用直接翻倍。100B 参数模型在 bfloat16 下，权重 200 GB，优化器状态再占 400 GB，合计 600 GB。仅仅是存储可训练状态，就需要大约 8 张 H100。
2. 激活值——前向传播需要把每层的激活值保存下来，供反向传播计算梯度用。对于 100B 参数的 Transformer，根据序列长度和 batch size 的不同，反向传播期间激活值可能再占 200–500 GB。
3. 梯度缓冲区——在执行优化器步骤之前，每个参数对应的梯度副本需要临时驻留在内存中。

以合理的 batch size 和序列长度跑一步 100B 模型的训练，总内存需求大约在 800 GB 到 1.2 TB。H100 只有 80 GB，光是所有数据就需要 10–15 张卡。如果单纯用数据并行——每张卡都复制完整模型，处理不同 batch，计算梯度后取平均——那每个副本就得用 15 张卡。如果需要两个副本做容错？那就30 张。要是还要流水线化 batch就要继续加。

Meta 等机构最终得出的洞察是：并不是所有内容都需要在 GPU 之间复制。模型参数可以切片；激活值只在某一设备的前向和反向传播期间存在，可以流水线处理；优化器状态可以独立于参数单独切片。这一切都要以最小化跨网络的数据流量为前提——GPU 内存再多，只要网络比算力慢，就是白搭。

数据并行：最简单的基准方案

理解数据并行为何在 100B 量级失效是进入后续话题的前提。

数据并行中，每张 GPU 保存完整的模型副本，前向传播时各自独立处理不同的 batch，反向传播后各自得到梯度，再通过 all-reduce 操作取平均，最后每张卡独立执行一次优化器步骤。参数量在几十亿以内、模型能放进单卡时，这套流程运行良好。

100B 参数下的问题在于，bfloat16 的权重（200 GB）需要复制到每一张卡。16 张卡就意味着 3.2 TB 的模型副本。每次优化器步骤都要同步梯度，对 100B 参数来说，这是 200 GB（bfloat16）数据在集群网络上的一次传输。即便走 NVLink（900 GB/s）这样的高速互联，每步也要约 0.22 秒。如果一次计算本身只需要 300 毫秒，光同步就吃掉了 40% 的周期时间。同步时间不随模型规模线性改善，模型越大越难看——到 500B，这个问题会彻底失控。

数据并行实现简单（PyTorch DDP 开箱即用）所以是起点，但在 100B 量级单独使用，几乎不是正确选择。

张量并行（Tensor Parallelism）：跨 GPU 切分权重矩阵

张量并行，也称层内并行（intra-layer parallelism），把每一层的权重矩阵切分到多张 GPU 上，参数本身不再集中在一张卡。

以 Transformer 的多头注意力（multi-head attention）层为例，将输入投影到各注意力头的权重矩阵形状为 [hidden_dim, num_heads * head_dim]。张量并行沿列方向把这个矩阵切开：GPU 0 持有第 0 到 N/2 列，GPU 1 持有第 N/2 到 N 列。前向传播时每张卡各自计算输出的一部分，再拼接结果；反向传播时梯度以相同方式切分。

通信开销体现为 all-gather 操作，而非数据并行中的 all-reduce。8 路张量并行下，每张卡只保存 1/8 的参数，不再需要复制完整模型——这是张量并行在 100B 量级优于数据并行的直接原因。但是每经过一个线性层都需要 all-gather 来重新拼装完整输出。若能与下一层的计算重叠进行，这部分开销是非阻塞的；若模型层数不够深，或并行度太高，同步时间就会超过计算时间。

Meta Llama 2 70B 的训练采用了 2 路张量并行（通过 NVLink 桥接），将通信开销控制在合理范围内。100B 以上的前沿实践通常用 4 路或 8 路，不会再高因为超过这个度，通信开销就会占主导。

流水线并行（Pipeline Parallelism）：按层切分模型

流水线并行的思路与张量并行截然不同。它不切分参数，而是把模型的层分配到不同 GPU：GPU 0 负责前 N 个 Transformer 块，GPU 1 负责接下来的 N 个，依此类推。前向传播时数据从左向右流经各卡，反向传播时从右向左流回。

100 层的 Transformer 4 路切分，每张卡只处理 25 层，内存大约降低 4 倍，好处是显而易见，但代价同样明显：流水线气泡（pipeline bubbles）。GPU 0 处理完一个 batch 并将结果传给 GPU 1 后，如果 GPU 1 还在处理上一个 batch，GPU 0 只能空等。典型前向传播过程中，任意时刻往往只有一张卡在做有效工作，其余都在等待，GPU 利用率可能低至 25% 甚至更低。

Gpipe 以及 vLLM 的变体通过在流水线中交错多个 batch 来缓解这个问题，但根本矛盾仍在：流水线并行是以延迟换内存，训练时延迟直接等于浪费掉的 GPU 周期和能耗。

实践最优点：张量并行 + 数据并行

大多数 100B 量级的团队最终收敛到同一种混合策略：节点内张量并行，节点间数据并行。

原因在于集群架构本身。节点内 GPU 通过 NVLink 互联，H100 上达 900 GB/s；节点间走集群网络，速度可能只有 400 Gbps 或 800 Gbps，低一到两个数量级。张量并行放在节点内，通信走的是快速通道；数据并行放到节点间，不需要在每层都同步，每个训练步骤只发生一次。

这种组合的效果是：模型参数在每张卡上放得下，张量并行的通信走高速 NVLink，数据并行的同步分散到较慢的集群网络，但每步只做一次。

梯度累积（Gradient Accumulation）：不增加内存地模拟大 batch

几乎每次 100B 训练都会用到梯度累积——不在每次前向-反向传播后立即更新，而是跨多个微批次（microbatch）累积梯度再统一执行优化器步骤。

这样做有两个原因：1、有效 batch size：batch size 越大，训练动态越稳定，但单卡内存有限时只能用小 batch；梯度累积允许用四个 128-token 的微批次模拟一个 512-token 的 batch，无需一次性分配 512-token 激活值的完整内存。2、使用流水线并行时，梯度累积能让流水线保持满载，多个微批次流过各阶段后再同步，减少空闲时间。

不过会更加耗时，因为每更新一次权重要做四倍的前向-反向传播，每步训练时间更长。但有效 batch size 更大，通常有助于收敛，也允许使用更大的学习率。

如何选择并行策略

训练一个 100B 参数模型之前，可以按以下决策树逐步确认：

1. 模型能放进一张 GPU 吗？（100B 量级极少，除非使用 int4 量化。）→ 数据并行。结束。
2. 模型能放进一个节点（8 张 GPU）吗？→ 节点内张量并行（如 8 路），节点间数据并行。
3. 切分到一个节点仍然太大 → 节点内 4 路张量并行 + 节点间 2 路流水线并行 + 更大范围集群的数据并行。这是 100B 以上的主流方案。
4. 模型层数很深且通信延迟低 → 可以考虑加入轻量级流水线并行，但务必仔细测量实际耗时，它往往适得其反。
5. 追求最大吞吐量且集群网络带宽充足 → 可以考虑 Megatron 风格的 3D 并行（张量 + 流水线 + 数据），但仅当集群规模足够大、节点间带宽良好（每链路 ≥400 Gbps）时才值得。

张量并行度越高，模型内存越小，节点内同步开销越大；流水线并行度越高，内存进一步降低，气泡也越多；数据并行度越高，每步网络开销越大，但实现成本最低。100B 量级没有唯一最优解，只有最适合当前硬件和模型架构的那个答案。

为什么通信会成为瓶颈

超过某个临界点后，网络速度就慢于算力，再多的 GPU 也改变不了这一点。一次 100B 参数的训练步骤中：

• 计算时间（前向 + 反向）：经过良好并行化的 8-GPU 配置上约 300 ms
• 梯度 all-reduce（数据并行同步）：典型集群网络上约 200 ms
• 张量并行 all-gather：约 50 ms（NVLink 上，与计算重叠进行）
• 流水线气泡：0–100 ms，取决于策略

各项相加，每步约 500–600 ms，而计算只占其中 50%。把计算速度提升 20%，每步时间只改善 10%，因为网络依然卡在那里。

这正是 Llama 2 70B 训练报告中指出"跨机器通信是主要瓶颈"的原因，也是 Meta 等机构在 ZeRO（零冗余优化器，zero redundancy optimizer）和梯度检查点（gradient checkpointing）上投入大量资源的原因——目的是压缩需要通信或驻留在内存中的状态量。

PyTorch FSDP 与 Megatron

张量并行和流水线并行不需要手动实现，目前有两个主流框架。

PyTorch FSDP（Fully Sharded Data Parallel）是 PyTorch 的原生方案，将优化器状态、梯度和参数切分到各 GPU，侧重内存效率。FSDP 通过自动切分层来达到接近手动张量并行的效果，使用上更简单，但若调优不当，可能引入更多同步开销。

Megatron-LM 是 NVIDIA 专为 100B 以上规模训练构建的框架，实现了显式的张量并行、流水线并行，以及序列并行（sequence parallelism——对序列维度做切分的变体）。Megatron 更复杂，但对并行策略的控制更精细。世界上规模最大的一批模型，大多是在 Megatron 或其变体上训练出来的。

典型的 100B 训练要么用 FSDP 配合选择性激活检查点，要么用 Megatron 风格的张量并行。具体选择通常取决于模型架构能否与框架的假设良好匹配。

总结

如果你读完关于 100B 模型训练的论文，几乎都会得出同一个结论：瓶颈不在于数据移动的速度，而在于内存里能存多少、以及在移动数据的同时能让 GPU 保持多忙。100B 参数模型不只是 1000 亿个等待算力处理的数字，它是 1000 亿个在处理期间必须有地方存放的数字，而"存放的地方"是稀缺资源。

并行策略的选择，决定了这个"存放的地方"的结构。数据并行在各卡间复制这个地方；张量并行把它切片分配到各卡，附带节点内通信；流水线并行让它跨时间和节点展开，以延迟换内存；梯度累积则通过顺序处理更小的片段来模拟一个更大的空间。

每种方案都是一个取舍判断：内存、网络带宽和计算时间，哪个更贵？三者的约束在 100B 量级同时压下来，这正是混合策略成为主流的原因。工程不是要找到某种巧妙方案绕过物理限制，而是接受这些限制，围绕它们做设计。

成功扩展 PyTorch 训练的团队，花力气的方向不是让训练跑得更快，而是让训练不必等待。

https://avoid.overfit.cn/post/9215b925239c4b4ab4bd60668977dd4e

by Shriom Tripathi