昇腾社区首页
中文
注册

简介

TorchAir(Torch Ascend Intermediate Representation)是为Ascend Extension for PyTorch(torch_npu)提供图模式能力的扩展库,支持用户使用PyTorch和torch_npu在昇腾设备上进行图模式的训练和推理。TorchAir对外提供昇腾设备的图模式编译后端,对接PyTorch的Dynamo特性,将PyTorch的FX(Functionalization)计算图转换为昇腾的GE(Graph Engine,图引擎)计算图,并提供在昇腾设备上启动GE计算图编译和执行的能力。

TorchAir继承了Dynamo的大部分特性,如动态shape图功能,在PyTorch的基础上又新增了如下功能:

  • 在当前版本中,TorchAir仅作为一项试用特性,此功能在后续版本中可能会有所调整或改进。请用户在使用过程中关注后续版本的迭代。
  • 对于experiment特性,一般适用于特定场景,如果用户希望进一步提升模型执行性能,可按需配置对应功能项。

使用约束

对于以下单算子场景,暂不支持图模式:

  • 单进程多线程多卡并发执行场景
  • 单进程单线程多卡执行场景