多GPU显存分配实战：从拓扑到通信关键路径的优化指南

多GPU机器到手、驱动装完、nvidia-smi能扫出四张卡之后，大部分人以为接下来就是写代码分配device的事。真跑过并行训练你会知道，显存分配才是真正让人头大的地方——分配不好，要么显存浪费一大半，要么训练隔几分钟OOM，要么看着占用不高但所有卡都在原地等待。多GPU环境下的显存分配策略之所以是优化并行计算的关键路径，是因为它把硬件连接、并行算法、运行时缓存和通信调度四件事绑在了一起，任何一个环节不合拍，整条链路都会卡住。

这篇文章我不打算讲教科书式的抽象理论，就以一台常见的8卡服务器为例，从物理拓扑开始，把单服务器多GPU怎么连、不同并行策略下显存往哪分、框架底层怎么分配和回收显存、最后怎么在通信与计算重叠中优化关键路径，一层层拆开来讲。多数内容是我在真实环境里反复调试过的经验，适合正在搭建多卡训练环境、或者被多卡OOM和通信瓶颈折磨的工程师和数据科学方向的学习者参考。

1. 先看物理拓扑：单服务器多GPU怎么连是显存分配的上限

1.1 为什么连接方式先于分配策略

先说一个多数人忽略的事实：你的显存分配策略，无论设计得多精巧，都跑不出物理互联的约束。

之前有同事在一台8卡服务器上跑并行训练，为了把大模型塞进显存，用张量并行把一层切成四份放进四张卡。他按顺序把0、1、2、3号卡放了一组，4、5、6、7号卡放了一组。看起来挺对称，但每次训练速度忽快忽慢。后来我用nvidia-smi topo -m一查，这台机器的8张卡并不是全互联的，0到3号之间走的是PCIe，4到5号之间才有NVLink高速链路。他这个"看起来公平"的分配方式，让一部分通信流量压在低带宽的PCIe上，训练时间直接受影响。

这里有个核心观念要先建立起来：显存分配不只是"这块数据放哪张卡"，还包括"这张卡和其他卡怎么交换数据"。 数据交换的带宽和延迟，由物理连接方式决定；而物理连接方式，在买机器、插卡的时候就基本定死了。所以优化显存分配的第一步，不是调代码，而是搞清楚服务器当前的GPU拓扑。

很多初学者会把"多GPU怎么连"理解成一个硬件问题，觉得把卡插上去、供电、装驱动就结束了。实际上这里说的"连"，至少有三层含义：

• 物理层面：卡插在哪些PCIe插槽，走哪条root complex路径，是否有NVLink桥或NVSwitch背板；
• 系统层面：操作系统和驱动识别出的拓扑顺序，nvidia-smi的GPU编号对应哪块物理卡；
• 逻辑层面：分布式框架里的rank编号和GPU编号怎么映射，通信组怎么建。

这三层任何一层错位，都会让显存分配策略失效。尤其是第三层，很多人不重视，结果就是明明高速链路就在那里，程序的通信流量却全部绕了远路。

1.2 nvidia-smi topo -m：一张你必须学会看的表

在任意一台多GPU机器上执行：

bash复制nvidia-smi topo -m

输出是一个GPU之间的连接矩阵，大概长这样：

code复制    GPU0    GPU1    GPU2    GPU3    GPU4    GPU5    GPU6    GPU7
GPU0   X      NV12    NV12    NV12    NV12    NV12    NV12    NV12
GPU1  NV12     X      NV12    NV12    NV12    NV12    NV12    NV12
...
GPU6  NV12    NV12    NV12    NV12    NV12    NV12     X      NV12
GPU7  NV12    NV12    NV12    NV12    NV12    NV12    NV12     X

但也有可能看到PCIe、PIX、PXB、SYS、CPU这些字样。简单说：

• NV#：两块卡之间有NVLink高速连接。NV12表示12条NVLink链路，NV4、NV6同理。
• PIX：两张卡挂在同一个PCIe交换机下面，通信走PCIe交换机，延迟相对低。
• PXB：经过PCIe根复合体和交换机，距离更远。
• SYS：需要跨CPU通信，通常是最慢的路径，这种连接一般出现在多路CPU服务器的跨CPU场景。

