RmsNorm算子使用指南

表1 RmsNorm算子基础信息

算子名称

RmsNorm

torch_npu API接口

torch_npu.npu_rms_norm(x, gamma, epsilon)[0]

支持的torch_npu版本

1.11, 2.0, 2.1

支持的昇腾产品

Atlas 推理系列产品Atlas A2 训练系列产品

支持的数据类型

float16,bfloat16,float

算子IR及torch_npu接口参数

模型中替换代码及算子计算逻辑

算子替换的模型中小算子

使用限制

Atlas A2 训练系列产品支持全泛化case,Atlas 推理系列产品当前仅支持gamma shape 大于等于32byte。

已支持模型典型Case