多GPU训练显存分配实战：从OOM到优化

多GPU训练大家聊得最多的是"怎么把模型跑起来"，但真正跑到大规模并行时，最先拦路的往往是显存分配。前两周我调一个多卡训练任务，8张A800，按理说显存总容量接近640GB，结果跑一个70B模型的LoRA微调，居然频繁OOM，单卡占用就崩了。排查了很久，最后发现不是模型太大，而是显存分配策略出了问题。这种事情在刚接触多GPU环境的朋友那里几乎天天发生，所以这次我就把显存分配这条线从头梳理一遍，从底层机制到实操调优，一次讲透。

这篇文章主要面向两类人：一类是从单卡迁移到多卡训练的算法工程师，另一类是负责训练平台和推理服务部署的运维/平台开发。你会看到多GPU环境下显存的分配方式、并行模式对显存的影响，以及一组可以直接落地执行的优化手段。理解了这些，你就不会再把"显存不够"简单归因于"卡不够多"。

1. 多GPU环境下的显存分配到底在解决什么问题

1.1 多卡并不是"显存翻倍"这么简单

很多人对多GPU训练的第一印象是：一张卡装不下的模型，那就把模型切到多张卡上，合起来的总显存总是够的。这个思路理论上没错，但实际落地时你会发现，多卡环境下的显存分配远远不是"加总"这么简单。原因在于数据并行、模型并行、流水线并行这些并行方式，对显存的消耗模式完全不同。

拿最常见的分布式数据并行（DDP）来说，它会把一个batch的数据切分到每张卡上，每个GPU各持有一份完整的模型参数、梯度、优化器状态。也就是说，如果你一张卡需要20GB显存来放模型相关数据，那8张卡一共需要160GB，而不是20GB。显存总量确实是翻倍了，但单卡需求一点没降。这也是很多人的误区——以为上了多卡，单卡压力就小了，实际上在DDP模式下单卡压力几乎没有变化。

而显存分配策略要解决的核心问题，是让有限显存在训练过程中被合理利用，避免碎片化、避免突发峰值、避免通信等待导致的内存堆积。它决定了你的训练是能稳定跑很久，还是每隔几小时就OOM重启。

1.2 显存分配不只是PyTorch的事

在动手调参前，我们必须建立一个全局视角。一块GPU的显存使用，至少包含几层：

• 驱动层显存管理：CUDA driver会为context、module、CUDA event等分配固定开销，这部分虽然不大，但确确实实存在。
• CUDA runtime缓存：CUDA在底层会缓存一些分配和释放的空间，不会立即归还给操作系统。
• 框架层显存分配器：PyTorch、TensorFlow都有自己的缓存分配器（caching allocator），它会向CUDA申请大块显存，然后内部按需切片分配，避免频繁调用cudaMalloc导致性能下降。
• 模型运行时张量：前向激活值、梯度、参数、优化器状态、临时缓冲区等。

实际排查显存问题时，我们通常看的是PyTorch层面的显存统计，但底层CUDA缓存的占用经常给人"惊喜"。这就是为什么你经常在nvidia-smi里看到进程占用了很多显存，但代码里统计不到对应大小的张量——因为它们被缓存住了，并没有真正释放回驱动层。

多GPU环境下，这个缓存机制还会被NCCL通信库放大。NCCL在初始化时会为每个通信组分配一定量的显存作为buffer，这些buffer平时不参与计算，但会一直占着地方。如果你并行组很多，或者每个rank都建了多个通信组，这些buffer加在一起也能吃掉不少显存。

2. 硬件拓扑：分配策略必须先看懂的说

2.1 单服务器多GPU卡之间的连接方式

要讲显存分配策略，就绕不开GPU之间的通信拓扑。因为多卡训练中，模型同步、梯度聚合、张量搬运全都依赖GPU间的互联，而互联方式直接决定了通信的带宽和延迟，进而影响你该选择哪种并行模式和显存分配方案。

单服务器多卡最常规的连接方式是PCIe。PCIe的带宽目前主流是PCIe 4.0 x16，单向带宽约32GB/s，双向约64GB/s。听上去不低，但在大模型梯度同步场景下，一个几百MB的梯度张量要广播到8张卡，PCIe会非常吃力。而且PCIe是树状拓扑，不同槽位的GPU之间通信可能要经过CPU或PCIe Switch，延迟和带宽都有明显差异。

更快的方案是NVLink。NVLink是GPU直连的高带宽互联总线，以A100为例，第三代NVLink单条链路带宽是50GB/s，一般有12条链路，总双向带宽达到600GB/s。8卡A100通过NVLink全互联（fully connected）后，任意两张卡之间都能以很高的带宽直接通信，不用经过CPU。这种拓扑下，数据并行通信开销会小很多。

所以你看硬件拓扑，本质上是在搞清楚一个问题：卡与卡之间"说话"快不快。通信越快，你越可以用激进的并行策略；通信越慢，你就越要想办法让数据少跨卡流动，或者在显存分配上预留出通信缓冲的量。

2.2 如何确认你的服务器拓扑

实操层面，上来别急着写代码，先确认机器拓扑。推荐用下面这些命令：

bash复制# 查看NVIDIA驱动和GPU状态
nvidia-smi

# 查看GPU之间的连接拓扑（包括PCIe和NVLink）
nvidia-smi topo -m

# 查看NVLink连接状态
nvidia-smi nvlink -s

nvidia-smi topo -m输出的矩阵里，NV#表示两个GPU之间有NVLink连接，PIX表示通过PCIe直连同一个交换机，PHB表示通过PCIe桥连接，NODE表示要通过跨CPU socket访问。这个矩阵基本决定了你的通信效率。

我调试的一台8卡机器，拓扑输出是4条NVLink一组，前4张卡之间全互联，后4张卡之间全互联，但第4张卡和第5张卡之间只有PCIe连接。这种情况下，如果你用8卡DDP，每次梯度同步时跨NVLink和PCIe混合通信，NCCL会自动选择最佳路径，但PCIe链路会成为瓶颈。更合理的做法是拆成两个4卡组，组内通信走NVLink，组间通信通过主卡中转，或者干脆用模型并行时把需要频繁通信的层放在同一组内。

这些拓扑信息直接影响了你的显存分配方案。如果通信走PCIe，NCCL的buffer分配策略可能需要调整；如果走NVLink，你可以更放心地使用张量并行，因为通信开销低，显存交换更频繁也不会太拖慢速度。

3. 三种并行模式下的显存分配逻辑

3.1 数据并行：显存是"复制"，不是"分割"

数据并行（Data Parallelism）是最常见的多卡训练方式，PyTorch中的DDP就是典型代表。它的做法是每张卡持有一份完整的模型副本，各自处理不同的数据子集，然后在反向传播后同步梯度。

从显存分配角度看，数据并行有三个特点：

• 每张卡都要完整放下一份模型参数、梯度和优化器状态。
• 前向计算产生的激活值按mini-batch大小均匀分割后分布在各卡上。
• 梯度同步阶段，NCCL会额外分配通信buffer。

所以DDP模式下，显存分配策略的核心是：如何让每张卡的显存占用尽量均衡，同时留出足够的梯度同步buffer空间。

实际调优时，我通常会关注两点。第一是batch size的切分，确保每张卡上的样本数大致一致，避免某张卡因样本过多导致显存高出其他卡一大截。第二是通信buffer的预留，如果你的模型比较大、梯度张量多，可以考虑设置NCCL_BUFFSIZE环境变量。默认值一般是16MB，但大模型场景下建议调大到64MB或128MB，避免小buffer频繁申请释放带来的性能抖动。

bash复制# 调大NCCL通信buffer的示例
export NCCL_BUFFSIZE=134217728  # 128MB

3.2 模型并行：切分参数的显存艺术

模型并行（Model Parallelism）是指把模型不同层放到不同GPU上，每张卡只负责模型的一部分计算。这种模式下，显存分配面临的核心挑战是负载均衡和通信依赖。

以最简单的按层切分为例，假设一个Transformer有16层，我用4张卡，每张卡放4层。理想情况下每张卡的显存占用应该基本一致，但实际情况往往不是这样。attention层和FFN层的激活值大小不一样，第一层的输入嵌入层可能很大，最后一层的输出头也可能很大，切分不均会导致某张卡显存溢出，其他卡却闲着。

更细粒度的做法是张量并行（Tensor Parallelism），把每一层的计算按矩阵维度切分到多张卡上。这种方式显存分配更均匀，但通信量巨大，每层前向反向都要做all-reduce，只有NVLink等高带宽互联才扛得住。

模型并行的显存分配策略要点是：先摸清每层激活值的峰值区间，再做切分，而不是简单按层数平均分。你可以用PyTorch Profiler跑一个小batch，导出每个层的内存占用曲线，再根据曲线做切分决策。

3.3 流水线并行：显存与吞吐的折中

