异构计算图调度实战:HeteroOpt如何用多目标优化分配算子
2026/9/9 15:49:43 网站建设 项目流程

如果你做过多模型部署,一定有这种感受:同一个模型在GPU上跑得飞快,换到NPU上却可能因为某个算子不支持而整体卡死;或者GPU已经排队排到冒烟,旁边的CPU却闲得发慌,可调度器就是不知道把一部分算子切过去。这个场景背后,就是异构硬件上的计算图调度问题。HeteroOpt这个名字听起来很学术,但它解决的事情非常具体:在一张深度学习计算图里,把每个算子分配到最合适的设备上,同时决定执行顺序,并兼顾时延、功耗、内存等多个互相打架的指标。我做AI推理部署这么多年,早期用过传统贪心调度,也试过纯规则策略和大量手工调参,后来才意识到,真正能扛住多设备场景的,是把它当成一个全局多目标优化问题来做。这篇内容适合正在做推理引擎、边缘计算部署或者编译器后端优化的同学参考,也可以当作一份异构调度框架的设计笔记来读。

1. 这问题到底难在哪:先想清楚为什么要做计算图调度

1.1 异构硬件的性能鸿沟:一个模型在不同设备上跑出两个世界

我们先从硬件说起。深度学习模型部署的目标设备很少只有一种,常见的有CPU、GPU、NPU、DSP、FPGA,甚至还有各类定制加速卡。它们不是“快慢”的区别,而是“擅长什么”的区别。CPU适合逻辑复杂但不重的任务,GPU适合大规模并行计算,NPU在卷积、矩阵乘这些结构规整的算子上效率极高,DSP则在信号处理类场景里有独特优势。

举个例子,一个卷积层在GPU上可能只需要0.2毫秒,在CPU上要3毫秒,但在NPU上可能只要0.08毫秒。而同一边缘设备上的另一个算子,比如一个动态shape的Resize或者一个条件分支,NPU甚至根本不支持,只能乖乖回到CPU执行。也就是说,每个算子的“最优设备”不是固定的,它跟算子的类型、输入shape、内存布局、设备状态都有关系。

我见过不少团队在单卡GPU上把模型调得很好,一换到边缘盒子就全线崩溃。原因往往是他们把所有算子硬塞到算力最强的芯片上,结果某个算子不支持,推理直接报错。这种场景用生活类比就是:一个项目组里有几个能力不同的人,有人擅长数学,有人擅长写作,有人速度慢但从不偷懒——你需要把一个项目拆成很多碎活发下去,每个人处理碎活的速度还不一样,而且碎活之间还有前后依赖,前一个人的输出是后一个人的输入。调度员的任务,就是决定谁在什么时候做什么活。

1.2 计算图调度到底在调什么

深度学习模型在编译和推理阶段会被表示成一张计算图。图中节点是算子,边是张量依赖。常见的有TensorFlow的GraphDef、PyTorch的torch.fx Graph、ONNX Graph,以及各家编译器内部的IR图。调度要做的事情,看起来就三件:顺序、设备、资源。

顺序,是指算子在同一个设备上的执行先后。拓扑排序只是底线,面对有分支、有汇合、有可并行的子图,不同的顺序会带来完全不同的执行时间。设备,是把每个算子指定到哪个计算单元上,这一步直接决定算子的执行耗时和数据搬运开销。资源,主要是内存和张量生命周期。两个算子如果生命周期不重叠,它们的临时buffer就可以复用;如果把算子切分到不同设备,中间结果要跨设备拷贝,内存和带宽都会受影响。

这三件事不是独立的。你把算子A放到GPU,把依赖它的算子B放到CPU,那么A的输出就得从GPU搬到CPU,这个搬运时间可能比A和B本身的执行时间还长。调度策略如果只顾着给每个算子挑最快的设备,很容易做出一个“单点最快、整体很慢”的方案。

1.3 为什么单单快不够:多目标之间的真实冲突

早期调度器最喜欢干的事,就是把目标函数定为“执行时间最短”。但在实际工程项目里,这个目标往往不够用。你只优化时延,算法会把所有算子都往GPU上塞,显存直接爆掉;你只优化内存,算子全塞到CPU,推理慢到用户无法接受;你只优化功耗,性能又可能掉一大截。

