☰
卷积计算图专用内存与Spill调度器设计详解
2026/10/18 19:39:02 网站建设 项目流程

一个针对特定计算图(尤其是卷积运算)的内存分配与溢出(Spilling)调度器,听起来像是框架内核人才会整天琢磨的东西。但如果你曾经把一个卷积网络的batch size一降再降,最后仍收到设备内存不足的报错,就会意识到这东西比想象中实用得多。显存是硬约束,而卷积计算图的中间张量又格外能占地方,所以与其靠“手动缩小batch”来妥协,不如在分配器层面做文章:把暂时不用的中间结果搬到宿主端,等用到时再搬回来,让计算图在更低的显存预算下继续跑。这篇文章我会把这类调度器的设计思路、底层原理、卷积特化处理,以及实际落地时容易踩的坑完整写出来,供做推理框架、边缘部署或模型裁剪的读者参考。

1. 为什么卷积计算图需要专门的内存与Spill调度

1.1 显存短缺下常见的两条弯路

很多人在显存不够时的第一反应是缩小batch size。这个办法简单,但代价很大:batch一旦变小,卷积的矩阵乘规模也跟着变小,设备计算单元很容易喂不饱,吞吐量断崖式下跌。另一条更常见的弯路是“干脆把所有张量都常驻显存,靠一套大内存池复用”。这套方案在模型小的时候没问题,可一旦网络深一点、特征图分辨率高一点,峰值驻留量会轻松突破预算。尤其是在降采样层、上采样层交错出现的卷积网络里,多分支相加会让某些中间结果的生命周期拉得很长,内存复用窗口被严重压缩,峰值显存几乎等于所有大尺寸特征图的和。

所以真正需要的是一个懂得“时间维”的分配器:它知道每个张量什么时候创建、什么时候最后一次被读取,并且允许在峰值压力较大的时间点,把某些低优先级张量主动转移出去。这本质上就是操作系统虚拟内存里的Spill思想,把它放到计算图调度层来做,就可以让模型在硬件物理显存小于计算图峰值需求的情况下正常运行。

1.2 计算图内存生命周期的本质

理解这个问题,先要抓住一个核心概念:liveness interval,也就是活跃区间。每个张量从被某个算子写出来的那一刻开始存活,直到被所有下游算子读完之后才死亡。在拓扑序确定的计算图中,每个张量都可以对应到一段时间区间,区间长度取决于它在网络中的位置。

举例来说,一个卷积层的输出会被后面的ReLU读取,ReLU的输出又会被池化层读取,而池化输出还会被另一个卷积读取。那么第一个卷积的输出张量,其活跃区间要从该卷积执行结束算起,一直到池化层完成才结束。在这整段时间里,无论它是否参与计算,它都必须占据显存。给一张包含大量分支和跳跃连接的卷积计算图做生命周期分析后你会发现,不同张量的活跃区间重叠程度非常高,这直接决定了内存的峰值。

很多传统分配器只做“空间复用”,即比较两个张量的尺寸是否一样,一样就共用一块内存。但这种方式忽略了空间复用能成立的前提是活跃区间不重叠。如果没有对区间做排序和检查,尺寸相同但生命周期重叠的张量会被错误地分配到同一块地址,轻则数据覆盖,重则产生随机错误结果。一个合格的调度器必须先把整张计算图的活跃区间算清楚,再做区间层面的复用判断,否则后面的Spill策略完全是空中楼阁。

1.3 卷积算子的三副“胃口”与其他算子完全不同

卷积算子对内存的消耗远不是“一个输入、一个输出”那么简单。我用实际算过的一些网络结构来分析,它至少有三块额外开销跟内存分配强相关。

第一块是常规的输入输出特征图。这个大家都理解,NCHW布局下,一张特征图动辄好几MB到几十MB,网络一旦有几十层,这部分天然是大头。

