你的大模型为什么训不快？大模型预训练技术精要

关注

原创最新推荐文章于 2026-06-08 21:57:58 发布 2k 阅读 AI 写同款· GEO 优化 ›

大模型训练显存优化推理加速方案当前的深度学习框架大都采用的都是fp32来进行权重参数的存储，比如Python float的类型为双精度浮点数fp64，pytorch Tensor的默认类型为单精度浮点数fp32。随着模型越来越大，加速训练模型的需求就产生了。在深度学习模型中使用fp32主要存在几个问题，第一模型尺寸大，训练的时候对显卡的显存要求高；第二模型训练速度慢；第三模型推理速度慢。其解决方案就是使用低精度计算对模型进行优化。本文主要讲解几种优化显存存储的方法。阅读详情

大模型训练用时可以拆解成两方面：一个是模型运行时间，这部分要想快就需要买更好的机器，研究fused kernel来加速，一般时间周期长，难度高，这部分后面文章会有专门讨论。第二部分是通讯时间，大模型训练过程中需要不断从其他GPU，其他节点收集梯度信息更新模型。千万不可小看第二部分通讯时间，笔者在默认设置下在自己环境上跑大模型并行，发现通讯时间竟然高达90%，模型训练异常低效，优化后占比降低到20%左右，显著提升了模型训练效率，并且吞吐率到能对齐甚至略高于llama2官方的吞吐率。

训模型常见思路

为啥要做并行方式验证和调优？

是基于你的训练GPU，训练集群的inter-gpu带宽，inter-node带宽，模型大小，训练数据，选择一个最优吞吐率的并行方案（人工加粗加黑，集群不同并行方案就不同）。别人在别人机器上的验证并行方式不一定适用你自己的机器。一般做好并行方式验证调优的话吞吐率可以提升30%-50%。通过并行方式验证和调优，我们就是找到适合自己机器通讯环境的最优模型并行方式，这样可以逼近集群运算能力的上限。

如何确定使用什么样的并行策略？

如果模型不是特别大7B，13B直接数据并行Data parallel(DP)搞起（zero0，zero1，zero2往上尝试），还塞不下然后考虑加Pipeline parallel(PP)。如果模型特别大70B，130B的时候，这时候考虑加Tensor parallel(TP)到2-4-8（一般先设置为2比较合适，优先增加PP），扩大PP直到能塞下模型，如果已经能塞的下模型了，再增大PP速度会降低速度，因为需要通讯量增加了，在笔者的8节点64 A800机器上增加BS塞满显存并不能抵消通讯量带来的速度损失。如果还是实在塞不下再考虑gradient chekpointing和offload等技术。

注意：如果用了Pipeline parallel(PP)，就不能用zero2和zero3。因为zero-2/3就将梯度也切分了，但在做梯度累计的时候，流水线并行的时候的需要从不同的nodes汇集梯度会需要巨量的通信开销，所以deepspeed直接禁用了zero2/3 + PP的组合。

一般单GPU batch size （micro batch size）设置为1就行了，大的batch size通过梯度累计实现，预训练过程是一个GPU计算密集型任务，GPU利用率一般都能跑满，所以加大batch size提速不大。更重要的是，显存应该用于存模型训练相关参数，比如模型参数，梯度，优化器参数等。这样可以减少整体通信开销，加快训练速度。有一种场景是可以增大micro batch size的，就是模型本来是塞不下的，我的PP从4增大到8以后可以塞下了，但空出来20%的显存，这时候我可以将micro batch size 从1变成2把多出来的20%显存占满，实验发现也可以提速5%左右，原理是GPU计算更大矩阵乘法效率会更高。

Zero3只适合机器比较少的时候玩一下，扩展到多机就很慢，因为zero3本质上其实是DP+TP，如果我们有64个GPU，就相当于64路TP了，通讯成本太高了限制了训练速度。

为啥优先用PP而不是TP（2以上）呢？

我们知道做Ring-AllReduce通讯量是和参数的通讯节点数量线性相关的。

具体可参考文章：https://zhuanlan.zhihu.com/p/617133971

这个问题和通讯量有关，PP简单来说就是横着切模型，模型切成很多块（chunk），但层还是完整的。这样每次通讯的话只需要在断点出进行块与块之间的通讯，因此PP无论切多少刀，每次只需要在两个节点之间通讯，并且不同chunk通讯是异步的，同一时间的通讯压力小。而TP是竖着切模型，每一层都切开了，通讯的话需要每一层不同块进行all reduce通讯，如果切得太多，则同一时间需要参与通讯的GPU数量就很多，训练就会慢。如果TP=2切一刀那么通讯量和PP是一样的，如果切成4份那么需要在4个节点之间进行同步通讯，通讯量直接翻倍。甚至TP如果切多了，比GPU个数还多了，就需要inter-node的通讯了，这个一般比Intra-node通讯慢很多，自然是无法承受的。

