别迷信“模型文件多大显存就占多少”，实测下来ESMFold这条线比你想的更“吃”显存。我的调度器当时只给任务留了16GB显存上限，结果模型一加载，nvidia-smi里直接冲到24GB，整整多了8GB。最后是用RTX 5090实测压出来的问题，发现这8GB不是模型权重，也不是序列太长，而是PyTorch缓存分配器、cuDNN workspace、CUDA context这些“隐形账单”在作怪。如果你也要在RTX 5090上跑ESMFold或类似Transformer推理，这篇文章的实测数据和处理方法应该能让你少踩一次OOM的坑。

1. 问题现场：调度器自信地说“够用”，两分钟后CUDA out of memory

先说背景。我在用ESMFold批量跑一批蛋白质序列的结构预测，任务丢到集群上，由调度器按每个作业声明的显存大小来分配GPU上的显存配额。因为我用的ESMFold权重是BF16格式，文件大约16GB，我自己粗算了一下：权重16GB、激活值给2GB、留点余量，总共申请20GB也合理。但考虑到同卡还要跑其他服务，我把配额压到了16GB。

结果作业一启动，模型加载到一半，直接报“CUDA out of memory”。我当时第一反应是“模型文件都16GB了，我给它16GB显存是不是刚好卡线？”，但仔细一看，模型文件16GB，参数加载进显存其实差不多也就16GB，理论上还没开始推理，怎么就没内存了？后来单独跑了一次抓显存曲线，才发现峰值稳稳地落在24GB，比申请的16GB高出一大截。

1.1 一次深夜跑批翻车记录

那次翻车的过程很典型。深夜提交了十几个任务的批量预测，第一个任务跑了不到两分钟就退出，日志最后一行是：

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 128.00 MiB (GPU 0; 31.86 GiB total capacity; 16.45 GiB already allocated; 14.95 GiB free; 16.95 GiB reserved in total by PyTorch)

注意看这里有个细节：GPU总容量31.86GiB，PyTorch reserved显存16.95GiB，free还有14.95GiB，但分配128MB却失败了。这说明调度器把显存配额锁在了16GB左右，PyTorch已经reserved了16.95GB，再往上分配就撞墙了。问题不在于显卡不够大，而在于“PyTorch预留的显存”已经超过了调度器给的上限。

我当时顺手查了同卡上的其他进程，确认没有别人占用。那么问题的核心就变成了：一个BF16权重16GB的模型，为什么PyTorch要reserved将近17GB？而且还没算推理阶段的激活值。这显然说明我对显存估算缺了一块。

1.2 为什么RTX 5090也会翻车

很多人觉得RTX 5090有32GB显存，跑模型绰绰有余。实际上，32GB只是物理显存总量，调度器如果通过某种方式给了你一个更小的“可见上限”，你的进程就只能在上限内玩。更关键的是，PyTorch的显存分配不是“用多少拿多少”，而是按块向驱动申请，申请到的显存会留在自己的缓存池里，就算暂时没用到也不还给系统。这个缓存池会把“实际占用”放大。

还有一个容易被忽略的因素：新显卡的驱动版本、cuDNN版本对workspace的默认策略越来越激进。在加载模型或者第一次跑attention的时候，cuDNN/cuBLAS可能会为了加速卷积和矩阵运算申请一块很大的workspace，这个workspace同样算在进程显存里。RTX 5090因为算力版本较新，PyTorch与之配套的CUDA后端默认行为也会有些差异，显存分配的“起步价”比旧卡更猛。

所以，真正的问题不是“5090能不能跑”，而是“调度器给作业的配额和PyTorch实际的内存行为没有对齐”。如果你也用类似方式管理显存，大概率也会遇到同样的8GB差额。

2. ESMFold的显存需求到底由什么决定

要搞清楚那8GB多在哪，得先把ESMFold的显存构成掰开看。ESMFold和传统AlphaFold不太一样，它不需要通过数据库搜索生成MSA（多序列比对），而是直接把单条氨基酸序列丢给一个大的蛋白质语言模型提取表征，再用结构模块折叠成三维结构。这个流程中，显存的消耗主要集中在两个部分：语言模型主体和结构模块。

2.1 先给不熟的朋友科普一下ESMFold在做什么

ESMFold的推理流程可以粗略分为两步：第一步，把氨基酸序列变成token序列，过一遍类似BERT的Transformer encoder，输出每个残基对应的embedding；第二步，把这些embedding交给structure module，通过迭代精化预测残基间的距离和坐标。

有点像一个翻译模型：第一步把“氨基酸语言”翻译成“深度特征”，第二步把“深度特征”翻译成“三维坐标”。因为少了MSA这个重头，ESMFold比AlphaFold快很多，但代价是它的语言模型部分非常庞大，参数数量动辄几亿到几十亿，这直接决定了模型权重的显存下限。

