4、显存分配机制：CUDA显存分配策略、PyTorch显存分配机制、TensorFlow显存分配机制

聊到压测中的显存瓶颈，咱们得先搞清楚一个底层问题：显存到底是怎么分配的？

说白了，GPU显存不像CPU内存那样随用随取。它有一套自己的分配逻辑。如果你不了解这套逻辑，压测时就会出现「明明显存还有几百MB，但程序就是报OOM」的诡异现象。我当年第一次遇到这个问题，差点把显卡拔下来重插——后来才发现，是分配策略在搞鬼。

4.1 CUDA显存分配策略

CUDA作为底层接口，它的显存分配策略直接影响上层框架的行为。我个人习惯把CUDA的分配策略分为两种模式：懒惰分配和预分配。

4.1.1 懒惰分配（Lazy Allocation）

这是CUDA的默认行为。你调用cudaMalloc时，驱动并不会立刻在物理显存上划一块给你。它只是给你一个虚拟地址，等到你真正往这块地址写数据时，才触发实际的物理显存分配。

嗯，这里要注意：如果你只是分配了显存但没有使用，nvidia-smi里可能看不到显存占用。我在项目中遇到过这种情况——压测脚本跑完一轮后，显存占用看起来很低，但下一轮突然就OOM了。为什么？因为上一轮分配的内存虽然没写数据，但虚拟地址空间已经被占用了。

4.1.2 预分配（Eager Allocation）

有些场景下，懒惰分配会导致性能抖动。比如你训练一个模型，每轮迭代都要分配新的显存，如果每次都走懒惰分配，那实际分配的开销会分散到计算过程中，导致每轮迭代的时间不稳定。

这时候可以设置环境变量CUDA_CACHE_DISABLE=0或者使用cudaSetDeviceFlags(cudaDeviceScheduleBlockingSync)来改变分配行为。但我个人建议：除非你明确知道自己在做什么，否则别轻易改这个。我曾经为了追求压测的稳定性，强行开启了预分配，结果显存占用直接翻倍，模型都加载不进去了。

4.2 PyTorch显存分配机制

PyTorch在CUDA之上做了一层缓存分配器。说白了，它不会每次需要显存都去调用cudaMalloc，而是自己维护一个显存池。你释放的显存，PyTorch先留着，下次再用。

4.2.1 缓存分配器的工作原理

PyTorch的分配器把显存分成一个个块（block）。当你申请显存时，它先从缓存池里找有没有大小合适的空闲块。如果有，直接复用；如果没有，才去CUDA申请新的显存。

这个机制的好处很明显：减少了cudaMalloc的开销。但坏处也很要命：显存碎片化。你想想看，如果你反复分配和释放不同大小的张量，缓存池里就会留下很多大小不一的空洞。最后可能总空闲显存还有2GB，但最大的连续块只有200MB，一个大张量就分配失败了。

4.2.2 显存监控与调试

PyTorch提供了一些工具来查看显存分配情况：

import torch

# 查看当前显存占用
print(torch.cuda.memory_allocated())
# 查看缓存分配器占用的显存
print(torch.cuda.memory_reserved())
# 查看详细的内存快照
print(torch.cuda.memory_summary())

我个人习惯在压测脚本的关键节点打印memory_summary()，这样能清楚地看到哪些张量占用了显存，哪些是缓存碎片。有一次我发现一个张量在循环中反复创建，导致显存泄漏，就是靠这个函数定位的。

4.3 TensorFlow显存分配机制

TensorFlow的显存分配策略和PyTorch完全不同。它默认是霸占型的——启动时就尽可能多地占用显存。我第一次用TensorFlow跑压测时，看到nvidia-smi显示显存占用直接飙到95%，还以为程序出bug了。

4.3.1 显存预占策略

TensorFlow默认使用GPUOptions中的per_process_gpu_memory_fraction来控制显存占用比例。默认值是1.0，也就是占满所有显存。这样做的好处是避免了后续分配的开销，但坏处也很明显：如果你在同一张卡上跑多个进程，它们会互相争抢显存。

4.3.2 动态增长模式

TensorFlow也支持动态增长模式，通过设置allow_growth=True来实现：

import tensorflow as tf

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

这个模式下，TensorFlow不会一开始就占满显存，而是按需分配。但要注意：它只增不减。也就是说，一旦分配了显存，即使不再使用，也不会释放回系统。我在压测长周期任务时发现，显存占用会随着时间推移缓慢增长，最后稳定在一个较高的水平。

4.4 三种机制的对比与选择

好了，咱们把三种机制放在一起对比一下：

我个人在压测时，如果是PyTorch项目，会重点关注缓存分配器的碎片化问题；如果是TensorFlow项目，第一件事就是设置allow_growth或者per_process_gpu_memory_fraction。至于CUDA原生的分配策略，一般只有在排查底层问题时才会直接操作。

这张图把三种分配机制的核心要点串起来了。你压测时遇到显存问题，先对照这张图看看：是CUDA层的分配策略问题？还是框架层的缓存碎片问题？定位到具体环节，才能对症下药。

想要更多资料，联系可可