我这里列一下真实部署时通常要同时盯的几个指标:

  • 时延:端到端推理时间,最核心的体验指标。
  • 功耗:边缘设备有散热和电池限制,数据中心则直接关系电费和制冷。
  • 峰值内存:尤其是大模型和移动端场景,内存可能比算力更稀缺。
  • 设备利用率:不能让某一个芯片空闲太久,否则整体资源浪费。
  • 编译时间:搜索调度方案本身就消耗时间,这个时间算不算成本也看场景。

这些目标天然冲突。最典型的就是“时延最优”和“内存最优”对着干:为了时延最小,你倾向于把多个算子并行放到不同设备上,中间结果全部驻留,内存峰值快速上涨;为了内存最小,你又倾向于串行执行、及时释放buffer,但总时延就上去了。传统做法是拍脑袋定一个加权公式,把多个目标压成一个分数。但权重怎么定?不同业务场景权重完全不一样,线上推理服务看重时延,无人机之类的东西看重功耗和散热,你不可能用一个固定权重覆盖所有场景。

2. HeteroOpt的思路拆解:多目标搜索与异构特征融合

2.1 把调度问题改写成搜索问题

HeteroOpt和传统规则式调度最大的区别,是它不试图用一组if-else规则覆盖所有情况,而是把调度问题明确地改写成“搜索问题”。在这个视角下,一个调度方案就是搜索空间里的一个点,搜索算法的任务就是在空间里找到足够好的点。

搜索空间的维度怎么定义?最简单的方式是:每个算子一个维度,每个维度的取值是该算子可用的设备编号。那一个可能的调度方案就是给每个算子分配一个设备。但如果仅仅做设备分配,还远远不够。真实推理中,算子的执行顺序、张量切分策略、显存复用计划、是否需要算子融合,都会影响最终性能,所以一个完整的搜索方案通常会包含好几类决策变量。

HeteroOpt这类框架的特征是“全局”两个字。它不是按层做局部优化,而是把整张计算图放进同一个搜索过程里,让设备分配、顺序调整、内存规划相互影响、统一评估。这种全局视角会带来一个副作用,就是搜索空间巨大。一个有100个算子、3种设备的图,不考虑顺序调整时,已经有3的100次方个组合了——你不可能枚举完,必须依赖高效的搜索算法和剪枝策略。

2.2 多目标怎么权衡:帕累托前沿不是空话

在数学上,“多目标优化”和“单目标加权”有本质区别。单目标加权是提前把多个指标揉成一个数,搜索时只朝一个方向走;而多目标优化保留“方案A在时延上优于方案B,但B在功耗上优于A”这类比较关系,最终输出的不是唯一解,而是一组帕累托前沿解。

帕累托前沿的意思,用大白话说:在候选集合里,如果一个方案在某个指标上更好,同时在所有其他指标上都不比另一个方案差,那它就是“支配”了另一个。所有不被别人支配的方案,构成了帕累托前沿。这些方案没有绝对的好坏,只是侧重不同。

我记得第一次把一个调度搜索从单目标加权改成帕累托多目标输出时,业务方给我的反馈是:太好了,我不需要提前定死功耗和时延的权重,我只需要拿到几个候选方案,然后在不同的机型上现场选。不同设备形态,选出来的方案不一样。这个价值在真实工程里非常明显,因为你永远无法预知客户会跑什么负载、什么温度环境、什么剩余电量。

所以,一个实用的多目标调度框架,通常不会替你拍板“应该用哪个调度方案”,而是给你一条“前沿曲线”,让你根据现场需求决策。这种设计思路虽然看起来“把选择困难抛给了用户”,但实际用起来,比黑盒给一个确定方案要灵活得多。

2.3 异构特征融合是让搜索“懂行”的关键

调度搜索能不能做好,很大程度上取决于你拿什么特征去描述算子和设备。如果只是拿算子类型和设备编号喂给搜索算法,那算法基本靠猜。HeteroOpt这类框架能跑出效果,背后靠的是异构特征融合。

这里说的“异构特征”,分三个层面:

第一是算子特征。包括算子类型、输入输出的shape、通道数、计算量、访存量、是否适合融合等。这些信息可以从计算图和shape推导里拿到。

第二是设备特征。包括算力峰值、内存带宽、内存容量、缓存大小、功耗系数,以及最重要的“算子支持矩阵”——某个设备能不能跑某个算子,能跑的时候性能如何。

第三是关系特征。包括张量依赖关系、数据量大小、设备间的通信带宽和延迟,以及算子之间是否存在可融合的候选路径。

