☰
计算图执行优化:缓冲区复用与关键路径调度
2026/9/26 6:18:37 网站建设 项目流程

先说结论:计算图这玩意儿,把节点连起来只是第一步,真正跑起来之后,内存占用和执行效率完全是另一回事。这个Python程序是我在调一个推理管线的运行时问题时动手写的——图里有两百多个算子,按默认的拓扑顺序直接执行,峰值内存逼近2GB,关键路径上一有重计算节点卡住,后面一堆无关节点全在空等。折腾了大半个月,最后发现瓶颈不在算子实现,而在缓冲区怎么分配、节点什么时候调度。于是就有了这个专门做计算图缓冲区管理和调度优化的程序,把原本盲目无序的执行过程,变成一套有统一内存管家、有动态优先级的执行框架。

1. 为什么要单独做一套缓冲区和调度层

很多用计算图的场景,执行方式比我一开始还粗暴:解析完整张图,按拓扑排序的节点顺序一个个跑下去,每个节点的输出缓冲区现场分配,用完等整个图结束再统一清理。能跑,确实能跑。只要图规模小、数据量小,这方案一点问题都没有。但一旦节点数上百、单节点输出buf从几MB涨到几十MB,三个问题就会接踵而来。

第一是内存无序爆炸。每个节点的输出缓冲从生成那一刻起,要到全图结束时才随整体清理释放。而实际上,大部分buffer的生命周期是短暂且错开的:一批节点的buffer早就可以回收了,只是因为没人做回收动作,就一直占着内存。最后的结果是,明明活着的buffer只有几百MB,实测峰值却吃掉了几个GB。

第二是分配开销被严重低估。Python里如果用numpy array现场分配大块内存,一次分配本身不慢,但次数一多就出问题。两百个节点的图跑一遍,产生上千次大块buffer分配和释放,malloc开销、页面触页开销叠加在一起,执行时间被白白拖长10%到20%。这个在机器监控里很难看,因为CPU看着不高,但墙钟时间就是下不去。

第三是调度没考虑依赖之外的客观耗时。拓扑排序保证的只是“每个节点执行时依赖都已就绪”,它不保证“整体等待时间最短”。实际跑起来以后,某个重计算节点还在执行,后方一大批相互不依赖的轻量节点已经完全就绪,却在排队等着。默认的静态顺序不会管这些,只会按序往下走,把本可以并行的节点一个一个处理完。

这个Python程序处理的就是这三件事:做一个BufferPool统一管buffer,按引用计数延迟释放、按大小分桶复用;做一个Scheduler动态决定当前时刻跑哪个节点,关键路径上的节点能插队、短任务能快速清场,并用调度窗口限制同时存活的buffer数量。核心目标只有三条:降峰值内存、升buffer复用率、缩整体执行时间。这套东西对写手写计算图执行器、或者想优化推理管线内存延迟的人最有用;如果用的是TensorFlow这类现成框架,内部已经有类似机制,但理解这套思路对读懂框架底层也很有帮助。

1.1 先明确这个程序到底处理什么问题

要讲清楚,先给计算图下个定义:节点代表一次运算,边代表数据依赖。应用场景比很多人想的多——不光是深度学习模型,还有ETL数据管道、流式计算、编译器的中间表示层。我的程序针对的是“一次性执行完一张静态图”的场景,也就是图结构固定、每个节点只跑一遍。动态图、循环图不适用,那几个场景需要完全不同的策略,后面我会讲到为什么。

插这一层之后,程序要回答的问题就非常明确了:当一个节点执行完毕,哪些buffer可以立刻回收?哪些节点已经具备启动条件?在并行执行资源的限制下,先跑哪一个能获得最佳的整体收益?调度器每做一次决策,既要问内存,又要问时间,合在一起就是这套运行时策略引擎的核心。

这也带出一个原则:调度器不能只看依赖是否满足,还要看buffer是否允许复用。比如同一块存储空间,只要前一个使用者已经不再需要,新的节点就可以把它当作自己的输出缓冲。这是缓冲区管理和调度编排能协同工作的根本前提。把这两件事分开做,效果的提升会非常有限。

1.2 为什么选Python而不是C++

