动态shape图分档执行功能

功能简介

当用户网络中仅Batch的少量维度发生了变化，且维度变化的范围是有限的、可被枚举的，可以开启动态Shape图分档功能。通过将一张图划分为不同档位的方式，使一张图能支持多Batch，提升网络执行性能。

此处以图1为例，阐述三种图执行方式的差异。假设用户图只有Add算子，存在input1、input2、output三个tensor。编译该图时传入的input1与input2的shape都为(2, 2)，而用户实际输入时input1与input2的shape有(2, 2)、(2, 2)与(4, 2)、(4, 2)两种。

图1 Add算子示意图

使用约束

使用方法

以图1为例，动态分档的关键操作如下：

  1. 在模型脚本中通过set_dim_gears接口设置档位。

    • 支持对同一个Tensor设置一个或者多个维度的档位。
    • 若图编译、执行时Tensor的shape不在设置的档位中，会导致编译或执行报错，请合理设置档位值。
    • 不支持对同一个tensor使用该接口设置两次不一样的档位。
    • 生成的总档位数量要不大于100。
    • 首次执行时对输入tensor设置档位即可，保证首次编译时能够获取到档位，后续执行时无需设置档位，避免因为设置档位的动作造成性能下降。
    1
2
3
4
5
    import torch, torch_npu, torchair
input1 = torch.ones(2, 2).npu()
input2 = torch.ones(2, 2).npu()
torchair.inference.set_dim_gears(input1, dim_gears={0:[2, 4]})
torchair.inference.set_dim_gears(input2, dim_gears={0:[2, 4]})

  2. （可选）在模型脚本中设置动态分档的组合模式。

    通过torchair.get_npu_backend中compiler_config配置，示例如下，参数介绍参见表1。

    1
2
3
4
5
6
    import torch_npu, torchair
config = torchair.CompilerConfig()
# 动态分档组合模式配置
config.inference_config.dynamic_gears_merge_policy = "zip"
npu_backend = torchair.get_npu_backend(compiler_config=config)
opt_model = torch.compile(model, backend=npu_backend)
    表1 参数说明

    参数名

    参数说明

    dynamic_gears_merge_policy

    指定动态分档组合模式。若采用zip模式配置档位繁琐时，可使用product模式配置。

    • zip（缺省值）：位置一一对应。
    • product：排列组合，可以组合成更多的档位。
    说明：

    举个例子，当input3输入的shape为(2, 10)，若想要对其划分6个档位(2, 10)、(3, 10)、(4, 10)、(2, 20)、(3, 20)、(4, 20)，两种组合模式的配置代码如下：

    • zip模式，档位配置复杂
      1
      torchair.inference.set_dim_gears(input3, dim_gears={0:[2, 3, 4, 2, 3, 4], 1:[10, 10, 10, 20, 20, 20]})
    • product，档位配置简单
      1
      torchair.inference.set_dim_gears(input3, dim_gears={0:[2, 3, 4], 1:[10, 20]})

    按步骤1中配置，不同组合模式的档位结果如下：

    • zip：input1、input2档位划分结果为[2, 2]、[4, 4]，即input1、input2的shape支持(2, 2)、(2, 2)与(4, 2)、(4, 2)两种档位。
    • product：input1、input2档位划分结果为[2, 2]、[2, 4]、[4, 2]、[4, 4]，即input1、input2的shape支持(2, 2)、(2, 2)；(2, 2)、(4, 2)；(4, 2)、(2, 2)；(4, 2)、(4, 2)四种档位。其中input1、input2的shape为(2, 2)、(4, 2)与(4, 2)、(2, 2)两种档位在图1中不符合相同shape相加的规则。

后续操作

在模型脚本中设置动态分档后，再次执行模型脚本，验证分档结果是否符合预期。

假设模型脚本定义如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
import torch
import torch_npu
import torchair
from torchair.configs.compiler_config import CompilerConfig

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x1, x2):
        return x1 + x2

config = CompilerConfig()
# 缺省为zip模式，也支持设置product模式
# config.inference_config.dynamic_gears_merge_policy = "zip"
npu_backend = torchair.get_npu_backend(compiler_config=config)
model = Model().npu()
# 必须整图编译
model = torch.compile(model, fullgraph=True, backend=npu_backend)
npu_input0 = torch.ones([2, 2]).npu()
npu_input1 = torch.ones([2, 2]).npu()
# 设置档位
torchair.inference.set_dim_gears(npu_input0, {0: [2, 4]})
torchair.inference.set_dim_gears(npu_input1, {0: [2, 4]})
# 首次编译+执行，shape为（2，2）、（2，2）
print(model(npu_input0, npu_input1))

# 再次执行，shape为（4，2）、（4, 2）在档位中，不会触发重新编译
npu_input0 = torch.ones([4, 2]).npu()
npu_input1 = torch.ones([4, 2]).npu()
print(model(npu_input0, npu_input1))

开启TorchAir C++层日志，会显示图被划分的档位信息：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
[INFO] TORCHAIR [concrete_graph/concrete_graph.cpp:123] Compiling concrete graph 0 with options:
[INFO] TORCHAIR [concrete_graph/concrete_graph.cpp:125]     ge.deterministic = 0
[INFO] TORCHAIR [concrete_graph/concrete_graph.cpp:125]     ge.dynamicDims = 2,2;4,4
[INFO] TORCHAIR [concrete_graph/concrete_graph.cpp:125]     ge.dynamicNodeType = 1
[INFO] TORCHAIR [concrete_graph/concrete_graph.cpp:125]     ge.exec.atomicCleanPolicy = 1
[INFO] TORCHAIR [concrete_graph/concrete_graph.cpp:125]     ge.exec.memoryOptimizationPolicy = MemoryPriority
[INFO] TORCHAIR [concrete_graph/concrete_graph.cpp:125]     ge.exec.outputReuseMemIndexes = 0
[INFO] TORCHAIR [concrete_graph/concrete_graph.cpp:125]     ge.exec.reuseZeroCopyMemory = 1
[INFO] TORCHAIR [concrete_graph/concrete_graph.cpp:125]     ge.featureBaseRefreshable = 0
[INFO] TORCHAIR [concrete_graph/concrete_graph.cpp:125]     ge.inputShape = arg1_1:-1,2;arg2_1:-1,2
[INFO] TORCHAIR [concrete_graph/concrete_graph.cpp:125]     ge.jit_compile = 2
[INFO] TORCHAIR [concrete_graph/concrete_graph.cpp:125]     ge.topoSortingMode = 1

对应支持的档位如下，符合用户预期的档位。

父主题： 更多功能