CANN/cann-learning-hub：第三方框架集成npugraph_ex指南

前置文章【免费下载链接】cann-learning-hubCANN 学习中心仓支持在线互动运行、边学边练提供教程、示例与优化方案一站式助力昇腾开发者快速上手。项目地址: https://gitcode.com/cann/cann-learning-hub《npugraph_exCANN aclGraph的图模式样板间》1 背景介绍npugraph_ex融合了ACLGraph的调度能力和亲和NPU的图优化能力在大模型推理场景下有效的消除了算子调度时延显著提升了大模型的推理表现。然而使能npugraph_ex由于需要对模型进行编译优化会引入较长的冷启动时间此时便需要使用编译缓存功能将编译优化后的产物保存到磁盘以便下次运行直接加载节省编译耗时。当前npugraph_ex不仅支持了第三方框架对于模型性能优化的集成也支持了编译缓存功能的集成端到端的降低模型推理耗时。2 原理介绍现在主流的AI框架都会支持作为torch.compile的后端进行模型图模式优化当直接使用npugraph_ex作为后端时torch.compile主要分为以下三个阶段1.TorchDynamo也称为图捕获Grapph Capture阶段将Python模型代码进行跟踪记录转换成FX图Pytorch图存储格式。2.AOT AutoGrad通过TorchDynamo捕获的FX图生成前向和反向的FX图。3.FX Compile对FX图进行编译优化。而FX Compile这个阶段就是npugraph_ex使能的关键其中也分为了四个阶段的优化1.FX Pass优化在FX图上执行特定变换操作优化图执行例如算子融合、常量折叠、内存优化、ReInplace优化等。2.Tiling UpdateTiling的动态刷新部分算子如FlashAttention Tiling随输入prompt长度变化需要在执行阶段动态刷新算子的Tiling信息。3.Static Kernel算子的静态编译提前将aclnn kernel编译缓存节省算子编译耗时。4.aclgraph capture推理时先将模型的执行任务全部捕获之后再统一下发并随着推理循环多次执行节省任务的下发调度耗时。对于第三方框架作为后端而言通常TorchDynamo这个阶段是一致的不同之处在于得到FX图后框架对于FX图的优化流程和方法因此第三方框架集成npugraph_ex的推理加速能力就是在输入FX图的前提下使能npugraph_ex FX Compile的四个阶段优化。npugraph_ex提供了_NpuCompiler类给定FX图和输入Tensor后便可以实现FX Compile阶段的优化并提供各个阶段编译产物的缓存功能加速程序二次运行的冷启动耗时。3 功能使用介绍3.1 torch.compile自定义后端集成npugraph_extorch.compile支持自定义后端要求后端函数具有(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]) - Callable的约定,返回的就是自定义后端利用输入FX图编译优化后的可执行对象要求(*args: torch.Tensor) - List[torch.Tensor]。TorchAir提供了get_compiler接口返回_NpuFxCompiler实例def get_compiler(compiler_config: CompilerConfig None): Retrieves the NPU compiler instance. Args: compiler_config (CompilerConfig, optional): Compiler configuration. Defaults to None. Returns: _NpuFxCompiler: NPU compiler instance. if compiler_config is None: compiler_config CompilerConfig() return _NpuFxCompiler(compiler_config)而_NpuFxCompiler实例可以利用后端函数的输入gm和example_inputs生成_CompiledFxGraph对象也就是npugraph_ex根据FX图编译优化后的NPU图可执行对象class _NpuFxCompiler: Main compiler class for converting FX graphs to NPU-compatible graphs. def __init__(self, compiler_config: CompilerConfig) - None: self.config compiler_config pretty_error_msg def __call__(self, gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]): Compiles the FX graph into an NPU-compatible graph. Args: gm (torch.fx.GraphModule): The FX graph module to compile. example_inputs (List[torch.Tensor]): Example inputs for compilation. Returns: _CompiledFxGraph: Runner wrapping the compiled graph. return self._get_compiled_gm(gm, example_inputs) ......_CompiledFxGraph满足了torch.compile后端函数返回对象的要求符合(*args: torch.Tensor) - List[torch.Tensor]。因此自定义后端既可以自行对TorchDynamo跟踪输出的FX图优化处理成新的FX图要求与AotAutoGrad处理后的FX图格式相同再通过get_compiler返回的_NpuFxCompiler实例生成_CompiledFxGraph可执行图实现后端函数也可以直接将_NpuFxCompiler实例作为AOTAutoGrad的参数输入完成后端函数实现以下给出后者的简易示例import os import torch from torch._functorch.aot_autograd import aot_module_simplified import torch_npu import torchair from torchair.configs.compiler_config import CompilerConfig class MM(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x, y): x x y return x def custom_backend(gm: torch.fx.GraphModule, example_inputs): compiler_config CompilerConfig() compiler_config.mode reduce-overhead compiler torchair.get_compiler(compiler_config) return aot_module_simplified(gm, example_inputs, fw_compilercompiler) torch.npu.set_device(0) x torch.ones([2, 2], dtypetorch.int32).npu() y torch.ones([2, 2], dtypetorch.int32).npu() model torch.compile(MM().npu(), backendcustom_backend, dynamicFalse) ret model(x, y) print(ret)参考资料1.get_compiler、_NpuFxCompiler和_CompiledFxGraph定义参考TorchAir开源仓:https://gitcode.com/Ascend/torchair/blob/master/python/torchair/npu_fx_compiler.py2.torch.compile自定义后端接入详情参考Pytorch官网手册https://docs.pytorch.ac.cn/docs/stable/torch.compiler_custom_backends.html3.2 编译缓存_CompiledFxGraph提供了dump_artifacts和load_artifacts方法共同实现编译产物的缓存落盘功能在模型首次编译时使用dump_artifacts方法获取内存中的编译产物并写入到磁盘指定路径中。模型二次编译时检测磁盘指定路径是否存在编译产物的落盘文件存在便读入内存使用load_artifacts方法加载直接生成编译优化后的图。示例如下程序首次执行将缓存编译产物到磁盘二次执行直接加载磁盘编译产物节省编译耗时import os import pickle import torch import torch_npu import torchair from torch._functorch.aot_autograd import aot_module_simplified from torchair.configs.compiler_config import CompilerConfig class MM(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x, y): x x y return x def custom_compiler(gm: torch.fx.GraphModule, example_inputs): cache_file ./compiled_model.pkl # 尝试从缓存加载 if os.path.exists(cache_file): print(发现缓存正在加载...) with open(cache_file, rb) as f: artifacts pickle.load(f) # 使用 load_artifacts 重建编译图 from torchair.npu_fx_compiler import _CompiledFxGraph compiled_graph _CompiledFxGraph.load_artifacts(artifacts) print(从缓存加载成功) return compiled_graph # 缓存未命中执行编译 print(缓存未命中开始编译...) compiler_config CompilerConfig() compiler_config.mode reduce-overhead # 必须设置为支持 codegen 的模式 compiler torchair.get_compiler(compiler_config) # 直接调用编译器获取 _CompiledFxGraph 实例 compiled_graph compiler(gm, example_inputs) print(type(compiled_graph)) # 保存到缓存 print(正在保存编译产物...) artifacts compiled_graph.dump_artifacts() with open(cache_file, wb) as f: pickle.dump(artifacts, f) print(f编译产物已保存到 {cache_file}) return compiled_graph def custom_backend(gm: torch.fx.GraphModule, example_inputs): return aot_module_simplified(gm, example_inputs, fw_compilercustom_compiler) def main(): torch.npu.set_device(0) # 准备数据 x torch.ones([2, 2], dtypetorch.int32).npu() y torch.ones([2, 2], dtypetorch.int32).npu() # 编译并执行 model torch.compile(MM().npu(), backendcustom_backend, dynamicFalse, fullgraphTrue) ret model(x, y) print(f结果: {ret}\n) if __name__ __main__: main()4 第三方框架集成npugraph_ex实例——XPU_GRAPH4.1 XPU_GRAPH介绍XPU_GRAPH是一个支持多种设备的自定义图编译器基于Pytorch FX图和Aten IR进行图编译优化通过集成npugraph_ex支持了昇腾设备存在以下特点通用图优化公共子表达式消除CSE死代码消除DCE算子折叠常量折叠以及更激进的常量传播。厂商自定义算子转换将效率较低通常会引发大量内存访问的算子转换为定制化融合算子。结构化模式匹配XPU_GRAPH 对常见结构模式进行抽象支持用户实现对应目标结构并将指定模式转换为用户定义的格式。后端编译器集成XPU_GRAPH 采用 “FX 图输入–FX 图输出” 设计因此可与 Inductor、GE 等其他 FX 图编译器兼容。同时支持推理与训练场景。4.2 XPU_GRAPH集成npugraph_exXPU_GRAPH在backends/__init__.py中定义了vendor_compile方法为自定义图编译器提供了统一的入口def vendor_compiler( gm: torch.fx.GraphModule, fake_inputs: list, target: Target, *, is_inference: bool False, is_backward: bool False, **config_dict: Dict[str, Any], ) - Callable: try: target_mod importlib.import_module(f.{target.value}, __package__) except Exception: logger.warning(f{target.value}_compiler not found, return gm) return gm compile_fn getattr(target_mod, f{target.value}_compile) logger.info(f{target.value}_compile start...) xpu_compiled compile_fn(gm, fake_inputs, is_inferenceis_inference, is_backwardis_backward, **config_dict) logger.info(f{target.value}_compile complete) return xpu_compiledcompile_fn便是不同的图编译器函数名称为***_compile返回编译优化后的图可执行对象这也正好方便了上层接入torch.compile。npugraph_ex对应的编译器函数是npu_compile定义在backends/npu.py中def npu_compile( module: torch.nn.Module, inputs, *, is_inference: bool False, is_backward: bool False, **config_dict: Dict, ) - torch.nn.Module: compiler config_dict.get(compiler, ge) if compiler ge: assert is_inference, Currently, we use ge only for inference. return ge_compiler(module, inputs, **config_dict) elif compiler device_graph: assert is_inference, Device graph capture/replay is intended for inference-style execution. return device_graph.device_graph_compiler(module, inputs, targetnpu, **config_dict) else: return inductor_compiler(module, inputs, is_inferenceis_inference, is_backwardis_backward, **config_dict)当使能npugraph_ex时会走到ge_compiler,其也定义在npu.py中def ge_compiler(module: torch.nn.Module, example_inputs, **config_dict: Dict) - torch.nn.Module: import torch.fx as fx import torch_npu torch.npu.set_compile_mode(jit_compileFalse) import torchair as tng import torchair.ge_concrete_graph.ge_converter.experimental.patch_for_hcom_allreduce from torchair.configs.compiler_config import CompilerConfig config CompilerConfig() recursive_set_obj(config_dict, config) if ( mode : config_dict.get( mode, ( max-autotune if compiler in config_dict else reduce-overhead ), # NOTE(liuyuan): If user specify the compiler, then we should consider it as GE instead of AclGraph. ) ) reduce-overhead: config.mode mode from torch import SymInt for ele in example_inputs: if isinstance(ele, SymInt): raise TypeError(ACL Graph does not support dynamic shape!!) if mempool : config_dict.get(use_custom_pool, None): config.aclgraph_config.use_custom_pool mempool npu_backend tng.get_compiler(compiler_configconfig) from torchair._utils import get_npu_default_decompositions module make_fx( module, decomposition_tableget_npu_default_decompositions(), tracing_modefake, record_module_stackTrue, )(*example_inputs) compiled_module npu_backend(module, example_inputs) if not has_triton_kernel(module): compiled_module NpuSerializableArtifact(compiled_module) return compiled_module代码中也是将config.mode设置为reduce-overhead,通过TorchAir的get_compiler接口获取NPU图编译器npu_backend并将模型torch.nn.Module)转换处理后的FX图通过npu_backend使能npugraph_ex的编译优化得到编译优化后的可执行图对象compiled_module。同时在不含有Triton算子时支持了编译产物的缓存功能将可执行图对象直接包装进NpuSerializableArtifact继承自SerializableArtifact不同编译器编译产物序列化的抽象基类class SerializableArtifact(ABC): def __init__(self, artifact): if isinstance(artifact, SerializableArtifact): return super().__init__() assert callable(artifact), fartifact must be callable, but got {type(artifact)} self._artifact artifact if getattr(artifact, _boxed_call, False): self._boxed_call True def __call__(self, *args, **kwargs): return self._artifact(*args, **kwargs) # NOTE(liuyuan): allow implicit no-conversion between subclasses of Serializable. def __new__(cls, artifact): if isinstance(artifact, SerializableArtifact): return artifact else: return super().__new__(cls) property def artifact(self): return self._artifact def __reduce__(self): return self.rebuild_from_bytes, (self.convert_to_bytes(),) abstractmethod def convert_to_bytes(self) - bytes: # TODO(liuyuan): For performance, try to make it as a zero copy byte strings. Convert artifact to bytes. The return value should be artifact_bytes, which can rebuild the artifact via rebuild_from_bytes. ... staticmethod abstractmethod def rebuild_from_bytes(byte_s: bytes): Rebuild artifact from bytes. The input is the artifact_bytes from convert_to_bytes. ...可以看到通过定义__call__方法继承SerializableArtifact的实例可以直接作为可调用对象执行编译优化后的图因此在ge_compiler中虽然根据是否存在Triton算子返回了不同的类型对象但是作为调用对象功能是一致的,不需要再做额外的区分处理。同时子类通过实现convert_to_bytes和rebuild_from_bytes完成编译产物的缓存和加载功能NpuSerializableArtifact中的实现如下class NpuSerializableArtifact(SerializableArtifact): def __init__(self, artifact): assert hasattr(artifact, dump_artifacts) super().__init__(artifact) def convert_to_bytes(self): # NOTE(liuyuan): Since tng_backend does not save any tenosr, would it be necessary? with temp_disable_tracing_envs(): return pickle.dumps(self._artifact.dump_artifacts()) staticmethod def rebuild_from_bytes(bytes): from torchair.npu_fx_compiler import _CompiledFxGraph # NOTE(liuyuan): Since tng_backend does not save any tenosr, would it be necessary? with temp_disable_tracing_envs(): return __class__(_CompiledFxGraph.load_artifacts(pickle.loads(bytes)))分别通过3.2中介绍的dump_artifact和load_artifacts方法简单实现了缓存和加载的抽象方法使得XPU_GRAPH既能使能npugraph_ex在NPU设备进行模型推理的编译优化也可以通过缓存功能加速冷启动时间。参考链接XPU_GRAPH仓链接https://github.com/XPU-Forces/xpu_graph/npugraph_ex接入PR链接https://github.com/XPU-Forces/xpu_graph/pull/4425.总结npugraph_ex提供了出色的模型推理加速功能在torch_npu和TorchAir开源的帮助下更进一步支持了第三方生态的接入XPU_GRAPH的成功集成便是其中的优秀典范。npugraph_ex也会持续地完善功能、优化性能并降低使用以及接入的难度开发者们可以关注TorchAir开源仓查看相关的最新技术动态。TorchAir仓链接https://gitcode.com/Ascend/torchair 推理的编译优化也可以通过缓存功能加速冷启动时间。参考链接XPU_GRAPH仓链接https://github.com/XPU-Forces/xpu_graph/npugraph_ex接入PR链接https://github.com/XPU-Forces/xpu_graph/pull/4425.总结npugraph_ex提供了出色的模型推理加速功能在torch_npu和TorchAir开源的帮助下更进一步支持了第三方生态的接入XPU_GRAPH的成功集成便是其中的优秀典范。npugraph_ex也会持续地完善功能、优化性能并降低使用以及接入的难度开发者们可以关注TorchAir开源仓查看相关的最新技术动态。TorchAir仓链接https://gitcode.com/Ascend/torchair【免费下载链接】cann-learning-hubCANN 学习中心仓支持在线互动运行、边学边练提供教程、示例与优化方案一站式助力昇腾开发者快速上手。项目地址: https://gitcode.com/cann/cann-learning-hub创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考