大模型训练硬件测评：GPU内存分块与并行策略

在AI和大型语言模型（LLM）这个圈子里，有一个普遍的共识：模型、数据和计算，这三驾马车是推动技术进步的核心动力。它们之间的关系盘根错节，缺一不可。就拿Llama 3来说，那个最大的模型参数超过了4000亿，是在16000块GPU上没日没夜地训练了好几周甚至好几个月才跑出来的。所以，优化计算，本质上就是想办法用更低的成本去训练更大的模型。

这篇文章，我们就来聊聊GPU那些最关键的特性，然后基于这些特性，探讨一下怎么设计出跑得更快的算法。

GPU 与 CPU 的核心差异

CPU的设计初衷是为了追求极致的单任务延迟，它擅长快速处理完一个任务，然后立刻转向下一个，这对于日常的通用计算来说，非常合理。但GPU的思路完全不同，它追求的是吞吐量，目标是在同一时间并行处理海量任务。打个比方你就明白了：CPU就像一个个人能力超强的全能工匠，而GPU则像是一群普通工人一起干活。在LLM训练这种大规模并行计算的场景下，GPU的架构显然有天然的优势。

我们再拿工厂来打个比方。可以把GPU想象成一个巨大的工业城镇。这个城镇里有好几个“工厂集群”（技术上叫流式多处理器，SM），每个集群里又包含多个工厂（流式处理器，SP）和一个小仓库（共享内存）。整个城镇还有一个巨大的中央仓库（DRAM），离各个集群比较远，但容量大得多。

这个比喻虽然简化了，但它揭示了GPU的一个核心关键：集群内部的小仓库（共享内存），访问速度比那个遥远的中央仓库（DRAM）要快得多，但代价是容量小得可怜。

至于DRAM的读取到底有多慢？一组来自20年间的数据会让你大吃一惊：硬件的浮点运算能力提升了60000倍，而DRAM的带宽只提升了100倍，互连带宽更是只有30倍。

这意味着，过去的瓶颈在计算，而现在，真正的瓶颈在于内存带宽。既然数据搬运才是最大的麻烦，那么，减少搬运的次数和数量，就是让GPU跑得更快的关键。下面这五个技巧，都是围绕这个核心思路展开的，它们源自著名的CS336课程。

技巧 1：低精度计算

在做矩阵乘法的时候，数字的精度是可选的。精度越高，存储一个数字需要的字节就越多，好比“9.327595”比“9.33”占的空间大得多。用低精度的数字，意味着搬运的“货物”更少，在拥堵的“数据通路”上花费的时间自然更短。

具体来说，就是用fp16代替fp32。不过，并不是训练的所有阶段都需要低精度，我们只需在数据搬运阶段用fp16就行。具体做法是：数据以fp16的格式传入，矩阵乘法本身在32位精度下完成（计算本身不是瓶颈，而且高精度能防止舍入误差的累积），最后输出时再降回fp16用于传输。

回到工厂那个比喻：道路太拥堵了（下图红线），所以进出工厂的箱子越小越好。而工厂内部空间很充裕，可以在大空间里完成高精度的加工，加工完再打包成小箱子运出去。

技巧 2：算子融合

假设工厂有三道工序：先把正方形变成圆形，再把圆形变成三角形，最后把三角形变成星形。如果每完成一步，都要把半成品送回仓库，再取回来做下一步，那来回搬运的次数可就太多了。

算子融合的思路，就是把多步操作干脆在工厂里一次性完成，省去中间产品的反复搬运。

实现的方式有两种：要么手写底层代码来控制融合的细节，要么直接用像`torch.compile`这样的工具来自动完成优化。

技巧 3：重计算

这个情况稍微复杂一些。假设工厂从仓库取了个正方形，依次加工成了圆形、三角形和星形。最后，星形被送回仓库供后续使用。但到了最终步骤，我们需要四种形状全用上——正方形、圆形、三角形、星形。工厂内部存不下东西，所有存储都得依赖仓库。

