深度学习显存优化指南：从OOM到DP、MP、PP并行策略实战

显存不够用这件事，几乎每个跑过深度学习训练的人都撞过墙。我第一次遇到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×4=8字节。

我拿一个10亿参数的模型举例，算一笔账：

也就是说，10亿参数的模型，光静态开销就要16GB左右。一张24G的卡，还没开始算激活值，就已经用掉三分之二。这就是为什么很多人发现模型明明不大，训练却跑不起来。

注意：上面这个16字节/参数的估算是一个常用的经验值，实际会因框架实现、是否使用ZeRO、是否开启梯度累积等有所浮动。但量级上可以作为快速估算的依据。

如果换成纯FP32训练、不用混合精度，权重4字节、梯度4字节、Adam两个状态各4字节，合计16字节，跟混合精度的账其实差不多，但混合精度在计算速度和激活值占用上有优势，所以现在基本是标配。

1.2 动态开销：激活值才是隐形杀手

激活值的显存占用没有固定公式，它跟网络深度、每层的输出尺寸、batch size都有关。粗略估算，激活值占用跟 batch_size × 序列长度 × 隐藏维度 × 层数 成正比。Transformer类模型里，注意力机制的中间结果还会额外放大这个数字。

我实际测过一个中等规模的Transformer，隐藏维度768、12层、序列长度512，batch size从8加到32，激活值占用从不到2GB涨到接近7GB，几乎是线性增长。这时候哪怕静态开销只有几个G，总显存也会被激活值顶爆。

激活值有一个特点：前向传播时逐层产生，反向传播时逐层消耗。理论上，算完某层的梯度后，那层的激活值就可以释放。但框架为了效率，通常会保留一部分，这就引出了**激活重计算（gradient checkpointing）**这个技巧——牺牲约30%的计算时间，换取激活值显存大幅下降。我实测下来，开启重计算后激活值占用能降到原来的三分之一甚至更低，代价是每步训练时间增加20%到40%，具体取决于模型结构。

1.3 那些容易被忽略的零碎开销

除了上面四大块，还有几项零碎但不可忽视的占用。CUDA上下文本身就要占几百MB，cuDNN的workspace在卷积网络里可能占1到2GB，通信缓冲区在分布式训练时也要预留空间。这些加起来，一张卡上轻松再吃掉1到3GB。

所以实际估算显存时，我习惯在理论值上再加20%到30%的余量。比如算出静态加激活大概20GB，那24G的卡基本就是极限操作，稍微改个batch size或者序列长度就爆。稳妥起见，留出3到4GB的缓冲比较安心。

2. DP、MP、PP三种策略的显存逻辑拆解

单卡塞不下的时候，就得考虑多卡。DP、MP、PP是三种最基本的并行维度，它们切分的东西不一样，显存账本也完全不同。很多人把它们混为一谈，实际上选错策略可能让你多花好几张卡却收效甚微。

2.1 DP数据并行：复制模型，分摊数据

DP（Data Parallelism）的思路最直观：每张卡上都放一份完整的模型副本，把batch的数据切成N份，每张卡算一份，然后同步梯度。它的显存特点是：每张卡的静态开销完全一样，都是完整模型的量。也就是说，DP不减少单卡的模型显存占用，它减少的是每张卡处理的batch大小，从而降低激活值占用。

DP的通信量在于梯度同步。每步训练结束后，所有卡的梯度要做all-reduce，通信量跟模型参数量成正比。卡越多，通信开销占比越大。我见过8卡DP训练一个小模型，通信时间占了每步的40%以上，加速比远达不到线性。

DP适合什么场景？模型能塞进单卡、但数据量大、想加快训练速度的时候。它的实现最简单，PyTorch的DistributedDataParallel几行代码就能跑起来。但如果模型本身就塞不进单卡，DP帮不上忙，因为每张卡还是要放完整模型。

2.2 MP模型并行：切分模型，分摊参数

MP（Model Parallelism）是把模型本身切开，不同层或者同一层的不同部分放到不同卡上。最典型的是张量并行（Tensor Parallelism），把一个大矩阵乘法拆成几块，每块在一张卡上算，结果再拼起来。Megatron-LM就是这套思路的代表。

MP的显存优势在于：每张卡只存模型的一部分，静态开销按切分比例下降。比如一个模型切成4份，每张卡的权重、梯度、优化器状态都降到四分之一。这对超大模型来说是刚需。

但MP的代价是通信频繁。张量并行在每层的前向和反向都要做all-reduce或者all-gather，通信量大且延迟敏感。所以MP通常只在节点内做，靠NVLink这种高带宽互联撑着，跨节点做MP性能会掉得很厉害。另外MP对模型结构有要求，不是所有层都能随便切，实现复杂度比DP高不少。

2.3 PP流水线并行：按层切分，分摊深度

PP（Pipeline Parallelism）是把模型按层切成若干段，每段放到一张卡上，数据像流水线一样依次经过各段。它的显存特点是：每张卡只存模型的一段，静态开销同样按段数下降。跟MP不同的是，PP切分的是层，通信只发生在段与段之间，通信量比MP小得多。

PP的难点在于流水线气泡。因为各段是串行依赖的，如果只喂一个batch，那同一时刻只有一张卡在干活，其他卡都在等。为了解决这个问题，需要把batch切成多个micro-batch，让各段尽量并行起来。但即使这样，流水线的首尾仍然会有气泡，气泡比例跟micro-batch数量和流水线深度有关。

我实际跑过一个4段PP的配置，micro-batch设为8的时候，气泡开销大概在15%左右，还能接受。如果micro-batch太少，气泡能占到一半以上，多卡还不如单卡。

2.4 三种策略的显存与通信对比

把三种策略放在一起对比，能看得更清楚：

实际生产中，这三种策略经常组合使用。比如大规模训练里常见的是TP+PP+DP的三维并行：节点内用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(f"step {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(backend="nccl") 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, shuffle=True) loader = DataLoader(dataset, batch_size=per_gpu_bs, sampler=sampler)

DP调优的关键在梯度同步的效率。几个实操要点：开启gradient_as_bucket_view=True能减少显存拷贝；用find_unused_parameters=False（确认没有未使用参数时）能减少开销；通信后端用nccl，别用gloo。另外，如果模型有BatchNorm，记得用SyncBatchNorm，否则各卡统计量不一致会影响收敛。

我踩过的一个坑是：DDP默认在每次反向传播时同步梯度，如果模型有多个输出分支、部分分支不参与loss，会报未使用参数的错误。这时候要么设find_unused_parameters=True（有性能损失），要么重构模型让所有参数都参与计算。后者更好，但需要改代码。

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, chunks=8) # chunks即micro-batch数

chunks参数就是micro-batch数量，设得越大气泡越小，但显存占用也越高，因为要同时保留多个micro-batch的激活值。我一般从4开始试，逐步加到8或16，观察显存和吞吐的平衡点。

MP的配置更细，需要指定张量并行的维度。Megatron风格的实现里，通常用tensor_model_parallel_size控制切几份。这个值不能超过单节点的卡数，因为TP通信太频繁，跨节点会拖垮性能。我实测过TP=8跨两个节点，吞吐比TP=4单节点还低，得不偿失。

提示：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最好判断，静态开销算一下就知道。训练几步后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替代FP16，BF16的动态范围更大，不容易溢出，虽然精度略低但训练稳定性好很多。现在新卡基本都支持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，出了问题很难定位是哪个环节的锅。这个习惯帮我省了不少调试时间。