流水线并行（Pipeline Parallelism）是介于数据并行和模型并行之间的方案。它把模型按层切分成多个stage，每个stage放在一张卡上，数据按micro-batch流式经过各stage。比如一个20层的模型切成4个stage，每个stage 5层，放在4张卡上。

流水线并行的显存分配特点是：每张卡只需要存放一个stage的模型参数和激活值，但需要额外的buffer来保存stage间传递的中间结果。这个中间结果buffer的大小，取决于你设置的micro-batch数量和激活值大小。

显存分配上，流水线并行最需要小心的是stage间传输buffer的峰值。如果你的micro-batch数量设置过大，这些中间结果的累积会占掉大量显存。通常建议micro-batch数量不超过stage数的整数倍，具体数值需要通过实测来确定。

4. PyTorch显存分配器的底层机制

4.1 缓存分配器：为什么不释放就是不释放

PyTorch的CUDA caching allocator是理解显存分配最关键的组件。它的工作逻辑是：向CUDA驱动申请一大块显存（比如2GB），然后在内部把这些显存切成小块，按需分配给张量使用。当一个张量被释放时，这块显存不会归还给CUDA，而是留在PyTorch的缓存池里，等待下次分配。

为什么要这么设计？因为cudaMalloc的调用成本非常高，频繁向驱动申请/释放显存会导致严重的性能下降。缓存机制用空间换时间，大幅降低了显存分配的开销。

但这带来了一个副作用：你在代码里看到某个张量被删了，torch.cuda.memory_allocated()指标下降了，但torch.cuda.memory_reserved()没有变化，nvidia-smi里的显存占用也纹丝不动。这属于正常现象，不是显存泄漏。

多GPU环境下，缓存分配器是per-device（每张卡）独立的。也就是说，每张卡都有自己的缓存池，不能跨卡共享。这带来一个很有意思的问题：如果你在代码中频繁地在不同GPU上创建和销毁张量，缓存池会各自膨胀，整体显存开销比理论计算值高不少。

4.2 显存碎片化是怎么形成的

显存碎片化是分配器最头疼的问题。它的形成过程是：显存被划分为很多不同大小的块，一些块被释放后，留下很多不连续的空隙。当你需要分配一个大块连续显存时，缓存池里每个空隙都不够大，分配器只能向驱动申请新的显存段，导致reserved显存不断上涨。

碎片化的典型场景是训练过程中动态调整batch大小，或者在不同大小的输入上反复做前向和反向。NLP自然语言处理场景尤其明显，因为句子长度不同，激活值大小差异巨大，频繁分配不同大小的张量，几个月跑下来显存占用会越来越离谱。

缓解碎片化的手段有几种：

• 设置PYTORCH_CUDA_ALLOC_CONF的max_split_size_mb参数，限制小块缓存的上限。例如设为128MB，小于128MB的显存块不会被单独缓存，而是合并到相邻空闲块中，能有效减少碎片。
• 在训练循环中尽量固定张量形状。如果输入会动态变化，用padding统一长度。
• 定期用torch.cuda.empty_cache()释放缓存池中的空闲块。注意这个方法只释放缓存中不使用的部分，不改变allocated值。

python复制# 设置max_split_size_mb的例子，通常在main.py开头设定
import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"

4.3 显存统计指标：你得知道看哪个数

排查显存问题时，首先得搞清楚几个指标的区别：

实际业务中，我们最关心的通常是memory_allocated，因为它代表真正的需求。但排查OOM时，memory_reserved和nvidia-smi的差值往往能提供重要线索。

比如有一次我排查一个推理服务的显存增长问题，memory_allocated稳定在8GB，但nvidia-smi显示的进程显存从9GB涨到了20GB。这就不是PyTorch张量导致的，而是推理框架底层CUDA缓存或者NCCL buffer在累积。这种情况下，单纯靠PyTorch层级的优化解决不了，需要检查底层库的显存管理行为。

5. 实操：多GPU显存分配的调优清单

5.1 第一步：摸清你的显存占用画像

任何优化都从测量开始。我的习惯是先在训练脚本里加上显存监控，每50步打印一次每张卡的allocated、reserved和nvidia-smi的占用。这一步能快速确定显存分配是否均衡。

python复制import torch
import pynvml

def print_gpu_memory(step):
    # PyTorch侧统计
    allocated = torch.cuda.memory_allocated()
    reserved = torch.cuda.memory_reserved()
    print(f"Step {step} | allocated: {allocated / 1024**3:.2f} GB | reserved: {reserved / 1024**3:.2f} GB")

    # nvidia-smi侧统计（需要pynvml库）
    pynvml.nvmlInit()
    for i in range(torch.cuda.device_count()):
        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
        info = pynvml.nvmlDeviceGetMemoryInfo(handle)
        print(f"  GPU {i}: used {info.used / 1024**3:.2f} GB, total {info.total / 1024**3:.2f} GB")

跑上几百步之后，你会得到一张表：哪个step显存峰值最高、哪个step的reserved出现异常跳变、各卡之间的占用是否均衡。基于这些数据，再决定下一步怎么调。

这里分享一个实测案例。我之前跑一个多模态模型的多卡训练，每张卡在step 200时的allocated只有14GB，reserved却有20GB，nvidia-smi显示已用22GB。一看就知道缓存池里有6GB空闲块没释放，而且NCCL buffer还占了一部分。这种情况不影响正常运行，但如果你后面要加载更大的验证集，就可能会出现显存不足。所以我会在验证阶段前主动torch.cuda.empty_cache()，把缓存池里的空闲块释放掉。

5.2 第二步：按需设置显存上限

如果你的服务是多进程共享GPU，比如一个服务器上同时跑了训练进程和推理进程，就必须考虑显存上限的设置。PyTorch提供了两个关键接口：

python复制# 设置进程可分配的最大显存（比例制，0到1之间）
torch.cuda.set_per_process_memory_fraction(0.8, device=0)

注意，这个接口是"硬限制"，当分配超过比例上限时，会直接报OOM。但它限制的是caching allocator从驱动层申请的显存总量，不包括CUDA context等固定开销，所以实际预留时要留一些余量。

另一个常用的手段是环境变量CUDA_VISIBLE_DEVICES，它决定进程能看到哪些GPU。在多进程场景下，每个进程只绑定自己需要的卡，可以避免显存分配互相干扰。

bash复制# 让当前进程只看到第2和第3张卡
export CUDA_VISIBLE_DEVICES=1,2

在多GPU显存分配中，这个变量还有一个隐藏作用：它会改变CUDA device的编号映射。比如你物理上有8张卡，用CUDA_VISIBLE_DEVICES=6,7后，这两个设备在代码中的编号变成0和1。如果不小心，容易在分布式训练中把rank和GPU对应错，导致多个进程绑到同一张卡上，直接OOM。

5.3 第三步：梯度检查点与混合精度，双管齐下

如果显存实在不够，最有效的两类方案是梯度检查点（gradient checkpointing）和混合精度训练。

梯度检查点的核心思想是：前向传播时不保存所有激活值，只保存部分关键节点，反向传播时重新计算需要的激活值。这是用计算换显存的典型手法，通常能减少70%左右的激活值显存占用，但会增加约30%的计算开销。在显存吃紧的场景下，这是非常值得的取舍。

python复制# Hugging Face Transformers中开启gradient checkpointing
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("your_model_path")
model.gradient_checkpointing_enable()

混合精度训练则是把模型参数、梯度和激活值从前向计算中部分改为FP16或BF16存储和计算，显存占用减半。在大规模并行训练中，混合精度已经成了标配。用PyTorch的话，最简单的写法是：

python复制from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
for batch in dataloader:
    with autocast():
        loss = model(batch)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

实际使用时要特别注意BN层和梯度裁剪在混合精度下的行为，有些操作需要特殊处理。比如梯度裁剪时，一定要先把梯度unscale出来再裁剪，否则梯度的scale因子会影响裁剪阈值。多卡训练里，如果不同卡上梯度scale不一样，同步时会出现数值问题。

5.4 第四步：数据加载与预取对显存的影响

这是一个容易被忽略的点。数据加载器（DataLoader）的num_workers和prefetch_factor，会直接影响CPU侧内存和GPU显存的配合。当num_workers过高、预取batch较多时，主板上的内存压力大，但更关键的是，如果数据加载过程中涉及GPU上的预处理（比如把图片搬到GPU再增强），那么预取的batch会在GPU显存中堆积，造成显存峰值升高。

我见过一个训练任务，跑几步就OOM，报错显示张量在cuda:0上溢出。排查后发现DataLoader设置了num_workers=32和prefetch_factor=8，每个worker都会往GPU上预取一批数据做增强，GPU显存里堆积了好几个batch的数据。把prefetch_factor降到2，OOM问题直接消失。

多GPU训练中建议：数据预处理尽量放在CPU侧完成，GPU只负责模型计算。如果必须在GPU上做预处理，要控制好预取batch的数量。