安排生产线有两条路：

• 选项 1： 在加工过程中，把圆形和三角形也送回仓库保管。到时候需要了，再从仓库取回来。
• 选项 2： 干脆不保存任何中间形状，做完一步就丢掉。需要的时候，再从正方形重新加工一轮。

选项1节省了重新计算的电，但增加了仓库的搬运量。选项2搬运量小，但要额外消耗算力。这本质上是一个内存与计算之间的权衡。

既然瓶颈在于“道路拥堵”而不是“车间产能”，那么重计算（选项2）显然是更合理的选择：重新计算的成本相对较低，但从仓库搬运的成本可能高出好几个数量级。用算力去换内存带宽，这笔账很划算。

技巧 4：内存合并访问

仓库有一个特点：货物是按板条箱整箱发出的。工厂如果请求其中任何一件物品，仓库都会把整个板条箱一起送过来。那么，优化的要点就在于：尽量把我们需要的物品，集中放在同一个箱子里。

假设每箱能装4件物品，工厂总共需要8件。如果这8件恰好集中在2个箱子里，那只需要取2箱就够了。但如果它们散落在8个不同的箱子里，那就得搬8箱——搬运成本翻了四倍。

技术上来说，DRAM是以“突发模式”读取的，每次读取都会返回一段连续的字节。即使处理器只需要其中一个地址的数据，整个突发段也会被送过来。当所有线程访问的地址都落在同一个突发段内时，只需要一次DRAM请求，这种情况就叫做“完全合并访问”。

由此可以得出一个直接的推论：把矩阵的维度（比如词汇表大小）对齐到64的倍数，会带来可观的速度提升。

原因很简单：分块操作（见下一个技巧）需要沿着突发段的边界来读取数据，如果分块的边界和突发段对不齐，读取次数就会急剧增加。

技巧 5：分块

分块的核心思想是：把大矩阵切成小块，加载到共享内存（也就是集群里那个小仓库）中，避免反复去访问那块慢悠悠的全局内存。

以两个4x4矩阵A和B的乘法为例，结果是另一个4x4矩阵C。要计算C的某几个元素，需要在A和B矩阵上多次跨行/跨列读取，而每一次读取都要访问全局内存。

分块的思路是，把A和B各自切成四块。这样，小块就可以完整地加载到共享内存里。首先加载红色块，计算出部分和（图中橙色部分）：

接着，再加载下一组块，继续累加部分和。总计算量并没有变，但每一步都在共享内存中完成，而不是反复去访问全局内存，节省的时间相当可观。

FlashAttention：站在技巧的肩膀上

有了上面这五个技巧做铺垫，现在可以来看FlashAttention了。它正是这些技巧的集大成者。

先简单回顾一反赌意力机制。权重矩阵会将隐藏向量投影为Q、K、V，然后对每个词的q和k向量求点积（这等价于一次Q × Kᵀ的矩阵乘法），得到原始注意力分数——也就是每个查询词对各个键词的关注程度。最后，对原始分数做一个softmax归一化，让它们的和为1。

不过这里有个数值稳定性问题。取指数之前，得先减去最大值。因为e¹²已经是162755了，超过了fp16能表示的上限65504，直接计算会溢出。减去最大值之后，计算结果不变，但能完美规避溢出问题。

归一化后的softmax分数，再与每个词的“值”向量相乘、求和，最终就得到了注意力输出。

现在回到FlashAttention上。Q和K相乘会产生一个N × N的矩阵，其中N是序列长度。当上下文窗口变得非常大时，这个矩阵根本放不进共享内存。

解决方案和技巧5一样：沿着N维度分块。比如，上下文窗口是1028，按64切块，每块就能放进共享内存了。这样，虽然得到的是完整点积结果，但我们是逐块填充结果矩阵的。

