做深度学习编译器或者推理引擎的朋友几乎都会撞上同一堵墙——显存OOM。输入尺寸稍微一调或者网络结构加深几层Caffe、PyTorch、TensorFlow跑起来就报“CUDA out of memory”然后只能灰溜溜地调小batch size。但换个角度想这真的是物理显存不够吗很多时候不是。是计算图里的内存没有被精细调度很多中间张量明明生命周期错开却各自占着一块独立显存有些峰值本来能削平却没有调度器去动态腾挪。我这次要记录的就是一个针对特定计算图、尤其是卷积运算设计的内存分配与溢出Spilling调度器。它解决的问题说白了就三件把张量的生命周期算清楚让该复用的内存复用在显存确实不够时把最不划算的张量Spill到主机内存或者干脆重算以及针对卷积运算这种“中间结果多、内存碎片重灾区”的场景做专门的适配。适合谁看呢写推理引擎、做神经网络编译器、或者被显存优化逼到想自己动手的人这篇应该能让你们少走不少弯路。1. 为什么需要专门的内存分配与溢出调度器1.1 深度学习内存管理的痛点先说一个我经常用来跟新人解释的比喻深度学习计算图有点像一条流水线每个算子在跑的时候都要从边上拿原材料输入张量然后产出新的部件输出张量这些部件可能还要被后面好几道工序继续加工。原始工厂的做法是每个工序都给自己新申请一块存储区用完了随手丢在那儿不管结果仓库越塞越满到后面明明还有好多空闲角落新工序却已经申请不到一块足够大的连续区域了。传统深度学习框架早期的内存管理就是这个状态——每个张量执行完后并不立即释放要等整个图的引用计数归零才回收。更糟糕的是不同迭代之间还存在默认的数据缓存导致你盯着一块巨大的显存占用却说不清到底是哪一层在消耗。即便框架后来引入了内存池也大多是按“需用即分用完放回池子”的思路缺少对全局生命周期的预判。真正的解法是什么是静态调度。在计算图真正执行之前先做一次全局的内存规划把每个张量的首次使用、最后一次使用的位置提取出来画成生命周期区间。两个区间没有重叠的张量就可以共用同一块显存。这本质上就是一个区间着色问题——把不冲突的变量染成同一种颜色每种颜色对应一块唯一的内存地址。这种方案在传统编译器的寄存器分配里已经用了五六十年但现在很多深度学习框架还没认真去做导致显存被白白浪费掉。1.2 卷积运算为何是特殊场景如果说普通网络结构的内存管理已经很粗糙那卷积运算就是粗糙中的粗糙。原因有三。第一卷积会产生大量中间张量。一个卷积算子内部按最朴素的做法要先把输入做im2col图像转列生成一个巨大的数据矩阵然后才能做GEMM。这个im2col矩阵的体积常常是原始特征图的几十到几百倍。另外如果使用Winograd变换还会额外生成变换域的输入、输出、核三个中间体每个都是额外包袱。这些中间张量在一个算子内部就存在更不用说整个网络里几十层卷积串起来。第二卷积的特征图天然带有维度信息NCHW、NHWC、NC1HW1等等。在分配内存时不能只考虑字节数还得考虑维度对齐。比如CUDA的shared memory和某些DSP的Local Memory都有严格的对齐要求你分配了mismatch的起始地址轻则性能下降重则直接报非法访问。而这些对齐约束恰恰会大幅度影响内存池的碎片率——你每隔一段就要padding几个字节。碎片率一旦上来即便总空闲空间够大也分配不出一块连续的大块。第三卷积层常常跟pooling、pad、激活、批归一化连在一起。如果调度器只看到一个个独立算子它就不会知道相邻算子之间其实可以融合融合后某些张量根本不用在全局内存里落地。这就是算子级调度的局限性。所以针对卷积的调度器至少要在做内存规划前先知道哪些算子可以被融合成一个大的“超级算子”再去考虑内部和外部的内存复用。我接触过不少模型从AlexNet到后来各种轻量网络只要把卷积部分的内存调度搞清楚整套网络的内存峰值能砍掉一半以上。这也是为什么我坚持要把卷积作为核心场景来设计调度器而不是泛泛做一套通用内存池。2. 调度器核心设计方案与权衡2.1 从生命周期到干扰图构建调度器拿到手的不是张量列表而是一个完整的计算图每个节点是算子每条边是张量。第一步工作是遍历图给每个张量标记它的“诞生时刻”和“死亡时刻”。我在实现的时候用的是一个二遍扫描法。第一遍按拓扑序正向遍历所有算子记录每个张量第一次被某个算子作为输入读取的时刻诞生时刻可以更早但是在该算子执行前必须已分配以及它被最后一个算子作为输入使用的时刻。第二遍倒序扫描可以更精确地判断“最后一次使用”从图末尾往回走某个算子的输出被谁消费最后一个消费者就是死亡边界。有了每个张量的 [出生, 死亡] 区间就可以构建干扰图。干扰图也是一个图节点是张量如果两个张量的生命周期区间有重叠它们之间就接一条边。注意这里重叠的定义是“只要它们的存活时间区间有交集”哪怕只有一毫秒的计算重叠也不能共用同一块内存。因为该时刻两者都确实是存在的、都可能被访问到。这里有一个特别容易被忽略的坑执行顺序不是严格按照拓扑序的。如果运行时支持多流并行、异步执行那“生命周期区间”就不再是一个数轴上的线段。比如两条并行分支分支A的某个张量明明在拓扑序上先结束但由于GPU上的流A还没执行完CPU侧可能已经去执行流B了此时这块内存究竟能不能被复用不能。因为你无法保证流A已经释放。我做调度器时遇到这个坑后来干脆采用保守策略把张量的生命周期扩大到“整个并行子图执行的起始到结束”。虽然损失一点复用率但能保证安全。2.2 内存分配策略选型图着色与区间规划构建好干扰图之后就该做真正的内存规划了。这里主流有两条路线。第一条是图着色。Kempe链式着色算法被传统编译器用来分配寄存器核心思路很简单把所有节点按度数排序度数大的优先选择颜色选择颜色时挑一个没有与其邻居冲突的最小编号。落在内存池里“颜色”就对应一块内存块。颜色总数就是需要的最大内存块数。图着色胜在通用、好实现而且能处理任意图结构缺点是它关注的是“冲突关系”而不是字节大小可能把一个8KB的张量一个512MB的张量涂成同一个颜色导致内存块大小不匹配浪费严重。第二条是区间背包。这个思路更贴合张量分配的本质把内存空间看成一条从左到右的地址轴所有张量都是需要占据一段长度的矩形我们按照它们生命周期内的时间顺序往里摆放。问题变成对于每个时间步把所有活着的张量当成背包里的物品去优化它们的最小总长度对应最大地址。这可以用贪心或者整数规划求解。区间背包的优点是非常贴近物理内存的连续分配碎片控制比图着色好得多缺点是计算复杂度高网络一大就容易卡住。我实际项目里采用的是混合策略先做一次粗粒度的图着色给张量分组每一组内的张量互不冲突然后对每个组的内容再按字节数和启动结束时间做一次best-fit的区间摆放。这样避开了图着色“只看冲突不看大小”的毛病又不会让区间背包一上来就面对几千个变量而崩溃。2.3 溢出策略换出到主机内存还是重算内存调度再精细也总有扛不住的时候。比如输入分辨率从224翻倍到448中间特征图的size变了四倍显存不够了。这种时候调度器必须做出Spilling决策——也就是选择一部分张量不在GPU显存里待着。业界有几种思路。第一种是交换Offload把张量在不需要时拷贝到CPU主机内存等后面要用时再拷贝回来。这种策略适合那种生命周期特别长、中间有一段很长时间都闲置的张量。举个典型例子主干网络的特征图要在后面一个很晚的层才被使用中间隔了几十个算子那这段时间它完全可以在CPU内存里躺着。代价是PCIe拷贝的带宽和时间。H2D和D2H的速度通常在12GB/s到60GB/s之间如果张量很大拷贝开销会相当可观。第二种是重计算Recompute/Checkpointing把某个中间张量的显存直接释放到要用时才重新跑一遍它依赖的算子。这个策略在反向传播里已经被广泛使用前向时只保留每经过几个节点才存储一次需要用梯度时再从最近的checkpoint重新计算那段子图。代价是计算时间翻倍但好处是几乎不占额外内存。在Pure推理场景如果网络中有一些诸如残差分支里的重复计算可以用到重计算通常比跨卡通信还要划算。第三种是淘汰Evict and discard干脆丢掉数据后面要用时从下游的其他缓存里再生成出来。比如某些网络的中间结果可以从下一层的输出去反推但这实现难度极高我一般不建议。具体怎么选我通常用一个成本函数overhead_cost offload_time * bandwidth / realtime_slack recompute_cost再加一个“金丝雀原则”先Spill那些生命周期长、活跃期短、体积大的张量不Spill生命周期短、活跃期长的小张量因为这种张量刚腾出去马上又要用来回拷贝纯属瞎折腾。3. 针对卷积运算的深度适配3.1 卷积算子的内存特征分析如果只把张量当成一块单纯的内存来做分配那等于白费了“针对卷积”这四个字。卷积算子的内存特征比普通张量要复杂得多。先看输入输出特征图的尺寸它不是一个流线型的线性关系。假设输入是 NCHW [1, 64, 112, 112]卷积核是 3x3stride 1padding 1输出通道128。输出就是 [1, 128, 112, 112]两者都是806KB左右float32。但真正跑底层实现时卷积会被展开成GEMM而GEMM的输入矩阵是 [CKhKw, NHW]也就是 [576, 12544] 的矩阵再乘上各类对齐padding实际存储可能是理论值的两倍。这就导致一个看起来只需要800KB内存的卷积真实运行时可能瞬间需要2MB以上。如果调度器按理论尺寸来做分配执行时极容易发生越界踩内存。第二卷积输出特征的写回方式是分块的。现代GPU kernel为了利用寄存器会按tile来一种种计算输出比如一个block处理4个输出通道上的8x8空间位置。这意味着算子内部除了输入、输出之外还有tile级别的局部buffer。这些局部buffer不需要在全局分配但在设计调度器的内存池时需要考虑卷积算子固定分配的临时缓冲否则kernel内部会因显存不足而失败。第三很多运行时喜欢把卷积权重和偏置做成常量内存constant memory或者专用缓存这类内存和普通数据内存不是同一片区域。调度器如果一视同仁地把权重也放进去共享分配就会破坏const数据的读取效率。所以至少要对张量分类常量、权重、激活、中间缓冲、临时buffer。分类不同分配的位置可以完全不一样。3.2 融合与分块对调度的改变做调度器之前必须先把图优化过。这是我在这个项目里最深的感触内存规划这件事不只是“分配方式”的优化更是“图结构”的优化。卷积算子的邻居常有pad、pooling、激活、BN。如果能先做算子融合把这些操作合并成一个超级算子原本pad生成的零填充张量、BN的scale/bias、激活函数的输出就都可以只存在于超级算子内部的寄存器/shared memory里压根不需要落在全局内存中。这样留给调度器的活就少多了内存峰值自然显著下降。所以调度器的第一步永远不是贪心分配而是问一句“这个算子到底是不是必须存在”真正做到底的时候要注意融合之后的算子本身的输出和输入可能会产生“微生命周期”——也就是中间结果存活时间极其短可能只存在于一个GPU block的整个执行周期内。这类中间张量在微软的宏观调度框架里被归为“瞬时张量”应该直接从全局内存池的候选清单中剔除。另外还有一种有趣的方案叫“分块执行Tiling”。如果整个特征图太大放不下可以把卷积按空间维度切成分块每个分块内部做完一整条完整的计算路径。这种情况下内存复用的模式就变成一个分块一个分块地单独加载、计算、卸载调度器需要按分块来循环分配内存而不是按整张特征图一次性分配。这个方案在Kaldi和某些端上推理引擎里用得很多效果特别好。3.3 与运行时协作的关键点调度器产出的结果不是一个静态脚本而是要真正交给执行引擎去跑。因此调度器跟运行时之间的协作协议就很重要。我设计的方案是“两阶段式API”第一阶段调度器会为每个张量分配一个虚拟内存IDmem_id并给每个算子打上额外的Allocate、Free、Spill_Out、Spill_In这四种标注。第二阶段运行时拿到这个标注序列在真正的执行循环中依次执行。为什么要多此一举搞虚拟ID因为物理地址在编译期可能还没有分配好尤其当还有多线程并发申请时直接给物理指针会让调度器的复用逻辑失效。通过虚拟ID运行时可以自由决定把ID映射到哪一块物理内存也便于调试时做一层抽象层的校验。协作的第二层是关于异步执行问题。调度器生成的Free操作必须插在最后一个使用该张量的算子提交之后而不是它完成之后。如果按完成之后插会白白浪费掉很多可复用内存但如果按提交之后插又得担心算子在GPU上还在读这块内存。这个矛盾的解决需要调度器知道运行时的“完成事件”机制。我的做法是给每个Free操作绑定一个事件Event运行时在事件触发时才真正把内存送回池子。4. 从零实现一个原型调度器4.1 调度器的数据结构和接口纸上谈兵说完直接看怎么落地。我砍掉了项目里很多工程上的边角料保留一个最小可用的原型。先定义几个核心结构。用Python伪代码来描述思路比语言重要。dataclass class TensorInfo: name: str shape: tuple dtype: str producer: int # 产生它的算子ID consumers: list # 使用它的算子ID列表 birth_step: int # 首次被写出的调度步 death_step: int # 最后一次被消费的调度步 size_bytes: int category: str # Weights / Activations / Intermediate / Temp alignment: int 256 # 对齐字节数 dataclass class AllocOp: type: str # alloc | free | spill_out | spill_in tensor_id: int mem_id: int # 复用的内存ID target: str # device | host event_id: int 0整个调度器的对外接口只需要三个方法class MemoryScheduler: def __init__(self, graph, runtime_meta): ... def analyze_lifetime(self): # 遍历图填充 birth/death ... def plan(self, memory_budget_bytes): # 生成AllocOp列表和每个tensor的mem_id ... def emit(self): # 转换成一串有序的调度指令交给runtime执行 ...这里有个经验alignment字段非常关键。很多显存分配的默认对齐是256字节但在某些推理设备上是64字节或者512字节。调度器做地址复用时必须让同一块内存上多个Tensor的起点都满足对齐要求否则kernel会报错。4.2 生命周期提取的完整流程生命周期提取是调度器的地基做得不够精确后面所有复用都会出错。我提供一个实在的算法大家可以照搬第一步是确定拓扑序号。用一个队列从没有依赖的算子开始做拓扑排序按排到的先后给算子编号0, 1, 2...。这一步需要保证计算图没有环有环的话先做环路解除比如把循环展开到固定次数。第二步是正向遍历。每遇到一个算子先查看它的每个输入张量如果该张量已经被某个算子消费就把消费者ID记录进consumers列表。同时对于该算子的每个输出张量记录producer为该算子ID。在这一步我约定张量的birth_step是它的producer算子的执行序号如果该张量本身是个模型入口的输入那birth_step 0。第三步是反向遍历。从图末尾的算子开始往前回推。对于每个算子把它所有输入张量的death_step更新为当前算子的ID注意如果某个张量被多个算子消费需要取最后一个被消费的算子ID所以反向遍历时不断取max即可。这个三步法的时间复杂度是O(VE)几乎可以忽略不计。真正麻烦的是处理控制流比如如果网络里有if-else分支或循环拓扑排序就不能简单做了。我在原型里只支持无控制流的静态图动态形状和动态分支暂不处理但实现上你可以把每个分支都拍平成独立子图再做合并判断。4.3 内存规划与溢出决策的伪代码拿到所有张量的生命周期后我采用一个基于事件堆的扫描算法。事件流里有两类事件allocate张量诞生和free张量死亡。按调度步数依次处理维护一个“当前存活张量”的集合。events [] for t in tensors: events.append((t.birth_step, alloc, t)) events.append((t.death_step, free, t)) events.sort(keylambda e: e[0]) mem_pool MemoryPool(budget_bytes) alloc_result {} for step, kind, t in events: if kind alloc: mem mem_pool.find_best_fit(t.size_bytes, t.alignment) if mem is None: # 显存不足触发Spilling策略 candidate selective_spill(t, alive_tensors, mem_pool) if candidate is None: # 无法腾出空间 - 报告失败 raise OutOfMemoryError() spill_out(candidate) mem mem_pool.find_best_fit(t.size_bytes, t.alignment) alloc_result[t.id] mem_pool.reserve(mem, t.size_bytes, t.alignment) elif kind free: mem_pool.release(t.size_bytes)这里的find_best_fit用的不是传统的最小空闲块而是“最小可容纳且对齐后浪费最少”的块。原因很简单显存池经过多次分配回收后会产生很多碎片如果只找最小块可能导致一个60KB的空闲块将来再也放不下任何大张量。用浪费最少相当于在时间维度上尽量保持大块的连续性。Spilling的selective_spill逻辑我写出来给大家参考def selective_spill(need_tensor, alive_tensors, mem_pool): # 1. 按张量大小降序排序越大越优先Spill # 2. 排除生命周期短且活跃期长的马上要用Spill无意义 # 3. 排除类别是Weight权重类Spill不划算 # 4. 优先挑剩余空闲最少、且可腾出空间不低于需求1.2倍的这个1.2倍余量是我个人经验。因为真正执行时可能出现微小的对齐padding或突发性临时buffer留出20%的安全余量能避免运行到最后差一口凑不上而整体报错。4.4 对接现有执行引擎有了AllocOp列表怎么让它成为可执行的代码最直接的做法是在调度器后增加一个代码生成器把每个算子的执行循环嵌进AllocOp指令中。举一个极小的例子指令流看起来是[alloc] mem_id3 tensorA [exec] opConv2D inputA outputB [alloc] mem_id7 tensorC [spill_out] tensorA targethost # 暂时不再用A [exec] opReLU inputB outputC [free] tensorB [exec] opConv2D inputC outputD [spill_in] tensorA targetdevice # 后面又需要A了 [exec] opConcat inputA,D outputE在实际项目中这套指令流会交给runtime里一个小循环来执行。runtime不需要知道调度器的内部逻辑只要照着指令推就行。调试的时候我习惯在每条alloc/free指令处打印当前显存池的使用峰值用来对比实际分配的峰值和调度器预估的峰值看是否一致。如果实际超过了预估那多半是张量生命周期分析漏掉了某些隐式读取比如算子的内部buffer没有纳入统计。另外如果运行时是C写的实现时建议把张量ID和mem_id分开。张量ID是图的逻辑标识mem_id是调度器给的内存复用标识。这样你在调试时可以直接打log查看同一mem_id在历史上分别对应哪些tensor复盘复用关系是否合理。5. 实际运行中的问题与调试技巧5.1 常见问题速查表一定会有读者跑完以后发现一堆诡异问题我把最常见的问题直接列成表格照着排查省时间现象可能原因排查与解决程序启动后直接非法内存访问张量生命周期重叠判断失误两个同时存活的张量被复用了同一块内存调大对齐打印干扰图人工检查重叠区间最大的几个张量显存峰值比调度器预估的高很多算子内部临时buffer没计入统计比如im2col腾出来的中间大矩阵为每个算子补充“内部峰值内存”字段在调度时作为额外约束偶尔出现数据错乱但复现困难异步流执行导致free时机过早前面的算子还没读完就回收了内存free指令后绑定事件事件触发后再真正释放做了spill之后性能反而下降选错了Spill对象频繁在设备与主存之间拷贝小张量增加一个“最小spill体积”阈值小于该阈值的张量宁可留在显存里程序报“无法对齐”错误分配块的起始地址不满足设备要求全局统一对齐字节数并在find_best_fit时先做对齐padding再判断大小某些张量尺寸动态变化动态shape下生命周期会变导致池内分配计划失效预留一个连续动态区域或者退化成按需分配并禁用复用5.2 三个值得铭记的踩坑记录第一个坑是关于im2col的隐性内存。我在第一次做调度器时自以为已经把张量的生命周期分析得很干净跑一个简单的卷积网络预估显存峰值只有实际峰值的60%。一查才发现底层卷积库在调用前要额外申请一块巨大的工作区workspace这个工作区可能比所有张量加起来还大。调度器完全不知道它的存在因此也就无法复用它。后来我把每个算子的“workspace_size”字段合并到一个special_reservoir里并单独规划出一块专用区域才把预估和实际对齐。第二个坑是分块执行与重叠。给特征图做Tiling之后调度器需要重新考虑每个分块的pixel操作对中间结果的依赖。有一次我为了优化内存把特征图按高度切成8个strip分别调度结果两个strip之间因为重叠的edge mask而需要重复读取同一块输入缓冲区。如果调度器不知道这层关系就会错误地认为第二个strip可以重用第一个strip输入的内存导致读取被覆盖。从此以后我明确了一个原则分块后的每个子图都必须保留原始图的“边界张量”生命周期同步信息绝不轻率合并。第三个坑关于host pinned memory。做Spill到主存时调度器如果直接调用cudaMemcpy会发现拷贝速度特别慢。原因是没有使用pinned memory页锁定内存。后来我在调度器里增加了一个HostBufferPool预先分配一批pinned的内存块拷贝速度提升到普通内存的2-3倍。而且更妙的是这个池子的大小也可以参与内存规划——如果溢出到主存的数据量小只需要一个很小的pool就能扛住。5.3 如何验证调度器是否真正有效验证点其实就三条。第一条是正确性第二条是峰值第三条是可复现性。正确性验证用最笨的办法把调度器开起来跑一遍标签正确的样例比对输出张量的数值和不开调度器时的数值误差必须小于浮点运算本身允许的范围。建议不要跑一遍完事要同时跑多个不同尺寸的输入特别是极端尺寸比如输入宽高不能被16整除的。峰值验证用显存监控工具在调度器注释的指引下把每个步骤的当前占用量打出来。比较调度器预期的峰值和实际的峰值误差应该控制在5%以内。如果误差太大先查生命周期再查算子内部buffer。可复现这一点最容易被人忽略。调度器的决策逻辑如果你用了随机化比如某些启发式算法那每次生成的AllocOp可能不同但仍然都应该满足内存约束。我的建议是在发版前固定一个随机种子跑一个专门的回归测试集确保同一模型产物每次都生成相同的指令流否则后续做内核融合、算子拼接都会变得不可预测。最后再分享一个我自己的小习惯调度器生成的结果我会单独导出一个JSON文件里面包含每个张量的存活区间、mem_id、spill记录。这样一旦线上模型出现了诡异的偶发错误我不用重新翻代码直接查JSON就能看到哪几个张量在哪个时间段共享同一块物理地址。因为很多时候问题根本不在调度算法而在下游某个算子没有严格遵守我们约好的内存生命周期约定。这个脚本从设计到落地我前后迭代了差不多一个多月最大的体会是调度器本身不难难的是它必须跟算子实现、内核库、运行时执行模式形成一套完整的约束契约。只要有一个环节不守规矩调度器再怎么精确都会白搭。希望这篇踩坑笔记能帮同路人少交一点学费。