CANN/ops-nn RMS动态量化归一化

cann_ops_nn.rms_norm_dynamic_quant【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn产品支持情况Ascend 950PR/Ascend 950DT 不支持Atlas A3 训练系列产品/Atlas A3 推理系列产品 支持Atlas A2 训练系列产品/Atlas A2 推理系列产品 支持Atlas 200I/500 A2 推理产品 不支持Atlas 推理系列产品 不支持Atlas 训练系列产品 不支持Kirin X90 处理器系列产品 不支持Kirin 9030 处理器系列产品 不支持功能说明接口功能融合 RMS Normalization 与对称动态量化INT8封装 aclnnRmsNormDynamicQuant。将输入 x 先做 RMSNorm 归一化可选加 smooth_scales 缩放和 beta 偏置再按行做动态 INT8 量化输出 INT8 量化结果和 FLOAT32 的 scale 因子。计算公式$$ y \operatorname{RmsNorm}(x)\frac{x}{\operatorname{Rms}(\mathbf{x})}\cdot gamma beta, \quad \text { where } \operatorname{Rms}(\mathbf{x})\sqrt{\frac{1}{n} \sum_{i1}^n x_i^2epsilon} $$若smooth_scales不输入则直接对rmsnorm输出做量化$$ scaleOutrow_max(abs(y))/max_val $$$$ yOutround(y/scaleOut) $$若输入smooth_scales则先做smooth缩放再量化$$ input y\cdot smooth_scales $$$$ scaleOutrow_max(abs(input))/max_val $$$$ yOutround(input/scaleOut) $$其中row_max代表每行求最大值。max_val在INT8时为127。函数原型torch.ops.cann_ops_nn.rms_norm_dynamic_quant(x, gamma, smooth_scalesNone, betaNone, *, epsilon1e-6, dst_type2) - (Tensor y, Tensor scale)参数说明参数名参数类型可选/必选描述数据类型维度(shape)xTensor必选标准化过程中的源数据张量公式中的 x。最后一维须 ≤ 8192。FLOAT16、BFLOAT16(..., D)gammaTensor必选标准化过程中的权重张量公式中的 gamma。shape 需与 x 最后一维一致。FLOAT16、BFLOAT16(D,)smooth_scalesTensor可选量化过程中使用的 smoothScale 张量公式中的 smooth_scales。shape 和 dtype 需与 gamma 一致。FLOAT16、BFLOAT16(D,)betaTensor可选标准化过程中的偏置项公式中的 beta。shape 和 dtype 需与 gamma 一致。FLOAT16、BFLOAT16(D,)epsilonfloat可选用于防止除 0 错误公式中的 epsilon必须大于零。默认值 1e-6。FLOAT32-dst_typeint可选输出 y 的数据类型枚举值。当前仅支持 INT8 量化默认值 2。INT-返回值说明输出名输出类型可选/必选描述数据类型维度(shape)yTensor必选量化输出公式中的 yOutshape 与 x 一致INT8(..., D)scaleTensor必选量化 scale 因子公式中的 scaleOutshape 为 x 去掉最后一维FLOAT32(...,)约束说明输入 x 的最后一维必须小于等于 8192。输入 x 的每一行不能全为零否则归一化后仍全为零量化 scale 为 0 导致除法异常。确定性计算默认支持确定性计算。调用说明单算子模式调用eagerimport torch import torch_npu import cann_ops_nn rows, hidden 16, 64 x torch.randn(rows, hidden, dtypetorch.float16).npu() gamma torch.ones(hidden, dtypetorch.float16).npu() y, scale torch.ops.cann_ops_nn.rms_norm_dynamic_quant(x, gamma) print(y: , y) print(scale: , scale)带可选输入smooth torch.ones(hidden, dtypetorch.float16).npu() beta torch.zeros(hidden, dtypetorch.float16).npu() y, scale torch.ops.cann_ops_nn.rms_norm_dynamic_quant( x, gamma, smooth_scalessmooth, betabeta, epsilon1e-5)图模式torchair调用import torch, torch_npu, cann_ops_nn torch.compile(backendnpu) def func(x, gamma): return torch.ops.cann_ops_nn.rms_norm_dynamic_quant(x, gamma) x torch.randn(4, 64, dtypetorch.float16).npu() gamma torch.ones(64, dtypetorch.float16).npu() y, scale func(x, gamma) print(y: , y) print(scale: , scale)【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考