npugraph_ex快速上手

npugraph_ex后端可以提供一个基于torch.compile的简单整图aclgraph加速方案,aclgraph又称为捕获模式,其实现原理请参考《CANN 应用开发指南 (C&C++)》中“运行时资源管理>基于捕获方式构建模型运行实例”章节。

本章将提供aclgraph图模式功能配置的快速上手示例,仅供参考。请根据实际情况自行修改脚本,支持配置的功能参见功能列表

使用约束

  • 主要面向在线推理场景,暂不支持反向流程capture成图、随机数算子capture。
  • npugraph_ex与torch.cuda.CUDAGraph原生接口(参见《PyTorch 原生API支持度》中的“torch.cuda”)功能类似,约束与其保持一致(如不支持stream sync、动态控制流等),此处不再赘述。

使用方法

通过配置torch.compile的backend="npugraph_ex"进行编译。

import torch
import torch_npu

# 自定义Model
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x, y):
        return torch.add(x, y)

model = Model().npu()
# 基于npugraph_ex backend进行compile
opt_model = torch.compile(model, backend="npugraph_ex", fullgraph=True, dynamic=False)

# 执行编译后的Model
x = torch.randn(2, 2).npu()
y = torch.randn(2, 2).npu()
opt_model(x, y)

功能列表

npugraph_ex提供多项功能,包括: 基础功能进阶功能DFX功能

其中基础功能通过torch.compile的options参数进行配置,torch.compile为PyTorch原生接口,接口详细介绍请参见官网,接口原型如下:

torch.compile(model=None, *, fullgraph=False, dynamic=None, backend='inductor', mode=None, options=None, disable=False)

torch.compile参数配置说明参见表1

表 1 torch.compile参数说明(aclgraph模式)

参数名 PyTorch原生参数说明 aclgraph模式下参数说明
model 必选参数。入图部分的模型或者函数。 与原生含义一致。
fullgraph 可选参数,bool类型。是否捕获整图进行优化。
False(缺省值):非整图优化。
True:捕获整图优化。
建议设置为True,要求将整个函数或模型捕获到一个单一的计算图中。如果编译器遇到无法追踪到该单一图中的代码时(即“图中断”),则会引发错误。
dynamic 可选参数,bool类型或None。是否启用动态Shape追踪。
None(缺省值):自动检测是否启用动态Shape追踪。
False:不启用动态Shape追踪。
True:启用动态Shape追踪。
与原生含义一致。
backend 必选参数,后端选择,缺省值为"inductor"。 需显式传入backend="npugraph_ex"。
mode 开销模式,内存开销模式选择,缺省值为None。 昇腾NPU暂不支持
options 优化选项,缺省值为None。 提供多种基础功能配置,具体参见基础功能
disable 可选参数,bool类型。是否关闭torch.compile能力。
False(缺省值):开启torch.compile能力。
True:关闭torch.compile能力,采用单算子模式。
与原生含义一致。

表 2 npugraph_ex基础功能

功能 功能说明
force_eager功能 图执行前是否使用Eager模式运行。
force_recapture功能 是否强制每次执行时重新捕获aclgraph。
FX图优化Pass配置功能 是否开启FX图优化优化能力。以减少计算过程中的内存搬运,从而提升性能。
FX图算子融合Pass配置功能 是否开启FX图算子融合Pass。该Pass基于已有Aten IR进行融合,从而提升性能。
aclgraph间内存复用功能 aclgraph间内存复用功能,支持多种模式。
静态Kernel编译功能 是否开启静态Kernel编译。
冗余算子消除功能 是否对冗余Kernel进行优化处理。
固定权重类输入地址功能 图执行时是否固定权重类输入地址。
重捕获次数限制功能 设置重捕获次数。
集合通信入图 实现集合通信算子Ascend Converter,调用torch.compile时默认已支持集合通信算子入图。
Cat算子消除功能 是否开启Cat算子消除优化以减少内存拷贝和临时张量分配,提升执行性能。
图捕获安全策略配置 控制图捕获过程中,对某些可能不安全操作(如分配device内存)的处理策略。

