异构计算图全局调度：多目标优化延迟、功耗与内存的实战解析

最近在调一条多卡异构的训练推理链路时我突然意识到一个挺尴尬的事实算子层面的优化已经快“卷”到头了。在一个计算图里哪怕你把每个算子都各自压到了理论峰值图与图之间的数据搬移、设备同步、内存峰值依然可能把整体性能拖回去一大截。HeteroOpt 这个项目就是我对“深度学习计算图在多类异构硬件上到底该怎么调度”这个问题的一次系统梳理。它面向的不只是某一块 GPU 或某一块 NPU而是 CPU、GPU、NPU、FPGA 混在一起时如何从计算图全局出发同时优化延迟、功耗和内存占用这几个互相打架的目标。如果你也是做 AI Infra、推理引擎或者端侧部署的同学这篇文章应该能给你一些启发。1. 为什么“全局图调度”在异构时代被重新抬上桌面1.1 单个算子的优化空间正在收窄早几年做性能优化思路很直接找一个热点算子用 CUDA 或者汇编把它榨干就能拿到非常明显的收益。这个过程确实很有成就感算子融合、内核调优、内存复用这些手段到今天依然是基本功。但随着 cuDNN、OneDNN、TensorRT 这类底层算子库越来越完善通用算子的优化已经接近手工调优天花板。你花一个星期手写一个 Matmul 的 kernel很可能也就比官方库快 1% 到 2%而且还未必能在所有 shape 上稳定胜出。这意味着什么如果还把精力全部押在“让单个算子跑得更快”上投入产出比会越来越难看。真正剩余的优化空间反而集中在算子与算子之间的连接方式上谁和谁该放在同一块设备上、谁应该先执行谁应该后执行、中间张量要不要跨设备搬、哪些子图可以合并成一个大算子再执行。这些问题的答案已经不是单独看某个算子能算出来的了。1.2 异构硬件让“任务分配”从判断题变成了排列组合题在单一种类硬件的时代计算图调度基本就是个顺序问题。最坏的情况是在同一张 GPU 上决定某个算子用哪个 kernel 实现本质上还是一个局部寻优。但异构环境完全不一样。数据中心里可能有 Nvidia GPU、Habana、各种国产 NPU、FPGA 加速卡端侧还有 ARM CPU、Mali GPU、NPU、DSP。每一类硬件在算力峰值、内存带宽、功耗曲线上都差异巨大。一个 Transformer 层里矩阵乘法可能适合放在 GPU 或 NPU 上LayerNorm 这种访存密集型算子放在 CPU 上反而能省掉 GPU 的 kernel 启动开销Concat 这种算子如果放在设备端做可能还得把数据从显存搬过去再搬回来。算子的可执行设备集合一旦多起来整个计算图的调度决策空间就会爆炸式增长。更麻烦的是设备之间还需要协同。GPU 算完交给 NPU中间就有一条跨设备通信边如果通信链路和另一条数据加载链路共用同一条 PCIe 总线还会互相抢带宽。这种耦合关系单靠人工去排基本不可能做到接近最优。1.3 一个反直觉的例子局部最快不等于全局最快我经常跟团队里的小伙伴举一个例子有一个计算图A 算子在 GPU 上执行要 1 毫秒在 CPU 上要 3 毫秒B 算子依赖 A 的输出它在 GPU 上要 2 毫秒在 CPU 上要 4 毫秒A 和 B 之间的中间结果非常大大概有几百 MB。如果只看单算子性能A 放 GPU、B 放 GPU 好像没问题。但如果 A 放 CPU、B 放 GPU 反而可能整体更快因为省掉了一次大规模跨设备张量传输。而如果 A 放 GPU、B 也放 GPU虽然两个算子都快中间传输可能只要一次 PCIe 或者 NVLink 的拷贝但问题在于 GPU 同时还要处理别的高优先级算子可能出现排队等待。这种“局部最优叠加成全局不佳”的现象在异构调度里非常常见。它本质上跟操作系统里的优先级反转有点像每个任务都觉得自己拿到的是最优资源但整个系统却卡在等待和搬运上。所以异构调度必须有全局视角也就是从整张计算图出发去做决策。2. 从“局部下钻”到“全局视图”HeteroOpt如何建模一张计算图2.1 用DAG而不是算子序列来看待训练和推理HeteroOpt 的第一步是把模型从“一层一层顺序执行”的惯性思维里拽出来真正用有向无环图DAG来表达。DAG 里每个节点是一个算子每条边表示张量依赖。这个表达方式和编译器的中间表示IR非常像图的拓扑序对应执行的先后顺序但 DAG 本身并不限制并行度。使用 DAG 有一个好处调度器可以清楚看到哪些子图在逻辑上是完全独立的可以并行放到不同设备上执行。例如一个双塔模型的两个塔本质上互不依赖如果只按顺序执行就白白浪费了异构设备的并行能力。我在实际做模型切分的时候发现很多同学意识不到这一点因为 PyTorch 代码写出来是一个 Module 接一个 Module天然暗示了串行。2.2 算子的融合与并行实际上会改变图拓扑全局调度还有一个容易被忽略的点调度过程会反过来改图。图融合和算子拆分并不是在调度之前就一成不变地给定而是调度器需要一并考虑的决策。比如 Conv 和 BN 可以融合成一个算子减少一次中间张量写回和读取但过度融合也可能导致单个算子占据 GPU 太长时间挡了旁边本来可以并行执行的小算子。又比如一个非常大的 GEMM可以按行拆成两个子 GEMM分到两块设备上并行执行但拆分之后必然增加同步开销。这些图变换与设备分配是一体的问题不能先定图再定设备更不能先定设备再改图。所以我在 HeteroOpt 的实现里把图重写作为调度搜索的一部分搜索过程中会生成候选的融合版本、拆分版本再加上设备放置信息一起送入评估器打分。这样调度结果不再是一份“谁在哪跑”的表格而是一份“整张图应该被重写成什么样、每个节点在哪跑、按什么顺序跑”的完整方案。2.3 我在建模时额外多加了三种边普通 DAG 只包含数据依赖边但真实硬件上还有通信依赖和竞争关系。HeteroOpt 在建模时会在原始计算图基础上补充三类边数据边张量生产者和消费者之间的依赖决定执行顺序的下界。通信边如果两个相连节点被分配到不同设备上就需要一条跨设备传输边这条边的开销取决于张量大小和传输带宽。竞争边两个节点虽然无数据依赖但会争用同一块设备上的计算资源或同一根总线的带宽在调度时不能完全无视彼此。加入这三种边之后调度器的目标就变成了找一个满足所有约束的拓扑执行序和设备分配让目标函数最小。这个建模方式虽然比纯 DAG 复杂但更贴近硬件真实行为。曾经有一次我用纯 DAG 模型算出某个分拆方案能省 30% 延迟实际一测反而变慢了原因就是两条数据搬运路径抢了同一根 PCIe 总线。后来把竞争边补上模型预估值才和实测对得上。3. 多目标优化的落地延迟、能耗与内存的三角博弈3.1 只优化延迟是不太够的大部分调度器一上来就只盯着 latency 优化这是最容易实现的但放到真实生产环境里往往不够用。数据中心场景虽然希望推理延迟低但功耗是硬约束。机房电费、散热、PUE 都会折算进运营成本。如果你调度的目标是“尽量把更多算子丢到 GPU 上”延迟可能很漂亮但 GPU 功耗会飙升整体能效反而下降。端侧场景更明显手机上跑模型最怕的不是慢是发热和掉电快NPU 之所以存在很多时候不是因为它绝对算得快而是因为它每瓦特性能高。所以 HeteroOpt 从设计之初就把延迟、功耗、内存峰值作为三个并列目标而不是把后两个当成延迟的附属约束。3.2 从“加权和”到“Pareto前沿”把多目标压成一个标量最直接的方法就是加权求和TotalCost α × Latency β × Energy γ × PeakMemory但权重的选择非常主观。交互式应用看重延迟数据中心看重功耗端侧部署则内存占用可能才是第一瓶颈。如果固定一组权重调度器就只能给出一组方案换一个场景就得重新调权重甚至重新搜索。更合理的做法是搜索出一组 Pareto 最优解也就是所谓的 Pareto 前沿在这组解里任何一个目标都不能在不牺牲其他目标的前提下继续变好。比如方案 A 比方案 B 延迟低 10%但功耗高 20%那 A 和 B 可能同时出现在 Pareto 前沿上因为不同用户会做不同取舍。线下把这些候选方案全部保存下来线上根据当前硬件条件或业务需求查表选择比线上临时调整权重要稳定得多。维度单目标调度多目标调度优化对象延迟最小化延迟、功耗、内存等多维目标输出形式一组最优解一组 Pareto 候选解适配场景场景固定、目标单一场景多变、云端与端侧并存决策复杂度低拿到就直接用需要二次筛选计算开销相对小搜索空间更大3.3 一个调度方案如何被多维打分在 HeteroOpt 内部任何一个候选调度方案都会经历一个多维打分过程。延迟估计分两部分计算延迟和通信延迟功耗估计则参考各设备的历史功耗曲线内存估计则是对图中所有存活张量做峰值扫描。打分并不是简单拍脑袋而是每项都做了归一化。延迟会除以一个“单块最强 GPU 跑完整图”的参考时间功耗除以设备热设计功耗TDP内存除以设备可用内存上限。这样三个目标都落在 0 到 1 之间Pareto 搜索时不会因为量纲差异导致偏好偏移。实际测试中我还会对每个方案做一次小规模实测校准把成本模型的系统性偏差修正掉再进入最终的 Pareto 排序。4. 调度器内部管线图解析、成本估计、搜索与决策4.1 静态成本模型不真正跑也能估出个八九不离十调度搜索最怕慢每个候选方案都真的放到目标设备上跑一遍是不现实的。所以 HeteroOpt 的第一层筛选依赖静态成本模型。成本模型的核心输入有三个算子的计算量FLOPs、算子的访存量、算子的输出张量大小。然后根据设备参数估算执行时间EstimatedTime Max(ComputeCost / DeviceThroughput, MemoryAccessCost / DeviceBandwidth)这个式子虽然很简单但能区分两类算子Conv、GEMM 这类计算密集型算子瓶颈在设备总算力而 Flatten、Concat、Transpose 这类访存密集型算子瓶颈在内存带宽。估算时不是只看占比还要看算子内部是否存在 Kernel 启动开销和并行不足的问题。比如一个很小的算子放 GPU 上kernal 启动时间可能比实际计算时间还长这时候成本模型如果只算 ComputeCost就会严重低估延迟。我会在成本模型里额外加一个固定开销项大概是 3 到 10 微秒用于模拟内核启动和同步屏障。这样小算子自动会被模型惩罚调度器会更倾向把多个小算子融合或者在 CPU 上执行而不是频繁唤醒 GPU。4.2 候选生成先排序、后组合有了成本模型下一步就是生成候选调度方案。直接对图中所有算子做笛卡尔积式的枚举搜索空间太大。HeteroOpt 采用的思路是两阶段先粗粒度切分再细粒度分配。切分阶段会把计算图划分成若干子图或超节点。划分依据包括算子间的强耦合关系、中间张量大小、是否适合融合等。例如一个包含 Matmul、Add、LayerNorm、Dropout 的 Transformer Block很可能被切成一个“计算密集子图”和一个“访存密集子图”。切分完之后调度器为每个子图罗列出可执行设备集合再通过约束传播和剪枝去掉明显不满足内存上限或带宽上限的组合。这个阶段不会直接给出最终答案但会大幅压缩搜索空间。比如某个计算密集子图适合在 GPU 或 NPU 上执行那 CPU 相关的组合就不需要进入下一轮。4.3 全局搜索为什么单靠启发式或单靠学习都不够到了真正的搜索阶段我采用的方法是“启发式骨架 强化学习策略”混合。具体展开就是先跑一轮贪心拓扑排序得到一个可行的基础调度保证后续搜索不会比这个方案差太多然后在这个基础上使用一个轻量级 GNN 编码计算图的特征输出每个子图的设备放置偏好和拓扑序调整方向。启发式在这里起着约束探索边界的作用。每次强化学习策略给出一个倾向后系统会在它周围再扩大采样一批邻居解交给成本模型评估如果这批邻居解里有明显更优的就把策略模型的训练信号反馈回去。这样训练出来的策略模型不会像纯 RL 那样冷启动策略太差也不会像纯启发式那样陷入手工规则的天花板。在搜索预算有限时我还会用 Beam Search 保留 Top-K 个候选而不是只取一个最优。最后从 Top-K 个候选里再做一次 Pareto 过滤把非支配解全部保留下来。这套组合方案在小数据集上验证时搜索速度比纯进化算法快了一个数量级以上同时基本稳定找到更优的 Pareto 前沿。5. 我踩过的坑设备差异、动态shape与可复现性5.1 同一个算子在不同硬件上的数值结果并不完全一致这是一个特别隐蔽的坑。不同厂商的算子库在卷积、归一化等算子的实现细节上差别很大有的用 TF32有的用 FP16 中间精度有的对累加顺序做了重排。同一个模型的计算图调度到不同硬件上之后精度表现可能差距明显。我在做对比实验时遇到过类似的情况同一个模型GPU 上跑 Top-1 精度能到 92.3%切到另一块加速卡上直接变成 91.7%。表面上看是设备精度问题实际是调度器把某些对精度敏感的算子放到低精度设备上执行数值误差被放大了。HeteroOpt 在成本模型里加入了“精度惩罚项”如果某个算子被分配到只支持低精度计算的设备而模型本身对这块算子又比较敏感调度器会基于距离中心或误差预算的估算给它加一个额外的权重。真实部署时建议在调度结果确定之后对关键模型跑一遍精度对比测试损失超过 0.2 个点就要回退到高精度设备组合。5.2 dynamic shape 会让“全局静态调度”很尴尬静态计算图调度能做得快是因为所有 shape 都提前确定成本模型才能准确估算。可现实里很多模型输入长度是动态的。NLP 场景里 batch 内句子长度不同视频模型里序列帧数不同甚至检测模型里图片尺寸也会动态变化。一旦 shape 变了算子的计算量、访存量、输出张量大小都会变之前搜索出来的最优设备分配很可能不再成立。HeteroOpt 的解决方案是分两层图级调度保持相对静态只针对子图切分和设备偏好做决策算子级调度则在运行时根据实际 shape 动态决定到底用哪个 kernel 实现。这样既保留了全局调度的大部分收益又不会因为在运行时遇到一个不同 shape 的输入而整套方案失效。5.3 多目标优化的可复现性问题我踩得最深刻的坑是同一个调度配置隔了几个星期再跑一遍实验结果完全不一样。原因可能是驱动版本更新了算子库的 kernel 选择策略变了也可能是机器上多跑了一个占用显存的服务导致内存峰值估计失真。为了保住实验结论我在 HeteroOpt 的评测管线里引入三个强制约定。第一所有搜索实验固定随机种子和确定性算法开关第二记录完整环境快照包括 GPU 驱动、算子库版本、BIOS 功耗策略第三每个调度方案至少运行 5 次取 P50 或 P95 而不是简单平均值。下面是几个常见坑与应对方式的速查表坑现象建议驱动版本更新同一方案延迟变化明显锁定驱动版本并记录快照算子库 kernel 选择差异同一算子时间波动大固定 cuDNN 或 OneDNN 的算法启发式选项多任务抢占显存功耗出现异常峰值评测前清理环境独占机器shape 变化静态调度方案失效分层调度或设置 shape 分组精度差异部署后指标下滑加入精度惩罚项并做回退测试6. 在真实模型上观察效果与评估方法6.1 评测集合的覆盖思路HeteroOpt 的实验评测没有只拿一个 ResNet 就跑因为计算图调度的难易程度跟模型结构高度相关。我这边最终采用了一套“三模型 两模式”的评估策略。一组是 CNN 类模型比如 ResNet50特征是计算图相对规整算子以卷积和池化为主结构上比较适合静态切分。一组是 Transformer 类模型比如 BERT-base 和一个 GPT 风格的 decoder特征是有大量矩阵乘法、Attention 和 LayerNorm图中还包含复杂的 skip connection 与多头并行结构。最后一组是 GNN 类模型特征是图结构本身不规则算子在 gather/scatter 和 GEMM 之间切换非常考验调度器对访存密集型算子的处理能力。两个模式分别是训练和推理。训练场景关注吞吐和显存峰值推理场景更关注延迟和功耗。这样跑下来基本能覆盖大多数深度学习模型在异构部署时会遇到的调度难点。6.2 指标怎么看Pareto前沿与超体积如果只给一个“平均加速比”很容易掩盖多目标调度真正想解决的问题。我会把结果画在二维坐标里比如横轴是归一化延迟纵轴是归一化功耗每个调度方案就是一个点Pareto 最优的那些点连成一条折线就是这一组实验的 Pareto 前沿。比较两个调度器的优劣除了肉眼看前沿是否更靠左下角还有一个更客观的指标叫超体积Hypervolume也就是“Pareto 前沿和参考点之间围成的面积”。超体积越大说明这一组方案在目标空间里的覆盖质量和最优性越好。HeteroOpt 在这个指标上跑出的结果普遍比单目标调度器高出 12% 到 20% 左右比纯贪心启发式高出 25% 以上。6.3 对比实验的三种基线对比基线至少得有三种否则说明不了问题。第一种是单目标最快基线把所有算子都丢到算力最强的设备上不管功耗和内存。这是一个非常强的延迟基线通常延迟最低但功耗和内存往往不是最优。第二种是贪心启发式按照“计算密集算子去 GPU/NPU、访存密集算子去 CPU”的手工经验分配。这是工程团队最常见的做法效果稳定但没有全局协调。第三种是分组专家策略比如把所有卷积放在 GPU所有 Attention 放在 NPU一切按层类型分配。这种策略看起来合理但在跨设备传输开销巨大时经常失灵。实测结论是HeteroOpt 在延迟上可能比单目标基线慢 3% 到 5%但功耗能低 20% 到 30%内存峰值能低 15% 左右。这种“牺牲一点点延迟换回能耗和内存双重收益”的结果对很多生产场景来说非常值得选。最后分享一个小实践体会多目标调度不一定每次都要在线实时搜索。对于结构和部署环境相对固定的模型我会把搜索得到的最优 Pareto 前沿缓存成一张表线上做查表决策调度开销几乎可以忽略。只有当输入 shape 范围变化、硬件占用变更或模型结构更新时才触发一次增量重搜。HeteroOpt 在我这里不只是一套搜索算法更像是一个帮你想清楚“该让哪块硬件干什么”的框架想清楚了之后怎么执行反而是一件可以做得非常轻的事。