不过权重只占一部分。真正麻烦的是Transformer encoder里的自注意力层，它对序列长度是二次复杂度。序列长度从256涨到512，激活内存不是翻倍，而是可能变成三到四倍。这就是为什么同一个模型，跑短序列和跑长序列，显存需求天差地别。

2.2 显存账单的四个主要项目

我在排查时把显存账单分成了四类，建议你也照着这个思路去估算：

• 模型权重：BF16精度下，1B参数约等于2GB显存。我的模型权重文件16GB，说明参数规模大约8B，这部分是固定的。
• 激活值和中间变量：模型前向传播时每一层都会保留中间结果供反向传播使用。推理模式虽然不需要反向传播，但实际框架仍然会保存一部分算子需要的中间张量。这部分随序列长度和batch size变化。
• CUDA context和库workspace：进程第一次调用GPU时，CUDA runtime要建立context，cuDNN/cuBLAS也可能分配workspace。这部分和模型大小无关，但绝对值并不小，新驱动下往往能达到500MB到1GB，复杂场景下几GB也是常事。
• PyTorch显存缓存：PyTorch使用caching allocator管理显存，申请方式是对着驱动要一大块，然后切成小块给张量用。当你释放一个张量时，显存不会立刻还给系统，而是留在缓存池里。这就导致你在任务管理器里看到的进程占用，往往大于torch实际分配的张量总和。

上述第三和第四项，就是我的调度器少算的那8GB的主要来源。

2.3 序列长度和batch size怎么推高显存

为了给你一个直观感受，我摘一组实测数字。同一个模型，BF16权重固定16GB，只做推理，batch size=1，不同序列长度下的进程显存峰值如下：

可以看到，序列从384涨到512（只增加三分之一），显存却从24.1涨到28.6，多了4.5GB。原因是注意力矩阵的大小和序列长度平方成正比，虽然这部分经过算子融合有所缓解，但架子搭在那，省不了多少。

batch size更是线性放大的变量。如果把384残基的序列从batch=1提到batch=4，激活值直接翻几倍，峰值轻松突破32GB。所以在RTX 5090上，单个大模型做长序列推理，显存基本是贴着天花板走，调度器给配额时必须留够空间。

3. 实测记录：RTX 5090上的完整显存画像

光说不练没意义。我专门腾出一张RTX 5090，没有限制显存配额，把模型完整跑了一遍，记录了加载阶段、预热阶段、推理峰值阶段的变化。测试尽量贴近日常使用：PyTorch 2.x，BF16混合精度，ESMFold的distilled变体权重，单张RTX 5090 32GB。

3.1 测试环境与怎么测才准

测试环境如下：

• GPU：RTX 5090，32GB GDDR7
• 驱动与CUDA：驱动为最近一个稳定版，CUDA 12.x
• PyTorch：2.x，使用CUDA后端
• 精度：模型参数和计算都用BF16，部分结构模块算子回落到FP32
• 模型：基于ESMFold架构的8B参数蒸馏版本，BF16权重约16GB

这里有个关键经验：测显存不要只看“任务管理器”或“nvidia-smi”的瞬时数值，而是要在代码里显式追踪。nvidia-smi每隔一秒钟刷新一次，很可能错过瞬间峰值。正确做法是在推理前后调用PyTorch的显存统计接口。

我用来抓峰值的模板代码很简单：

import torch torch.cuda.reset_peak_memory_stats() torch.cuda.empty_cache() # 模型加载与推理代码 output = model(sequence) # 你的前向推理 # 推理结束后 peak_allocated = torch.cuda.max_memory_allocated() peak_reserved = torch.cuda.max_memory_reserved() print(f"Allocated peak: {peak_allocated / 1024**3:.2f} GB") print(f"Reserved peak: {peak_reserved / 1024**3:.2f} GB")

注意区分两个概念：max_memory_allocated是PyTorch实际发给张量的显存峰值；max_memory_reserved是PyTorch向CUDA驱动申请的显存峰值，后者一定大于等于前者，而nvidia-smi看到的进程显存更接近后者。我们要做调度器配额决策时，应该按reserved峰值来算，因为这块内存消耗是真实发生在GPU上的。

3.2 不同长度序列下的显存实测数据

我用上面的方法跑了一组序列长度从128到512的推理，记录下max_memory_reserved数值，同时用nvidia-smi -l 1在后台抓进程占用。数据如下：

三个数值之间的关系很清楚：allocated是理论上的“有用占用”，reserved是PyTorch实际跟驱动要的，nvidia-smi看到的进程占用又比reserved稍高一点，因为还有CUDA context这些底层开销。

