显存又炸了。这个场景我见过太多次模型参数明明不大一个batch喂进去占用曲线却一直在涨最后直接OOM。很多人第一反应是上大显存的卡、砍batch、换混合精度但很少有人意识到真正决定显存峰值高低的是框架内部那套计算图内存调度系统。它管着每一块显存什么时候分配、什么时候释放、能不能复用几乎所有的“不合理高占用”都跟它有关。这篇文章想把计算图内存调度系统拆开讲透。它是什么、靠什么工作、常见实现里有那些坑、实际调优时又该怎么利用它。适合正在做推理引擎优化、训练框架二次开发、或者被模型显存问题反复折磨的部署工程师。我会把原理和工程经验混在一起写这样看完全文至少你能对着profiling工具自己把问题定位到具体是哪几个张量把自己的显存顶爆了。1. 显存瓶颈的真实面貌内存调度系统解决什么问题1.1 一次显存OOM背后藏着什么先说一个我调过的典型案例。某个推理服务上线前压测模型本身也就几百MB的权重batch从8提到16直接OOM。当时几乎所有同事都说“模型快饱和了建议换大卡”代价太高。我没急着下结论先导了一份显存分配记录结果发现真正的问题模型在某个多分支汇合阶段产生了一批中间结果这批张量生命周期高度重叠同时存活的中间张量加起来超过了3GB而模型权重才不到1GB。这种情况非常典型。计算图的峰值显存不是“最大单张量的大小”而是“所有同时存活的张量大小之和”。一批张量像集装箱一样叠在同一个时间窗口里占用自然就冲上去了。如果框架没有做图级内存复用那么每个中间结果都会独立占用一块显存峰值就是所有生命周期有交叉的张量累加——这是最坏、最原始的内存分配方式。1.2 计算图内存调度到底做什么要理解调度系统先看计算图是什么。深度学习模型在框架内部是一个有向无环图每个节点是一个算子比如矩阵乘、激活、归一化每条边是算子之间流动的张量。内存调度系统要回答三个问题每个张量什么时候分配内存、它被映射到哪一块物理显存上、什么时候可以回收并让给下一个张量复用。把计算图想成一个仓库调度场景。一批货张量有明确的进库时间和出库时间调度员的工作是提前知道所有货的进出计划然后把“不会同时出现的货”放到同一个货架上这样仓库总面积才能压到最小。计算图内存调度的本质就是这个“货架复用”问题只是目标函数从仓库面积换成了显存峰值。1.3 为什么不直接交给系统分配器可能有人会问直接调系统的分配器或者底层驱动接口逐块分配不行吗当然行但代价极大。系统分配器只能看到单次请求看不到全局生命周期所以无法跨张量复用。而且逐块申请会带来两方面的浪费一是频繁小粒度分配导致头部开销和地址碎片化二是一旦张量释放地址是否归还系统也不受模型控制。工程上成熟的做法是“先分析、后申请”。调度器在计算图执行前把全部张量的生命周期解析出来生成一个显存复用计划执行时先一次性向系统申请一大块显存池然后根据计划在池内做手工切分。这样既绕开了系统分配器的碎片问题又把复用率拉高了。下面就从“分析”这步讲起。2. 生命周期分析让每块显存都有明确的“存活区间”2.1 张量生命周期从哪来内存调度的核心输入是每个张量的生命周期区间。在一个静态计算图里这个区间是可以精确计算的一个张量由某个算子产生被一个或多个下游算子消费当最后一个消费它的算子执行完毕后该张量就可以被释放。因此一个张量的生命周期就是从生产者算子结束开始到最后一个消费者算子结束为止的时间区间。伪代码级的计算逻辑大概是这样的for op in topological_order(op_graph): for output_tensor in op.outputs: alloc[output_tensor.id] allocator.acquire(size, alignment) for input_tensor in op.inputs: ref_count[input_tensor.id] - 1 if ref_count[input_tensor.id] 0: allocator.release(input_tensor.id)这里依赖一个前提计算图必须按拓扑序执行。拓扑序保证每个算子的输入都已就绪输出在算子结束后可用。沿拓扑序扫描一遍为每个张量累计引用计数计数归零的时刻就是该张量生命周期终止的时刻。2.2 离线引用计数与运行时计数的分工生命周期分析有两条路线离线静态分析和运行时动态追踪。静态图模式下比如推理引擎里常见的图优化阶段张量生命周期在拿到完整计算图后就能全部算出来稳定性极高调度器可以提前排布好显存规划执行时基本就是“按计划走”。动态图模式下图的执行路径是运行时才知道的没法提前做全局规划。这种场景通常退而求其次在每次算子执行时更新引用计数计数归零就把内存块归还给内存池。虽然做不到全局最优但至少能达成局部复用。实际工程中两者经常混合使用静态规划做粗粒度分配运行时计数做细粒度回收。2.3 生命周期标注如何决定复用窗口生命周期标注的价值在于当两个张量的存活区间不重叠它们就可以共享同一块物理显存。这就把内存调度问题转化成了一个经典的区间重叠问题。假设模型里有四个中间张量 A、B、C、D生命周期区间如下表张量存活区间大小A算子1 ~ 算子3256MBB算子2 ~ 算子4128MBC算子4 ~ 算子6512MBD算子5 ~ 算子7128MBA 和 C 的生命周期不重叠可以复用同一块显存B 和 C 重叠必须单独分配D 和 C 也有重叠但和 A、B 的关系不同。画一条时间轴就能发现同时存活的最大张量集合决定了理论峰值。调度器的目标就是通过复用策略把实际峰值尽量逼近这个理论峰值。3. 分配策略与复用算法峰值最小化的关键取舍3.1 显存池分桶与大小对齐有了生命周期区间之后接下来要解决“如何复用”的问题。直接的做法是每来一个张量请求就在可复用列表中找一个大小合适的内存块。但频繁比较大小成本很高所以工程上普遍引入内存池分桶机制。系统启动时分配若干大块内存池池内按固定大小分桶比如 256B、1KB、4KB、32KB、256KB、1MB、4MB……张量请求会向上对齐到最近的桶规格。好处是查表很快、管理简单坏处是会产生内部碎片——一个 1.05MB 的张量可能占掉一个 2MB 的桶白白浪费。实际调优时桶规格设计要根据模型的张量大小分布来配而不是随便拍一组数。3.2 四种常见分配策略的对比在内存池内做复用选择时不同策略对显存占用和执行性能的影响差异很大。业界常用的策略大致如下策略匹配逻辑优点缺点最先适配找到第一个足够大的空闲块查找快可能把大块切到不可用最佳适配找最小可容纳的空闲块省空间查找慢容易出小孔洞图级贪心复用按生命周期排序后复用复用率高依赖静态分析固定桶分配大小对齐后映射到桶最快最稳内部碎片偏大我的经验是训练场景优先保证速度推理场景优先压峰值。训练时生命周期分析不完全可靠动态图常见最佳适配反而会拖慢速度推理引擎都是静态图图级贪心复用能把显存吃得非常干净值得多花一点分析时间。3.3 区间复用的近似策略与理论边界如果所有张量大小一样区间复用问题就等于区间图着色按起始时间排序后贪心即可得到最优解。但现实里张量大小不等问题变成了装箱和着色的混合体复杂度显著上升业界普遍采用近似策略。最常见的近似思路是先按张量大小降序排序大张量优先分配然后对每个大张量在空闲块中挑生命周期不冲突且大小最接近的块。这块要注意一个坑贪心分配后如果发生冲突不要直接增大池子而是先尝试挪动生命周期端点的算子执行顺序。算子重排序在少数场景可行比如两个不相关的分支节点之间的调度顺序是灵活的换一下可能就把峰值降下来。这也是图优化器喜欢做“节点重排”的原因之一。4. 工程实现中的隐形雷区就地操作、内存对齐与碎片4.1 in-place操作如何打乱生命周期假设理论上的生命周期分析很完美工程实现时第一个雷就是就地操作in-place。比如某个激活函数会覆盖输入张量如果输入还被后面的算子引用那么逻辑上输入张量的生命周期还没结束但底层物理内存已经被写坏了。调度器如果不感知这种别名关系轻则计算结果错误重则随机崩溃。处理这个问题的常规方案是算子注册时就声明“本算子是就地操作”调度器在给该算子分配输出内存时必须检查输入张量的引用计数只有当引用计数为1且确认没有其他消费者时才允许输出复用输入的地址。否则就要强制为输出单独分配一块新内存。很多框架早期版本为了省显存盲目开启就地模式结果出现各种诡异错误。这里我特别提醒一句分析内存问题时如果发现“偶尔算错、难以复现”先怀疑就地操作的别名管理。4.2 内存对齐与张量连续性问题第二个雷是内存对齐约束。设备侧的内存访问有强对齐要求地址通常需要按16字节、64字节甚至128字节对齐不满足对齐的地址在某些算子上是根本无法执行的。调度器分配内存时必须把请求大小向上对齐比如 100MB 的张量可能要占 100MBpadding一个张量padding几个字节不可怕成百上千个中间张量累积起来也不小。比对齐更隐蔽的是张量连续性问题。计算图里一个算子的输出可能被切片、转置等操作引用产生非连续视图。调度器如果单纯按视图逻辑大小分配内存等算子真正执行时才发现数据不连续就只能额外做一次拷贝。表面上看省了显存实际上白白消耗了带宽。所以好的调度器会分析算子的内存布局偏好必要时“为保证连续性多分配一点显存避免触发额外拷贝”。4.3 碎片整理与高水位控制内存池用久了碎片问题会重新浮现。虽然调度器按生命周期复用已经比系统分配器好很多但由于张量大小差异大小块长期保留、大块被切分池内可用地址被割裂成很多碎片。此时总剩余空间够却没有一块连续空间能满足某个大张量请求。针对碎片工程上有两类做法。一是按张量用途区隔池子把权重、优化器状态这类长期驻留张量放进专用池临时中间张量放进临时池不让长期张量“污染”临时池。二是定期重启内存池训练场景下可以每隔若干轮将池子压缩重整重新分配一次。推理场景下如果模型支持多batch动态请求建议在请求低谷期主动做池重整能有效压住高水位。这里没有银弹碎片控制靠的是对自身模型结构的理解。5. 从Profiling到端到端调优一个真实优化复盘5.1 先看生命周期峰值表接手一个显存占用异常的模型时我的第一件事就是导出一份生命周期峰值表。做法不复杂先打开框架的profiling接口记录每个张量的分配时间点、释放时间点和大小然后按执行时间轴叠加成一条占用曲线。峰值曲线的最高点位置基本就是参与优化最重要的一组张量。拿到峰值表后只看三个东西峰值位置发生在哪个算子附近、峰值那一时刻同时存活的最大几个张量是哪些、这些张量各自生命周期长短。多数情况下问题会集中在某个“汇集点”比如多头注意力里的拼接、残差结构里的叠加、或者多分支网络的特征合并。这些位置的中间张量不仅大而且大量同时存活。5.2 三板斧内存复用、激活重计算、数据交换定位到问题张量后优化手段基本是三类。第一板斧内存复用。调整计算图节点的执行顺序让生命周期冲突的大张量错峰分布或者改造算子实现使输出可以复用某些不再使用的输入。这是侵入性最小、收益最直接的优化也是我优先用的方法。第二板斧激活重计算。把生命周期特别长的大张量按需在反向或推理的后期重新算一遍而不是一直驻留显存。代价是增加计算量适合“显存比算力更贵”的场景。在训练时梯度回传和推理时多层缓存场景都很实用。判断标准很简单如果这个张量的计算代价远小于它占掉的显存带宽成本就可以重计算。第三板斧数据交换。把暂时用不到的中间结果搬去其他存储介质等需要时再搬回来。这条能兜底但同步开销非常大尤其当模型频繁在设备端和其他介质之间读写时很容易把延迟拉高。我一般只在“退出内存复用后仍然超限”时才考虑它。5.3 一次端到端的优化过程以一个实际测试过的推理模型为例。模型是某Transformer结构batch8初始峰值显存5.2GB。跑完生命周期分析后我看到两个问题。一是某两路分支的中间结果都在汇合点附近存活重叠区间达到上百毫秒拖高了峰值二是模型中大量小块张量散落在临时池里导致大块请求需要二次扩容。我的操作分三步。第一步将其中一个分支的算子与汇合算子做指令重排让该分支的中间结果提前释放实测峰值降到4.4GB。第二步把激活函数全部改为基于引用计数安全才走就地路径减少了一份大张量的新增占用峰值降到4.0GB。第三步调整临时池分桶粒度把1MB到4MB区间的桶细分为更多档位碎片率下降池总容量不用再预留那么多最终峰值稳定在3.6GB分配次数下降了约三成。整个过程没有改模型结构收益全部来自计算图内存调度的细节优化。这类优化最考验的不是算法多高深而是能不能把张量生命周期看得足够清楚。最后分享一点个人体会。我做了很多显存优化之后发现最高效的路径永远是先弄清楚“谁和谁同时活着”再谈分配策略和复用算法。很多框架的默认内存调度已经不错但面对具体模型总有不合理之处。如果你正在排查一个显存占用反常的模型不要急着改模型或者换卡先导一次生命周期峰值表把时间窗口对齐到具体算子问题往往会自己现出原形。另外检查分析输出时优先看生命周期最长的那几个张量它们通常才是显存压力的真正源头。