训练或推理AI视频生成模型时，显存报错并不总是因为真实显存不足。很多情况下，nvidia-smi显示的占用率只有七八成，但程序却抛出了CUDA out of memory。如果仔细查看PyTorch的报错信息，有时会出现“CUDA out of memory. Tried to allocate X MiB”后面跟着一段缓存分配器的统计数据，其中free与reserved的差值特别显眼。这个差值通常就是碎片浪费掉的显存。PyTorch的CUDA缓存分配器为了减少cudaMalloc和cudaFree的系统调用开销，会自己维护一个显存池：释放的张量显存并不会立即归还给CUDA驱动，而是留在池子里等待复用。但当分配请求的大小与池中空闲块大小不匹配时，空闲块无法被使用，只能向CUDA驱动申请新的显存，碎片便由此产生。
对于AI视频生成这类负载，碎片问题尤其突出。视频模型经常需要处理不同分辨率、不同时间步的中间张量，这些张量大小差异很大，并且会在前向与反向传播中频繁创建和销毁。例如一个扩散模型在采样初期使用较大尺寸的潜变量，随后逐步裁剪或插值，每一次尺寸变动都可能导致分配器内部出现大小不一的空闲块。如果缓存分配器的策略不当，这些空闲块既无法合并，又找不到足够大的连续空间满足下一个分配请求，于是CUDA驱动被迫分配新的显存块，最终把物理显存耗尽，而池子里却还躺着大量细碎的空闲块。要解决这个问题，就需要从PyTorch内存池的管理策略下手，而PYTORCH_CUDA_ALLOC_CONF环境变量就是那把钥匙。
PyTorch CUDA缓存分配器的工作机制
PyTorch的CUDA缓存分配器并不是直接把每个张量的显存请求转给cudaMalloc。它会先在自己的缓存池中查找是否有大小合适的空闲块，如果找到就直接复用；如果找不到，则向CUDA驱动申请一块更大的显存，然后切出一部分给张量使用，剩余部分继续留在缓存池中。这个池子由多个不同大小的块组成，块的来源可能是之前释放的张量显存，也可能是新申请的大块切分后的剩余部分。分配器会维护一个统计表，记录当前已保留的显存总量（reserved）和实际分配给张量的显存量（allocated），两者的差值就是池中空闲但尚未归还驱动的显存。
这种设计在单次分配大小相对固定的场景下非常高效，比如图像分类模型的卷积层，张量形状基本不变。但在AI视频生成中，张量形状变化剧烈，分配器频繁地遇到“找不到匹配的空闲块”的情况。比如池中有一个50MiB的空闲块，但现在需要分配一个80MiB的张量，这个50MiB的块没法用，分配器只能再去申请新的显存。如果下一次又需要分配一个30MiB的张量，那个50MiB的块虽然可以用，但剩余20MiB又会成为新的碎片。随着迭代推进，大量无法合并的小块堆积，导致可用连续显存越来越少。
值得注意的一点是，PyTorch缓存分配器默认会在一个已经保留的显存段内进行块管理。这个段的大小通常由初始分配决定，后续的分配如果超出该段容量，需要申请新的段。段与段之间是独立管理的，跨段的空闲块无法合并。这种“段”的机制进一步限制了碎片整理能力。因此，要让分配器更好地适应视频生成负载，就需要调整段的行为和块拆分的粒度。这正是PYTORCH_CUDA_ALLOC_CONF所控制的内容。
PYTORCH_CUDA_ALLOC_CONF关键配置项详解
PYTORCH_CUDA_ALLOC_CONF是一个以逗号分隔的键值对字符串，用于在PyTorch初始化CUDA缓存分配器时传入自定义配置。目前支持的主要选项有backend、max_split_size_mb、roundup_power_two_divisions、garbage_collection_threshold以及expandable_segments。其中对AI视频生成优化最直接的是expandable_segments和max_split_size_mb两个参数。
expandable_segments可以取True或False，默认在较新版本的PyTorch中通常为False。当设置为True时，分配器不再使用固定大小的段，而是让段可以根据需要动态扩展。这意味着分配器可以把新申请的内存直接追加到已有的段上，从而减少段的数量，并且允许相邻的空闲块在同一个段内进行合并。对于反复创建和销毁不同大小张量的视频任务，启用可扩展段可以显著降低碎片率。设置方法如下：
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
需要说明的是，expandable_segments并不是银弹。它改变了段的管理模型，在某些PyTorch版本或特定的算子实现中，可能引起额外的内存映射开销。但就AI视频生成常见的推理与训练负载而言，将其设置为True通常能带来明显的显存利用率提升。
另一个常用参数是max_split_size_mb，它限制分配器在拆分大块时的最小保留块大小。默认情况下，分配器会尽量把大块拆成刚好能满足请求的大小，剩余部分作为新的空闲块保留。如果每次拆出来的剩余块都很小，这些小碎片很快就会失去复用价值。通过设置max_split_size_mb为一个合适的值，例如64或128，分配器会避免产生小于该尺寸的碎片块，转而将整个大块分配给请求，超出的部分由张量自行管理（实际上不会发生，这里只是说明分配策略的调整）。更准确地说，max_split_size_mb控制的是当一个空闲块被拆分成两个部分时，较小那一部分的最小尺寸。如果拆分后的小块小于该阈值，分配器就不会进行拆分，而是直接使用更大的块，从而减少碎片数量。示例：
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
在实际调优时，可以同时设置多个选项，用逗号分隔，注意逗号前后不要加空格，因为环境变量的解析对空格比较敏感。例如：
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:128
此外，backend选项可以切换底层的分配器实现。PyTorch原本支持native和cudaMallocAsync两种后端，cudaMallocAsync利用CUDA的流序内存分配器，在部分场景下可以进一步降低碎片。不过该后端对CUDA版本和驱动有要求，并且与某些自定义算子可能存在兼容性问题。对于AI视频生成任务，如果expandable_segments已经生效，通常不需要切换到cudaMallocAsync。建议先观察native后端加可扩展段的效果，再决定是否尝试其他后端。
针对AI视频生成场景的调优实践与验证
视频生成模型在推理时最常见的显存压力来自两个方面：一是扩散模型的多步采样，每一步都涉及编码器、解码器和UNet的多次前向；二是时间维度的注意力机制，需要存储大量的中间激活。这两类操作恰好都会产生大量不同形状的临时张量。因此，在调整PYTORCH_CUDA_ALLOC_CONF之前，需要先确认当前的显存碎片到底有多严重。执行以下代码可以查看分配器的统计信息：
import torch
# 模拟视频生成中的张量分配与释放
for i in range(100):
if i % 3 == 0:
a = torch.randn(1, 4, 64, 80, 80, device="cuda")
elif i % 3 == 1:
a = torch.randn(1, 8, 32, 40, 40, device="cuda")
else:
a = torch.randn(1, 16, 16, 20, 20, device="cuda")
del a
print(torch.cuda.memory_summary())
运行这段代码后，注意看输出中的“reserved memory”和“allocated memory”两行，以及下面的“Segments”和“Blocks”统计。如果reserved远大于allocated，且空闲块数量很多、平均块大小很小，就说明碎片严重。此时可以尝试设置环境变量并重新运行，对比前后的数据。
一个更贴近真实视频生成的测试是加载一个预训练的视频扩散模型，例如基于Stable Video Diffusion的变体，在固定batch size下进行一次采样，记录峰值显存和采样耗时。先在默认配置下跑通并记录基线，然后设置PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True再跑一次。很多情况下，启用可扩展段后，峰值显存会下降10%到30%，采样耗时几乎不变甚至略有下降。如果同时设置了max_split_size_mb，还需要注意该值不能设得过小，否则反而增加分配器查找空块的次数。建议从128开始尝试，如果显存仍不够，可以逐步减小到64或32，观察碎片率和分配器查询开销之间的平衡。
需要特别提醒的是，该环境变量必须在PyTorch首次初始化CUDA上下文之前设置，否则不会生效。如果使用的是Python脚本，最好的方式是直接在shell中export，而不是在Python代码里通过os.environ动态修改，因为后者的时机很难保证。对于使用多进程数据加载或分布式训练的场景，每个进程都需要独立继承该环境变量，否则子进程可能使用默认策略，导致整体显存行为不一致。
还有一种情况值得关注：如果你的模型使用了torch.compile或某些第三方CUDA扩展，expandable_segments可能会与底层内存管理产生交互。如果遇到奇怪的显存分配失败或性能回退，可以先关闭expandable_segments，只保留max_split_size_mb进行测试，逐步定位问题来源。总之，PYTORCH_CUDA_ALLOC_CONF提供了一个非常灵活且低成本的调节入口，理解它的作用机制，结合AI视频生成负载的特征反复实验，就能有效缓解显存碎片带来的假性OOM困扰。
PyTorch内存池显存碎片PYTORCH_CUDA_ALLOC_CONF修改时间：2026-10-02 20:59:10