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

现在训练一个 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 内存再多，只要网络比算力慢，就是白搭。

 

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