3、显存分配机制：CUDA显存分配策略、PyTorch缓存分配器、TensorFlow显存管理

显存分配，说白了就是GPU怎么给你“划地盘”。

很多同学跑模型时报OOM（显存溢出），第一反应是“显存不够大”。其实不一定。很多时候是分配策略没用好，导致显存碎片化严重，或者缓存没释放。我调过不少案例，把分配策略优化一下，同样的模型能省下20%-30%的显存。

3.1 CUDA显存分配策略

CUDA底层提供了两套分配接口：cudaMalloc 和 cudaMallocManaged。前者是显式分配，后者是统一内存（UM）。

cudaMalloc：直接向GPU要一块连续显存。优点是速度快、延迟低。缺点是必须手动管理生命周期，容易产生碎片。

cudaMallocManaged：让CPU和GPU共享同一块虚拟地址空间。数据会自动在CPU和GPU之间迁移。听起来很美好，对吧？但我个人建议慎用。为什么呢？因为自动迁移有性能开销。我在项目中遇到过，用了UM之后，训练速度直接掉了30%。

CUDA还提供了显存池（Memory Pool）机制。从CUDA 11.2开始，你可以用cudaMemPool来管理显存。它的核心思想是：预先分配一大块显存，然后按需切分。这样可以减少cudaMalloc的调用次数，降低碎片。

// 创建显存池
cudaMemPool_t pool;
cudaMemPoolCreate(&pool, 0);

// 从池中分配
void* ptr;
cudaMallocAsync(&ptr, size, pool);

// 释放回池
cudaFreeAsync(ptr, stream);

嗯，这里要注意：显存池不是万能的。如果你的模型动态变化很大（比如NLP里的变长序列），池子里的碎片还是会累积。我一般建议每训练1000步，手动重置一下池子。

3.2 PyTorch缓存分配器

PyTorch的显存管理，核心是缓存分配器（Caching Allocator）。它不直接调用cudaMalloc，而是自己维护一个显存缓存池。

为什么这么做？因为cudaMalloc是内核态调用，很慢。PyTorch的做法是：先向GPU要一大块显存（比如2GB），然后切成小块给张量用。张量释放后，显存不还给OS，而是留在缓存池里，给下一个张量复用。

你可以用torch.cuda.memory_summary()查看当前缓存状态。我个人习惯在训练脚本里每100步打印一次，看看缓存利用率。

import torch

# 查看显存分配详情
print(torch.cuda.memory_summary())

# 清空缓存（不释放给OS，但内部整理）
torch.cuda.empty_cache()

# 重置峰值统计
torch.cuda.reset_peak_memory_stats()

举个例子，我调过一个BERT训练任务，显存占用8.5GB，但实际模型参数只有3GB。剩下的全是缓存碎片。加上expandable_segments:True之后，直接降到5.5GB。

# 启动时设置
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128,expandable_segments:True
python train.py

你想想看，这省下来的3GB，是不是可以再塞一个batch？

3.3 TensorFlow显存管理

TensorFlow的显存管理，和PyTorch思路不太一样。它默认是贪心分配：启动时就把所有显存都占上。你跑一个单卡模型，它直接吃掉全部显存。很多新手第一次看到都吓一跳。

其实这是故意的。TF认为“占满”可以减少后续分配的开销。但如果你要同时跑多个模型，或者和PyTorch共享GPU，那就得改配置。

import tensorflow as tf

# 方式一：按需增长
gpus = tf.config.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

# 方式二：限制最大显存
tf.config.set_logical_device_configuration(
    gpus[0],
    [tf.config.LogicalDeviceConfiguration(memory_limit=4096)]  # 限制4GB
)

我个人更推荐按需增长模式。它不会一开始就占满，而是随着计算逐步分配。但要注意：这个模式有性能损耗，因为每次分配都要和驱动交互。如果你的模型显存需求很稳定，不如直接设一个固定上限。

TF 2.x还引入了PluggableDevice机制，允许自定义显存分配器。不过说实话，我很少用。大部分场景下，set_memory_growth就够用了。

3.4 三种分配策略对比

说白了，没有银弹。PyTorch适合研究阶段，灵活但碎片多。TF适合生产环境，稳定但不够灵活。CUDA原生适合你写自定义算子的时候用。

嗯，最后说一句：不管用哪个框架，监控显存都是基本功。我习惯用nvidia-smi配合gpustat实时看显存变化。遇到OOM，先看碎片率，再看缓存占用，最后才考虑是不是真的显存不够。