第二块是算法相关的workspace。卷积的实现路径非常多,常见有直接卷积、隐式GEMM、FFT、Winograd等。不同路径对临时工作空间的需求差别巨大。有的算法需要额外的缓冲区来存放变换后的输入矩阵,有的需要额外的空间做分块累加。一个大规模卷积的workspace可能达到数十MB甚至几百MB。这部分内存不是张量,但分配器必须同样考虑进去。

第三块是表示转换的中间结果。比如把输入按卷积窗口展开成im2col矩阵,再用GEMM计算时,im2col结果本身就是一个巨大的中间张量,它的大小是输入特征图的kernel_size的平方倍。如果网络里有好几个紧挨着的卷积层,这类展开结果一叠加,显存压力立刻翻倍。

正因为卷积相关临时内存的种类多、尺寸跨度大,一个“一刀切”的通用内存分配器很难处理好。通用分配器通常只处理固定尺寸的张量池,对workspace、im2col这类不规则临时buffer的适配很差,要么提前分配过大空间导致浪费,要么频繁申请释放导致碎片化。所以标题里特别强调“特定计算图”是有道理的:只有针对卷积计算图的特征做特化,调度器才能真正产生收益。

2. 内存分配的底层设计:活跃区间、内存池与碎片治理

2.1 如何把张量生命周期“剪”成可供调度的区间

在调度器内部,张量生命周期通常用step来刻画。每个step对应计算图中的一个算子执行位置,从0到N编号。一个张量的创建step是产生它的那个算子的执行位置,死亡step是最后一个读取它的算子的执行位置。于是这个张量就变成了一个半开区间[start, end]。所有张量的区间合集,构成了一张时间轴上的区间集合。

有了区间表,才能做无冲突分配。最简单可靠的方法是事件扫描:把所有区间的开始和结束事件按时间排序,然后使用一个维护空闲内存块的分配器模拟整个过程。每当某个区间开始,就从空闲块里切一块;每当区间结束,就把对应块释放回空闲池。这样得到的分配方案天然满足“重叠区间不共用同一块内存”的约束。

这里有一个容易被忽略的细节:区间分析必须考虑同一个算子内部产生的多个临时结果。比如一个卷积算子既产生了中间累加buffer,又产生了输出张量,如果它们的创建step都标成同一个数字,有的实现会误以为它们生命周期完全相同,从而拒绝复用,实际却可以更精细地错开。更合理的做法是给算子内部也划分sub-step,把一次算子执行拆成“输入读取、内部计算、输出写入”等多个更细的时点,这样才能把复用窗口挖到极致。

2.2 图着色分配与连续内存池的取舍

明确了区间以后,分配策略通常有两种流派。一种是基于区间图着色的方法:把张量看成图的节点,区间重叠的两个张量之间连一条边,然后对图着色,相同颜色的节点可以共享同一块内存。图着色法的优点是复用逻辑清晰,最小化峰值直观;缺点是图着色本身是NP难的,大图上只能靠贪心近似,实现起来复杂度偏高。

另一种是连续池分配:维护一整块地址连续的大缓冲池,按最佳适配或首次适配策略给新区间切块。这种方案工程上更好落地,也更容易支持动态shape,因为不需要预先算出全局着色方案,只要在模拟过程中记录空闲偏移即可。我在实际项目中采用的是后者,只是额外加了一条规则:同尺寸区间优先复用最近释放的块。这条规则能显著提高缓存命中率,减少内存地址跳变带来的TLB开销。

工程实现上,内存池并不需要真的“malloc一块整的大内存”,可以直接基于硬件的显存分配接口预留一块大区域,然后内部用空闲链表管理。分配器内部还会维护一个“按大小分桶”的加速索引,方便快速找到满足size且地址对齐的空闲块。只要不涉及API层面的真实申请释放,这套方案的开销就非常低,可以在图编译阶段完成全部模拟。

2.3 卷积特征图的对齐要求与碎片治理