这是项目一开始被问得最多的一个问题。如果目标是极限性能,那一定是C++或者Rust更合适。但这个程序的定位是“运行时策略引擎”,不是被反复调用百万次的算子内核。真正的算子计算可以是numpy向量化、C扩展、外部子进程,Python层只负责决策和调度。在这个前提下,Python的开发效率就是最大优势。

实际使用中,Python做这种工作的短板也确实存在,主要是GIL和单线程吞吐。我的解决方案是把调度器里所有纯计算部分——比如关键路径深度、就绪队列优先级——全部收敛到几个自包含函数里,用numba编译或直接C扩展加速;执行部分用线程池,重算子本身不在Python循环里空转,实测8个线程下加速比能到4到5倍,这个后面实测章节会讲。

我的建议是:如果你的算子也是numpy或C扩展这种能释放GIL的,可以放心用Python做调度层;如果算子本身就是慢吞吞的Python循环,那并行优化基本没救,不如先优化算子实现。这一点在设计架构的时候就要想清楚,不然吃力不讨好。

2. 架构拆解:三个模块怎么协同

程序整体分成三个模块:Graph负责图的解析和静态信息计算,BufferPool负责所有缓冲区的分配、复用、回收,Scheduler负责生成动态执行顺序。三者各干各的活,边界非常清晰,调试问题的时候能快速定位故障层。

2.1 从节点定义到依赖关系

节点用dataclass来表示,这是最常见的做法。每个节点需要记录四类信息:名字、算子类型、输入依赖列表、输出buffer大小。算子的预估耗时一开始没有,我会让它带上一个可选的est_duration属性,从上一轮运行的滑动平均里取值。

@dataclass class Node: name: str op_type: str inputs: list[str] output_size: int est_duration: float = 0.0

图本身就是一个nodes列表加一个邻接表。依赖关系的解析只需要一遍拓扑排序:用Kahn算法逐层剥离入度为0的节点。但这里有个关键点:拓扑排序的结果只用于确认图结构合法和计算关键路径,真正的执行顺序不归它管。如果谁把拓扑序直接当成调度策略,那这篇文章前面说的调度优化就白搭了。

2.2 缓冲区池的模块边界

BufferPool的定位是整个图生命周期内唯一的内存管家。它对外暴露的接口就三个:acquire申请一块输出缓冲、acquire_for_consumer给下游节点登记一次读引用、release将一个buffer还给空闲池。

模块边界上有一条硬规矩:池不负责buffer内容的初始化。只有上层明确要求清零时才做fill_zero,池本身只判断大小和引用状态,不碰数据。这样省掉了大量无意义的memset,也让池的逻辑足够单纯。池最外层统一销毁,内部不逐块去GC,避免了反复分配释放的开销。

另外池内部要维护的东西比表面多不少。一是空闲分桶,二是忙碌表,三是每个buffer的引用计数。池的峰值上限也要在这里控制,否则就会出现“为了省内存反而撑爆内存”的怪圈。这个我会在第三节详细展开。

2.3 调度器的模块边界

Scheduler要回答的是:当前时刻,哪些节点可以开始执行,哪个节点优先级最高。它不关心算子的具体实现,也不负责执行算子,只负责把“就绪且最高优先级”的节点交给Executor。

每个调度决策点出现在两类时机:一是有节点执行完毕,二是有新的buffer被释放回池。第一种时机很直观,一个节点跑完以后,它的下游节点入度可能会清零。第二种时机很容易被忽略:某个大buffer被释放后,一个等待了半天的节点可能因为“目标输出buffer终于有空间”而变得可执行。这类节点本质上是受了内存限制,不是受依赖限制,调度器必须把这种约束也纳入判断。

这也就是为什么我有意把Scheduler和Executor解耦。执行器只负责跑节点,调度器只负责选节点。否则项目一复杂,两边逻辑缠在一起,出了问题根本不知道应该查调度决策还是查执行实现。

3. 缓冲区池实现:分配、复用、回收的完整闭环

缓冲区管理是这套程序里收益最直接的部分。很多图的节点输出buffer大小是有规律分布的,只要把生命周期错开的buffer复用起来,峰值内存能省一半以上。

