RMSNorm¶
- class torchtune.modules.RMSNorm(dim: int, eps: float = 1e-06)[source]¶
Implements Root Mean Square Normalization introduced in https://arxiv.org/abs/1910.07467.
Reference implementation (used for correctness verification) can be found here: https://github.com/facebookresearch/llama/blob/main/llama/model.py
- Parameters:
- forward(x: Tensor) Tensor [source]¶
- Parameters:
x (torch.Tensor) – input tensor to normalize
- Returns:
The normalized and scaled tensor having the same shape as
x
.- Return type: