FEATURED · 精选文章

如何用 ONNX Runtime 的 session.name_based_layer_assignment 在不改模型的情况下把指定层分配到 GPU 或 CPU

发布时间 / 2026/9/14 1:54:37
来源 / 创域科博编辑部
栏目 / 资讯中心
如何用 ONNX Runtime 的 session.name_based_layer_assignment 在不改模型的情况下把指定层分配到 GPU 或 CPU 如何用 ONNX Runtime 的 session.name_based_layer_assignment 在不改模型的情况下把指定层分配到 GPU 或 CPU【免费下载链接】onnxruntimeONNX Runtime: cross-platform, high performance ML inferencing and training accelerator项目地址: https://gitcode.com/GitHub_Trending/on/onnxruntime当你已经有一个 ONNX 模型想把其中一部分层放到 GPU、另一部分留在 CPU但又不想改动模型文件本身时ONNX Runtime 提供了会话配置项session.name_based_layer_assignment。它直接对模型中每个节点的名字做子串匹配把命中的节点预先分配到指定设备之后再走正常的 EP 能力划分流程。适用前提是模型的节点名带有结构化的层级信息——大多数 HuggingFace 导出和 PyTorch 导出的 ONNX 模型都满足这一点节点名形如/model/layers.0/self_attn/q_proj/MatMul其中layers.0、layers.10这样的子串可以直接用来定位层。如果节点名里根本没有可识别的层级字符串这个配置项就不适用。匹配规则写配置之前先理解语法session.name_based_layer_assignment的值使用如下格式device1(pattern1, pattern2, ...); device2(pattern3, pattern4, ...)与它配套的行为规则来自文档见 PartitioningWithAnnotationsAndMemoryConstraints.md子串匹配pattern 只要出现在节点名的任意位置就算命中。例如layers.0/能匹配/model/layers.0/self_attn/q_proj/MatMul。最长匹配优先多个 pattern 命中同一个节点名时最长的 pattern 生效。例如对名为/model/layers.10/...的节点layers.10/会胜过layers.1/。不支持前缀annotation 语法里的精确匹配限定符pattern在这里会被直接拒绝并报错所有 pattern 一律按子串处理。设备部分与session.layer_assignment_settings相同支持的设计符包括设备设计符含义可能绑定的 EP 示例cpu任意 CPU 设备CPUExecutionProvidergpu被发现并标记为 GPU 的设备CUDA EP、ROCm EP、DirectML EPcudaNVIDIA CUDA GPU等同gpu:nvidiaCUDAExecutionProviderdmlDirectX 兼容 GPUDMLExecutionProvidernpu神经网络处理单元Qualcomm QNN EP、Intel NPU EPfpgaFPGA 加速器自定义插件 EPaccelerator任意非 CPU 设备任意厂商 EPgpu:vendor厂商特化gpu:nvidia、gpu:amd、gpu:intelgpu:index指定设备索引GPU 0、GPU 1多个设备规则之间用;分隔。第一步确认节点名里有哪些可匹配的字符串写 pattern 之前先确认自己模型的节点名长什么样。文档建议使用 Netron 查看模型的节点名从中找出适合做子串的片段。这一点决定了整条操作路径能否走通如果节点名里找不到能区分各层的字符串就无法用名称匹配来表达层级。以 HuggingFace 风格的节点名为例文档给出的结构是/model/layers.0/self_attn/q_proj/MatMul /model/layers.0/self_attn/k_proj/MatMul /model/layers.15/mlp/gate_proj/MatMul /model/embed_tokens/Gather /model/norm/LayerNormalization对这类命名layers.N/就是天然的分层子串。第二步编写 pattern注意尾随斜杠文档特别提示pattern 末尾要带上/写layers.1/而不是layers.1否则layers.1作为子串会错误地命中layers.10、layers.11等层。虽然最长匹配优先规则在一定程度上能兜住这种情况但带路径分隔符是从源头消除歧义的做法。第三步在 Python 中设置会话配置并创建会话按文档中的示例把前 8 层分配到 GPU、后 8 层分配到 CPUpattern 需按你自己模型的层数改写import onnxruntime as ort opts ort.SessionOptions() # Assign layers 0–7 to GPU, layers 8–15 to CPU based on node names opts.add_session_config_entry( session.name_based_layer_assignment, gpu(layers.0/, layers.1/, layers.2/, layers.3/, layers.4/, layers.5/, layers.6/, layers.7/); cpu(layers.8/, layers.9/, layers.10/, layers.11/, layers.12/, layers.13/, layers.14/, layers.15/) ) session ort.InferenceSession(model.onnx, opts, providers[CUDAExecutionProvider, CPUExecutionProvider])上面示例假设环境已安装带 CUDA EP 的 onnxruntime Python 包model.onnx是待加载的模型文件。配置项本身的定义可以在 onnxruntime_session_options_config_keys.h 中的kOrtSessionOptionsNameBasedLayerAssignment处找到。可选把规则指向某一块具体的 GPU多卡机器上可以用gpu:index指定设备索引opts.add_session_config_entry( session.name_based_layer_assignment, gpu:1(encoder, decoder); cpu(postprocess) )这里有一个文档明确的限制ONNX Runtime 当前只允许每个 EP 类型在每个 session 中出现一个实例所以不能在单个 session 里把不同 annotation 拆到多块 GPU 上——gpu:index只是选择已注册的那个 EP 落在哪块卡。注意gpu:1这个例子中cpu(postprocess)属于 annotation 语法的写法在name_based_layer_assignment里前缀会报错实际使用时应写成不带的子串形式。如何判断配置是否按预期生效文档没有提供逐节点打印分配结果的命令能依据的判断信号有以下几个配置错误会直接报错pattern 里带前缀会收到错误提示所有 pattern 都按子串处理请去掉前缀如果session.name_based_layer_assignment与session.layer_assignment_settings同时设置会返回错误两者互斥见下文限制。未匹配的设备设计符会告警如果配置里的设备设计符例如npu匹配不到 session 里注册的任何一个 EPONNX Runtime 会记录一条警告并跳过该条规则被跳过规则覆盖的节点不做预分配回落到正常的 EP 能力划分。回退行为是确定的没有任何 pattern 命中的节点会走正常的 EP 能力划分文档说明通常是 CPU。另外分配表达的是偏好而非保证——如果目标 EP 没有该节点的已注册 kernel比如某个数据类型/opset 组合在 CUDA EP 未实现该节点不会放到指定设备上而是落到 provider 列表里下一个能处理它的 EP。限制与常见坑与 annotation 方式互斥session.name_based_layer_assignment和session.layer_assignment_settings不能同时设置同时设置会返回错误。模型带layer_ann元数据时用前者对应的 annotation 方式模型未改动、只有节点名时用名称方式。需要把某个单独节点强制到 CPU 这类细粒度例外时把该节点名的 pattern 加进名称配置即可不要两种方案混用。没有子图继承annotation 方式下未标注的子图节点会继承父节点的设备分配名称方式则每个节点独立按自己的名字匹配因为节点名本身就是密集的、几乎每个节点都带结构位置信息不需要继承。图变换产生的新节点按文档说明预划分阶段的融合等变换通常会用原节点名构造新节点名子串匹配仍能命中但个别融合会生成通用名如Attention这类节点需要额外加对应子串 pattern 才能覆盖。节点名稳定性文档指出节点名在不同导出器之间不保证稳定所以用前缀/子串匹配而不是精确全名。可选进阶与内存预算配合使用名称分配可以和 CUDA 的容量感知划分组合使用名称规则表达哪些层想放哪session.resource_cuda_partitioning_settings提供的内存预算作为安全网超出预算的节点会被移出 GPU 分配。opts ort.SessionOptions() # Name-based assignment (no model modification needed) opts.add_session_config_entry( session.name_based_layer_assignment, gpu(layers.0/, layers.1/, layers.2/, layers.3/); cpu(layers.4/, layers.5/, layers.6/, layers.7/) ) # Memory budget as a safety net opts.add_session_config_entry( session.resource_cuda_partitioning_settings, 4194304,node_memory_stats.csv ) session ort.InferenceSession(model.onnx, opts, providers[CUDAExecutionProvider, CPUExecutionProvider])内存预算的统计文件需要先跑一次采集配置session.collect_node_memory_stats_to_file获得具体流程见同一文档中 Capacity-Aware Partitioning 一节。这一步是可选的只在你需要约束 GPU 显存时才使用。更完整的背景与语法细节可继续阅读 docs/annotated_partitioning/PartitioningWithAnnotationsAndMemoryConstraints.md 中 Name-Based Layer Assignment (No Model Modification) 一节。【免费下载链接】onnxruntimeONNX Runtime: cross-platform, high performance ML inferencing and training accelerator项目地址: https://gitcode.com/GitHub_Trending/on/onnxruntime创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
RELATED — 相关阅读

相关资讯

LATEST — 最新资讯

最新发布

TODAY — 本日精选

新闻

WEEKLY — 本周精选

新闻

MONTHLY — 本月精选

新闻