功能

RMS归一化处理。

图1 NORM

约束

所有输入输出Tensor的最后一维大小相等并且小于或等于8192。

定义

struct RmsNormParam {
    enum RmsNormType : int {
        RMS_NORM_UNDEFINED = 0,
        RMS_NORM_NORM,
        RMS_NORM_PRENORM,
        RMS_NORM_POSTNORM,
    };
    struct NormParam {
        QuantType quantType = QUANT_UNDEINFED;
        float epsilon = 1e-5;
        double layerNormEps = 1e-5;
        bool rstd = false;
    };
    struct PreNormParam {
        QuantType quantType = QUANT_UNDEINFED;
        float epsilon = 1e-5;
        bool hasBias = false;
    };
    struct PostNormParam {
        QuantType quantType = QUANT_UNDEINFED;
        float epsilon = 1e-5;
        bool hasBias = false;
    };
    RmsNormType layerType = RMS_NORM_UNDEFINED;
    NormParam normParam;
    PreNormParam preNormParam;
    PostNormParam postNormParam;
};

成员

成员名称	描述
layerType	归一化类型，参数如下： RMS_NORM_NORM RMS_NORM_PRENORM RMS_NORM_POSTNORM。
normParam	NORM参数。
preNormParam	PRENORM参数。
postNormParam	POSTNORM参数

表1 NormParam成员
参数	描述
quantType	量化类型。
epsilon	Epsilon，归一化时加在分母上防止除零。
layerNormEps	Epsilon，默认为1e-5，暂时不使用。
rstd	默认为False，设置为true时会使用训练的rmsnormforward算子。

表2 PreNormParam成员
参数	描述
quantType	量化类型。
epsilon	Epsilon，归一化时加在分母上防止除零。
hasBias	是否叠加偏置。默认为False，当需要输入beta时设置为True。

表3 PostNormParam成员
参数	描述
quantType	量化类型。
epsilon	Epsilon，归一化时加在分母上防止除零。

输入

表4 RMS_NORM_NORM输入
参数	维度	数据类型	格式	描述
x	[-1,…,-1,n] -1表示当前维度的大小没有约束。	float16/float/bf16	ND	当quantType不为QUANT_UNDEINFED时，仅支持数据类型float16。仅当rstd为true时，支持数据类型float。
gamma	[1,...,1,n] / [1,n]	float16/float/bf16	ND	最后一维的大小需进行32字节对齐。 quantType = INT8时, 维度为[1,n]
beta	[1,n]	float16/float/bf16	ND	quantType = INT8时输入。最后一维的大小要32字节对齐。
scale	[1]	float16	ND	当quantType为QUANT_INT8时才输入
offset	[1]	int8	ND	当quantType为QUANT_INT8时才输入
注1：所有Tensor最后一维的n大小需保持一致。注2：-1表示当前维度的大小没有约束。注3：仅Atlas 800I A2推理产品支持输入为bf16数据类型。

表5 RMS_NORM_PRENORM（quantType = QUANT_UNDEINFED）输入
参数	维度	数据类型	格式	描述
x	[-1,…,-1,n] -1表示当前维度的大小没有约束。	float16	ND	最后一维的大小要32字节对齐。
beta	[1,n]	float16	ND	当hasBias为True时输入最后一维的大小要32字节对齐。
residual	[-1,…,-1,n] -1表示当前维度的大小没有约束。	float16	ND	最后一维的大小要32字节对齐。
gamma	[1,n]	float16	ND	最后一维的大小要32字节对齐。
注1：所有Tensor最后一维的n大小需保持一致。

表6 RMS_NORM_PRENORM（quantType = QUANT_INT8）输入
参数	维度	数据类型	格式	描述
x	[-1,…,-1,n] -1表示当前维度的大小没有约束。	float16	ND	最后一维的大小要32字节对齐。
residual	[-1,…,-1,n] -1表示当前维度的大小没有约束。	float16	ND	最后一维的大小要32字节对齐。
gamma	[1,n]	float16	ND	最后一维的大小要32字节对齐。
beta	[1,n]	float16	ND	最后一维的大小要32字节对齐。
注1：所有Tensor最后一维的n大小需保持一致，n ≤ 8192。

表7 RMS_NORM_POSTNORM输入
参数	维度	数据类型	格式	描述
x	[-1,…,-1,n] -1表示当前维度的大小没有约束。	float16	ND	当前支持quantType = QUANT_UNDEINFED。最后一维的大小要32字节对齐。
beta	[1,n]	float16	ND	当hasBias为True时输入。最后一维的大小要32字节对齐。
residual	[-1,…,-1,n] -1表示当前维度的大小没有约束。	float16	ND	最后一维的大小要32字节对齐。
gamma	[1,n]	float16	ND	最后一维的大小要32字节对齐。
注1：所有Tensor最后一维的n大小需保持一致。

输出

表8 RMS_NORM_NORM输出
参数	维度	数据类型	格式	描述
result	[-1,…,-1,n] -1表示当前维度的大小没有约束。	float16/float/int8/bf16	ND	quantType=INT8时输出int8。
rstd	[-1,…,-1,n] -1表示当前维度的大小没有约束。	float16/float/bf16	ND	-

表9 RMS_NORM_PRENORM输出
参数	维度	数据类型	格式	描述
output	[-1,…,-1,n] -1表示当前维度的大小没有约束。	float16	ND	-
resOut	[-1,…,-1,n] -1表示当前维度的大小没有约束。	float16/int8	ND	quantType=INT8时输出int8。

**表10** RMS_NORM_POSTNORM输出
参数	维度	数据类型	格式	描述
output	[-1,…,-1,n] -1表示当前维度的大小没有约束。	float16	ND	-

RmsNormOperation

功能

约束

定义

成员

输入

输出