固定权重类输入地址功能

功能简介

推理场景下，Parameter类型（权重类）的图输入内存地址通常保持不变，可以开启本功能缩短图下发时间，提升下发性能。

对于PyTorch v2.6.0及以上版本，通过torch._dynamo.mark_static_address接口标记的内存地址不变的图输入Tensor（如LLM模型的kv_cache）也可以开启本功能。

该功能适用于ChatGPT、LLaMA等开源大模型，请根据自身实际情况开启。

使用约束

本功能仅适用于GE图模式场景。

使用方法

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

import torch_npu, torchair
config = torchair.CompilerConfig()
# 固定权重类输入地址开关
config.experimental_config.frozen_parameter = True
npu_backend = torchair.get_npu_backend(compiler_config=config)
opt_model = torch.compile(model, backend=npu_backend)

表 1 参数说明

参数名	说明
frozen_parameter	图执行时是否固定权重类输入地址。False（默认值）：不固定权重类输入地址。True：固定权重类输入地址。

特殊场景

PyTorch的to算子转换示例如下：

import torch
import torch_npu
import torchair

config = torchair.CompilerConfig()
config.experimental_config.frozen_parameter = True
npu_backend = torchair.get_npu_backend(compiler_config=config)

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x, y, z):
        return torch.add(x, y*z)
 
model = Model()
# 正确转换方式：先将Tensor转换为NPU Tensor，再转换为Parameter类型，转换后in1是Parameter类型
in1 = torch.nn.Parameter(torch.randn(4, 1).float().npu())
# 错误转换方式：先转换为Parameter类型，再将Tensor转换为NPU Tensor，转换后的in1不是Parameter类型
# in1 = torch.nn.Parameter(torch.randn(4, 1).float()).npu()
in2 = torch.randn(4, 4).float().npu()
in3 = torch.randn(4, 4).int().npu()
model = torch.compile(model, backend=npu_backend, dynamic=True)
graph_result = model(in1, in2, in3)