卷积算子的内存对齐要求比普通张量严格。通道维度的对齐尤其常见:很多底层卷积库要求输入通道数按4、8或16对齐,某些特定指令集还要求基地址按64字节或更高对齐。如果分配器忽略了这些约束,直接在原地址上偏移分配,后续算子可能根本没法直接使用这块内存,只能额外拷贝一份对齐后的数据,既费时间又费空间。

我在设计分配器时,特意把所有区间的大小都先做“向上取整到对齐粒度”,然后再进行池分配。这样虽然会浪费一点点尾部空间,但换来的是整个执行过程里零额外拷贝。对卷积网络来说,这个是稳赚不赔的。另一个务实做法是大小分桶:把卷积层常见的特征图尺寸分别设置固定大小的专用槽位,其他不规则的中间张量走通用池。专用槽位的复用路径简单、地址稳定,碎片率大幅降低。

碎片治理还有一个小技巧:如果发现某些卷积层反复出现“上一轮释放的block刚好比下一轮需要的block小几十KB”的情况,可以在池分配模拟时增加“相邻块合并”机制,并在分配前优先检查左右邻居是否空闲。卷积网络里这种尺寸微差非常常见,合并机制能有效减少碎片造成的无效峰值。

3. Spill调度器的核心决策:换哪个张量、何时换、何时取回

3.1 候选集筛选:不是所有张量都适合搬出去

Spill的本质是用一次搬移数据的开销,换取一段显存空间的使用权。如果这个张量刚搬出去马上又要用,那一次搬移不仅没换来多少“喘息期”,还白白增加了两次拷贝延迟。因此必须先筛出值得Spill的候选集合。

按我的实践,候选张量至少要满足以下条件:生命周期跨度足够长,一般要求活跃区间覆盖超过若干个算子;占用的空间足够大,否则Spill节省的内存和它引入的复杂度不成正比;生命周期内没有太多下游算子紧挨着读取,最好可以在下一次使用之前有足够时间完成异步搬回;同时该张量没有被外部持续引用,比如网络的输出层结果就不能随便搬走。筛选完成后,剩下的张量进入真正的Spill规划阶段。

候选条件可以写成一张快速检查表:

判断项推荐基准反例
生命周期长度覆盖大于3个算子以上仅被下一个算子立即读取
占用大小大于总内存预算的10%几十KB级别的小张量
后续使用位置距离当前较远,可异步预取下一个算子立刻消费
外部引用状态无别名/视图引用被多个分支引用且引用结束的时间不一致
重计算代价重新生成代价很高可以几百次运算轻松重建的中间值

这张表用于启发式评分时,可以避免把调度器变成一个“什么都往外搬、结果卡顿更严重”的憨憨工具。

3.2 代价模型:不能只算带宽,还得算启动延迟和重叠

当一个张量被选中Spill时,需要估算它带来的时间开销。很多人的第一反应是看带宽:数据量除以带宽,得出搬移耗时。但在真实设备上,一次搬移的开销还包括启动延迟、同步等待、内存分配器加锁等固定成本。所以单次搬移时间更接近:

搬运耗时 = 数据量 / 设备-主机有效带宽 + 固定启动开销

固定启动开销在小数据量时格外明显。比如一次几十MB的块搬移,固定成本可能只占百分之几;但一次几百KB的搬移,固定开销可能占一半以上。这就是为什么候选筛选时我会把“小张量”直接排除,因为Spill它们的启动开销会成倍放大总延迟。

比单次耗时更重要的是搬移能不能被隐藏。如果reload操作能在上一个算子还在计算时异步发起,那么搬移时间不一定出现在关键路径上。调度器设计里会把reload点尽量提前:在目标算子开始前若干个step就发起异步取回,等到真正需要值时,数据已经在设备端了。只有当取回时间超过从发起点到使用点之间的计算总时长时,这部分额外等待才会被计入整体延迟。