我提交给调度器的配额是16GB，但即使最短的128残基序列，reserved峰值也到了18.7GB，比申请多了2.7GB；而我要跑的批量序列平均长度在384上下，reserved峰值24.2GB，比申请的16GB整整多了8.2GB，基本对应标题里的8GB。换句话说，不是调度器写错了，是我申请时根本没把reserved和context算进去。

3.3 那8GB到底藏在哪

看到数据后，我做了个简单的减法：模型权重16GB，384残基序列的allocated峰值20.6GB，reserved峰值24.2GB。allocated比权重大了4.6GB，这部分是激活值和中间张量；reserved比allocated又大了3.6GB，这部分是PyTorch缓存和CUDA context。加上nvidia-smi里额外的一点开销，总共额外8GB就出来了。

还有一个隐藏大户是cuDNN workspace。在模型第一次执行attention里的矩阵乘时，cuDNN/cuBLAS自动选择算法，会保留一块workspace。如果环境变量CUBLAS_WORKSPACE_CONFIG没有限制，默认可能申请一份较大的空间。再加上几个过时的缓存块，PyTorch的reserved就被推高了。

在实测中，我试着设置了CUBLAS_WORKSPACE_CONFIG=:4096:8，发现workspace变小，但性能稍微下降了一点。对于RTX 5090这种算力很足的卡，性能损失可接受，显存余量却多出好几GB。这算是8GB里可压缩的关键空间。

4. 调度器为什么会少算：常见估算公式的漏洞

调度器本身没有错，错的是提交任务的人——也就是我。我低估了显存需求，但低估的原因很典型：只用“模型文件大小+一点缓存”来估算，没有考虑PyTorch真实的内存策略。这一节我把估算漏洞和修正方法写清楚。

4.1 标准估算为什么不靠谱

很多人（包括之前的我）会这样估算推理显存：

所需显存 = 模型权重大小 + 序列激活估计 + 10%余量

这个公式在推理小模型时勉强能用，遇到ESMFold这种大模型就不准了。原因有三点：

第一，模型权重大小不完全等于加载后的显存。PyTorch在加载权重时会经历临时张量、转换精度、重组state_dict的过程，这个过程本身可能产生一份或多份临时拷贝，峰值会短暂超过最终权重大小。模型越大，这个瞬时峰值越高。

第二，激活值很难拍脑袋估。Transformer的注意力矩阵是平方复杂度，不同序列长度下激活值差别极大。比如384和512，看起来长度只差33%，但激活显存增量能到5GB以上。用固定“2GB”或者“5GB”去估所有场景，必然翻车。

第三，最容易被忽略的就是PyTorch缓存分配器和CUDA workspace。这两项是进程级的，不仔细做实验根本不知道多大。尤其在新驱动和cuDNN版本下，workspace的默认申请可能很激进，直接吃掉数个GB。

4.2 用三个API把显存账算清楚

要告别瞎猜，可以直接在代码里把账算清楚。我最后用下面三个API组合统计显存：

allocated = torch.cuda.memory_allocated() # 当前给张量用的显存 reserved = torch.cuda.memory_reserved() # 当前缓存池占用的显存 max_reserved = torch.cuda.max_memory_reserved()# 运行至今的reserved峰值

在模型加载前后、推理不同batch时分别打点，把数据记录下来，就能得到一条显存曲线。重点看max_memory_reserved，它是整个进程生命周期里“最紧张时刻”的申请量。这一步做完，调度器配额直接按这个数去申请，就不会再因为显存被打爆。

我建议你做一个一次性profiling实验：把真实要跑的最长序列、最大batch，在一个不受限制的GPU上跑一遍，用reset_peak_memory_stats抓max_memory_reserved，然后把这个数乘以1.05到1.1作为作业申请的显存。这比任何经验公式都可靠。

4.3 给调度器写预留建议公式

结合实测，我给自己的调度器写了一版新的显存申请公式：

作业申请显存 = max_memory_reserved(最长序列, 最大batch) * 1.1

其中max_memory_reserved必须通过实测获得，而不是猜测。如果没有实测条件，可以用更保守的公式：

估算值 = 模型权重(GB) + 激活保留(序列长度/256 * 8GB) + 8GB固定开销

这里的8GB固定开销就是我从RTX 5090上实测得到的经验值，覆盖CUDA context、workspace、PyTorch缓存池以及权重复制瞬时峰值。对8B参数左右的模型，这个值实际相当接近。如果你的卡是显存更大的A100，固定开销可能少一些；但考虑到新驱动和新cuDNN的习惯，宁多勿少。

