RMSNorm#

General#

RMSNorm (Root Mean Square Layer Normalization) operation performs normalization on the input tensor using the root mean square statistic.

The RMSNorm operation performs the following transformation of the input tensor:

\[y = \gamma \cdot \frac{x}{\sqrt{\text{RMS}(x) + \epsilon}},\]

where

\[\text{RMS}(x) = \sqrt{\frac{1}{n} \sum_{i=1}^{n} x_i^2}\]

Attribute Name	Description	Value Type	Supported Values	Required or Optional
epsilon	The constant to improve numerical stability.	f32	Arbitrary positive f32 value, `1e-5` (default)	Optional
begin_norm_axis	`begin_norm_axis` is used to indicate which axis to start RMS normalization. The normalization is from `begin_norm_axis` to last dimension. Negative values means indexing from right to left. This op normalizes over the last dimension by default, e.g. C in TNC for 3D and LDNC for 4D.	s64	[-r,r-1],where r=rank(src). -1 is default	Optional

The inputs and outputs must be provided according to below index order when constructing an operation.

Index	Argument Name	Required or Optional
0	`src`	Required
1	`gamma`	Optional

Note

gamma is scaling for the normalized value. gamma shape should be broadcastable to the src shape.

Index	Argument Name	Required or Optional
0	`dst`	Required

RMSNorm operation supports the following data type combinations.