3.3 两层策略:静态规划和运行期动态决策

Spill调度最稳妥的做法是“静态规划为主,动态决策为辅”。静态阶段,调度器在拿到完整计算图后,按拓扑序模拟一遍内存分配。每遇到显存超出预算的时刻,就从候选池里挑一个策略评分最高的张量标记为Spill,并把这个张量的下次使用位置标记为reload点。这样得到的是一次确定性计划,运行期只要严格照做就行。

动态决策主要用于两种情况:一是图里有数据相关的动态shape,导致张量大小在运行期才能确定;二是存在条件分支,个别路径可能根本不消费某个张量。对于这两种情况,我不会在静态阶段就把Spill标记写死,而是在运行期保留一个决策缓存:同一shape组合只计算一次,后续直接复用上次的Spill选择。这个折中让实际执行的调度开销保持在极低水平,又保留了对动态输入的适应能力。

这里有一个心得:千万不要在运行期每走一个算子就重新算一遍完整Spill规划。那会让调度器本身成为性能瓶颈,而且不同决策之间还可能抖动,导致同一个算子前后的内存布局不稳定。静态计划 + 缓存化动态修正,是工程上最划算的组合。

4. 卷积计算图专属优化:Workspace、im2col与重计算的协同

4.1 Workspace的“临时借出”与优先级保护

卷积算法的workspace有一个显著特点:它不是张量,却在分配器看来是必须解决的内存需求。很多卷积库在图上做算法选择时,会试探多个卷积实现算法,每个都返回一个建议workspace大小,最终选择时通常要保证该大小在内存预算内。

我的处理方法是给workspace单独预留一个“全局最大workspace块”。因为不同卷积层不会同时执行,它们的workspace可以复用同一块物理内存。这块内存在所有卷积算子之间共享,优先级高于一般候选Spill张量。换句话说,当显存吃紧时,宁可先把某个中间特征图Spill出去,也不能压缩workspace。原因是workspace不足时卷积只能退回更慢的算法,性能损失远大于一次中间张量搬移。

但workspace本身也可能成为峰值压力来源。如果某个卷积算法的workspace需求特别大,占用甚至超过了整个预算的一半,调度器就应该抛弃这个算法,强制选择较小workspace的替代实现,哪怕它的计算速度慢一点。这本质上是一个“时间换空间”的权衡,需要把该层的计算耗时也纳入Scoring。我通常会把workspace也作为Spill候选,但优先级被调到最低,因为搬走它往往意味着算法锁定失败,代价远高于普通张量。

4.2 im2col展开矩阵:大而短命的特殊中间张量

在卷积转GEMM的实现路径里,im2col展开矩阵是典型的“大而短命”张量。它的创建瞬间非常大,但在GEMM计算完成后立刻失去读取需求,活跃区间很短。这种特性让它看起来很适合Spill——反正很快就不用了,搬出去节省的空间非常可观。但问题是它马上就会用到,如果搬出去再搬回来,两次耗时几乎全部落在关键路径上,反而得不偿失。

更合理的处理是“分块生成、分段计算”:把im2col矩阵按行或按输出通道切成几个block,每生成一个block就立刻做对应那部分GEMM,生成下一个block时释放前一个。这样完整im2col矩阵从未真正驻留过设备端,峰值自然被压下来。相比Spill,这是一种更彻底的形态优化。调度器在识别到卷积层存在im2col展开选项时,可以优先尝试分块展开,再决定是否将展开结果纳入全局Spill规划。

当然分块实现会带来额外复杂度:GEMM的切分会影响矩阵乘的tile大小,如果不小心可能让计算效率降低。实践中我会做一个简单判断:如果展开矩阵总大小超过当前空闲内存,优先分块展开;否则维持原逻辑。这个启发式在多数卷积网络上表现稳定。

4.3 与重计算(Activation Checkpointing)的分工与协作