3.1 从“分配即新建”到“复用即分配”

最朴素的写法是这种:

buf = np.zeros((1024, 1024), dtype=np.float32) # 用完就丢,等GC

问题在于,对一个有几百个节点的图来说,这句话会被执行几百次,每次分配的都是新的内存,用完之后引用消失,又被垃圾回收。下次再有节点需要同样大小的buffer,就再来一轮分配、回收。分配的开销在Python里比想象中大,因为numpy还要做数组对象初始化和类型检查。

buffer pool的版本是反过来设计:

class BufferPool: def __init__(self): self._free_buckets = defaultdict(list) self._busy = {} self._refcount = defaultdict(int) self._bucket_limits = (4096, 65536) def _bucket_for(self, size: int) -> int: for limit in self._bucket_limits: if size <= limit: return limit return size def acquire(self, owner: str, size: int) -> np.ndarray: bucket = self._bucket_for(size) if self._free_buckets[bucket]: buf = self._free_buckets[bucket].pop() else: buf = np.empty(size, dtype=np.float32) self._busy[owner] = buf self._refcount[owner] = 0 return buf def release(self, owner: str): if owner not in self._busy: return buf = self._busy.pop(owner) bucket = self._bucket_for(buf.nbytes) self._free_buckets[bucket].append(buf) self._refcount.pop(owner, None)

acquire的时候优先从空闲桶里拿现成buffer,没有才新建。release的时候不销毁buffer,直接扔回空闲桶。这样同一块内存可以被图中不同节点在生命周期错开的情况下反复使用。

这里有个细节要特别说明:为什么用np.empty而不是np.zeros。因为很多算子根本不在乎输出缓冲的旧内容,直接覆盖写,用了np.zeros等于白白多一次memset。只有那些确实需要初值零的算子,上层才自己调用buf.fill(0)。这种按需清零的策略,在一张图动辄几百个buffer时,省下来的时间非常可观。

3.2 引用计数与提前回收

如果图是一根直线,节点执行完立刻释放buffer没问题。但实际图里到处都是分支结构,一个输出buffer可能要喂给多个下游节点。最安全的回收时机是:所有依赖它的节点都已经读完。

我给每个buffer单独维护一个consumer_count,初始值等于它的下游节点数。每个下游节点在开始执行前,通过acquire_for_consumer登记一次读引用,读完一次就减一。只有引用计数归零,buffer才真正进入空闲池。

def attach_consumer(self, producer: str): self._refcount[producer] += 1 def mark_consumed(self, producer: str): self._refcount[producer] -= 1 if self._refcount[producer] == 0: self.release(producer)

这个“延迟释放”我是踩过坑之后才重视的。早期版本图省事,节点跑完就直接release,结果就是同一个buffer在某个下游节点还没读完的时候被重新分配给了别的节点,产生脏数据。当时排查了很久,最后靠打印每个buffer的分配时间戳才发现问题。从那时候起,我坚持所有回收操作都必须走引用计数,宁可多写几行代码,也不为图省事埋炸弹。

3.3 大小分桶与峰值内存控制

如果只是简单地把空闲buffer丢进一个list,运行一段时间就会发现一个问题:空闲list里积压了大量size各异的buffer,而真正有用的却拿不到。比如一个1MB的buffer释放后,接下来所有节点只要不是恰好1MB,就不会复用它。最后池子里养着一堆“看起来能用、实际用不上”的空间,峰值内存还是高。

所以我学内存分配器的思路,搞了大小分桶。把buffer按size分成几档:小于4KB的进小桶,4KB到64KB的中桶,更大的进大桶。申请时向上取整到自己所在的桶,释放时按实际大小入桶。这样小块内存不会被大buffer长期占用,大buffer也不会被零碎的小请求反复折腾。

峰值内存控制则靠一条硬规则:整个池子累计分配出去的buffer字节数不能超过max_pool_bytes。超过之后,release的buffer不入空闲池,直接丢弃交还给系统。这样内存峰值被硬生生顶在一个上限,不会出现“省内存省到反而撑爆内存”的荒诞局面。max_pool_bytes这里怎么取?我的经验是先按所有节点buffer的一起生命周期重叠部分估算,通常取总和的一半到一个三分之一,再跑两轮实测微调。

