RmsNorm算子样例

概述

样例概述：本样例介绍了利用ATVOSS实现RmsNorm单算子并完成功能验证

• 算子功能：RmsNorm算子是大模型常用的归一化操作，相比LayerNorm算子，其去掉了减去均值的部分
• 调用方式：Kernel直调

样例支持的产品

• Ascend 950PR/Ascend 950DT

算子描述

• 算子数学计算公式：

  $$\mathrm{RmsNorm}(x_i)=\frac{x_i}{\mathrm{Rms}(\mathbf{x})}{g}_i,\text{ where }\mathrm{Rms}(\mathbf{x})=\sqrt{\frac{1}{n}\sum _{i=1}^n{x}_i^2}$$

• 算子规格：

参数名	输入/输出/属性	描述	数据类型	数据格式
in1	输入	表示进行归一化计算的输入。公式中的`x`。	float	ND
in2	输入	表示进行归一化计算的缩放因子（权重），公式中的`g`。	float	ND
out	输出	表示进行归一化后的最终输出，公式中的`RmsNorm(x)`。	float	ND

规格说明：

• 当前只支持二维输入
• 总的输入Shape(M, N)要满足：
  • M < 8160，N <= 7168
  • N需要32元素对齐
• Tile块的Shape(m, n)，要满足n = N，m * n <=7168
• 目前只支持float类型

目录结构

文件名	描述
rms_norm.cpp	RmsNorm样例算子代码实现
CMakeLists.txt	RmsNorm样例算子的编译构建文件
README.md	RmsNorm样例算子的说明文档

RmsNorm样例算子的编译和运行

• 编译 在代码仓根目录下执行：

bash scripts/build.sh -DSOC=ascend950 rms_norm

• 运行 在代码仓目录下执行：

output/bin/rms_norm --help // 查看帮助
output/bin/rms_norm --shape=16,32 // 运行样例