（beta）torch_npu.npu_rms_norm

接口原型

torch_npu.npu_rms_norm(Tensor self, Tensor gamma, float epsilon=1e-06) -> (Tensor, Tensor) 

功能描述

RmsNorm算子是大模型常用的归一化操作，相比LayerNorm算子，其去掉了减去均值的部分 。其计算公式为

参数说明

• self：Tensor类型，支持float16、bfloat16、float32，输入shape支持2-8维。
• gamma：Tensor类型，数据类型需要和self保持一致，输入shape支持2-8维，通常为self的最后一维。
• epsilon：float数据类型，用于防止除0错误。

输出说明

共两个输出，格式为： (Tensor, Tensor)

第1个输出为Tensor，计算公式的最终输出y；

第2个输出为Tensor，rms_norm的reverse rms中间结果，用于反向计算。

约束说明

输入数据类型仅支持float16、bfloat16和float32。

调用示例

>>> x = torch.randn(24, 1, 128).bfloat16().npu()
>>> w = torch.randn(128).bfloat16().npu()
>>> out1 = torch_npu.npu_rms_norm(x, w, epsilon=1e-5)[0]
>>> out1
tensor([[[ 0.2852,  0.2988, -0.6055,  ..., -0.3730,  1.1484,  0.3613]],

        [[-0.3516,  1.6641, -0.2344,  ..., -0.3555,  1.2891, -0.2168]],

        [[-0.0041, -2.2188,  1.7031,  ...,  0.4141, -0.6992,  0.1621]],

        ...,

        [[ 0.3633,  0.9336, -0.1855,  ..., -0.0549, -0.6797,  0.2988]],

        [[ 0.1562, -0.5156,  1.1172,  ...,  0.5625, -0.5742, -0.8203]],

        [[-0.0444, -1.7578, -0.6875,  ..., -0.8281, -0.4902,  0.1406]]],
       device='npu:0', dtype=torch.bfloat16)