把这些特征融合成统一的向量或者图表示,搜索算法才能做出有依据的决策。比如一个卷积算子的计算量是10GMACs,设备的访存带宽只有20GB/s,那算法就能判断瓶颈可能在访存而不是算力;如果两个相邻算子都是逐元素操作,算法会优先考虑它们能不能融合到同一个设备上执行。

现在比较新潮的做法,是用图神经网络对计算图做嵌入,让每个算子节点学到周围上下文的信息。但这不意味着必须上GNN才能干活。最简单的方案是构造一个特征矩阵,再配合规则约束,也能让遗传算法或者贝叶斯优化跑出还不错的结果。特征融合的目的是压缩信息、引导搜索,不是炫技。

3. 从零复现一个简化版HeteroOpt调度器

3.1 环境准备:先把模型导成一张干净的计算图

我默认你已经有一台能跑PyTorch的机器,CUDA环境能正常用就行。如果环境还没配好,先花点时间把驱动、CUDA、PyTorch装好再回来,这一步不难但很碎,别急。

要复现HeteroOpt的思路,我们不需要把整套框架从论文里搬过来,只需要把三个模块串起来:图解析、代价评估、搜索算法。第一步,把模型导成计算图。PyTorch里最简单的方式是用torch.fx:

import torch import torch.fx as fx model = torchvision.models.resnet18(pretrained=True) model.eval() graph_module = fx.symbolic_trace(model) graph = graph_module.graph

跑完这一步后,graph.nodes里就是一张完整的计算图。每个node自带op类型、targetargs这些元信息,配合一个假输入做一次 shape 推导,就能拿到每个中间 tensor 的 shape。这一步很关键,因为后续算通信量、内存占用,全都要靠 shape 信息。

input_tensor = torch.randn(1, 3, 224, 224) graph_module(input_tensor) shape_map = {} # node_name -> tensor shape for node in graph.nodes: if hasattr(node, "meta") and "tensor_meta" in node.meta: shape_map[node.name] = node.meta["tensor_meta"].shape

我这里用meta拿 shape,是PyTorch 2.x常见的做法,老版本可能需要靠profiling或者模拟执行来补全。拿到节点列表和shape之后,把边关系也导出来,就得到了一个干净的图数据结构。图不大,完全可以用networkx或者igraph来存,后面遍历比较方便。

3.2 代价评估函数:给每种算子打分

调度搜索需要一把“尺子”,用来量一个调度方案的各项指标。这一步没有标准答案,必须贴近你手上的硬件来定。我给出一个非常实用的思路:先构造一个“算子-设备耗时表”作为经验值,然后用公式估算通信和内存开销。

# 简化版:算子类型在不同设备上的耗时估算函数 # 实际项目中,这张表应该用本机benchmark跑一次,把结果填进来 op_cost_table = { ("conv2d", "cuda"): 0.2, # ms per 1 GMACs-ish ("conv2d", "cpu"): 3.0, ("conv2d", "npu"): 0.08, ("add", "cuda"): 0.01, ("add", "cpu"): 0.05, ("add", "npu"): 0.03, # ... } def estimate_op_time(op_type: str, device: str, flops: float): base = op_cost_table.get((op_type, device), 1.0) return base * max(flops, 1e-6) def estimate_transfer_bytes(shape, dtype=torch.float32): numel = 1 for dim in shape: numel *= dim return numel * dtype.itemsize if hasattr(dtype, "itemsize") else numel * 4

设备之间的搬运开销可以单独算。比如把张量从GPU搬到CPU,搬的是device_to_device,时间大致等于数据量除以实际带宽。注意实际带宽不是理论带宽,我一般按理论带宽的50%~70%估算,宁高勿低,宁可调度结果保守一点。

内存峰值怎么算?我的做法是模拟执行一遍计算图:按照拓扑序逐个执行算子,算子输出变成活跃张量,等依赖它的所有算子都执行完了,这个张量才能释放。活跃张量总量会随着设备分配方案的不同而变化,这个模拟器虽然只有几十行,但对调度结果的影响非常大。