我在不同型号的服务器上见过完全不同的拓扑。有些4卡工作站是两两NVLink，剩余路径走PCIe；有些8卡平台因为有NVSwitch，所有GPU之间都是NVLink，矩阵看起来非常整齐。这两种机器的显存分配策略，起点就不一样。

另一个很有用的命令是：

bash复制nvidia-smi -q | grep -E "GPU Name|Memory Size|Link" 

能快速看每张卡的显存容量、驱动信息和拓扑链接状态。多卡环境下，确认所有卡的显存容量一致是前提，不然做对称分配时小的那张会变成瓶颈。

1.3 带宽数字背后，藏着什么样的显存交换代价

为了理解拓扑为什么影响这么大，几个数字要记住：PCIe Gen4 x16的双向带宽大约是64GB/s（单向32GB/s），PCIe Gen5翻倍。而NVLink 3.0单条单向25GB/s，一块A100上12条NVLink加起来单向300GB/s、双向600GB/s，是PCIe Gen4的大约10倍。NVSwitch则让多张卡之间的通信不必绕CPU，任意两卡之间都能以接近全带宽直接通信。

这个差异落在显存分配上是什么概念？比如数据并行训练时，每张卡算完一个阶段就要做梯度AllReduce。假设模型有1GB梯度数据，在NVLink上可能只要几毫秒到十几毫秒，走PCIe就得几十毫秒甚至上百毫秒。训练迭代多了以后，通信时间在总时间里的占比会大得离谱。你显存分得再均匀，通信路径慢，整体照样拉垮。

所以，拿到机器的第一件事就是跑nvidia-smi topo -m，把手里的卡按连接关系分组。把需要频繁通信的并行进程放在同一组高速互联的卡上，把跨组通信降到最低。这一步看起来和"显存分配"没有直接关系，实际上决定了后面所有分配策略的可行性。

2. 从并行策略切入：不同并行方式的显存分配逻辑

2.1 数据并行：显存分配的重点是"对称冗余"

数据并行（DP/DDP）是最容易理解，也最容易踩坑的并行方式。它的思路是：每张卡都放一份完整的模型副本，训练数据切成多个批次分给各卡，每卡独立算梯度，最后通过AllReduce把所有卡的梯度加总，再各自更新参数。

这种模式下，显存分配策略的核心关键词是"对称"。每张卡的显存占用结构是完全一样的：一份模型参数、一份优化器状态、一份当前批次的激活值、以及一套梯度缓冲区。于是总显存需求约为单卡需求乘以卡数。如果机器上插了4张24GB的卡，模型稍微大一点就能利用近96GB的"名义容量"。

但这里有个细节常被忽略：数据并行并不降低单卡显存压力。如果你的模型本身单卡放不下，单纯加卡是救不了你的。这时需要的是模型切分，而不是数据并行。所以设计分配方案时，先判断瓶颈是"显存装不下"还是"算力不够"。显存装不下，就要考虑下面的张量并行、流水线并行或FSDP类方案；算力不够，数据并行就够。

我见过不少团队拿着8×80GB的机器跑DDP，却在一个单卡就能跑得动的小模型上浪费大量通信时间——这是对显存资源和通信资源的双重浪费。数据并行对显存分配相对无压力，但对通信压力大，所以它对拓扑的要求反而最苛刻。rank进程之间的梯度聚合频率很高，如果两张卡的AllReduce链路走了PCIe，整个训练的实际吞吐会明显下降。

2.2 张量并行：把单层切开之后，显存分配在张量维度上展开

当模型单层都大得夸张（比如超大embedding或者超大FFN），数据并行解决不了，就要把一层内的矩阵运算拆到多张卡上。这就是张量并行。最常见的做法是把一个Linear层的权重按行或按列切分，两张卡各存一半，计算时通过AllGather交换中间结果。