6. 常见问题与排查技巧实录

6.1 CUDA Out of Memory报错

这是最经典的错误，但原因各不相同。我的排查思路是分三步：

• 看memory_allocated，确认张量实际占用多少。
• 看memory_reserved，确认缓存池是否异常膨胀。
• 看nvidia-smi，检查是否有其他进程或底层库占用显存。

如果allocated没到上限，reserved却逼近上限，说明缓存池碎片化严重。先用PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:64重启任务试试，这个参数对碎片化场景有奇效。

如果3个值都逼近上限，那就真的是模型或batch太大。这时候优先开gradient checkpointing，其次尝试把batch再切小，用梯度累积（gradient accumulation）来做替代。

多GPU场景还有个特殊情况：显存分配不均衡导致某张卡OOM，其他卡闲得很。原因通常是代码里有些张量固定在某张卡上创建，比如loss计算、指标统计，或者某个模型层没有正确用device映射到多卡。

6.2 推理进程显存不断上涨

多卡服务器上经常会有多个推理进程并存，这类进程的显存增长很多时候不是内存泄漏，而是CUDA caching allocator和NCCL buffer的累积效应。

解决方案：在推理进程里，每隔一段时间调用torch.cuda.empty_cache()。但要注意，这个操作有一定性能开销，频繁调用会影响推理延迟，所以要控制频率，比如每1000次推理清理一次。

如果清空了缓存但显存还在涨，那就得检查是不是有张量在GPU上堆积。可以用torch.cuda.memory_snapshot()导出内存快照，看看哪些分配点占用最多。这是PyTorch自带的排查工具，信息量很大。

6.3 NCCL通信卡死或异常

NCCL卡死通常表现为训练进程卡在一个step上，显卡利用率下降到0。这往往与显存分配也有关系——NCCL通信buffer可能需要额外分配显存，但显存不足时分配失败，导致通信初始化异常。

排查步骤：

• 用torch.distributed.barrier()在各个rank上打印日志，确认卡在哪个阶段。
• 设置NCCL_DEBUG=INFO环境变量，查看NCCL的初始化细节。
• 检查网络拓扑，确认NCCL_P2P_LEVEL是否需要手动指定。

bash复制export NCCL_DEBUG=INFO
export NCCL_P2P_LEVEL=NVL  # 限制P2P通信只在NVLink连接之间启用

有些机器上，PCIe连接的两个GPU之间P2P通信不稳定，禁用P2P后反而更稳定。代价是通信走共享内存或CPU内存，带宽会下降，但能保证训练不卡死。

6.4 多进程频繁切换导致显存碎片化

如果你在一个GPU上同时跑多个训练任务（比如用小batch做超参搜索），进程频繁创建销毁，CUDA缓存很容易碎片化。我建议优先用MPS（Multi-Process Service）来管理这种场景。MPS可以把多个CUDA进程的上下文合并在一个MPS server里，显存分配更集中，碎片化会明显减少。

启动MPS的方式：

bash复制# 设置MPS模式
export CUDA_MPS_PIPE_DIRECTORY=/tmp/mps_pipe
export CUDA_MPS_LOG_DIRECTORY=/tmp/mps_log
nvidia-cuda-mps-control -d

不过MPS有自己的限制，比如要求所有进程在同一个GPU上，而且对PyTorch的多进程分布式支持不算完美。如果你的任务需要跨GPU通信，MPS反而会成为瓶颈。所以这个方案更适合同一物理GPU上多个小任务并存的场景。

7. 多GPU显存调优的几个核心心得

• 显存分配不是一锤子买卖，需要贯穿整个训练流程。我现在的习惯是每次训练任务启动前都会写一段显存监控代码，跑前几百步观察内存曲线，发现问题立刻处理。
• 优先使用框架层工具而不是手工管理。PyTorch的caching allocator已经做得很好了，手动torch.cuda.empty_cache()和set_per_process_memory_fraction用错了反而会降低性能。只有遇到具体问题时，才需要干预。
• 一步一个指标。调优的时候不要同时改多个参数，每次只改一个，对比性能指标和显存指标的变化，不然出了问题你根本不知道是哪一步改坏的。

多GPU环境下的显存分配，说到底是"硬件拓扑 + 并行模式 + 框架机制"三者的平衡。硬件拓扑决定了通信的底线，并行模式决定了显存的需求模型，框架机制决定了你能在多大程度上精细控制。把这三层都理解透了，遇到显存问题才不会只想到"加卡"这个笨办法。