def simulate(execution_order, device_assign, nodes, shapes): alive_tensors = {} peak_memory = 0 total_time = 0.0 total_energy = 0.0 device_ready_time = {"cpu": 0.0, "cuda": 0.0, "npu": 0.0} for node in execution_order: dev = device_assign[node.name] ready = max(device_ready_time[dev], max((alive_tensors[src] for src in node.args if isinstance(src, str) and src in alive_tensors), default=0.0)) op_time = estimate_op_time(node.target.split(".")[-1], dev, shapes.get(node.name, 0)) start = ready if dev != "cpu" else ready # 简化处理 end = start + op_time alive_tensors[node.name] = end device_ready_time[dev] = end total_time = max(total_time, end) total_energy += op_time * 5.0 if dev == "npu" else op_time # 释放不再使用的tensor used_by_future = set() # 根据依赖关系判断是否还有后续算子需要该tensor # ... return total_time, peak_memory, total_energy

这个模拟器虽然简化,但已经能反映很大的问题:同一个模型,不同调度方案之间的总耗时会差出好几倍。后面真正要做的,就是在这个评估函数的基础上去做搜索,找到更优的方案。

3.3 用NSGA-II风格的遗传搜索找帕累托候选集

调度搜索的算法选型,很多人第一反应是“用强化学习”。但我的经验是,先用一个足够好的启发式搜索算法把baseline打出来,再考虑要不要上更复杂的东西。NSGA-II是经典的多目标遗传算法,实现不复杂,效果却很稳,特别适合这种离散组合优化问题。

搜索编码可以这样定义:一条染色体就是一组设备编号,长度等于算子数,每一位的取值是0/1/2分别对应cpu/cuda/npu。执行顺序在这次简化版本里先固定为拓扑序,把顺序调整留到下一轮迭代再做。

import random def random_chromosome(nodes, device_num=3): return [random.randrange(device_num) for _ in range(len(nodes))] def crossover(a, b): cut = random.randrange(len(a)) return a[:cut] + b[cut:], b[:cut] + a[cut:] def mutate(chromosome, prob=0.1): for i in range(len(chromosome)): if random.random() < prob: chromosome[i] = random.randrange(3) return chromosome

多目标选择的核心是非支配排序。对每个个体,判断它被多少个体支配,然后一层一层地剥出帕累托层。这个算法写出来没几行,但要注意细节:当两个个体互不支配时,要保留多样性,不能全丢掉。NSGA-II的实现可以参考论文里的伪代码,网上也有不少开源实现,但自己写一遍之后对调度问题的理解会深很多。

每一轮迭代,用上一轮产生的帕累托前沿个体做父代,经过交叉变异生成下一代,再合并做非支配排序,优胜劣汰。跑上几百代之后,最终输出的就是一个帕累托前沿集合。我实际跑下来,100个算子左右的图,种群规模设200,迭代300代,在我笔记本上大概十几分钟能出结果,搜索时间完全可接受。

3.4 结果怎么验证:仿真指标与真机校准

搜索出的方案不能只活在模拟器里,要拿到真机上去验证。我建议至少验证三组对比:所有算子放CPU、所有算子放GPU、搜索出的帕累托方案。看一下时延、功耗、内存这三项的真实数据,跟模拟器预估对比,误差通常在20%~50%之间,原因是估算表本身不够准。

校准方式也很粗暴但有效:先用profiler把每个算子在每个设备上的真实耗时跑出来,回填到op_cost_table里,然后重新跑一遍搜索。多迭代几次,你会发现搜索结果越来越接近实际最优。这一步就是很多论文里说的“评估-搜索-实测闭环”,放到工程里就是最实用的tuning流程。

这里有一个很容易忽略的点:很多设备在并行执行多个算子时,耗时不能简单累加。GPU上的kernel是流水线式并行的,CPU上多核也能并行,所以如果只是简单把执行时间加起来,你会低估并行收益,也会高估串行时间。我在模拟器里加了一个粗粒度的并行估计:同一设备上如果有多个算子的时间窗重叠,就把总占用时间按重叠系数压缩。这个系数需要靠跑真实样例去标定,不同设备差异挺大。

4. 我踩过的坑:搜索爆炸、模型失真与真机打脸

4.1 搜索空间太大,跑不动怎么办

第一次完整跑通整个框架,我兴冲冲地拿一个Bert-base的图跑了NSGA-II,结果发现算到天荒地老。后来一查,问题出在两处:一是算子粒度太细,一个模型拆出几千个节点,每个节点都是搜索维度;二是设备分配和顺序调整混在一个搜索里,解空间成倍膨胀。

应对策略其实就一句话:先粗后细,分层搜索。第一轮搜索按“层”为单位做设备分配,把几十个算子合并成一个搜索单元;确定了大方向之后,再对关键路径上的算子做更细粒度的调度。搜索空间立刻小了一个数量级。另外可以加剪枝:每个算子的候选设备不是所有设备,而是先用规则筛出top2或者top3,搜索时根本不考虑明显不行的选项。

