aclgraph间内存复用功能

功能简介

本功能为试验特性，后续版本可能存在变更，暂不支持应用于商用产品中。

reduce-overhead模式下，由于aclgraph本身不支持动态shape，因此执行过程中可能会存在多张aclgraph。

• 当torch.compile使用dynamic为False编译时，每一种shape将对应一张FX graph，同时也对应一张aclgraph。
• 当torch.compile使用dynamic为True编译时，多种shape只要满足规则都只会对应一张FX graph，但不同的shape会对应不同的aclgraph。

这可能导致多张aclgraph占用多份内存，从而在模型执行时引发 OOM（Out Of Memory）问题。为解决这一问题，TorchAir提供两种内存复用模式，可按实际需要选择。

• 模式一：用户可在config中显式指定要使用的内存池，后续使用该config的graph均使用该内存池，从而达成多个aclgraph内存复用的效果。
• 模式二（默认开启）：torch.compile编译过程中默认将一张FX graph里的多种shape捕获成多张aclgraph，实现同一张FX graph内多张aclgraph间的内存复用。如需关闭，需显式设置。

此外，由于aclgraph是基于固定内存地址执行，因此前次执行的输出Tensor内存会被后续执行覆盖。如需将输出结果长时间保存和使用，可能会引入精度问题。针对该问题，可使用Clone接口对指定输出结果做克隆。

使用约束

• 目前仅适用于如下产品：
  • Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件
  • Atlas A3 训练系列产品/Atlas A3 推理系列产品
• 本功能仅支持reduce-overhead模式，约束与reduce-overhead模式配置一致。
• 注意，开启模式一时，模式二功能会自动关闭。

使用方法

该功能通过torchair.get_npu_backend中compiler_config配置，示例如下，仅供参考不支持直接拷贝运行，参数介绍参见表1。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

import torch_npu, torchair
config = torchair.CompilerConfig()
# 配置图执行模式
config.mode = "reduce-overhead"
# 开启模式一
config.aclgraph_config.use_custom_pool = torch.npu.graph_pool_handle()
# 关闭模式二
config.debug.aclgraph.disable_mempool_reuse_in_same_fx = True
# 主动开启输出内存拷贝
config.debug.aclgraph.enable_output_clone = True
npu_backend = torchair.get_npu_backend(compiler_config=config)
opt_model = torch.compile(model, backend=npu_backend)