分块本身是标准操作，真正棘手的地方在于softmax以及后续的值向量加权求和。通常情况下，计算softmax需要看到一整行的数据才能做归一化，而访问一整行就意味着要去全局内存取数据。FlashAttention的突破点在于它实现了“在线softmax”——softmax的计算和值向量加权求和可以在块内一次性完成，而无需看到全行数据。关键在于，最终的加权求和操作，给了我们逐块修正的数学基础。

下面用一个例子来说明。假设QK矩阵乘法产生了六个原始分数，代表某个查询词对六个其他词的关注度。常规做法是，一次性对这六个分数做softmax，再与六个值向量加权求和，得到结果A。

但遍历整个长度为N的序列，在块内放不下。于是，我们用“在线”的方式来做：把这六个分数分成三个块（每块2个元素），然后逐个处理。在第一个块里，我们只有两个原始分数，先基于这两个值开始计算：

这一步先不做归一化。虽然可以用当前的和（1+0.0082）来归一化，但后续的块会改变总和，到头来还得修正。所以，更好的做法是先记录下归一化分母的累积值，最后一步再做统一归一化。

接着处理第二个块。目标是得到一个和一次性看到所有四个值时相同的结果。为此，我们需要把第一个块的最大值传递过来，同时，累积的加权和与归一化分母也需要一并传递。

到目前为止，如果只有四个值，那么用A₁₊₂除以归一化总和1.3098就能得到最终结果。

最后一个边界情况是：新块里出现了更大的最大值。第三个块的最大值从12变成了13。但之前的A₁₊₂是按最大值12来计算的。为了让最终结果和一次性看到全部六个值完全一致，我们需要进行修正——把之前所有的旧指数乘以e⁻¹（即e⁽¹²⁻¹³⁾），来补偿最大值的变化。

我们不需要逐个回去修正每个指数值，只需将A₁₊₂和它的归一化分母整体乘以e⁻¹即可：

最后，用累积的分子A₁₊₂₊₃除以更新后的分母1.4955，就得到了结果。整个过程从未回访之前的块：只要一直跟踪最大值和归一化分母，就能逐块完成softmax。所有这些计算都在共享内存中完成，完全不需要频繁访问全局内存。

效果如何？FlashAttention的原始论文显示，在GPT-2模型上，注意力计算的耗时减少了数倍。

处理大规模模型时，如何巧妙地在共享内存里完成计算，对整体性能的影响，远比我们想象的要大。

并行计算的挑战与策略

以上所有讨论，都局限在一块GPU上。小模型这么跑没问题，但现代大型LLM，一块GPU根本装不下。就像Llama 3，它用了16000多块GPU。核心问题就变成了：如何把训练计算分配到这么多的机器上，然后再把结果汇总起来？

在展开不同的并行策略之前，我们先回顾一下标准的训练流程。以一个2层神经网络、batch size为16、使用Adam优化器的训练为例。

数据并行

拆分计算的第一种方式，是拆分数据。假设有效的batch size是16，但每块GPU的内存只够放4个样本。如果是单GPU，我们就需要跑4轮前向传播来累积梯度，再做一次反向传播，这叫做梯度累积。

数据并行的做法则是：把16个样本分给4块GPU，每块拿4个样本，然后各自并行执行前向传播。问题在于，如何聚合梯度？

一种方法是把所有激活值汇集起来，计算平均loss再求梯度。但更聪明的做法是，让每块GPU各自计算自己那4个样本的梯度，然后再求和——这样搬运的数据量更小，而且数学上完全等价。最后，把求和后的梯度传回每台机器，各自更新本地的模型。

这个操作在技术术语里叫做 all-reduce：每台机器贡献出自己的梯度，合并后，每台机器都能拿到最终的聚合结果。虽然示意图里画了一个“聚合器”（灰色方框），但实际实现中，all-reduce通常是通过环形传递来完成的——GPU之间互相传递梯度，最终每块GPU都拿到了平均值。