4. 调度器实现:从拓扑排序到关键路径优先

调度是整套程序里最考验设计能力的部分。它不只是把合法顺序排出来,而是在“合法”之外再优化“高效”。

4.1 拓扑排序只是起点,不是终点

Kahn算法做拓扑排序,本质是“哪个节点入度先清零就先输出哪个”。这个顺序合法,但是很傻:它不知道节点之间的耗时差距,也不知道哪条链对整体延迟影响最大。

我做一个最简单的假设场景来说明问题:一条关键链上有一个耗时80ms的重计算节点A,旁边有一百个0.1ms的轻节点B都处于就绪状态。拓扑排序如果先把A排到后面,那所有依赖A的节点都要多等80ms;如果先把A排前面,轻节点在后面跑,整体等待时间就会短很多。这还只是单核情况。如果做成并行,调度错乱导致的等待浪费更明显。

调度器要做的是对每个节点计算两个值,一是关键深度,也就是从它的下游链路中最长的那条链到终点的耗时估算;二是最晚开始时间,基于全图的总关键路径长度反推。关键路径上的节点,也就是那些“一旦延迟就会让整图延迟”的节点,必须获得最高优先级。

4.2 动态就绪队列:关键路径优先,短作业前置

整个调度循环用事件驱动方式实现,核心简化版本是这样:

ready_queue = PriorityQueue() def on_ready(node): # 第一优先级:是否关键路径节点;第二优先级:预估耗时短的先跑 ready_queue.put((0 if node.on_critical else 1, node.est_duration, node.name)) while not all_done: if executor.has_idle_slot() and not ready_queue.empty(): _, _, name = ready_queue.get() node = nodes[name] # 先挂上游buffer的引用,再执行 for inp in node.inputs: pool.attach_consumer(inp) executor.submit(node, pool) else: # 等待执行完成事件或者buffer释放事件 event = wait_event() if event.type == NODE_COMPLETED: for dep in out_edges[event.node_name]: dec_in_degree(dep) if in_degree[dep] == 0: on_ready(dep) elif event.type == BUFFER_RELEASED: for waiter in waiting_for_memory: if pool.has_space(waiter): on_ready(waiter)

为什么优先级是“关键路径节点优先 + 短作业优先”?关键路径优先很好理解,它直接压缩了整图的理论最短时间。短作业优先则是为了让就绪队列快速腾出位置,让更多节点进入可并行状态,避免一个0.1ms的小任务卡住后面一堆节点。这两者组合在一起,能明显提升调度窗口内的任务饱满度。

4.3 调度窗口:卡住内存和线程池的边界

调度窗口这个参数的用处,很多人一开始想不到。它不是时间窗口,而是“同时执行的最大节点数”,我通常设为CPU核数的1到1.5倍。窗口有两个作用:第一,限制同时存在的buffer数量,配合BufferPool的max_pool_bytes一起卡峰值内存;第二,避免一次性往线程池提交上百个节点,让任务全在排队等调度,造成无意义的时间分片。

窗口大小怎么设?我会先跑一遍profiling,把每个节点的预估耗时和并行阶段并发数算出来。如果某个阶段平均并发只有3,窗口设8就是浪费;如果平均并发已经到7,窗口设4就会人为串行化,得不偿失。我推荐的流程是:先设一个偏小的窗口,比如4,跑一遍,然后逐步往上加,观察墙钟时间和峰值内存的变化。一般来说,在这个数据到达拐点之前停止增加就够了。我自己的压测图里,8是最合适的值,再往上内存飙升,时间却几乎不变。

还有一个坑要提醒:调度窗口如果做得太激进,比如窗口内都是重节点,可能出现“窗口被占满但在等一个还没就绪的节点”的假死状态。我的处理方式是,窗口的空位不固定分配给某个节点,而是每次做决策时重新评估,如果当前就绪队列里最好的节点预估耗时太长,而其他可等待任务能在窗口内穿插完成,就允许插入少量轻节点。本质上是贪心策略,但加了关键路径这个强制优先级来防止本末倒置。

5. 实测数据与排查实录:这些坑必须提前知道

