固定权重类输入地址功能

功能简介

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

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

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

使用约束

本功能支持的产品型号参见使用说明。

使用方法

该功能通过npugraph_ex的options配置，示例如下，仅供参考不支持直接拷贝运行，参数说明参见下表。

import torch
import torch_npu

opt_model = torch.compile(model, backend="npugraph_ex", options={"frozen_parameter": True}, dynamic=False, fullgraph=True)

表 1 参数说明

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

特殊场景

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

import torch
import torch_npu

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="npugraph_ex", options={"frozen_parameter": True}, fullgraph=True, dynamic=False)
graph_result = model(in1, in2, in3)