RmsNormParam

Attribute	Type	Default Value	Description
layer_type	torch_atb.RmsNormParam.RmsNormType	torch_atb.RmsNormParam.RmsNormType.RMS_NORM_UNDEFINED	This default type is unavailable. You need to set this parameter.
norm_param	torch_atb.RmsNormParam.NormParam	-	-
pre_norm_param	torch_atb.RmsNormParam.PreNormParam	-	-
post_norm_param	torch_atb.RmsNormParam.PostNormParam	-	-

RmsNormParam.RmsNormType

Enumerated items:

RMS_NORM_UNDEFINED
RMS_NORM_PRENORM
RMS_NORM_NORM
RMS_NORM_POSTNORM

RmsNormParam.NormParam

Attribute	Type	Default Value	Description
quant_type	torch_atb.QuantType	torch_atb.QuantType.QUANT_UNQUANT	Quantization is not performed.
epsilon	float	1e-5	-
layer_norm_eps	double	1e-5	-
rstd	bool	False	-
precision_mode	torch_atb.RmsNormParam.PrecisionMode	torch_atb.RmsNormParam.PrecisionMode.HIGH_PRECISION_MODE	-
model_type	torch_atb.RmsNormParam.ModelType	torch_atb.RmsNormParam.ModelType.LLAMA_MODEL	-
dynamic_quant_type	torch_atb.DynamicQuantType	torch_atb.DynamicQuantType.DYNAMIC_QUANT_UNDEFINED	-

RmsNormParam.PrecisionMode

Enumerated items:

HIGH_PRECISION_MODE
HIGH_PERFORMANCE_MODE

RmsNormParam.ModelType

Enumerated items:

LLAMA_MODEL
GEMMA_MODEL

RmsNormParam.PreNormParam

Attribute	Type	Default Value	Description
quant_type	torch_atb.QuantType	torch_atb.QuantType.QUANT_UNQUANT	Quantization is not performed.
epsilon	float	1e-5	-
has_bias	bool	False	-

RmsNormParam.PostNormParam

Attribute	Type	Default Value	Description
quant_type	torch_atb.QuantType	torch_atb.QuantType.QUANT_UNQUANT	Quantization is not performed.
epsilon	float	1e-5	-
has_bias	bool	False	-

Example

import torch
import torch_atb  
import numpy as np

def rms_norm():
    rms_norm_param = torch_atb.RmsNormParam(layer_type = torch_atb.RmsNormParam.RmsNormType.RMS_NORM_NORM)
    rms_norm_param.norm_param.rstd = True
    rms_norm = torch_atb.Operation(rms_norm_param)
    shape=[8, 8, 8]
    shape_gamma=[8]
    x = torch.from_numpy(np.random.uniform(low=0, high=100, size=shape).astype(np.float32))
    gamma = torch.from_numpy(np.random.uniform(low=0, high=100, size=shape_gamma).astype(np.float32))
    in_tensors = [x.npu(), gamma.npu()]
    print("in_tensors: ", in_tensors)

    def rms_norm_run():
        rms_norm_outputs = rms_norm.forward(in_tensors)
        return rms_norm_outputs

    outputs = rms_norm_run()
    print("outputs: ", outputs)

if __name__ == "__main__":
    rms_norm()

Parent topic: OpParam