1 Star 0 Fork 152

陈睿敏 / torchair_5

forked from Ascend / torchair 
加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
克隆/下载
贡献代码
同步代码
取消
提示: 由于 Git 不支持空文件夾,创建文件夹后会生成空的 .keep 文件
Loading...
README
BSD-3-Clause

torchair实现将torch的FX图转换为GE计算图,并提供了GE计算图的编译与执行接口。

编译准备

配置编译环境,执行配置命令

配置只需要进行一次,用于获取pytorch的编译选项(如当前的torch是否开启ABI)及Ascend sdk的安装路径(如果需要在本地CPU上进行调试)。

bash ./configure

默认情况下,执行上述命会弹出如下的交互式会话窗口

您的会话可能有所不同。

Please specify the location of python with available torch 2.1.x installed. [Default is /usr/bin/python3]
(You can make this quiet by set env [TARGET_PYTHON_PATH]):

此时,要求您输入安装了 Torch 2.1 版本的python解释器路径,如果默认路径是正确的,直接回车,否则请输入正确的 python 解释器路径。

您可以通过设置 TARGET_PYTHON_PATH 环境变量,来抑制交互式窗口弹出,但是要确保路径是有效的,否则,仍然会要求您输入正确的 python 解释器路径。

键入后,会耗费几秒钟以确保您的输入是有效的,接着,会弹出下面的交互式窗口

Specify the location of ascend sdk for debug on localhost or leave empty.
(You can make this quiet by set env [ASCEND_SDK_PATH]):

如果您不需要在本地CPU调试,可以直接回车跳过。否则,需要输入昇腾处理器开发套件的安装路径(需指定至opensdk/opensdk目录)。

您可以通过设置 ASCEND_SDK_PATH 环境变量指定 SDK 目录 或 设置 NO_ASCEND_SDK 环境变量指定不需要 SDK 来抑制交互式窗口弹出。

键入后,等待配置完成。

编译与安装

mkdir build
cd build
cmake ..
make torchair -j8

编译完成后,会在build/dist/dist/目录下生成名为torchair-{version}-py3-none-any.whl的安装包文件。

您可以直接使用pip安装该安装包,或者使用make命令安装至您configure时指定的python环境中。

make install_torchair

执行测试

如果您在配置时未指定Ascend sdk的安装路径,则无法执行CPU上的调试,需要在NPU环境上进行测试。

CPU调试时,需要设置LD_LIBRARY_PATH到生成的fake so文件目录以及sdk目录

tools/env.sh会根据配置生成对应的LD_LIBRARY_PATH(如果您在配置时指定了Ascend sdk安装路径)

source tools/env.sh

python3 examples/example.py

关于CI编译

工程目录下的build.sh用于对接CI编译,但是您也可以通过其在本地执行编译和UT/ST测试。

查看帮助

./build.sh -h

编译安装包

./build.sh -c

编译完成后,会在output/目录下生成名为torchair-{version}-py3-none-any.whl的安装包文件。

执行UT测试

本地执行UT测试需要设置环境变量ASCEND_CUSTOM_PATH,将其指定至Ascend sdk的安装路径(指定至ai_cann_x86目录)

./build.sh -u

执行ST测试

本地执行ST测试需要设置环境变量ASCEND_CUSTOM_PATH,将其指定至Ascend sdk的安装路径(指定至ai_cann_x86目录)

./build.sh -s

查看覆盖率报告

UT或ST执行通过后,会在coverage目录下生成覆盖率文件coverage.info,如果您想要查看覆盖率报告,可以执行如下命令

genhtml coverage.info -o coverage_report

基于torchair调试pytorch模型

在合适的位置添加torch.compile,并指定npu后端

import torchair
# torchair提供了一些额外的编译选项,您可以在此处指定,也可以设置为None使用默认选项
npu_backend = torchair.get_npu_backend(compiler_config=None)
torch.compile(model, backend=npu_backend)

注意,如果您基于torch_npu执行,可以使用torch_npu提供的"npu"后端。

torch_npu目前正在集成torchair作为其torch.compile后端。

torch.compile(model, backend="npu")

确定需要补齐的converter

torchair提供配置项config.debug.fx_summary开关来确定FX图中涉及需要补齐的converter,您可以通过如下方式来配置

config = torchair.CompilerConfig()
config.debug.fx_summary.type = "csv"
npu_backend = torchair.get_npu_backend(compiler_config=config)

torch.compile(model, backend=npu_backend)

执行后,会生成summary_{timestamp}.csv文件,您可以通过excel等工具来查看。 通过导出的csv文件,您可以看打当前模型fx图中涉及的所有converter

  • 对于支持状态未支持的converter,您需要在ge_concrete_graph/ge_converter目录下对应文件中补齐实现

当前已经为aten下的op提供了一个壳子实现(固定抛出未支持的异常),您应当在此基础上补齐实现

  • 对于支持状态未注册的converter,您需要在ge_concrete_graph/ge_converter目录下对应文件中新增注册并补齐实现

  • 对于支持状态部分支持的converter,您需要查看位于ge_concrete_graph/ge_converter目录下的实现,并根据输入数据列决定是否需要补齐场景实现

实现converter

您可以参考ge_concrete_graph/ge_converter目录下的实现,实现对应的converter。 我们以torch.ops.aten.add.Tensor的converter实现为例说明converter实现时的一些细节:

@declare_supported([
    Support(F32(2, 2), F32(2, 2)),
    Support(F32(2, 2), F32(2, 1)),
    Support(F32(2, 2), F16(2, 1)),
    Support(F32(2, 2), F16(2, 2), alpha=2),
    Support(F32(2, 2), 2.0),
    Support(F32(2, 2), 2),
    Support(F32(2, 2), 2, alpha=2.0),
])
@register_fx_node_ge_converter(torch.ops.aten.add.Tensor)
def conveter_aten_add_Tensor(
        self: Tensor,
        other: Tensor,
        *,
        alpha: Union[Number, Tensor] = 1,
        meta_outputs: Union[TensorSpec, List[TensorSpec]] = None):
    """ NB: aten::add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor """
    if not isinstance(alpha, Tensor) and alpha == 1:
        # just for better permance
        self, other = dtype_promote(self, other, target_dtype = meta_outputs.dtype)
        return ge.Add(self, other)
    else:
        self, other, alpha = dtype_promote(self, other, alpha, target_dtype = meta_outputs.dtype)
        return ge.AxpyV2(self, other, alpha)

声明converter支持的场景

您应该在开头声明converter需要支持的全部场景,支持场景应该穷举aten.add.Tensor所支持的全部传参方式, 需要注意,不应当有不支持的场景,如果您发现有无法支持的传参方式,需要在converter实现中抛出异常。

@declare_supported([
    Support(F32(2, 2), F32(2, 2)), # 支持基础的torch.add()
    Support(F32(2, 2), F32(2, 1)), # 支持f32类型间的广播
    Support(F32(2, 2), F16(2, 1)), # 支持f32与f16类型的广播
    Support(F32(2, 2), F16(2, 2), alpha=2), # 支持带alpha入参场景
    Support(F32(2, 2), 2.0), # 支持与浮点常量的加法
    Support(F32(2, 2), 2), # 支持f32类型与整型常量的加法
    Support(F32(2, 2), 2, alpha=2.0), # 支持带alpha入参场景
])

当您实现了您的converter后,我们会根据您声明支持的场景,生成对应的测试用例,您可以通过如下方式来测试您的converter是否正确:

需要在真实NPU环境测试,并确保已经实现了对应的converter及正确安装了torchair

python3 smoke/converter_test.py

可以通过入参控制只测试满足某个prefix的converter,如下所示:

python3 smoke/converter_test.py aten.add.Tensor

converter的函数签名

from torchair.ge_concrete_graph.ge_graph import Tensor, TensorSpec

@register_fx_node_ge_converter(torch.ops.aten.add.Tensor)
def conveter_aten_add_Tensor(
        self: Tensor,
        other: Tensor,
        *,
        alpha: Union[Number, Tensor] = 1,
        meta_outputs: Union[TensorSpec, List[TensorSpec]] = None):
    """ NB: aten::add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor """

我们逐行解释下上述代码片段:

from torchair.ge_concrete_graph.ge_graph import Tensor, TensorSpec

文件的开头从GE graph的文件中导入了Tensor和TensorSpec,这两个类分别表示GE图上的Tensor和TensorSpec。 需要特别注意的是,Tensor和TensorSpec绝不是运行时的真实数据,你只能从Tensor上获得dtype和rank信息,而不能从Tensor上获得shape和数据。

@register_fx_node_ge_converter(torch.ops.aten.add.Tensor)

表示要为torch.ops.aten.add.Tensor注册converter函数,实现将fx图上的aten.add.Tensor节点,转换为GE图上的节点。对应的实现函数则是conveter_aten_add_Tensor。

self: Tensor,
other: Tensor,

表示fx图上的aten.add.Tensor节点的两个输入self和other对应的GE图输入,这两个输入都是GE图上的Tensor类型。

*,

这是python3的语法,表示后面的参数都是关键字参数,即必须使用参数名来传参。

