RmsNormOperation

功能

RMS归一化处理。

图1 NORM

约束

所有输入输出Tensor的最后一维大小相等。

定义

struct RmsNormParam {
    enum RmsNormType : int {
        RMS_NORM_UNDEFINED = 0,
        RMS_NORM_NORM,
        RMS_NORM_PRENORM,
        RMS_NORM_POSTNORM,
    };
    enum PrecisionMode : int {
        HIGH_PRECISION_MODE = 0,
        HIGH_PERFORMANCE_MODE,
    }
    enum ModelType : int {
        LLAMA_MODEL = 0,  
        GEMMA_MODEL,   
    };
    struct NormParam {
        QuantType quantType = QUANT_UNDEFINED;
        float epsilon = 1e-5;
        double layerNormEps = 1e-5;
        bool rstd = false;
        PrecisionMode precisionMode = HIGH_PRECISION_MODE;
        ModelType modelType = LLAMA_MODEL;
        DynamicQuantType dynamicQuantType = DYNAMIC_QUANT_UNDEFINED;
    };
    struct PreNormParam {
        QuantType quantType = QUANT_UNDEFINED;
        float epsilon = 1e-5;
        bool hasBias = false;
    };
    struct PostNormParam {
        QuantType quantType = QUANT_UNDEFINED;
        float epsilon = 1e-5;
        bool hasBias = false;
    };
    RmsNormType layerType = RMS_NORM_UNDEFINED;
    NormParam normParam;
    PreNormParam preNormParam;
    PostNormParam postNormParam;
};

成员

成员名称	描述
rmsNormParam	RMSNORM参数。 RMS_NORM_UNDEFINED：默认值，未定义。 RMS_NORM_NORM：NORM参数。 RMS_NORM_PRENORM ：PRENORM参数。 RMS_NORM_POSTNORM：POSTNORM参数。
layerType	归一化类型，参数如下： RMS_NORM_UNDEFINED：默认值，未定义。 RMS_NORM_NORM：NORM参数。 RMS_NORM_PRENORM：PRENORM参数。 RMS_NORM_POSTNORM：POSTNORM参数。
normParam	NORM参数。
preNormParam	PRENORM参数。
postNormParam	POSTNORM参数。

表1 NormParam成员
参数	描述
quantType	量化类型。当前支持以下类型。 QUANT_UNDEFINED QUANT_INT8
epsilon	Epsilon，归一化时加在分母上防止除零。
layerNormEps	Epsilon，默认为1e-5，暂时不使用。
rstd	默认为False，设置为true时会使用训练的rmsnormforward算子。仅Atlas 800I A2推理产品支持该算子。不支持和“precisionMode”，“modelType”同时设置。量化场景下不支持使用“rstd”。
precisionMode	默认为HIGH_PRECISION_MODE。支持参数如下： HIGH_PRECISION_MODE：默认值，中间计算使用fp32类型。 HIGH_PERFORMANCE_MODE：中间计算使用fp16类型。不支持和“rstd”，“modelType”同时设置。量化场景下不支持使用“precisionMode”，该场景下配置该参数将返回报错ERROR_INVALID_PARAM。
modelType	默认为LLAMA_MODEL，设置为GEMMA_MODEL时使用gemma模型的rmsnorm计算公式。支持参数如下： LLAMA_MODEL：默认值， Llama的rms norm计算公式。 GEMMA_MODEL：Gemma的rms norm计算公式。不支持和“rstd”，“precisionMode”同时启用。量化场景下不支持使用“modelType”，该场景下配置该参数将返回报错ERROR_INVALID_PARAM。
dynamicQuantType	动态量化类型。默认为DYNAMIC_QUANT_UNDEFINED非动态量化。当前版本暂不支持非对称动态量化。

8.0.RC2及后续版本normParam不再支持quantInputScale、quantInputOffset量化参数。

表2 PreNormParam成员
参数	描述
quantType	量化类型。当前支持以下类型。 QUANT_UNDEFINED QUANT_INT8
epsilon	Epsilon，归一化时加在分母上防止除零。
hasBias	是否叠加偏置。默认为False，当需要输入beta时设置为True。量化场景下不支持使用“hasBias”，该场景下配置该参数将返回报错ERROR_INVALID_PARAM。

