2.3 图分区与子图委派

2.3 图分区与子图委派

本节摘要：Graph Partition 按 EP GetCapability 切分子图。分区质量决定 CPU↔GPU 拷贝次数——异构推理的性能关键。

分区：异构执行的隐藏开关

图优化做完之后，整张图还只是一个逻辑实体。真正决定"谁在哪块硬件上跑"的，是图分区：ORT 拿着每个 EP 的 GetCapability() 声明，把图切成若干连续的子图，每个子图交给最匹配的 EP。分区不是按算子一对一切，而是尽量切出连续、大块、内部自洽的子图——因为子图边界如果跨 device，就要插入内存拷贝，拷贝次数与数据量直接决定异构推理的收益。

GetCapability() 声明的内容包括：能接管的节点列表、每个节点的约束（输入类型、shape 要求）、子图的输入输出接口。ORT 的分区器综合所有 EP 的能力，生成一个以最小化边界拷贝为目标的划分。这个决策在 Session 创建时一次性完成，之后 run 就按既定计划执行。

拷贝边界是异构瓶颈

每两层 CPU/GPU 来回拷贝，GPU 再快也被 PCIe 吃光。PCIe Gen4 x16 的理论带宽约 32GB/s，而一个 4K×4K 的 FP32 特征图是 64MB，单次拷贝就要 2ms 量级——若模型有几十个这样的边界，总开销轻松盖过 GPU 的算力优势。所以分区的第一原则是：让每个 EP 尽可能连续地接管大块子图，把边界数量压到最少。

分区质量的诊断信号很简单：Profile 里 Memcpy 类节点占比高，或运行日志里出现大量 host/device 交替，就是边界过多。

代际演进与分区策略

新代际让 EP 拥有"子图编译"能力：TensorRT 拿到整段子图后做层融合与 kernel 自动调优，分区的粒度从"单个算子"变成"一段子图"。这也解释了为什么同一模型在新版 ORT 上更快——不是单个 kernel 快了，而是分区变好了。

控制分区结果的代码

分区结果主要由 providers 顺序和 provider 选项决定，代码层能做的三件事：

import onnxruntime as ort

sess = ort.InferenceSession("model.onnx", providers=[
    "CUDAExecutionProvider",    # 优先接管
    "CPUExecutionProvider",     # 兜底
])
print(sess.get_providers())
# 查看每个节点的实际执行 device
for n in sess.get_session()._sess_options.__dict__:
    pass   # 结构化信息建议用 get_provider_details 或日志
​

# 开启 verbose 日志看分区决策
ORT_LOG_LEVEL=VERBOSE python run.py
# 日志里会打印 subgraph 划分，检查边界数量与落点
​

控制分区的另一手段是 set_providers 动态调整：先用全部 EP 建 Session，运行时根据设备状态切换。对多模型服务，把"GPU 只跑 backbone、前后处理固定 CPU"作为默认策略，能同时保住吞吐与稳定性。

边界拷贝的量化评估

判断分区是否健康，先量边界成本，再动配置。一个直接的方法是统计模型中跨 device 的边，估算总拷贝字节数：

import onnxruntime as ort

sess = ort.InferenceSession("model.onnx", providers=[
    "CUDAExecutionProvider", "CPUExecutionProvider"])

# 用 verbose 日志或 graph_partition 相关 API 拿到分区信息
# 简化做法：Profile 里统计 Memcpy 节点
import json
# 开启 profiling 后读取 json，统计类型为 Memcpy 的节点总耗时
with open("onnxruntime_profile.json") as f:
    prof = json.load(f)
memcpy_ms = sum(ev["dur"] / 1000 for ev in prof
                if ev.get("cat") == "Node" and "Memcpy" in ev.get("name", ""))
print("memcpy total ms:", memcpy_ms)
​

若 Memcpy 总耗时超过总延迟的 10%，分区就是第一怀疑对象。处置按优先级：调整 providers 顺序让 GPU 接管更大块 → 在导出侧把 CPU-only 算子包进前处理/后处理子图 → 必要时用 Custom Op 让 CPU 算子留在 GPU 端。

分区与数据依赖

分区器必须保证子图间的数据依赖正确。一个子图的输出若被另一个 device 的子图消费，ORT 会自动插入 Memcpy 节点并保证同步语义——但这也意味着该边界上的张量必须经过一次完整搬运。数据依赖越密集，可切分的大块就越少。

上图的切分是健康的：两个 Memcpy 边界把三段逻辑完全隔离。反之，若 backbone 内部每隔几层就跳回 CPU（比如混入一个 CPU-only 的 Gather），边界数量就会爆炸。分区审计时，数一下 backbone 内部是否有 device 交替，比看总节点数更有价值。

判断直觉与常见误区

⚠️ 只注册 CUDA——单个 Gather 不支持导致频繁回 CPU。GPU 模型里混入 CPU-only 算子时，分区器会把图切得很碎，每次边界都拷贝。先查 get_providers()，再看日志确认 Gather 落在哪个 EP。

💡 让 EP 连续接管 backbone，前后处理固定 CPU。预处理（resize、normalize）与后处理（argmax、nms）本就 CPU 友好，留在 CPU 还能避免把 host 端逻辑也搬上 GPU 排队。

分区与第1章的混合执行是同一件事的两面：混合执行是"每个节点都有 kernel 可用"的保底，分区是"每个节点被分配给哪个 EP"的最优解。前者保证能跑，后者决定多快。当边界拷贝成为瓶颈时，优先从模型结构上消除跨 device 的算子混排，而不是盲目换 EP。

要点串联

• Partition = Capability 驱动切分
• 拷贝边界 是异构瓶颈，PCIe 带宽吃光 GPU 收益
• 子图编译 让分区粒度变粗，性能随之提升
• 控制手段：providers 顺序、日志审计、set_providers
• 默认策略：GPU 跑 backbone，CPU 留前后处理
• Memcpy 占比高 → 优先检查分区

下一章：EP 内部机制与选型——分区的下游执行者。