工具做得好不好,最终要看实测。我给自己造了一个模拟推理图,尺寸不小:240个节点,3条并行主链(每条链40个节点),中间穿插一些共享节点和分支结构,节点buffer从几KB到几十MB不等。所有buffer需求总和有5.2GB,但因为生命周期重叠没那么高,同时存活的理论峰值约1.9GB。测试机是8核CPU,16GB内存。

5.1 三个策略的量化对比

我把三种情况放在同一台机器上跑:

执行方式峰值内存总耗时
拓扑序直接执行,每节点即时分配1.87GB328ms
只加缓冲池,不改调度1.02GB301ms
缓冲池 + 关键路径优先调度0.84GB214ms

这个结果有几个点值得细看。缓冲池的收益非常扎实,因为生命周期不重叠的buffer大量复用之后,峰值直接从1.87GB砍到1.02GB,这个效果是“把大块内存反复用”带来的,和算子本身快慢无关。调度优化的收益更隐蔽:它把总耗时压低了接近35%,但这个收益不是凭空来的,它是把关键路径上的重计算节点从“被轻节点插队”的困境中解救出来,让整图的理论最短时间真正兑现了。

我也试过不设关键路径优先级、只开缓冲池加并行窗口的组合,总耗时只到272ms,说明优先级策略对整个结果贡献很大。只做内存优化不做调度优化的结果则说明,内存搞好了,如果执行顺序还是一团糟,时间上的浪费一样很可观。

5.2 常见问题排查表

实操中一定会有问题,我把踩过的和帮别人看过的坑整理成了一张速查表:

现象根因排查/解法
某个节点的输出被后续节点改写,数据变得莫名其妙缓冲区提前释放,引用计数漏加检查每个input是否在上游启动前attach_consumer,打印buffer分配时间戳定位
内存不降反升,池子越养越大空闲分桶队列里积压了长期不用的buffer给空闲buffer加时间戳,超过阈值定期销毁
多线程调度时偶发卡死就绪队列空转,条件变量等待丢失用事件驱动循环重写,wait加上超时重查
性能比拓扑序还慢关键路径计算有误,或窗口设太大导致资源被低价值任务抢占打印关键路径节点列表核对,窗口减半重测
图有环时程序死循环没做环检测构造阶段用Tarjan算法判环,遇到环直接给出节点列表报错
大buffer一直没有被复用分桶粒度不合理,比如把1MB和100MB放同一档细化桶的档位,或者按2的幂次分桶

5.3 几个值得注意的工程细节

多线程执行节点时,Python的GIL影响比想象中大。如果节点算子是纯Python循环,那开八个线程也拿不到实际并行度,反而线程切换拖慢单节点执行。我用的是numpy向量化和C扩展,算子在执行期间能释放GIL,所以线程池才有效果。这是设计这套程序时就埋下的前提,如果你的算子做不到这一点,整个并行策略要推翻重来。

关于节点预估耗时的获取,不要在真实环境里临时profiling,开销太大。我为每个算子类型维护了一个滑动平均时长:oatype -> ema_duration,每执行完一个节点就用真实耗时更新一次。调度器的优先级依赖这个数据,所以估得越准,调度越聪明。第一轮跑的时候数据还是空的,我会用默认值,第二轮开始基本就收敛了。

还有一个很容易被忽略的点:输出缓冲区的大小必须按实际值上报。有的算子输出大小依赖输入的形状,比如padding、concat、slice,预估值很可能不对。我在节点执行完的回执里显式带上实际buffer大小,BufferPool按真实尺寸登记,这样后面释放和复用才不至于错位。

最后说一个我个人体会比较深的地方:缓冲区管理和调度优化这类工作,很容易被当成性能优化的“边角料”,但实际它在工程里的收益往往比打磨单个算子更大。这个程序后续想扩展的话,方向也挺明确——把分桶策略改成TLSF算法、把调度器接入异步IO事件循环、或者把池内buffer换成分布式共享内存,都是可以深挖的点。但不管怎么改,核心思路都是那句:内存要复用,调度要动态,关键路径要优先生跑。希望这篇记录能帮正在做计算图执行器的人少走几个来回。

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

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

立即咨询