整体说明
接口概述
TorchAir是为Ascend Extension for PyTorch提供图模式能力的扩展库，支持用户使用PyTorch在昇腾设备上进行图模式训练和推理。
本章详细介绍了TorchAir提供给外层使用的Python API接口，包括在昇腾设备上使能、调测图模式的相关接口，以及图模式下使用自定义算子涉及的相关接口说明。
接口总览
接口分类 |
接口名 |
接口说明 |
所属Python模块 |
|---|---|---|---|
基本接口 |
用于导出由TorchAir生成的离线图（air格式）。 |
torchair |
|
获取能够在NPU上运行的图编译器。 |
|||
获取能够在NPU上运行的图编译后端npu_backend，可作为backend参数传入torch.compile。 |
|||
将模型中的权重weight转成TorchAir定义的内部私有格式。 |
|||
将自定义算子注册到TorchAir框架中。 |
|||
针对PyTorch 2.1版本中不支持入图的集合通信算子提供的补丁函数，实现部分集合通信算子入图。 |
|||
GE相关接口 |
算子converter中的构图元素，表示一个Const节点，即图中的常量值。 |
torchair.ge |
|
算子converter中的构图元素，表示一个Cast节点，即图中Tensor的类型转换方法。 |
|||
基于算子原型（IR）实现算子conveter函数，完成PyTorch IR与GE IR的转换，方便自定义算子入图。 |
|||
inference接口 |
开启模型缓存编译功能时需要调用该接口实现模型缓存编译。 |
torchair.inference |
|
开启模型缓存编译功能时需要调用该接口读取封装的func函数缓存文件compiled_module，并以可读文件（格式不限，如py、txt）格式呈现。 |
|||
开启动态shape图分档执行功能时需要调用该接口设置图被划分的档位。 |
|||
ops接口 |
在图执行过程中，打印执行脚本中目标tensor值。 |
torchair.ops |
|
LLM-DataDist接口 |
通过一串Device地址创建PyTorch在NPU上的tensors。 主要使用场景是用于创建大模型中的kv cache tensors，所有kv cache的shape和dtype都一致。 |
torchair.llm_datadist |