这时候的显存分配逻辑完全不同：同一层的数据，被切到多个GPU上，每个GPU只持有1/N的权重；但前向传播过程中，激活值经常需要被完整地重建出来。比如列并行后要得到完整输出，就需要AllGather把两块结果拼起来，这部分显存是临时的、突发的，峰值会比平均占用高不少。

张量并行是显存分配比较精细、比较考验规划的一种方式。你需要清楚每个算子的输入输出形状在哪个阶段会扩大、哪个阶段会缩小，从而决定哪些张量放本卡、哪些要广播出去。如果不做这样的规划，框架默认的分配行为往往会让一块卡某个瞬间爆掉，另一块卡还闲着。

工作实践中的一个经验是：张量并行的通信频率极高，几乎每个算子前后都有集合通信。因此，张量并行的GPU组必须放在NVLink或NVSwitch全互联的拓扑内，千万别把张量并行的两个rank放到只能走PCIe的卡对上——那会让每次前向传播都像在搬一座山。

2.3 流水线并行：阶段之间的缓冲区同样要占显存

流水线并行是把模型按层切成多个stage，每个GPU负责其中一段。表面上看，每张卡只需保存自己那一段模型，显存压力小很多。但实际操作中，stage之间需要传递中间激活，这些激活要放在接收端的显存缓冲区里，而且因为流水线存在"冒泡"（bubble）和微批次调度，缓冲区不止一份。

比如经典的1F1B调度中，每个stage可能要保存多个微批次的中间结果，才能保证前后向交错执行时不冲突。这些缓冲区的显存开销经常被当成"模型显存"统计，导致设计分配方案时少算一块。

流水线并行对通信带宽的要求不如张量并行那么极致，但对通信延迟敏感。因为stage之间的数据依赖是顺序的，一次握手慢，整条流水线都要等。所以stage划分要尽量让相邻stage之间在拓扑上靠近，把跨交换机通信放在stage边界之外。

2.4 混合并行：一张卡上的显存具体由什么构成

真实大模型训练几乎不会只用一种并行。通常的做法是：先张量并行（TP）把超大层切开，再流水线并行（PP）按层切段，再用数据并行（DP）在每组之间做复制。这就是俗称的3D并行，也是Megatron-LM这类框架的默认思路。

在这样的混合并行下，一张GPU上的显存构成大体是：

如果把这五项列出来，你会发现显存分配的本质，就是在这几类需求之间做规划：权重和优化器状态是静态的，激活和通信缓冲区是动态的，碎片和缓存则是隐形的。做分配方案时，先算清楚静态部分占多少、动态部分峰值会有多高，才能决定该开多大batch size、要不要开检查点、分桶参数设多少。

3. 深挖运行时分配机制：缓存池、预分配和碎片问题

3.1 PyTorch默认分配器的缓存行为

很多人第一次排OOM时都会困惑：我明明释放了张量，nvidia-smi里的显存占用却没有降下来。这不是泄漏，而是深度学习框架的缓存分配器在起作用。

以PyTorch为例，它的CUDA缓存分配器会向CUDA驱动申请一大块显存，之后在内部维护一个空闲块列表（block pool）。tensor释放时，分配器不立刻把显存还给驱动，而是先放回自己的缓存池里，供后续分配复用。这么做的原因很现实：每次cudaMalloc都有驱动层开销，频繁申请释放会严重影响性能。缓存池能让相同大小的张量分配在毫秒级完成。

理解这个机制对显存分配有两个直接意义。

第一，程序里看到的"已占用"不等于"活跃使用中"，很多是分配器预留的缓存。如果在训练中途想塞一个大的评估张量，大概率会撞上缓存池已满而触发OOM，即使真正活跃的张量并不多。

第二，用torch.cuda.empty_cache()可以把缓存池中的空闲块释放回驱动，但注意它只能释放空闲块，不能移动正在使用的块，所以对峰值帮助有限。它更适合在训练和评估模式切换之间手动整理一次，而不是在训练循环里频繁调用。