4块GPU并行工作，有效的batch size仍然是16，但处理速度却快了很多。不过，这里有一个效率问题：每块GPU都在做完整的模型更新，意味着每块GPU都要维护所有参数的Adam状态（一阶矩和二阶矩）。

内存充裕的时候这不成问题，但实际上，每块GPU里都复制了完整的模型参数、梯度、主权重，以及Adam优化器状态。Adam的状态量本身是模型参数量的两倍，内存占用非常大。

对于大模型来说，内存就成了硬瓶颈。ZeRO（Zero Redundancy Optimizer）就是针对这个问题提出的：它是一组内存优化技术，能在保持数据并行的前提下，大幅减少每块GPU的内存占用。

ZeRO Stage 1

核心思想是让每块GPU只负责更新一部分参数。比如，把每层参数分成四份，GPU 1负责Part 1，GPU 2负责Part 2，以此类推。

流程是这样的：数据仍然拆分到四块GPU上，每块GPU基于自己看到的4个样本计算完整的梯度——到这里还是标准的数据并行。但梯度汇总后，不再发回给所有人，而是按参数分片发送：每块GPU只收到自己负责的那部分参数的梯度。这个过程在术语上叫reduce-scatter——每人只拿到合并结果的一个切片。

然后，各GPU只更新自己负责的那部分参数，相应地，也只保留该部分的优化器状态。更新完成后，各GPU再把自己的参数切片分享出去，拼接成完整的模型。这个过程叫all-gather——每人贡献一个切片，每人拿到完整的拼接结果。

这个过程可以概括为两阶段：第一阶段按数据维度拆分，各GPU算出全参数梯度后再汇总；第二阶段按参数维度拆分，各GPU只更新自己负责的参数切片，最后再拼接出完整模型。

这样做的效果是，每块GPU只保留了一小部分优化器状态，内存节省非常可观。从计算量上看，reduce-scatter加上all-gather的总通信量，和朴素数据并行中的all-reduce是等价的，没有额外的开销。

ZeRO Stage 2

ZeRO Stage 2则更进一步——不仅优化器状态分片，梯度本身也要分片。

关键在于，反向传播是逐层进行的。每一层的梯度算完后，立刻将不属于自己管辖的部分发送给对应的GPU，然后丢弃。这样，在任何时刻，我们都不需要存储全部层的完整梯度。

与ZeRO Stage 1的流程相比，要改变的是这一步：

现在，我们改为逐层处理梯度。红框内的部分变成了这样：

算完第2层的梯度后，把不负责的部分发出去、丢弃，然后处理第1层，重复同样的步骤。层数很多的LLM能从中获益巨大——我们不再需要同时存储所有层的梯度。代价是，逐层通信会带来一点额外开销。

ZeRO Stage 3（完全分片数据并行）

ZeRO Stage 3把分片思路推到了极致——连模型权重都只存各自负责的那部分。这意味着，前向传播的过程也会受到影响。

流程同样是逐层进行的。处理到第1层时，先执行all-gather，各GPU拿出自己的权重切片，拼出完整的第1层。然后，每块GPU用完整的第1层权重和各自的数据计算激活值，算完之后，立刻丢弃不属于自己的权重切片。第2层同理。

反向传播的过程和ZeRO Stage 2类似，但多了一步：在计算每层的梯度之前，要先执行all-gather，把完整的权重拼出来（因为本地没有完整权重），算完之后再丢弃非本地的切片。

本质上，我们是在按需、逐层地从各GPU拼出模型，任何时候，都没有一块GPU持有全部权重。通信开销确实增加了，但内存节省是巨大的。对于给定的GPU配置，ZeRO Stage 3能训练的模型规模，远超前两个阶段。

CS336课程提供了一组数据：在8块A100 80GB GPU上，不同策略能训练的最大模型尺寸，差异非常明显。