还有一招是限制时间预算。遗传算法这类方法有个特点,前几十代提升飞快,后面几百代几乎是在精雕细琢。如果业务要求秒级出结果,可以把迭代次数压到50代以内,配合更聪明的初始化策略,效果不会差太多;如果离线编译允许跑几小时,那就放心加大种群规模和迭代次数。

4.2 代价模型失真,调度方案明显反直觉

有一次我搜出来的调度方案,竟然把一堆简单算子全放到了NPU,把卷积放到了CPU,怎么看都觉得不对。检查之后发现,是我的op_cost_table里NPU的add算子耗时填得太低,卷积的CPU耗时填得太低,导致算法产生了一个在现实里不存在的荒谬方案。

代价模型失真,是这类框架最容易翻车的地方。解决方式只有一个:用真机数据填表,不要拍脑袋。你可以先写一个benchmark脚本,把模型里出现的所有算子类型分别放到每个设备上跑一遍,记录耗时、峰值内存、功耗,然后把这些真实数据喂给搜索器。注意跑benchmark时要用真实的输入shape,因为同样一个卷积,跑512x512和跑64x64,在不同设备上的相对快慢可能完全不一样。

另外,设备不是孤立工作。多个算子同时跑在一个NPU上时,存在争抢算力和带宽的问题,真实耗时不是单算子耗时的简单相加。我后来在代价模型里加入了“设备负载系数”:当某个设备上同时活跃的算子数超过阈值时,单算子耗时按比例放大。这个修正让调度结果在真机上更稳健。

4.3 搜索结果不稳定,换台机器就变味

同一份调度方案,在开发机器上实测很好,客户机器上一测又慢又卡。这个“换台机器就变味”的问题,几乎每个做调度的人都会遇到。原因是每台机器的CPU频率、NPU固件版本、GPU驱动、内存带宽都不一样,甚至同一个型号的机器,因为散热差异,稳定后的算力都有波动。

我的经验是:不要把调度方案做成全局固定值,而是做成“候选集+运行时选择”。离线搜索阶段不要只输出一个最优方案,输出5到10个帕累托候选;在线部署时用一个轻量级runtime,根据当前设备属性、温度、剩余内存等状态,在候选集里挑一个最合适的。这个思路和计算机系统的动态电压频率调节有点像,静态调度负责找到优质候选,动态决策负责现场适配,各干各的活。

另外,固定随机种子、增加重复实验次数,也是保证结果可复现的基本操作。遗传算法本身有随机性,同一份代码跑两次结果不一样很正常,但如果你连实验都不固定种子,后面排查问题会非常痛苦。

4.4 真机部署被“最后一公里”坑到

调度方案在模拟器里完美,真机部署还是可能崩,最常见的坑是算子不支持。你搜索时假设某个算子可以放到NPU,但实际部署时这个算子的某一种输入shape没有对应的kernel实现,推理直接报错。解决方案是在搜索前就把算子支持矩阵做扎实,不支持的算子直接过滤掉,别让搜索器浪费时间。

另一个坑是隐藏开销。kernel启动时间、内存分配锁竞争、D2D拷贝和D2H拷贝的同步延迟,这些在模拟器里往往被忽略,但真机上它们可能占推理总时延的30%以上。尤其是小算子的情况,kernel启动开销比计算本身还大,这时候调度搜索的意义就不大了,优先做算子融合反而收益更高。

我的建议是:无论离线调度搜得多好,最后都要保留一个运行时profiling的兜底接口。上线之后真机数据回流,再迭代一轮代价模型和新候选方案。调度框架不是一个一劳永逸的离线工具,它是一个持续进化的系统。

回到HeteroOpt这个思路本身,它的价值不在于某个具体算法的花哨程度,而在于它提供了一个可扩展的框架:把复杂的异构硬件差异,收敛成“图解析-多目标评估-全局搜索”的标准流程。我实际用下来,最大的体会是,做调度的核心不是把优化算法调得多精通,而是把代价模型和特征表示做得足够贴近真实机器。这两件事做好了,哪怕只用最简单的遗传算法,也能出很好的效果。如果你也在被异构设备折腾,建议先别急着上重型解决方案,把你手上那张代价表校准一遍,跑一个全局搜索baseline,很可能就已经比你现在的规则调度好出一大截了。

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

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

立即咨询