TensorFlow那边也有类似的机制，默认会抢占几乎全部显存，通过set_memory_growth或者per_process_gpu_memory_fraction可以改成按需增长。多卡环境下如果用TensorFlow，建议显式设置每卡可分配上限，防止单卡把所有显存占光。

3.2 显存碎片：为什么有空闲却仍OOM

这是多GPU环境下最隐蔽的OOM来源。显存碎片和内存碎片的道理一样：分配器手里的空闲块是分散的，而新请求需要的是连续的大块。训练过程中，前向传播不断产生不同形状的激活，这些张量申请和释放的尺寸差异很大。一会儿一个512MB的块，一会儿一个64KB的块，反复分配释放之后，大块内存被切碎，后面突然来了一个1GB的请求，分配器找不到连续空间，就报OOM。

PyTorch提供了一批可调参数，其中最常用的环境变量是：

bash复制export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

max_split_size_mb控制分配器在多大尺寸以上的块不再拆分。调大这个值可以减少碎片，却也可能导致大块被长期占用、浪费显存；调小则碎片更严重但更省。我自己的经验是从默认值开始，如果OOM日志里出现"block"相关的提示，说明碎片可能是元凶，试着调到256或者更大。

还有一种更有效的配置：

bash复制export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

开启后，分配器会尽量使用可扩展段来减少碎片化，显著改善训练过程中激活尺寸波动导致的碎片问题。代价是某些版本下可能有少量性能开销，但对长稳训练非常值得。

3.3 显存峰值怎么测，才能指导分配决策

优化显存分配，前提是知道内存都去哪了。很多人只盯nvidia-smi的进程占用，这其实很粗糙。PyTorch里提供了更精确的接口：

python复制torch.cuda.reset_peak_memory_stats()
# 训练若干步
peak = torch.cuda.max_memory_allocated()
print(f"峰值: {peak / 1024**3:.2f} GB")

max_memory_allocated()返回的是分配器实际分配出去的峰值，配合torch.cuda.memory_summary()可以查看当前占用的详细构成，包括每个块的大小和用途。想要更细粒度地定位是哪个阶段导致峰值飙升，可以在每个训练子阶段前后记录一次快照，对比增长出现在前向、损失计算、反向还是优化器更新。

这里我要提醒一句：如果你开了混合精度，峰值显存不等于模型权重乘个系数，激活值的峰值往往才是大头。用torch.profiler做一次profile，能看到每个算子的显存申请事件，这是做精确分配优化最好的工具。

3.4 激活重计算和混合精度：用计算换显存的典型策略

在多卡环境下，显存不够时最直接的解法不是换更大的卡，而是调整分配策略。

激活重计算（gradient checkpointing）是使用最广泛的方法：前向传播时只保存少数关键激活值，其他全部丢弃；反向传播需要时再重新计算一遍。这样能把激活显存从与层数线性相关降到接近开平方级别，代价是大约增加30%~50%的计算量。在数据并行场景下，如果单卡激活太高，全部开重计算甚至可能比加一张卡更划算。

混合精度（FP16/BF16）的收益也很直接：权重和激活从FP32变成BF16，显存直接砍半；如果使用FP32主权重维护优化器状态，则可以通过FSDP或ZeRO状态切分进一步分散到多卡。注意BF16的指数位和FP32一致，在大模型训练中数值稳定性更好，也是我现在的默认选择。

这两个策略的共同点是：都在"用计算换显存"。在做多卡分配时，不要只想着把数据铺到更多卡上，也要考虑每一张卡上能不能通过重计算和精度调整把峰值压下来。有时候压掉30%的激活峰值，比多插一张卡还管用。

4. 优化关键路径：通信、计算与显存访问的重叠

4.1 关键路径到底在哪：多数情况下瓶颈不是算力

题目里说的"关键路径"，我理解为一条训练迭代里从数据进来到梯度更新完成之间，最串行、最不可压缩的那条链路。单卡训练时，这一般是前向计算和反向计算的串行链路，显存只是用来放中间结果的仓库。多卡训练时，链路里多出了通信环节，关键路径就变成了"计算 + 等待通信"的复合体。