在同样的硬件配置下，使用ZeRO Stage 3能训练的模型，比前两个阶段大了很多。

不过，数据并行有一个天然的约束：batch size。batch size不能小于GPU数量——你没法给一台机器分配“半个”样本。而batch size越大，收益往往越低：大batch可以降低数据噪声的方差，但超过某个阈值后，边际收益就接近于零了。因此，batch size的“自然上限”直接限制了数据并行的扩展规模。

模型并行

除了按数据维度拆分，我们还可以按模型的维度进行切分，这就是模型并行。这里介绍两种主要形式：流水线并行和张量并行。

模型并行：流水线并行

流水线并行是沿着模型的深度方向进行切分，比如把一层分配给一块GPU。这样做的问题在于，前向和反向传播都是逐层串行的——每一层都需要前一层的输出才能开始计算。这就导致了一部分GPU在等待输入时处于空闲状态，形成了所谓的“气泡”。

缩小气泡的有效方法是引入mini-batch级别的流水线：当第二块GPU在处理某个mini-batch的第二层时，第一块GPU已经可以开始处理下一个mini-batch的第一层了。

流水线并行的优势在于两点：一是内存节省显著，每个设备只需要存储一层的参数；二是通信模式极其简单，只需将激活值从一层传递到下一层。这种简单的通信特性，使它特别适合部署在跨集群等带宽较低的网络链路上。

张量并行

张量并行则是沿着模型的宽度方向进行切分，把单层内的矩阵乘法分配到多块GPU上并行执行。每块GPU各自得到部分结果后，再跨GPU进行求和。这个概念上有点像之前提到的分块运算，但区别在于：分块是串行处理各个块，而张量并行是并发处理的。

这种方法的通信量很大——每一层都要同步激活值。好在节点内部的NVLink带宽能达到600-900 GB/s，而跨节点的互连则要慢上10到20倍。实践经验表明，当张量并行扩展到超过8块GPU时，收益会急剧衰减。所以，通常会将张量并行限制在单个节点（最多8块GPU）内部。

张量并行有一个独特的优势：它不依赖于batch size。batch size是数据并行和流水线并行都共享的约束资源，而张量并行与之正交，可以叠加使用，不会消耗batch size这项宝贵的资源。

组合不同形式的并行

我们讨论的这几种并行策略，分别沿着不同的维度拆分计算：数据维度、模型深度维度、模型宽度维度。在实际训练中，往往是多种策略的组合使用，才能达到最好的效果。

一个经验法则很简单：先解决内存问题，确保模型能装进GPU。如果装不下，就用流水线并行、张量并行、ZeRO Stage 3这些技术来节省内存。等模型能装下之后，再考虑用数据并行等手段堆算力，加快每一个batch的处理速度。

附录：Softmax 计算的细节

Softmax函数的作用，是把一组原始的分数，转换成一个加和为1的概率分布：对每个分数取指数，然后除以所有指数的和。

以三个分数（12， 7.2, 9.1）为例：

问题在于，指数函数的增长非常快。e¹²已经达到了162755，超过了fp16能表示的最大值65504。理论上的计算结果虽然正确，但在计算过程中会直接溢出。解决办法很简单：将分子和分母同时除以e^(max)，这等价于从所有原始分数中减去一个最大值：

数学上，结果完全一致，但完美地避免了溢出。可能会出现下溢（值太接近零），但在这种情况下，0已经是一个足够好的近似值。这个数学技巧，几乎被所有LLM中的softmax实现所采用。

总结

这篇文章从GPU的底层架构讲起，一直聊到多GPU的并行策略，涉及的这些都是将模型从玩具规模拉到生产规模所必须面对和解决的工程问题。在专业的AI团队中，训练一个无法放入单块GPU的LLM是常态，优化训练成本也是日常工作中很重要的一部分。真正理解底层硬件的工作原理和并行机制，是做好这一切工作的前提。