alpha: Union[Number, Tensor] = 1,

表示fx图上的aten.add.Tensor节点的alpha入参,其类型为Number(字面值)或者GE图上的Tensor,其默认值为1。

meta_outputs: Union[TensorSpec, List[TensorSpec]] = None):

可以发现conveter_aten_add_Tensor函数的入参与torch.ops.aten.add.Tensor几乎完全一致,但是多一个名为meta_outputs的入参外完全一致。 meta_outputs是GE graph下定义的TensorSpecList[TensorSpec](aten节点有动态输出时)类型,TensorSpec是对一个节点输出的描述信息。meta_outputs是在指示converter,最终你应该输入一个什么样的Tensor。meta_outputs由原始fx节点输出的aten.Tensor转换而来,包含dtype和rank信息。 什么时候应该使用meta_outputs呢,典型的场景包括:

  • 通过meta_outputs可以确定输出数量,有些aten的节点的输出数量是不确定的,比如aten.split,会根据输入shape的不同得到不同的输出数量。
  • 确定输出的dtype,用于类型提升,比如aten.add,会根据输入的dtype的不同得到不同的输出dtype,这时候通过meta_outputs上的dtype可以精确地确定应该把输入提升成何种类型。

实现converter

alpha的输入可能是Tensor或者Number,我们需要对其进行判断,如果是Number且为1,则可以直接使用ge.Add来实现,否则需要使用ge.AxpyV2来实现,下面的代码片段展示了aten.add.Tensor的converter实现:

@register_fx_node_ge_converter(torch.ops.aten.add.Tensor)
def conveter_aten_add_Tensor(
        self: Tensor,
        other: Tensor,
        *,
        alpha: Union[Number, Tensor] = 1,
        meta_outputs: Union[TensorSpec, List[TensorSpec]] = None):
    """ NB: aten::add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor """
    if not isinstance(alpha, Tensor) and alpha == 1:
        # just for better permance
        self, other = dtype_promote(self, other, target_dtype = meta_outputs.dtype)
        return ge.Add(self, other)
    else:
        self, other, alpha = dtype_promote(self, other, alpha, target_dtype = meta_outputs.dtype)
        return ge.AxpyV2(self, other, alpha)

主要关注其中的dtype_promote(类型提升)函数,这个函数接收任意个输入,并返回等量的输出。其作用是将传入的多个输出,提升为target_dtype指定的类型。 例如测试用例TestInput(F32(2, 2), 2),将一个f32类型与一个int类型的常量相加,如果不执行dtype_promote,最终生成的图上,会出现一个ge::Add节点, 其两个输入分别为F32和INT64类型,当前这种图无法编译通过(图编译不支持类型提升)。 dtype_promote(类型提升)的关键,在于converter实现时,需要根据算子语义,将需要保证类型一致的GE图节点输入进行类型提升,以保证最终生成的GE图可以编译通过。

注意事项

需要特别注意,所实现的converter必须支持动态shape,不应该试图从输入的Tensor上获取任何shape信息,Tensor也不会提供任何shape信息。

如果您的converter依赖shape才能工作,这通常意味着实现错误,或者没有选择正确的Ascend IR映射。

导出gegraph

torchair提供配置项config.debug.graph_dump开关来导出gegraph,您可以通过如下方式来配置

config = torchair.CompilerConfig()
config.debug.graph_dump.type = 'txt' # ['txt', 'pbtxt', 'py']
npu_backend = torchair.get_npu_backend(compiler_config=config)

model = torch.compile(model, backend=npu_backend)

执行后,会生成dynamo_{timestamp}.{graph_dump.type}文件,当前支持 ['txt', 'pbtxt', 'py'] 三种导出方式

  • 导出的txt文件是cann最终接收到的torchair的构图结果,为protobuf格式,您可以通过vscode等查看
  • 导出的pbtxt文件是可以被tensorboard读取的构图结果,您可以通过tensorboard等查看
  • 导出的py文件是torch代码经由converter转化后的GEIR代码,支持运行,您可以通过vscode等查看
BSD 3-Clause License Copyright (c) 2023, Huawei Technologies Co., Ltd All rights reserved. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: * Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. * Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

简介

暂无描述 展开 收起
Python
BSD-3-Clause
取消

发行版

暂无发行版

贡献者

全部

近期动态

加载更多
不能加载更多了
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
Python
1
https://gitee.com/rmch/torchair_5.git
git@gitee.com:rmch/torchair_5.git
rmch
torchair_5
torchair_5
master

搜索帮助

344bd9b3 5694891 D2dac590 5694891