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。

• 调用示例

  import torch
  import torch_npu
  x = torch.randn(24, 1, 128).bfloat16().npu()
  w = torch.randn(128).bfloat16().npu()
  ​
  out1 = torch.npu_rms_norm(x, w, epsilon=1e-5)[0]
  print(out1)
  tensor([[[-0.1123,  0.3398,  0.0986,  ..., -2.1250, -0.8477, -0.3418]],
  ​
          [[-0.0591,  0.3184, -0.5000,  ...,  1.0312, -1.1719, -0.1621]],
  ​
          [[-0.1445,  0.3828, -0.3438,  ..., -0.9102, -0.5703,  0.0073]],
  ​
          ...,
  ​
          [[-0.1631, -0.3477,  0.4297,  ...,  0.9219,  0.1621,  0.3125]],
  ​
          [[-0.1387,  0.0815,  0.0967,  ...,  1.7109,  0.1455, -0.1406]],
  ​
          [[ 0.0698,  1.3438, -0.0127,  ..., -2.2656, -0.4473,  0.3281]]],
         device='npu:0', dtype=torch.bfloat16)