8.0.RC2及后续版本PreNormParam不再支持quantInputScale、quantInputOffset量化参数。

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

RMS_NORM_NORM输入输出描述

参数	维度	数据类型	格式	描述
x	[d_0,…,d_n-1,d_n]	float16/float/bf16	ND	float数据类型仅在rstd为true时支持，且rstd为true时只能使用float数据类型。所有Tensor最后一维的d_n大小需保持一致，其他维度的大小没有约束。仅Atlas 800I A2推理产品支持输入为bf16数据类型。
gamma	[1,d_n] / [d_i,...,d_n] (0 <= i <= n)	float16/float/bf16	ND	最后一维的大小需进行32字节对齐。 rstd为true时，维度数需要大于0，并小于x的维度数，gamma的维度从最后一维向前，每一维都需要和x保持一致。rstd为false时,维度为[1,d_n]
beta	[1,d_n]	float16/float/bf16	ND	quantType = INT8时输入。最后一维的大小要32字节对齐。
scale	[1]	float16/bfloat16	ND	当quantType为QUANT_INT8时才输入。
offset	[1]	int8	ND	当quantType为QUANT_INT8时才输入。
output	[-1,…,-1,d_n]	float16/float/int8/bf16	ND	quantType=INT8时输出int8。
rstd	[d_0,..,1]	float16/float/bf16	ND	rstd为true时输出。第i维大于等于x的维度数减去gamma的维度数时，大小为1，否则和x中对应维度大小相等。
所有Tensor最后一维的n大小需保持一致。仅Atlas 800I A2推理产品支持输入为bf16数据类型。 quantType = INT8时，最后一维小于或等于8192。

RMS_NORM_NORM DYNAMIC_QUANT

参数	维度	数据类型	格式	描述
x	[d_0, ..., d_k, n]	float16	ND	最后一维n的大小要32字节对齐，且小于等于12288。
gamma	[1, ..., 1, n]	float16	ND	最后一维n的大小要32字节对齐，且小于等于12288。
beta	[1, ..., 1, n]	float16	ND	最后一维n的大小要32字节对齐，且小于等于12288。
output	[d_0, ..., d_k, n]	int8	ND	最后一维n的大小要32字节对齐，且小于等于12288。
scale	[d_0, ..., d_k]	float	ND	当quantType == QUANT_INT8且dynamicQuantType != DYNAMIC_QUANT_UNDEFINED时输出。
offset	[d_0, ..., d_k]	float	ND	当quantType == QUANT_INT8且dynamicQuantType == DYNAMIC_QUANT_ASYMMETRIC时输出。当前版本暂不支持。

RMS_NORM_PRENORM quantType = QUANT_UNDEFINED输入输出描述

参数	维度	数据类型	格式	描述
x	[-1,…,-1,n] -1表示当前维度的大小没有约束。	float16/bf16	ND	最后一维的大小要32字节对齐。
beta	[1,n]	float16/bf16	ND	当hasBias为True时输入最后一维的大小要32字节对齐。
residual	[-1,…,-1,n] -1表示当前维度的大小没有约束。	float16/bf16	ND	最后一维的大小要32字节对齐。
gamma	[1,n]	float16/bf16	ND	最后一维的大小要32字节对齐。
output	[-1,…,-1,n]	float16/bf16	ND	输出tensor，维度数与x一致。
resOut	[-1,…,-1,n]	float16/bf16	ND	输出tensor，维度数与x一致。
所有Tensor最后一维的n大小需保持一致。

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字节对齐。
scale	[1]	float16	ND	量化scale参数。
offset	[1]	int8	ND	量化offset参数。
outputQuant	[-1,…,-1,n]	int8	ND	输出tensor，维度数与x一致。
output	[-1,…,-1,n]	float16	ND	输出tensor，维度数与x一致。
所有Tensor最后一维的n大小需保持一致。

RMS_NORM_POSTNORM输入输出描述

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

父主题： atb/infer_op_params.h