很多团队优化多卡训练时只盯着GPU利用率，结果发现利用率上不去，不是因为显存不足，而是因为通信在关键路径上没有被隐藏。通讯时GPU核心在原地等数据，显存里明明有数据却传不过来，这就是关键路径变长的典型例子。

4.2 梯度分桶：让通信尽早开始

PyTorch DDP的一个关键优化是梯度分桶。反向传播计算梯度时，不会等每个参数梯度都算完再AllReduce，而是把相邻参数分成一个个bucket，每个bucket累积到一定大小（默认25MB）就发起一次AllReduce。这样通信可以和后续参数的反向计算重叠进行，而不是等所有梯度算完才通信。

python复制bucket_cap_mb = 25  # 默认值

有经验的工程师会根据网络状况调整bucket_cap_mb。NVLink上25MB的bucket通信很快，可以调大到50~100MB，减少AllReduce次数，降低通信启动的开销；PCIe环境下则要考虑小一点，让通信更早开始，尽量躲开计算峰值。这个参数没有绝对最优值，我一般跑一二十步看总耗时变化来定。

另外一个容易被忽略的点是：torch.nn.parallel.DistributedDataParallel相比早期DataParallel，重要的不仅是分布式设计，更是因为DDP把梯度规约下沉到了每个bucket，而DP是在每个看得见的同步点做梯度的多卡合并，通信开销和显存峰值完全不同。所以新项目不要再用DataParallel，用DistributedDataParallel。

4.3 通信计算重叠：还能在哪些环节做

除了梯度桶，通信和计算重叠还有几个常见切入点。

一是数据加载。数据从CPU侧搬到GPU显存的过程，也可以和前向/反向重叠，关键是尽量用pin_memory=True把CPU侧内存固定为锁页内存，让H2D拷贝走更快的DMA通道。如果不锁页，数据加载会隐式地做一次复制，显存分配和释放就会频繁发生，给全局分配器增加压力。

二是通信集合的异步执行。在多机多卡场景下，torch.distributed里的集合通信默认是同步的，但可以通过async_op=True把AllReduce变成异步，把下一步不需要通信结果的计算先调度起来。这个技术在需要多级并行时尤其有用。

三是把梯度裁剪、权重更新挪到通信结果的子集上。大模型训练里，参数更新往往可以按层划分，第一层参数的梯度AllReduce完成后就能立刻更新，不必等全部梯度同步完。框架层面的参数通信重叠、部分ZeRO实现都类似思路。显存分配此时要留意，异步通信持有的临时缓冲区不能过早被新张量覆盖，否则容易出现数据竞争。

4.4 显存分配与通信的多阶段重叠

再做更深一层。数据并行、张量并行、流水线并行混在一起时，显存的分配时机和通信时机高度耦合。比如，数据并行的AllReduce和张量并行的AllGather可能同时进行，如果分配器没有提前把本轮要用的缓冲区留好，等待分配的时间就会插入通信关键路径，导致不可预测的卡顿。

我习惯的做法是：在训练正式开始前，做一次"显存预埋"，把已知的通信缓冲区（比如梯度桶的最大累计容量、AllGather需要的分片拼接区）用持久化张量提前分配好，尽量让训练循环里不再做大的显存申请。这样虽然会稍稍抬高基线占用，但换来的是更平滑、更可预测的执行时间。对于追求极致吞吐的线上训练，这是值得的。

可以这样理解：显存分配和内存管理的道理类似，更像一套仓库管理系统。频繁向驱动申请和释放，相当于每次从仓库外面搬箱子，开销巨大；而提前把常用箱子放到手边，哪怕多占一点空间，换的是效率。工程上我们愿意用一点点空闲显存换关键路径的确定性。

5. 我实测中踩过的坑与一套可复用的优化思路

5.1 四个典型的"假OOM"

多卡调试中遇到的OOM，相当一部分不是显存总量不足，而是分配策略不当。我总结出四个高频场景。

场景一：分配器缓存导致假性不足。 训练中途切换评估模式，评估张量比较大，但训练阶段已经把缓存池塞