★ FEATURED ARTICLE 1. 为什么“能跑起来”和“能训得动”是两码事很多人第一次接触大模型训练脑子里想的都是“我先把环境搭好把数据喂进去跑起来再说”。这个思路在小模型时代没什么问题一个几百兆参数的模型哪怕显存管理粗糙一点、学习率拍脑袋定一个顶多就是效果差一些不至于直接崩。但大模型完全是另一回事——它像一个脾气极其敏感的精密仪器你稍微有一个地方没弄对它不会给你“效果差一点”的中间态而是直接给你一个Loss曲线飞天或者原地不动让你连排查的方向都找不到。我见过不少团队模型代码是从开源仓库里扒下来的数据是自己业务场景里攒的硬件是好不容易申请到的几张卡然后一跑就出问题。出问题之后就开始到处问“为什么我的Loss不下降”“为什么训练到一半突然爆了”“为什么多卡比单卡还慢”。这些问题看起来五花八门但根子上往往就集中在三个地方数值稳定性、显存管理、并行策略。这三个地方一旦出问题轻则训练效率打对折重则直接把你几天的训练成果清零。这篇手记不打算给你讲一遍Transformer的论文也不打算把PyTorch的文档复述一遍。我想做的是把这三个“会炸”的地方一个一个拆开告诉你它们为什么会炸、炸之前有什么征兆、以及最关键的——怎么在训练开始之前就把这些雷排掉。不管你是刚接手大模型训练任务的新手还是已经跑过几个小规模实验想往大了做的开发者这些内容应该都能帮你省下不少盯着Loss曲线发呆的时间。2. 第一个会炸的地方数值精度与梯度稳定性2.1 从一次Loss突然变成NaN说起先描述一个非常典型的场景。你搭好了一个几十亿参数的模型数据管道也通了前几百步Loss下降得很漂亮你心里正美滋滋结果去泡了杯咖啡回来发现Loss变成了NaN。你检查数据没问题检查代码没改过检查学习率也不算大。那问题出在哪这种情况十有八九是数值精度在作祟。大模型训练中数值精度的选择直接决定了你的训练能不能稳定跑下去。现在主流做法是混合精度训练也就是前向和反向用FP16或BF16但优化器状态和主权重用FP32。这个方案本身没问题问题在于很多人只知道“要开混合精度”却不知道混合精度背后有一堆需要配合的机制。FP16的动态范围很窄能表示的最小正数大概是6e-8最大也就65504。在反向传播过程中梯度值很容易小于这个最小正数直接变成0这就是下溢。一旦某一层的梯度全部下溢这一层的参数就再也更新不了了相当于“死”了。更糟糕的是如果梯度值超过65504就会变成Inf然后经过几层传播就变成NaN整个模型直接崩掉。BF16的动态范围和FP32一样宽所以不容易出现上溢和下溢的问题但它的尾数只有7位精度比FP16的10位要低。这意味着BF16在表示小数的时候误差更大对于一些对精度敏感的层可能会影响收敛效果。所以选FP16还是BF16不是拍脑袋决定的要看你的硬件支持情况和模型结构。2.2 损失缩放到底在做什么解决FP16下溢问题的标准方案是损失缩放。原理很简单在计算Loss的时候乘上一个很大的系数比如1024这样反向传播出来的梯度也会被同比放大就不会下溢了。然后在优化器更新参数之前再把梯度除回原来的系数。但这里有一个很多人忽略的细节动态损失缩放。静态损失缩放就是你手动设一个固定值比如固定乘1024。问题是训练过程中梯度的量级是会变化的前期梯度大后期梯度小。如果固定用1024后期可能梯度本身就很小乘完还是下溢或者前期梯度本来就大乘完直接上溢。动态损失缩放就是让框架自动调整这个系数。它的逻辑是先设一个初始值然后每隔若干步检查一下梯度有没有出现Inf或NaN。如果没有就慢慢增大缩放系数如果出现了就跳过这一步更新同时减小缩放系数。PyTorch的torch.cuda.amp.GradScaler就是干这个的。我踩过的一个坑是用了GradScaler但忘了在优化器更新之前调用scaler.unscale_(optimizer)结果梯度没有被除回来相当于每次更新都放大了上千倍模型直接飞了。这个错误很隐蔽因为代码不会报错只是Loss曲线会以一种非常诡异的方式上升。2.3 梯度裁剪不是万能的但没它万万不能梯度裁剪是另一个必须配上的安全措施。它的作用是在反向传播之后、优化器更新之前检查所有梯度的范数如果超过某个阈值就等比例缩小所有梯度让它们的范数不超过这个阈值。为什么需要这个因为大模型训练中偶尔会出现某个batch的数据特别“刺激”导致梯度突然变得非常大。如果不裁剪这一大步更新可能会把参数推到一个很糟糕的区域甚至直接导致后续Loss爆炸。裁剪阈值一般设在1.0左右但这不是绝对的需要根据模型规模和任务特点来调。这里有一个经验性的判断方法你可以在训练前几百步的时候把每一步的梯度范数打印出来看看。如果大部分时候梯度范数都在0.1到1.0之间那阈值设1.0就合适。如果经常跑到5.0甚至10.0那要么是学习率太大了要么是数据有问题需要先解决根因而不是简单把裁剪阈值调大。还有一个容易忽略的点裁剪是在所有参数上一起做的还是分层做的。PyTorch默认的clip_grad_norm_是把所有参数的梯度拼成一个向量算全局范数。这对于大多数情况是OK的但如果你的模型里有Embedding层和输出层这种参数量极大但梯度稀疏的层全局裁剪可能会让这些层的梯度被过度压制。有些实现会采用分层裁剪或者按参数组裁剪这个要根据具体情况来定。2.4 初始化与归一化层的隐藏陷阱数值稳定性的问题不只是在训练过程中初始化阶段就埋下了种子。大模型的参数初始化不是随便来个正态分布就完事了。如果初始化方差太大前向传播的激活值会逐层放大到最后一层可能已经溢出如果太小激活值会逐层缩小到深层就趋近于零梯度也传不回去。现在主流做法是用缩放初始化比如GPT系列用的就是标准差为0.02的正态分布同时把残差路径上的输出投影层额外缩小一个系数。这个系数的计算跟网络层数有关层数越多残差路径上的缩放就要越小否则残差累加会导致激活值爆炸。归一化层的位置也有讲究。Pre-LN和Post-LN在稳定性上差别很大。Post-LN就是把LayerNorm放在残差相加之后这是原始Transformer的做法但它在深层网络中非常容易不稳定需要小心的学习率预热。Pre-LN是把LayerNorm放在残差分支里面残差路径上是干净的恒等映射梯度可以无障碍地传回去稳定性好很多。现在的大模型基本都用Pre-LN或者它的变体如果你还在用Post-LN建议换掉。3. 第二个会炸的地方显存墙与激活值管理3.1 显存都去哪了一笔明白账很多人第一次算显存需求的时候只算了模型参数占多少。比如一个70亿参数的模型FP16存储也就14GB看起来一张卡就能放下。但实际训练的时候发现一张80GB的卡都不够用。为什么因为训练时的显存占用远不止模型参数。粗略来说主要有四块模型参数、梯度、优化器状态、激活值。用Adam优化器的话每个参数需要存一份FP32的主权重、一份FP32的一阶动量、一份FP32的二阶动量再加上FP16的模型参数和FP16的梯度每个参数总共需要2244416字节。70亿参数就是112GB已经超过一张80GB卡的容量了。这还没算激活值。激活值是前向传播过程中每一层的输出需要保留到反向传播时计算梯度用。激活值的显存占用跟batch size、序列长度、隐藏层维度都成正比。在长序列场景下激活值可能比模型参数本身还大。所以当你发现OOM的时候第一件事是搞清楚到底是哪部分占的显存最多。PyTorch提供了torch.cuda.memory_summary()可以看详细的内存分配情况。一般来说如果模型参数和优化器状态占了大头那需要上ZeRO或者模型并行如果激活值占了大头那需要上梯度检查点或者减小batch size。3.2 梯度检查点用时间换空间的经典操作梯度检查点的思路很直接前向传播的时候不保存中间激活值只保存少数几个检查点的激活值。反向传播的时候从最近的检查点重新做一次前向传播把需要的激活值算出来用完就扔。这样做的好处是显存占用大幅降低代价是计算量增加。因为反向传播时多了一次前向传播整体计算量大概增加30%左右。但显存可能降低到原来的三分之一甚至更少对于显存紧张的情况非常划算。不过梯度检查点有几个实操中的坑。第一不是所有层都适合做检查点。比如LayerNorm层本身计算量小、激活值也小做检查点反而增加开销。一般是对Transformer Block做检查点因为它的激活值占用最大。第二用了检查点之后反向传播的计算图会变得复杂如果跟混合精度、梯度累积一起用需要确保框架版本支持这些组合否则可能出现梯度计算错误。还有一个容易忽略的点检查点的粒度。你可以对整个Block做检查点也可以只对Block里的Attention部分做检查点。粒度越细显存节省越明显但计算开销也越大。需要根据实际显存瓶颈来权衡。3.3 ZeRO的三级划分与选择逻辑ZeRO是微软提出的显存优化方案核心思想是把模型参数、梯度、优化器状态这三样东西分散到多张卡上而不是每张卡都存一份完整的。它分三个阶段阶段分散内容显存节省通信开销ZeRO-1优化器状态约4倍低ZeRO-2优化器状态梯度约8倍中ZeRO-3优化器状态梯度参数约N倍N为卡数高ZeRO-1是最安全的通信开销增加很少基本不影响训练速度。ZeRO-2在大多数场景下也能保持不错的效率。ZeRO-3虽然显存节省最明显但通信开销大幅增加如果卡间带宽不够训练速度可能下降得很厉害。选择哪个阶段取决于你的瓶颈在哪。如果只是优化器状态占显存太多ZeRO-1就够了。如果激活值也很大可能需要ZeRO-2加上梯度检查点。ZeRO-3一般是在单卡显存实在放不下模型参数的时候才用比如训练百亿参数以上的模型。这里有一个实操经验ZeRO-3的通信开销跟模型结构有关。如果模型里有大量的AllReduce操作ZeRO-3的额外通信可能会让训练速度降到原来的几分之一。所以在决定用ZeRO-3之前最好先在小规模上测一下通信开销别上来就在大集群上跑浪费时间。3.4 激活值重计算与Offload的取舍除了梯度检查点还有激活值重计算和Offload两种手段。激活值重计算跟梯度检查点类似但更激进它把更多的中间结果都扔掉反向传播时全部重算。Offload是把暂时不用的数据从显存搬到内存甚至硬盘上需要的时候再搬回来。Offload听起来很美好但实际用起来限制很多。首先是PCIe带宽显存和内存之间的传输速度远低于显存内部带宽如果频繁Offload训练速度会断崖式下降。其次是CPU内存容量大模型的优化器状态动辄上百GB如果内存不够也Offload不了。我的建议是优先用梯度检查点不够再考虑ZeRO最后才考虑Offload。Offload适合那种对训练速度不敏感、但对显存极度受限的场景比如在少量卡上微调一个超大模型。如果是正常的多卡训练尽量通过并行策略来解决显存问题而不是依赖Offload。4. 第三个会炸的地方并行策略与通信瓶颈4.1 数据并行、模型并行、流水线并行的本质区别当单卡放不下模型或者单卡训练太慢的时候就需要上多卡。多卡训练有三种基本策略数据并行、模型并行、流水线并行。数据并行是最简单的每张卡上都放一份完整的模型但每张卡处理不同的数据batch。反向传播后所有卡把梯度做一次AllReduce求平均然后各自更新参数。这种方式实现简单但要求单卡能放下整个模型。而且当卡数增多时AllReduce的通信开销会线性增长因为每张卡都要跟其他所有卡交换梯度。模型并行是把模型的不同层放到不同的卡上。比如前10层在卡0中间10层在卡1后10层在卡2。前向传播时数据依次经过各张卡反向传播时梯度反向依次传回去。这种方式解决了单卡放不下模型的问题但卡之间的通信非常频繁而且容易出现负载不均衡——有的卡计算量大有的卡计算量小导致快的卡等慢的卡。流水线并行是模型并行的一种改进。它把模型切成多个阶段每个阶段放在不同的卡上然后把一个batch切成多个micro-batch让不同micro-batch在不同阶段上像流水线一样重叠执行。这样可以在一定程度上缓解负载不均衡的问题提高卡利用率。但流水线并行需要仔细设计micro-batch的数量和调度策略否则会出现“气泡”——有些卡在等数据白白浪费算力。4.2 张量并行的切分维度与通信模式张量并行是模型并行的一种细粒度形式它不是按层切而是把单个层内部的矩阵运算切到多张卡上。比如一个大的矩阵乘法可以把权重矩阵按列切或者按行切分到不同的卡上计算然后再把结果拼起来。张量并行的通信模式跟切分方式有关。按列切的话每张卡算出部分结果需要做一次AllGather把结果拼起来按行切的话每张卡算出部分和需要做一次AllReduce把部分和加起来。这两种通信在每一层都要做所以对卡间带宽要求很高。一般张量并行只在单机内部用因为机内NVLink带宽远高于机间网络。张量并行的并行度不是越高越好。当并行度超过单机GPU数量时通信就要走机间网络延迟大幅增加训练速度可能不升反降。所以通常的做法是机内用张量并行机间用流水线并行或数据并行。比如8台机器每台8张卡可以设张量并行度为8机内流水线并行度为8机间数据并行度为1。或者张量并行度为8数据并行度为8流水线并行度为1。4.3 通信原语AllReduce、AllGather、ReduceScatter多卡训练里的通信操作主要有三种AllReduce、AllGather、ReduceScatter。理解它们的区别对于分析通信瓶颈很重要。AllReduce是把所有卡上的数据加起来然后把结果发给所有卡。数据并行里的梯度同步就是AllReduce。AllGather是把所有卡上的数据收集起来每张卡都得到一份完整的拼接结果。张量并行里按列切后的结果拼接就是AllGather。ReduceScatter是先把所有卡上的数据加起来然后把结果分散到各张卡上每张卡只拿到一部分。ZeRO里的梯度分散就用到ReduceScatter。这三种操作的通信量跟数据量和卡数有关。AllReduce的通信量是2倍的数据量乘以卡数减一除以卡数当卡数很大时趋近于2倍数据量。AllGather和ReduceScatter的通信量各是1倍数据量。所以ZeRO-3里用AllGather拿参数、用ReduceScatter传梯度总的通信量比纯数据并行的AllReduce要大。4.4 并行策略组合的实操建议实际训练中通常是多种并行策略组合使用。组合的原则是优先用通信开销小的策略把通信开销大的策略限制在带宽高的范围内。具体来说数据并行的通信量相对较小可以跨机使用。流水线并行的通信量也不大但需要处理气泡问题也适合跨机。张量并行的通信量最大最好限制在单机NVLink范围内。一个常见的配置是张量并行度等于单机GPU数流水线并行度等于机数数据并行度根据总卡数除以这两者来定。比如64张卡8机8卡可以设张量并行8、流水线并行8、数据并行1。如果觉得数据并行度太低影响吞吐可以减小流水线并行度比如张量并行8、流水线并行4、数据并行2。还有一个经验并行策略的选择要跟模型结构匹配。如果模型是标准的Transformer各层结构一致流水线并行很容易切分。如果模型里有复杂的跨层连接或者动态路由流水线并行就很难做可能只能靠数据并行加张量并行。5. 训练前的排雷清单与监控指标5.1 开跑之前必须确认的七件事在按下训练启动按钮之前我习惯过一遍下面这个清单。每一条都是踩过坑之后总结出来的少检查一条就可能多浪费一天。精度配置确认混合精度用的是FP16还是BF16损失缩放是动态的还是静态的GradScaler的初始值和增长策略是否合理。梯度裁剪确认裁剪阈值设了多少是所有参数一起裁还是分组裁裁剪是在损失缩放之前还是之后。初始化确认参数初始化方案特别是残差路径上的缩放系数是否跟层数匹配。归一化确认用的是Pre-LN还是Post-LN如果是Post-LN学习率预热必须做足。显存估算用公式或者工具估算一下模型参数、梯度、优化器状态、激活值各占多少确认并行策略能放得下。并行配置确认张量并行、流水线并行、数据并行的度数和组合方式确认通信组初始化正确。数据管道确认数据加载不会成为瓶颈确认数据格式和预处理跟模型输入匹配。5.2 训练初期必须盯住的五个指标训练启动之后前几百步是最关键的。这个时候模型还没进入稳定状态任何配置问题都会很快暴露出来。我一般会盯这几个指标Loss值看是否在下降下降的斜率是否合理。如果Loss完全不降检查学习率和数据如果Loss震荡剧烈检查batch size和梯度裁剪。梯度范数看是否在合理范围内。如果经常超过裁剪阈值说明学习率可能太大如果一直很小说明学习率可能太小或者初始化有问题。学习率如果用了预热和衰减确认学习率按照预期变化。显存占用看是否接近上限如果接近上限后续可能会OOM。吞吐量看每秒处理的token数或样本数跟预期对比如果远低于预期检查通信或数据管道。5.3 那些不会报错但会要命的问题有些问题不会让程序崩溃但会让训练效果大打折扣甚至完全无效。这类问题最危险因为它们很隐蔽。一个是数据重复。如果数据管道有bug导致同一个batch被反复训练Loss会下降得很快但模型实际上只记住了那几个样本泛化能力极差。检查方法是定期打印batch里的样本ID确认没有重复。另一个是梯度累积与学习率的配合。梯度累积是在多个batch上分别计算梯度累加起来再更新一次参数。这相当于增大了等效batch size。如果增大了batch size但不调整学习率训练动态会发生变化。一般来说batch size翻倍学习率也要相应增大但具体增大多少需要实验。还有一个是位置编码的溢出。如果序列长度超过了位置编码的最大长度有些实现会静默截断或者循环使用位置编码导致模型看到的位置信息是错的。这个问题在长序列场景下特别容易发生一定要确认序列长度没有超过限制。6. 从炸到不炸一些个人体会训练大模型这件事说到底是一个不断跟各种“炸”作斗争的过程。我刚开始做的时候觉得把模型结构搭对、把数据准备好就完事了后来才发现真正的功夫都在这些“看不见”的地方。数值精度、显存管理、并行策略这些东西在论文里可能就一句话带过但在实际工程中每一个都够你折腾好几天。我的体会是不要等到炸了再去修要在炸之前就把它防住。上面那些清单和指标看起来繁琐但比起训练到一半崩掉重来的代价花半个小时检查一遍是非常划算的。另外从小规模开始验证。不要一上来就在最大的配置上跑先用小模型、小数据、少卡数把流程跑通确认没有数值问题和通信问题再逐步放大。这样即使出问题排查的范围也小得多。还有一个心态上的建议把每一次“炸”都当成学习的机会。大模型训练中的很多问题文档里不会写教程里不会讲只有自己踩过才知道。每次炸了之后搞清楚根因记录下来下次就能提前避开。积累多了你就有了别人没有的经验这才是真正值钱的东西。
阅读完成 · 觉得有帮助？