预训练常见技术：红榜

【Flash Attention】 可以提速+降显存，尤其是llama2这种序列较长的降显存非常明显，在Zero2+flash attention可以在单机8卡上放下7B模型，并且训练速度可以提速大约75%，不使用gradient checkpointing可以提速25%，相比于Zero2+gradient checkpointing的组合快了一倍。

具体可参考文章：https://zhuanlan.zhihu.com/p/645238961

【Zero优化器】 显存占用：zero0>zero1>zero2>zero3 显存降低大概每一级降低10%左右（具体看模型，以及卡的个数，总之不会多特别多），速度上来看的话zer0>zero1>zero2>zero3 每一级降低5%左右（速度上感觉不到特别明显的差别）。所以huggingface 推荐直接可以上zero2，因为速度没有慢多少，但可以降低显存比较多。从tflops来看会比较明显 zero 0 > zero 1 > zero 2 。因为优化器中的参数是最多的，一般在megatron框架中加入–use-distributed-optimizer就是zero1阶段就能升下不少显存，实测速度也不会降低。

【Pipeline Parallel(PP)和Tensor Parallel(TP)】 在保证能塞得下模型的前提下，PP * TP 应该越少越好，因为多了每个GPU显存降下去了，但通讯时间延长了。但PP和TP怎么配置也是讲究的，TP应该比较小，如果从2开到4吞吐量下降13%左右，2还行明通讯瓶颈一般到2了。PP开的如果太小的话 bubble就会比较多，所以提速不明显。64卡训练13B模型比如 pp2 + tp 1 和pp1 + tp2速度差不多，但pp4 + tp1 比 pp1 + tp4快20%。所以笔者发现13B模型64卡A800训练最优设置的PP 4 + TP 1，同时PP2 + TP 2 效果也非常接近，这个也是因为降低了PP的bubble。一般建议TP大小为Hidden size / TP > 1024。并且笔者的机器是A800机器，机内带宽是比较低的，PP 和 TP 的先后不存在明显的差异。同时TP一般会搭配sequence-parallel进行使用，可以在降低显存的情况下，不增加任何通讯成本。

具体可参考文章：

https://zhuanlan.zhihu.com/p/646406772

https://zhuanlan.zhihu.com/p/626553071

在使用PP的时候我们会发现各个节点的显存占用不一样，比如PP为4的时候显存占用Rank0>Rank1>Rank2>Rank3>Rank4。这个是因为PP不同Rank激活值存储时间不一样。比如下图中的Rank0-F0节点需要存储激活值的时间要比F1/2/3要长，这部分存储的激活值增加了显存占用。同时embedining层也在Rank0节点，比中间节点多占用了一些显存。

【Selective Activation Recomputation】 按照huggingface的描述可以在仅仅增加2.7%的FlOPs overhead的情况下对于GPT3可以降低70%的显存，是一个大杀器。

【梯度累积 Gradient Accumulation】 可以打印出来每次forward时间backward时间，梯度做传播的时间，如果发现做梯度传播的时间占整体训练时间1/5以上，那么可以考虑增加global batch size提升gradient accumulation，这样可以减少梯度在各个节点通讯的次数，从而显著提升吞吐量，笔者在256卡A800测试中发现，global batch size增加四倍后吞吐量增加25%，进一步提升训练效率，同时增加global batch size，需要适应调整学习率，一般可以线性增加，比如global batch size增加4倍，LR也要增加四倍，要不然训练的iteration太少，模型更新太慢的话，走不到最优点。

预训练常见技术：黑榜

【Full Activation Recomputation】 为虽然可以降显存，但是速度会慢25%，能不用的时候尽量不用。降显存完全可以用PP 和 TP，尤其是PP对通讯要求不是特别高，并且设置的得当的话还会加速。这个技术更像是一个显存严重不足场景下不得已为之的方法。

【Offload Optimizer/parameters】 使用后训练速度会慢非常多，基本慢了一倍。这个技术更像是一个显存严重不足场景下的不得已为之的方法。

一些调参的takeways

如何选择常见的训练策略

不同并行策略如何选择

如何选择合适的Zero优化器

以及Deepspeed相关训练策略

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述