RMSNorm

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

Root Mean Square Normalization in fp32.

See: https://pytorch.org/docs/stable/generated/torch.nn.RMSNorm.html

Parameters:

• dim (int) – embedding size

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

forward(x: Tensor) → Tensor

Parameters:

x (torch.Tensor) – input tensor to normalize

Returns:

The normalized and scaled tensor having the same shape as x.

Return type:

torch.Tensor