表 3 npugraph_ex进阶功能

功能 功能说明
模型编译缓存功能 在推理服务和弹性扩容等业务场景中,使用编译缓存可有效缩短服务启动后的首次推理时延。
多流表达功能 大模型推理场景下,对于一些可并行的场景,可划分多个stream提升执行效率。
AI-Core和Vector-Core限核功能 提供Stream级核数配置,可调整最大AI Core数和Vector Core数,避免算子执行并行度降低。
自定义FX图优化Pass功能 传入自定义FX Pass函数,该配置可控制自定义Pass在框架内置Pass执行前/后生效。

表 4 npugraph_ex DFX功能

功能 功能说明
图编译Debug信息保存功能 通过复用原生DEBUG环境变量TORCH_COMPILE_DEBUG开启日志打印和文件Dump。
算子Data-Dump功能 是否开启数据dump功能。
多流并发死锁检测功能 是否开启多流并发死锁检测功能。

问题定界

使用npugraph_ex后端提供的功能时,如果遇到异常场景,可参考如下问题定界方法,如仍无法解决,请单击联系技术支持。

  1. 定界是否是用户脚本问题。

    npugraph_ex是基于torch.compile扩展aclgraph的功能,因此用户脚本必须先经过PyTorch社区的compile验证才能进行NPU编译,即用户脚本必须先通过下面的尝试后才能继续使用npugraph_ex。

    torch.compile(backend="aot_eager", fullgraph=True)
    

    backend="aot_eager"是PyTorch社区为用户自定义后端做一个简单使用Eager模式直接运行fx.graph的方式。如果这一步出错,那么算子或用户脚本本身就不满足backend="npugraph_ex"的使用条件。一般来说,此时脚本直接使用torch.npu.graph会报错。

  2. 定界是否是aclgraph问题。

    由于npugraph_ex提供的所有功能均是基于fx.graph做的aclgraph增强体验,因此如果npugraph_ex出现问题的话,第一步需要定界是aclgraph的问题还是npugraph_ex增强的功能存在问题,可使用force_eager功能辅助定界。

    torch.compile(backend="npugraph_ex", fullgraph=True, options={"force_eager": True})
    

    如果用户脚本使用force_eager运行异常,则可能是npugraph_ex存在问题,否则可能是aclgraph的Runtime底层问题。

  3. 定界是否是图捕获问题。

    当确认问题与aclgraph相关时,可使用force_recapture功能辅助定界是图捕获执行问题还是对输入输出内存处理问题。

    torch.compile(backend="npugraph_ex", fullgraph=True, options={"force_recapture": True})
    

    启用force_recapture后,每次执行都会重新捕获aclgraph。通过对比重捕获前后执行效果,如果force_recapture后存在问题,说明问题是出在图捕获执行阶段。如果force_recapture后功能正常,说明问题是出在对输入输出内存处理阶段。

功能拓展

若您希望额外增加一些自定义的FX图优化功能,可通过torch.npu.npugraph_ex.compile_fx接口自定义compiler和backend,然后传入torch.compile进行编译,示例如下:

import torch
from torch._functorch.aot_autograd import aot_module_simplified
import torch_npu

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x, y):
        x = x + y
        return x

# 构建自定义的compiler
def custom_compiler(gm: torch.fx.GraphModule, example_inputs):
    test_options = {
        "clone_input": False
    }
    compiled_graph = torch.npu.npugraph_ex.compile_fx(gm, example_inputs, test_options)
    return compiled_graph

# 构建自定义的backend
def custom_backend(gm: torch.fx.GraphModule, example_inputs):
    return aot_module_simplified(gm, example_inputs, fw_compiler=custom_compiler)

x = torch.ones([2, 2], dtype=torch.int32).npu()
y = torch.ones([2, 2], dtype=torch.int32).npu()
model = torch.compile(Model().npu(), backend=custom_backend, fullgraph=True, dynamic=False)
ret = model(x, y)