注意：8GB不是精确常数，它是针对我这个模型和环境的实测值。你可以把它当作初始参考，真正上线前一定要用自己模型跑一遍。

5. 排查实录与修正后的配置

知道问题在哪里之后，剩下的就是改配置和验证。这一节记录我的完整排查步骤，以及最终采用的调度器配置和环境变量。

5.1 从OOM到精准定位的排查过程

我第一次看到OOM时，先做的是把模型文件大小和显存加起来对照，但怎么都对不上。于是我把调度器限制去掉，直接在前台运行同一个任务，同时每隔0.5秒用nvidia-smi记录显存。

结果发现，在模型加载阶段，显存会先猛涨到约17GB，然后稍微回落到16GB；进入第一次前向时，又突然跳到20GB以上，最后在第一次attention计算结束后达到峰值24GB。这个“两段式跳变”让我意识到，加载阶段的临时拷贝和工具库workspace各占一部分。

随后用torch.cuda.max_memory_reserved重跑，确认了reserved峰值24.2GB，而allocated峰值只有20.6GB，差距3.6GB。我立即把这两个数给调度器的同事看，才终于定位到问题：不是代码跑错了，是作业申请显存时的“附加开销”没算进去。

5.2 我最后调整的调度器配置

在调度器层面，我把这个任务的显存申请从16GB调整到24GB，同时为了让同卡的其他进程不受影响，我还在任务脚本里加了PyTorch进程内存上限：

# 调度器脚本片段 #SBATCH --gpus=1 #SBATCH --mem-per-gpu=24G export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128 export CUBLAS_WORKSPACE_CONFIG=:4096:8

max_split_size_mb:128是用来控制PyTorch缓存池里大块内存的拆分策略的。把拆分块调小一点，可以降低显存碎片，提高缓存利用率。CUBLAS_WORKSPACE_CONFIG限制了cuBLAS workspace的大小，效果立竿见影，但注意设置后第一次矩阵运算可能会变慢，换来的显存余量在RTX 5090上很划算。

如果你用的是容器方式限制显存，比如通过nvidia-container-toolkit给容器设置--memory=24G，也要注意这个内存上限是针对进程的，和PyTorch的显存概念不是同一回事。建议直接在作业里声明GPU显存配额，并让调度器开启显存隔离，否则容器内PyTorch看到的仍是整张卡的显存，照样超卖。

5.3 一组可以复制的环境变量

最后给出我测下来对显存控制比较有用的环境变量组合，你可以直接加在脚本里：

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:128 export CUBLAS_WORKSPACE_CONFIG=:4096:8 export NVIDIA_TF32_OVERRIDE=0 # 如果需要更精确的数值，可以关闭TF32

expandable_segments:True是PyTorch 2.3以后支持的新选项，它能让缓存段按需扩展，减少碎片化，对显存紧张的场景尤其有用。但注意，这个选项在部分模型和旧驱动上可能导致初始化变慢，如果遇到问题可以先去掉。

实测中，加了这两个环境变量后，相同的384残基序列推理，reserved峰值从24.2GB降到了22.7GB，省出的1.5GB虽然不算巨量，但在32GB卡上，足以让你把batch从1提到2。更关键的是，运行过程不再出现瞬时暴涨，nvidia-smi看起来平稳很多。

6. 常见问题速查表

为了让你排查时不用再翻文档，我把这次踩坑遇到的问题整理成一张速查表，方便对照。

还有一个容易被忽略的老问题：如果你在同一进程中反复加载和卸载模型，不清理缓存，reserved峰值会越来越高。最常见的处理办法是在加载新模型之前调用一次torch.cuda.empty_cache()，把之前的缓存池释放一部分。但注意这个函数只是提示CUDA归还当前未占用的缓存块，并不一定会全部归还，别再把它当成万能药。

7. 一点实测心得

这次跑ESMFold的经历让我彻底改掉了一个习惯：以前我看模型文件多大就申请多少显存，现在凡是上模型，第一件事就是在空跑环境里抓一次max_memory_reserved，再把结果写到调度器的配置备注里。对于ESMFold这种大模型，权重大只是一个下限，PyTorch缓存和算子workspace带来的隐性开销才是调度器少算的那8GB的真身。

如果你也要在RTX 5090或者类似显存的卡上跑ESMFold，我的建议很明确：先别急着上任务，花十分钟跑一遍不同序列长度下的峰值显存，把数据和模型一起固化下来。调度器本身不聪明，但它会忠实执行你指定的数字。你给的数字够准，后面的所有任务都会顺畅很多。最后再分享一个小技巧，把显存峰值测试写成一个简单脚本，每次换模型或换驱动后都跑一遍，你会发现所谓“玄学OOM”其实都是可以提前发现的数学问题。