torch_npu.npu_dynamic_quant

功能描述

算子功能：为输入的张量进行pre-token对称动态量化。
计算公式：
- 如果smooth_scales不存在：
  $\text{[math]}$
- 如果smooth_scales存在：
  $\text{[math]}$

接口原型

torch_npu.npu_dynamic_quant(Tensor x, Tensor? smooth_scales=None) ->(Tensor, Tensor)

参数说明

x：Device侧的Tensor类型，需要进行量化的源数据张量，必选输入，数据类型支持FLOAT16、BFLOAT16，数据格式支持ND，支持非连续的Tensor。

smooth_scales：Device侧的Tensor类型，对x进行scales的张量，可选输入，数据类型支持FLOAT16、BFLOAT16，数据格式支持ND，支持非连续的Tensor。

输出说明

两个Tensor类型的输出：

第一个输出为Tensor，量化后的输出y，shape和输入x一致，数据类型支持INT8。
第二个输出为Tensor，用于后面反量化的scale，shape为x剔除最后一维，数据类型支持FLOAT32。

约束说明

输入x的维度必须大于1，针对Atlas A2 训练系列产品，x的最后一个维度大小不超过11264。
输入smooth_scales必须是关键字传参，即调用接口时需要显式地写smooth_scales=xxx。

支持的PyTorch版本

PyTorch 2.1

支持的型号

Atlas A2 训练系列产品

调用示例

只有一个输入x

import torch
import torch_npu
 
x = torch.rand((3, 3), dtype = torch.float16).to("npu")
output, scale = torch_npu.npu_dynamic_quant(x)
print(output)
print(scale)

使用smooth_scales输入

import torch
import torch_npu
 
x = torch.rand((3, 3), dtype = torch.float16).to("npu")
smooth_scales = torch.rand((3,), dtype = torch.float16).to("npu")
output, scale = torch_npu.npu_dynamic_quant(x, smooth_scales=smooth_scales)
print(output)
print(scale)

父主题： torch_npu