Spill是把数据搬到更慢的存储介质,而重计算是直接把数据丢掉、用时再生成。两者方案目标相同,但代价模型完全不同。在一个完整的调度器里,应该同时支持这两种策略,并针对每个候选张量选择更优者。

简单说,可以按三条规则快速分流:

  • 如果张量重新生成的耗时明显小于搬移耗时,且生成它的算子依赖链不复杂,那么直接选择重计算,完全不占用显存。
  • 如果张量的重生成代价较高,但重新生成的输入数据在设备端仍然存在,那么Spill更划算。
  • 如果重生成既麻烦、体积又大,那就只能Spill,而且这时候要特别谨慎,因为这部分搬移开销基本无法被隐藏。

重计算在卷积网络里特别适合处理ReLU、池化、简单拼接这类轻量算子产生的中间结果。它们重新生成的计算量小,但数量多、尺寸大,如果每个都Spill,总线带宽会被打爆。调度器做全局规划时,我会先跑一遍重计算候选识别,把轻量算子产生的中间张量直接排除在Spill候选外,再针对剩余张量做代价对比。这样两者协同后,整个计算图的显存曲线会平滑很多。

5. 工程落地时最容易踩的三个坑

5.1 别名与视图引用会让Spill直接破坏数据

第一个坑来自张量视图。卷积网络里的slice、reshape、transpose操作会产生共享底层缓冲区的新视图。如果调度器只记录原始张量的活跃区间,而忽略了其视图的存活时间,就可能出现“原始张量已经Spill,但某个视图还在被使用”的严重错误。

解决办法是在构建内存图时维护每个底层缓冲区的引用计数,任何视图操作都会增加引用计数。只有当引用计数降为零时,该缓冲区才允许进入Spill候选池。这个规则听起来简单,但实现时很容易漏掉两个细节:一是视图创建的时点可能晚于基础张量的创建,从视图创建到基础张量死亡这段时间才是真正的活跃区间;二是某些算子会返回“输出即输入视图”的语义,比如padding后的slice,如果分配器不知道这个关系,极低概率触发的错位问题会非常难排查。所以调度器内部必须维护一个独立的“缓冲区依赖图”,所有spill决策之前都要先查这张图。

5.2 Spill搬移与计算的重叠度没有打满

第二个坑是搬移开销虽然被“隐藏”了,但实际延迟仍然不理想。理想状态下reload应该在计算节点开始前就完成,可实际执行时经常出现“发起太晚”的问题。原因是有些调度器把reload事件简单地排在当前计算流后面,相当于串行等拷贝结束才开始计算。

我的做法是使用双缓冲队列。当前一个算子还在执行时,调度器就把它后续需要的所有Spill张量的reload请求全部提交到独立的异步拷贝队列中。只有当消费这些张量的算子真正开始时,调度器才检查队列里的完成事件是否满足条件。如果所有reload请求都在算子启动前完成,那么搬移时间完美隐藏;如果没有,那就需要调度器在静态规划阶段把reload点往前调整。这一块是收益最明显也最容易做好的优化,我在历次落地项目里通常能减少三成以上的整体额外延迟。

5.3 带宽模型校准被当作“最后一公里”

第三个坑是把带宽模型当成一个拍脑袋估算值。理论带宽和实际带宽之间差距非常大,尤其是宿主端内存与设备端之间的传输,数据量大了以后实际受限于PCIe链路性能、主机内存分配策略、DMA通道冲突等多重因素,远非一个标称带宽可以描述。

我建议在调度器内部内置一个微型profiler:启动编译阶段时,对一组覆盖不同大小的数据块(从256KB到256MB)做真实的搬移计时,然后用线性分段模型拟合“数据量-耗时”曲线。这个模型要认真处理启动延迟和斜率两个参数,否则Spill代价评分会系统性偏差。实测下来,经过校准的模型在真实推理场景的预测误差可以控制在10%以内,而未经校准的模型误差动辄超过50%,足以改变Spill决策结果。

