1. 从“模型不大”到“显存爆掉”，ESMFold 的显存到底去了哪里？
事情起因很简单：我要在一台刚配好的工作站上跑 ESMFold 批量预测，机器用的是一张RTX 5090 32GB。按理说，32GB 显存跑 ESMFold 这种规模的模型，怎么想都绰绰有余。结果调度器报出的预估显存是“约 10GB”，我照着这个数去排任务，把单卡并发拉到了 3 个进程，然后眼睁睁看着任务一个接一个地CUDA OOM死掉。
一开始我以为是显存被打满了，但nvidia-smi一看，3 个进程加起来确实超出了物理显存，可单独跑一个进程时，占用峰值就远不止 10GB。调度器少算了大概8GB，而且这 8GB 恰恰是最容易被人忽略的那部分。
这篇就记录一下我实测出来的 ESMFold 显存占用明细，以及调度器预估显存时最容易漏掉的那几块。如果你也在用 RTX 5090、4090 这类大显存卡跑 ESMFold，或者自己写了个调度脚本按模型参数量去估算显存，这篇应该能帮你少踩一次坑。
先给结论：ESMFold 在 RTX 5090 上，单进程峰值显存大约17GB 到 18GB，其中真正的模型权重只占不到 4GB，激活值、计算图中间量、CUDA context 和 cuDNN 的工作区才是大头。如果你的调度器只按“参数量 × 精度字节数”去估，那必然会少算，因为 ESMFold 是个单序列（single-sequence）推理的模型，它的显存峰值和输入长度、回环数（recycling iterations）高度相关，而不是简单和参数规模成正比。
2. 调度器为什么少算了 8GB——显存估算的三个盲区
2.1 只算权重，不算激活值：这条最容易翻车
很多调度器在预估显存时，默认会拿“模型参数量 × 每个参数的字节数”来当作显存需求。ESMFold 的 ESM-2 主干是 650M 参数，FP32 推理时权重约 2.6GB，加上结构模块（structure module）的几百万参数，总共不到 4GB。于是调度器给出“约 10GB”这个数，已经算上了一些余量。
但 ESMFold 在推理时并不是“过一遍就完事”。它有一个recycling 过程，也就是把上一轮的输出再接回输入端，反复迭代多次，让模型逐步精修结构预测。每一次迭代都会在前向传播过程中产生大量的中间激活值（activations），这些激活值在反传时可能不需要保留（推理模式），但在前向的每一层之间仍然要存在显存里。尤其是 attention 层和 Evoformer 模块，激活值的尺寸和输入 token 数量的平方成正比。
我实测了一个 512 个残基（residues）的蛋白序列，单进程跑完一个 recycling 循环，峰值显存在 17GB 左右。如果序列更长，比如 800 残基，峰值直接逼近 24GB。这就是为什么单单评估权重远远不够——激活值才是 ESMFold 显存的大头。
提示：ESMFold 的 recycling 默认是 4 次，每一次迭代都会重新计算一遍激活值。虽然这些激活值在迭代之间会被释放，但峰值往往出现在一次完整前向传播的中间时刻，这时候显存使用量可远高于模型权重本身。
2.2 CUDA context 和 cuDNN 工作区：看不见的“基础开销”
第二个容易漏掉的是CUDA context。只要你在 PyTorch 里第一次执行 CUDA 运算，驱动就会为当前进程分配一块固定的显存，用于存放 CUDA context、cuBLAS/cuDNN 的句柄、kernel 编译缓存等。这块开销和你的模型大小无关，几乎是固定值。
在 RTX 5090 上，一个空的 PyTorch CUDA 进程，光初始化 CUDA context 就会占用大约 800MB 到 1.2GB 显存。如果你用torch.backends.cuda.matmul.allow_tf32 = True或者启用 cuDNN benchmark，还会额外产生一些自动调优（autotuning）所需的工作区显存。
cuDNN 的工作区（workspace）有时也容易被忽略。ESMFold 的结构模块里大量使用卷积和 LayerNorm，这些算子在某些 cuDNN 算法下会申请临时缓冲区。我观察到这部分在 RTX 5090 上大约要占用 1GB 到 2GB，取决于输入尺寸和算法选择。
这些开销加起来，即使在真正的模型计算开始之前，一个进程就已经默默占用了约 2GB 到 3GB 显存。如果调度器把这些一律当成“零”，那自然会少算。
2.3 PyTorch 显存分配器的“惰性缓存”行为
还有一个很多人都踩过的坑：PyTorch 的显存分配策略，并不是“用完就还”，而是默认使用缓存分配器（caching allocator）。比如你分配了 8GB 的中间张量，算完之后这些显存并不会立刻还给驱动，而是留在进程的缓存池里，下次要分配时可以复用。
这就导致一个现象：你用nvidia-smi看某个进程的显存占用，发现它居高不下，但实际“活跃”的显存可能只有一半。从调度器的角度看，这种情况很难精确预估——因为它没法预测你的代码会在什么时候触发新的分配，以及旧的缓存会不会被复用。
从我的实测来看，ESMFold 在单次推理结束后，进程的显存占用仍会维持在 12GB 左右，而不是回落到基础值。这就是 PyTorch 缓存分配器的“惰性”叠加了 ESMFold 中间张量尺寸变化大带来的后果。如果你的调度系统是“任务结束后统计峰值显存”，那没问题；如果你是“运行前预估”，那这 8GB 缺口就是这么来的。
3. 实测环境与测量方法：RTX 5090 上的真实数字
3.1 测试环境
- GPU：RTX 5090 32GB（驱动版本为较新的稳定版，CUDA 12.x）
- 软件：PyTorch 2.x + ESMFold 官方推理脚本（基于
openfold的权重） - 输入：一个 512 残基的蛋白质序列，FASTA 格式
- 测量工具：
nvidia-smi和 PyTorch 的torch.cuda.max_memory_allocated()
这里特别说一句，为什么用 RTX 5090 而不是之前的 4090：因为 5090 的显存带宽更高，跑 ESMFold 这种带循环迭代的模型时，整体吞吐比 4090 提升明显，但显存行为其实非常相似。也就是说，这段经验对 4090、4080 甚至 A100 都同样适用，只是具体数值略有差异。
3.2 测量方式与结果
我分别用了两种方式统计显存：
nvidia-smi的--query-gpu=memory.used --format=csv实时轮询，记录 VI 进程的峰值物理显存占用。- 在 ESMFold 推理循环外面包裹
torch.cuda.reset_peak_memory_stats()和torch.cuda.max_memory_allocated()，统计 PyTorch 分配器视角下的峰值。
结果如下（512 残基序列，FP32 推理，默认 recycling 次数）：
| 项目 | 显存占用 |
|---|---|
| CUDA context 基础开销 | 约 1.0GB |
| cuDNN/cuBLAS workspace | 约 1.5GB |
| 模型权重（FP32） | 约 3.6GB |
| 激活值与中间张量（峰值） | 约 10GB 到 11GB |
| 合计峰值 | 约 16GB 到 17.5GB |
| 推理结束后常驻显存 | 约 12GB 到 13GB |
所以，那个“约 10GB”的预估数，几乎就是把模型权重和基础开销加到了一起，而完全忽略了激活值的峰值。少算的 8GB，基本就是激活值那部分的低估。
细节：
torch.cuda.max_memory_allocated()统计的是 PyTorch 分配器记录的最大活跃分配量，它不会包含 CUDA context 等底层开销；而nvidia-smi的memory.used是物理显存占用，包含了缓存分配器尚未归还的缓存。两者之差可以作为估算 PyTorch 缓存和 CUDA 后端开销的参考。
3.3 不同序列长度下的显存变化
我又用几条长度不同的序列做了快速测试，结果符合预期：ESMFold 的显存峰值随序列长度近似平方增长（主要来自 attention 部分），但增长的起点并不是零。
| 序列长度（残基） | 峰值显存（实测） | 调度器预估（典型） |
|---|---|---|
| 128 | 约 4GB | 约 3GB |
| 256 | 约 7.5GB | 约 5GB |
| 512 | 约 17GB | 约 10GB |
| 800 | 约 25GB 以上 | 约 12GB |
你没看错，512 到 800 残基，显存从 17GB 飙升到 25GB 以上。如果你的实际任务里有大量长序列，按“平均长度”去预估显存，那翻车概率极高。
4. 实操解决方案：调度器的三种正确估算思路
4.1 方案一：经验公式法（最通用）
既然 ESMFold 的显存主要由 CUDA context、workspace、权重、激活值四部分组成，那我们可以用一个简单的经验公式来预估：
显存预估 ≈ 基础开销 + 权重大小 + 激活值估算其中：
- 基础开销：3GB 到 4GB（取 RTX 5090 实测值，包含 CUDA context、cuDNN workspace、PyTorch 缓存碎片）
- 权重大小：ESMFold 约 3.6GB（FP32）；如果使用半精度，可降至约 1.8GB
- 激活值估算：以 512 残基为“中间值”，约 10GB；可按序列长度的平方关系粗略外推
把这个公式翻译成实际调度规则：
- 序列长度 ≤ 256：单卡并发 4 个进程以内没啥问题
- 序列长度 256 到 512：单卡最多并发 2 个进程，且要保证显存余量
- 序列长度 > 512：建议单卡单进程，甚至考虑用显存更大的卡
我当时把调度器里的“每任务显存预估”从 10GB 改成了 18GB，单卡并发数从 3 降到 2，之后任务队列就再也没出现过 OOM。
提示：经验公式的误差在 ±2GB 左右，所以如果你拿不准，宁可把预估值往上调一档。调度器多估一点显存，损失的只是并发度；估少了，整批任务失败重跑，浪费的时间远大于多等的那几分钟。
4.2 方案二：动态采样法（更精确但复杂）
如果你们有自己的任务调度平台，可以做一个更聪明的逻辑：在每次任务启动前，先拿一条和真实任务长度相近的序列跑一个极短的 ESMFold 前向过程（比如只跑 1 个 recycling iteration），用torch.cuda.max_memory_allocated()直接测量峰值显存。
然后把测量结果加上一个固定余量（比如 1GB 到 2GB），作为该任务的实际显存需求传给调度器。这种方式比较适合“任务队列里序列长度差异很大”的场景——因为如果所有序列都是 256 残基左右，一个固定预估值就够了；但如果从 100 到 900 残基都有，动态采样才能真正避免“长短任务互相踩爆”。
唯一要注意的是，采样用的序列长度要和真实任务尽量一致，不然测出来的峰值没参考价值。同时，采样时要关闭 cuDNN benchmark 以外的其他无关显存开销，保证测量值干净。
4.3 方案三：直接改用半精度和梯度检查点
如果你的 GPU 是 4090 或 5090 这种消费级显卡，其实还有一个无脑操作：在 ESMFold 推理脚本里把模型切成FP16 或 BF16 推理。上面的实测数据是 FP32，如果切到半精度，显存峰值能直接砍半还多，512 残基序列的峰值可以降到 8GB 到 9GB，那调度器按 10GB 预估反而就够了。
不过这里有个坑：ESMFold 某些模块对数值精度比较敏感，直接切半精度可能导致结果微小偏差。我自己在测试中对比过，对大多数蛋白结构预测任务来说，半精度和 FP32 的结构输出差异很小（TM-score 相差不到 0.01），但如果你的下游任务需要非常精确的结构，还是建议保留 FP32。
5. 我在实际操作中踩过的坑和留下的经验
5.1 不要盲信官方脚本的默认并发逻辑
ESMFold 官方推理脚本本身没有内置显存预估功能，很多人的调度器是自己写的，或者基于第三方工具。我用过的几个开源调度工具里，有的确实按照“模型参数量 + 固定余量”来预判，有的甚至还包含一个“模型大小选择器”，但它们对 ESMFold 这类 recycling 模型的适配都不好。
如果你和我一样在写调度器，建议先把 ESMFold 的显存行为当成一个独立的“特殊模型”来处理，而不是套用普通的 CNN 或 Transformer 分类模型逻辑。老的公式只适合那些“前向传播一遍就结束”的模型，不适合有循环迭代、激活值暴涨的结构预测模型。
5.2 memory fragmentation 也是一个隐藏杀手
在跑批量任务时，我注意到一个现象：即使单进程的峰值显存只有 17GB，但当多个进程同时启动、显存分配交错进行时，PyTorch 的缓存分配器会产生更多的碎片，导致实际需要更多显存才能完成同一任务。
这个事很难定量测量，但它确实存在。所以我们在设调度余量时，不能只按“峰值相加等于总显存”来算，最好再加 10% 到 15% 的空闲余量，避免碎片导致的隐性 OOM。
5.3 一个小工具：快速打印当前显存开销
如果你也想自己测 ESMFold 在你的卡上的真实显存，可以在推理脚本里加这么一段：
import torch torch.cuda.reset_peak_memory_stats() # ... 执行完整的 ESMFold 前向预测 ... peak_memory = torch.cuda.max_memory_allocated() reserved_memory = torch.cuda.memory_reserved() print(f"PyTorch 峰值分配: {peak_memory / 1024**3:.2f} GB") print(f"PyTorch 保留缓存: {reserved_memory / 1024**3:.2f} GB")配合nvidia-smi实时查看物理显存占用，就能得到一个“PyTorch 视角 vs 驱动视角”的差异值。这个差异值正好就是调度器预判时容易漏掉的 CUDA context 和缓存部分。
我个人在实际操作中的体会是，像 ESMFold 这种模型，显存已经不再是“模型参数有多小=占用多少显存”的时代了。调度器少算 8GB，表面上是公式问题，本质上是对模型推理机制理解不够。如果你发现自己跑的模型也带 recycling、迭代优化或者自回归生成，记得把激活值和 CUDA 基础开销一起加到你的显存公式里。最后再分享一个小技巧：任务开始前先单独执行一个空的torch.zeros(1).cuda()来预热 CUDA context，可以让后续显存统计更干净，少一些“忽高忽低”的干扰。
