lite_boost.layers.rms_norm

lite_boost.layers.rms_norm(x, gamma, eps=1e-6)[source]

Applies per-row RMS normalization over the last dim of x.

Computes \(y = x / RMS(x) * gamma\), where \(RMS(x) = \sqrt{mean(x^2)}\) and \(mean(x^2)\) is the mean of squared values of x over the last dim. Usually used to replace the expanded F.normalize(x, dim) * sqrt(dim) * gamma chain in Wan-series VAE RMS_norm layers.

Parameters:

x (Tensor) – 2-D input tensor with shape \((N, C)\), the normalized dim is the last dim. Supported dtypes are float16, float32; on A2 and other SoCs bfloat16 is additionally supported.
gamma (Tensor) – Per-column scale with shape \((C,)\) and the same dtype as x, matching the last dim of x.
eps (float, optional) – Value added to the denominator for numerical stability. Default: 1e-6.

Returns:

Tensor, the normalized result with the same shape and dtype as x.

Raises:

ValueError – If x is not 2-D, or gamma is not 1-D or does not match the last dim of x.
ValueError – If on 300I Duo x is bfloat16, or the last dim of x is smaller than 16.

Supported Platforms:: Ascend

Examples

>>> import torch
>>> import torch_npu
>>> from lite_boost.layers import rms_norm
>>> torch.npu.set_device(0)
>>> x = torch.full((2, 16), 2.0, device="npu")
>>> gamma = torch.ones(16, device="npu")
>>> y = rms_norm(x, gamma)
>>> print(y.shape)
torch.Size([2, 16])