6. 用伪代码看一遍完整的调度循环

6.1 运行期调度器的骨架逻辑

说了这么多理论,落到代码层面其实可以浓缩成一个小型循环。下面是一个简化版运行期调度循环,完全用同步逻辑呈现核心思路,实际工程里再换成异步事件队列即可。

def run_with_memory_scheduler(graph, budget): # 静态阶段:计算每个张量的活跃区间和liveness状态 intervals = compute_liveness_intervals(graph) # 候选池:记录可在超预算时被Spill的张量 spilling_pool = build_spill_candidate_pool(graph, intervals) mem_allocator = ContiguousPool(alignment=256) async_copy_engine = CopyEngine() for node in topo_order(graph): # 1. 预取当前节点需要的被Spill张量 for tensor in node.inputs: if tensor.is_spilled: async_copy_engine.reload(tensor) tensor.is_spilled = False # 2. 模拟分配当前节点的输入输出和workspace required = node.required_bytes() allocated = mem_allocator.allocate(required) # 3. 如果超预算,选择最佳受害者并异步Spill while mem_allocator.peak_usage() + required > budget: victim = pick_best_victim(spilling_pool, intervals, graph) if victim is None: raise MemoryBudgetError("cannot fit even after spilling") async_copy_engine.spill(victim) mem_allocator.free(victim) victim.is_spilled = True spilling_pool.remove(victim) # 4. 等待当前节点真正需要的拷贝事件完成 async_copy_engine.wait(node.required_ready_events()) # 5. 执行当前算子 execute(node) # 6. 释放死亡区间对应的缓冲 for tensor in node.output_and_dead_inputs(): if not tensor.is_spilled: mem_allocator.free(tensor)

这里的核心思想有两个:一是“先预取,后分配”,确保当前节点需要的张量早已回到设备端;二是“边分配边Spill”,把超预算的压力拆成多个小步处理,避免一次性大起大落。pick_best_victim函数就对应前面第三部分讲的评分策略,实际执行时返回的是现在Spill代价最低的那个张量,同时要保证它不会和当前节点的输入重复。

6.2 小规模模拟验证与收益指标

我在一个模拟的卷积计算图上跑过这个小骨架。图结构是“输入卷积分支卷积融合”,包含三个分支汇聚,中间有六个大型卷积层和一个跳跃连接。不使用调度器时,峰值驻留需要约2.1GB;把预算压到1.2GB后,Spill总量约为900MB,整体额外耗时占总计算时长的11%;预算压到800MB时,额外耗时占比上升到29%。

显存预算峰值驻留Spill体积额外时间占比
2.1GB2.1GB0MB0%
1.5GB1.5GB380MB4%
1.2GB1.2GB900MB11%
800MB800MB1.6GB29%

从这组数据能明显看到一件事:Spill调度不是无代价的,预算压得越狠,额外耗时越高。它的价值在于让“本来跑不动的模型”在较小显存上能跑起来,而不是让模型跑得比原来更快。所以做这类调度器时,千万不要为了追求“峰值越低越好”而牺牲掉大部分性能。关键指标不是某个瞬间的显存占用,而是“总算时间 / 总吞吐”是否还能让业务接受。这个度,是每个落地项目需要反复调参才能找到的。

最后再分享一点个人体会:这类调度器做的不是精妙算法堆叠,而是把“时间轴上的资源复用”这一个核心问题想透。我每次在新框架里动手写内存管理,都会先写一个只处理线性区间、只做简单贪心Spill的粗糙版本,跑通典型卷积网络后再逐步加上workspace特化、im2col分块和成本模型校准。一上来就设计完整的动态决策加多优先级队列,结果往往是被各种角落case拖死,连最基本的显存占用都压不下来。先让简单版本暴露问题,再针对性优化,这条路走起来反而最快。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询