显存不够用这件事几乎每个跑过深度学习训练的人都撞过墙。我第一次遇到OOM的时候盯着那行CUDA out of memory发呆明明模型参数量算下来才几百兆怎么24G的卡说爆就爆。后来才慢慢搞明白显存占用远不止模型权重那么简单优化器状态、梯度、激活值、临时缓冲区每一项都在悄悄吃你的显存。而当你把单卡塞满之后下一步就是分布式训练DP、MP、PP这三个缩写会反复出现在你的视野里但它们的显存账本完全不同选错了策略可能卡越多反而越慢。这篇内容我打算把显存占用的账从头算一遍然后拆解DP、MP、PP三种并行策略各自怎么花显存、适合什么场景、落地时要注意什么。适合刚被OOM折磨过的朋友也适合正在纠结多卡方案怎么选的同行。所有数字和结论都来自实际训练中的观察和推算不是纸上谈兵。1. 先把显存这笔账算清楚很多人对显存的认知停留在“模型多大就占多大”这个误解是OOM的根源。实际训练时显存占用可以拆成四大块模型权重、梯度、优化器状态、激活值。前三个是静态开销跟模型参数量直接挂钩激活值是动态开销跟batch size、序列长度、网络结构强相关。推理时只有权重和少量激活训练时四块全上差距可能是好几倍。1.1 静态开销权重、梯度、优化器状态先拿一个参数量为Φ的模型来说。权重本身占多少显存取决于精度。FP32下每个参数4字节FP16/BF16下2字节。但训练时通常不会只用一种精度混合精度训练里会同时存在FP16的权重副本和FP32的主权重副本。梯度的大小跟权重一致每个参数对应一个梯度值精度通常跟权重同级别。优化器状态是大头Adam优化器会为每个参数维护一阶矩和二阶矩两个状态每个状态跟参数同精度。如果用FP32的Adam那就是每个参数额外2×48字节。我拿一个10亿参数的模型举例算一笔账组成部分精度每参数字节10亿参数占用模型权重FP16FP1622 GB主权重副本FP32FP3244 GB梯度FP16FP1622 GBAdam一阶矩FP32FP3244 GBAdam二阶矩FP32FP3244 GB合计1616 GB也就是说10亿参数的模型光静态开销就要16GB左右。一张24G的卡还没开始算激活值就已经用掉三分之二。这就是为什么很多人发现模型明明不大训练却跑不起来。注意上面这个16字节/参数的估算是一个常用的经验值实际会因框架实现、是否使用ZeRO、是否开启梯度累积等有所浮动。但量级上可以作为快速估算的依据。如果换成纯FP32训练、不用混合精度权重4字节、梯度4字节、Adam两个状态各4字节合计16字节跟混合精度的账其实差不多但混合精度在计算速度和激活值占用上有优势所以现在基本是标配。1.2 动态开销激活值才是隐形杀手激活值的显存占用没有固定公式它跟网络深度、每层的输出尺寸、batch size都有关。粗略估算激活值占用跟 batch_size × 序列长度 × 隐藏维度 × 层数 成正比。Transformer类模型里注意力机制的中间结果还会额外放大这个数字。我实际测过一个中等规模的Transformer隐藏维度768、12层、序列长度512batch size从8加到32激活值占用从不到2GB涨到接近7GB几乎是线性增长。这时候哪怕静态开销只有几个G总显存也会被激活值顶爆。激活值有一个特点前向传播时逐层产生反向传播时逐层消耗。理论上算完某层的梯度后那层的激活值就可以释放。但框架为了效率通常会保留一部分这就引出了**激活重计算gradient checkpointing**这个技巧——牺牲约30%的计算时间换取激活值显存大幅下降。我实测下来开启重计算后激活值占用能降到原来的三分之一甚至更低代价是每步训练时间增加20%到40%具体取决于模型结构。1.3 那些容易被忽略的零碎开销除了上面四大块还有几项零碎但不可忽视的占用。CUDA上下文本身就要占几百MBcuDNN的workspace在卷积网络里可能占1到2GB通信缓冲区在分布式训练时也要预留空间。这些加起来一张卡上轻松再吃掉1到3GB。所以实际估算显存时我习惯在理论值上再加20%到30%的余量。比如算出静态加激活大概20GB那24G的卡基本就是极限操作稍微改个batch size或者序列长度就爆。稳妥起见留出3到4GB的缓冲比较安心。2. DP、MP、PP三种策略的显存逻辑拆解单卡塞不下的时候就得考虑多卡。DP、MP、PP是三种最基本的并行维度它们切分的东西不一样显存账本也完全不同。很多人把它们混为一谈实际上选错策略可能让你多花好几张卡却收效甚微。2.1 DP数据并行复制模型分摊数据DPData Parallelism的思路最直观每张卡上都放一份完整的模型副本把batch的数据切成N份每张卡算一份然后同步梯度。它的显存特点是每张卡的静态开销完全一样都是完整模型的量。也就是说DP不减少单卡的模型显存占用它减少的是每张卡处理的batch大小从而降低激活值占用。DP的通信量在于梯度同步。每步训练结束后所有卡的梯度要做all-reduce通信量跟模型参数量成正比。卡越多通信开销占比越大。我见过8卡DP训练一个小模型通信时间占了每步的40%以上加速比远达不到线性。DP适合什么场景模型能塞进单卡、但数据量大、想加快训练速度的时候。它的实现最简单PyTorch的DistributedDataParallel几行代码就能跑起来。但如果模型本身就塞不进单卡DP帮不上忙因为每张卡还是要放完整模型。2.2 MP模型并行切分模型分摊参数MPModel Parallelism是把模型本身切开不同层或者同一层的不同部分放到不同卡上。最典型的是张量并行Tensor Parallelism把一个大矩阵乘法拆成几块每块在一张卡上算结果再拼起来。Megatron-LM就是这套思路的代表。MP的显存优势在于每张卡只存模型的一部分静态开销按切分比例下降。比如一个模型切成4份每张卡的权重、梯度、优化器状态都降到四分之一。这对超大模型来说是刚需。但MP的代价是通信频繁。张量并行在每层的前向和反向都要做all-reduce或者all-gather通信量大且延迟敏感。所以MP通常只在节点内做靠NVLink这种高带宽互联撑着跨节点做MP性能会掉得很厉害。另外MP对模型结构有要求不是所有层都能随便切实现复杂度比DP高不少。2.3 PP流水线并行按层切分分摊深度PPPipeline Parallelism是把模型按层切成若干段每段放到一张卡上数据像流水线一样依次经过各段。它的显存特点是每张卡只存模型的一段静态开销同样按段数下降。跟MP不同的是PP切分的是层通信只发生在段与段之间通信量比MP小得多。PP的难点在于流水线气泡。因为各段是串行依赖的如果只喂一个batch那同一时刻只有一张卡在干活其他卡都在等。为了解决这个问题需要把batch切成多个micro-batch让各段尽量并行起来。但即使这样流水线的首尾仍然会有气泡气泡比例跟micro-batch数量和流水线深度有关。我实际跑过一个4段PP的配置micro-batch设为8的时候气泡开销大概在15%左右还能接受。如果micro-batch太少气泡能占到一半以上多卡还不如单卡。2.4 三种策略的显存与通信对比把三种策略放在一起对比能看得更清楚维度DPMPPP切分对象数据模型层内模型层间单卡静态开销完整模型按切分比例下降按段数下降单卡激活值按batch比例下降部分下降按micro-batch下降通信频率每步一次all-reduce每层多次通信段间点对点通信量与参数量成正比大较小实现复杂度低高中适合场景模型能进单卡单层超大层数多、模型深实际生产中这三种策略经常组合使用。比如大规模训练里常见的是TPPPDP的三维并行节点内用TP切张量节点间用PP切层最外层用DP做数据并行。这样每一维都发挥自己的优势把显存和通信都控制在合理范围。3. 实操从单卡OOM到多卡策略落地理论讲完落到实操。这一节我按“先诊断、再选型、后调优”的顺序把整个流程走一遍附上我实际用过的配置和踩过的坑。3.1 第一步定位显存到底花在哪OOM的时候别急着加卡先搞清楚显存花在哪。PyTorch提供了几个好用的工具。最直接的是torch.cuda.memory_summary()能看到当前分配、峰值、缓存等细节。我习惯在训练循环里每隔若干步打印一次观察显存随训练步数的变化。import torch def print_memory(step): allocated torch.cuda.memory_allocated() / 1024**3 reserved torch.cuda.memory_reserved() / 1024**3 peak torch.cuda.max_memory_allocated() / 1024**3 print(fstep {step}: allocated{allocated:.2f}GB, reserved{reserved:.2f}GB, peak{peak:.2f}GB)allocated是实际被张量占用的显存reserved是PyTorch向驱动申请的显存池包含缓存peak是历史峰值。如果allocated远小于reserved说明有大量缓存没释放可以调torch.cuda.empty_cache()但别频繁调会影响性能。更细粒度的工具是torch.profiler能按算子级别看显存分配。我一般先用memory_summary粗看发现异常再上profiler。另外如果怀疑是激活值的问题可以临时把batch size调到1跑一步看显存降多少就能判断激活值占比。3.2 第二步根据模型规模选策略诊断完之后按模型规模选策略。我总结了一个粗略的判断流程模型静态开销按16字节/参数估小于单卡显存的60%激活值也能放下单卡就行或者上DP加速。静态开销接近或超过单卡显存必须切模型考虑MP或PP。单层特别大比如超大隐藏维度的Transformer优先TP。层数很多、单层不大优先PP。两者都有TP和PP组合。举个具体例子。假设手上有8张24G的卡要训一个70亿参数的模型。静态开销按16字节/参数算约112GB单卡肯定放不下。8卡DP也不行因为每张卡还是要放112GB。这时候必须切模型。如果模型是Transformer结构隐藏维度很大可以用TP切成8份每卡静态开销降到14GB加上激活值和零碎开销24G的卡能跑。如果隐藏维度不算特别大但层数多可以用PP切成8段每段静态开销也是14GB左右但PP的通信量比TP小跨节点也能接受。实际选型时还要考虑硬件拓扑。同一节点内的卡有高带宽互联适合TP跨节点带宽低适合PP和DP。这个原则在组网的时候就要想清楚。3.3 第三步DP的落地配置与调优DP的落地最简单PyTorch的DDP几行代码就能跑。核心是初始化进程组、包模型、配DataLoader的sampler。import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP dist.init_process_group(backendnccl) local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) model MyModel().to(local_rank) model DDP(model, device_ids[local_rank]) # DataLoader 要用 DistributedSampler sampler DistributedSampler(dataset, shuffleTrue) loader DataLoader(dataset, batch_sizeper_gpu_bs, samplersampler)DP调优的关键在梯度同步的效率。几个实操要点开启gradient_as_bucket_viewTrue能减少显存拷贝用find_unused_parametersFalse确认没有未使用参数时能减少开销通信后端用nccl别用gloo。另外如果模型有BatchNorm记得用SyncBatchNorm否则各卡统计量不一致会影响收敛。我踩过的一个坑是DDP默认在每次反向传播时同步梯度如果模型有多个输出分支、部分分支不参与loss会报未使用参数的错误。这时候要么设find_unused_parametersTrue有性能损失要么重构模型让所有参数都参与计算。后者更好但需要改代码。3.4 第四步MP和PP的组合配置MP和PP的实现比DP复杂通常借助现成框架。张量并行可以参考Megatron-LM的实现思路核心是把Linear层按列或按行切分前向时做all-gather或all-reduce。流水线并行可以用PyTorch的torch.distributed.pipeline.sync.Pipe或者DeepSpeed的pipeline实现。以PyTorch的Pipe为例需要把模型按层切成多个stage每个stage放到不同设备上然后定义micro-batch数量from torch.distributed.pipeline.sync import Pipe # 把模型切成4段分别放到4张卡上 model nn.Sequential( nn.Linear(1024, 1024), nn.ReLU(), nn.Linear(1024, 1024), nn.ReLU(), ) # 实际切分需要按层边界来这里仅示意 model Pipe(model, chunks8) # chunks即micro-batch数chunks参数就是micro-batch数量设得越大气泡越小但显存占用也越高因为要同时保留多个micro-batch的激活值。我一般从4开始试逐步加到8或16观察显存和吞吐的平衡点。MP的配置更细需要指定张量并行的维度。Megatron风格的实现里通常用tensor_model_parallel_size控制切几份。这个值不能超过单节点的卡数因为TP通信太频繁跨节点会拖垮性能。我实测过TP8跨两个节点吞吐比TP4单节点还低得不偿失。提示MP和PP组合时先确定TP的切分份数受节点内卡数限制剩下的卡用PP切层最外层再用DP扩数据并行。这个层级顺序不要搞反否则通信会成为瓶颈。3.5 第五步激活重计算与显存碎片治理不管用哪种并行策略激活重计算都是省显存的利器。PyTorch里用torch.utils.checkpoint包一下需要重计算的层from torch.utils.checkpoint import checkpoint def forward(self, x): x checkpoint(self.layer1, x) x checkpoint(self.layer2, x) return x注意checkpoint对输入有要求输入需要支持requires_grad且函数内部不能有随机性操作否则重计算时结果不一致。我踩过的坑是在checkpoint包裹的层里用了Dropout重计算时mask不一样导致梯度有偏差。解决办法是把Dropout移到checkpoint外面或者用固定随机种子。显存碎片是另一个隐形问题。长时间训练后reserved可能远大于allocated这就是碎片。PyTorch有环境变量PYTORCH_CUDA_ALLOC_CONF可以调分配器行为比如设expandable_segments:True能缓解碎片。我实测在长序列训练里开启这个选项后峰值显存能降5%到10%。4. 常见问题与排查实录实际训练里遇到的问题五花八门这一节我把踩过的典型坑整理成速查表附上排查思路和解决方法。4.1 OOM排查速查表现象可能原因排查方法解决方向启动就OOM静态开销超限算16字节/参数切模型或换大卡训练几步后OOM激活值累积或碎片打印memory_summary开重计算、清缓存验证时OOM验证batch太大单独设验证batch减小验证batch多卡比单卡更早OOM通信缓冲区占用看reserved变化调小buffer或换策略显存缓慢增长张量未释放profiler看引用检查是否有list累积启动就OOM最好判断静态开销算一下就知道。训练几步后OOM更隐蔽常见原因是某些中间结果被意外保留比如在list里累积了loss或者中间激活。我遇到过一次代码里为了画图把每步的attention map存进list跑了几百步显存就爆了。这种问题用profiler一看引用链就清楚了。验证时OOM也很常见因为验证通常不计算梯度但batch size可能跟训练一样大而验证的激活值不需要保留理论上更省。如果验证OOM先确认是不是没加torch.no_grad()加了还OOM就减小验证batch。4.2 分布式训练的通信瓶颈排查多卡训练没加速甚至变慢八成是通信瓶颈。排查思路先看每步耗时里通信占多少用torch.profiler或者NCCL的日志。如果通信占比超过30%就要考虑优化。DP的通信优化前面提过bucket view和减少未使用参数检查。MP的通信优化空间小主要靠硬件互联所以TP尽量放节点内。PP的通信量小但气泡是另一回事要调micro-batch数量。我遇到过一个典型问题8卡DP训练加速比只有3倍。查下来是每步的all-reduce耗时太长因为模型参数量大梯度同步成了瓶颈。后来改成梯度累积每累积4步才同步一次通信开销摊薄加速比提到6倍多。梯度累积的代价是等效batch变大可能需要调学习率。4.3 精度与显存的权衡陷阱混合精度能省显存但不是所有地方都能用FP16。我踩过的坑某些归一化层在FP16下数值不稳定loss会突然变NaN。解决办法是用BF16替代FP16BF16的动态范围更大不容易溢出虽然精度略低但训练稳定性好很多。现在新卡基本都支持BF16我默认用BF16。另一个陷阱是loss scaling。FP16训练需要动态loss scaling来防止梯度下溢如果scale设得不对梯度会变成0或者inf。PyTorch的AMP模块自带动态scaling一般不用手动调但如果发现梯度异常可以打印scale值看看。BF16不需要loss scaling这也是我偏好BF16的原因之一。4.4 策略选型的几个反直觉结论最后分享几个我实际踩出来的反直觉经验。第一卡多不一定快。PP的气泡、MP的通信、DP的同步都会吃掉多卡带来的收益有时候4卡配置调好了比8卡默认配置还快。第二显存省了不一定划算。激活重计算省显存但费时间如果显存本来就够开重计算纯属浪费。第三TP不是越多越好。TP切分份数超过节点内卡数后跨节点通信会让性能断崖式下跌宁可少切几份。还有一个容易被忽略的点不同并行策略对batch size的敏感度不同。DP下每卡batch可以很小因为梯度会同步PP下micro-batch不能太小否则气泡占比高MP下batch大小受切分方式影响。调参的时候要结合策略来定batch不能照搬单卡的经验。实际训练里我习惯先用小规模配置跑通流程确认显存和通信都在预期内再逐步放大。一上来就堆满卡和最大batch出了问题很难定位是哪个环节的锅。这个习惯帮我省了不少调试时间。