RMSNorm

class torchtune.modules.RMSNorm(dim: int, eps: float = 1e-06)

Implements Root Mean Square Normalization introduced in https://arxiv.org/pdf/1910.07467.pdf.

Reference implementation (used for correctness verfication) can be found here: https://github.com/facebookresearch/llama/blob/main/llama/model.py

Parameters:

• dim (int) – embedding size

• eps (float) – small value to avoid division by zero. Default: 1e-6

forward(x: Tensor) → Tensor

Parameters:

x (Tensor) – input tensor to normalize

Returns:

The output tensor after applying RMSNorm.

Return type:

Tensor