RmsNormOperation

功能

RMSNorm是一种归一化方法，它通过对数据的root mean square（RMS）进行归一化，避免了均值的使用。RMSNorm通过对数据的均值和方差进行调整，旨在优化深度学习模型的性能。RMSNorm不仅可以简化计算，还可以提高模型的收敛速度和稳定性。

当前支持Norm，PreNorm，PostNorm三种归一化操作。

图1 NORM

使用场景

通常使用在注意力机制attention前和attention与前馈网络ffn之间。

图2 RmsNormOperation算子上下文

计算公式

• RmsNorm

  

  beta在rmsnorm quant量化场景下会使用。

• add+RmsNorm(PreNorm, PostNorm)

  

  

  其中gamma为可训练参数，epsilon是一个取值范围为(0, 1]的浮点值，用于避免除零。

• 量化

  

定义

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;
};

参数列表

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

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

RMS_NORM_NORM DYNAMIC_QUANT

规格约束

• 除scale，offset，rstd tensors外，所有输入输出tensor的最后一维大小相等，需要32字节对齐。
• rstd仅Atlas 800I A2 推理产品/Atlas A2 训练系列产品支持，不能和precisionMode，modelType同时设置，不支持量化场景使用。
• precisionMode不能和rstd，modelType同时设置，不支持量化场景使用。
• modelType不能和rstd，precisionMode同时设置，不支持量化场景使用。
• dynamicQuantType目前只支持对称动态量化DYNAMIC_QUANT_SYMMETRIC。动态对称量化时，输入的最后一维大小需要小于等于12288。

功能列表

• RmsNorm
  • 非量化
    • 训练rmsnormforward算子，设置rstd为true开启。
    • precisionMode默认HIGH_PRECISION_MODE，中间计算使用float类型。HIGH_PERFORMANCE_MODE，在最后做gamma（weight）的乘法计算的时候使用float16类型。
    • modelType默认LLAMA_MODEL，使用llama模型的rmsnorm公式。GEMMA_MODEL，使用gemma模型的rmsnorm公式。区别在乘gamma（weight）前，会将gamma+1再做乘法。
  • 量化
    • rmsnorm量化。
    • dynamicQuantType默认值DYNAMIC_QUANT_UNDEFINED（普通量化）DYNAMIC_QUANT_SYMMETRIC 动态对称量化。
• add+RmsNorm融合（Prenorm，PostNorm）在做rmsnorm前，将x和residual相加
  • 非量化
    • PreNorm返回add+residual的结果，PostNorm不返回
    • hasBias参数。为true时，可以输入bias tensor。为false时，不输入bias tensor。
  • 量化

    只有Prenorm支持。