"""RMSNorm API for TransformerEngineNPU PyTorch"""
from typing import Optional, Union, Iterable, Any
import warnings
import torch
import torch_npu
import transformer_engine.pytorch.ops as ops
class RMSNorm(ops.RMSNorm):
r"""Root Mean Square Layer Normalization
Applies Root Mean Square Layer Normalization over a mini-batch of
inputs as described in the paper
`Root Mean Square Layer Normalization <https://arxiv.org/abs/1910.07467>`__
.. math::
y = \frac{x}{\sqrt{\mathrm{Var}[x] + \varepsilon}} * \gamma
where
.. math::
\mathrm{Var}[x] = \frac{1}{n}\sum_{i=0}^n x_i^2
:math:`\gamma` is a learnable affine transform parameter that
matches the inner-most dimensions of the input tensor.
Parameters
----------
normalized_shape : int or iterable of int
Inner dimensions of input tensor
eps : float, default = 1e-5
A value added to the denominator for numerical stability
sequence_parallel : bool, default = False
If ``True``, the weight parameter is marked for sequence parallel
all-reduce. This is custom logic for Megatron-LM integration.
zero_centered_gamma : bool, default = False
If ``True``, the :math:`\gamma` parameter is initialized to zero
and the calculation changes to
.. math::
y = \frac{x}{\sqrt{\mathrm{Var}[x] + \varepsilon}} * (1 + \gamma)
hidden_size : int, optional
**Deprecated.** Use ``normalized_shape`` instead.
params_dtype : torch.dtype, optional
**Deprecated.** Use ``dtype`` kwarg instead.
"""
def __init__(
self,
normalized_shape: Union[Iterable[int], int, None] = None,
eps: float = 1e-5,
sequence_parallel: bool = False,
zero_centered_gamma: bool = False,
hidden_size: Optional[int] = None,
params_dtype: Optional[torch.dtype] = None,
**kwargs,
) -> None:
super().__init__(
normalized_shape,
eps=eps,
zero_centered_gamma=zero_centered_gamma,
**kwargs,
)
if normalized_shape is None:
if hidden_size is None:
raise RuntimeError(
"Neither `normalized_shape` nor `hidden_size` (deprecated) args are provided"
)
warnings.warn(
"`hidden_size` arg has been renamed to `normalized_shape` "
"for compatibility with `torch.nn.LayerNorm`.",
DeprecationWarning,
stacklevel=2,
)
normalized_shape = hidden_size
elif hidden_size is not None:
raise RuntimeError(
"Both `normalized_shape` and `hidden_size` (deprecated) args are provided"
)
if params_dtype is not None:
if "dtype" in kwargs:
raise RuntimeError(
"Both `dtype` and `params_dtype` (deprecated) kwargs are provided"
)
kwargs["dtype"] = params_dtype
if not isinstance(normalized_shape, Iterable):
normalized_shape = (normalized_shape,)
else:
normalized_shape = tuple(normalized_shape)
self.normalized_shape = normalized_shape
self.eps = eps
self.sequence_parallel = sequence_parallel
self.zero_centered_gamma = zero_centered_gamma
dtype = kwargs.get("dtype", torch.get_default_dtype())
device = kwargs.get("device", torch_npu.npu.current_device() if torch_npu.npu.is_available() else "cpu")
self.weight = torch.nn.Parameter(torch.ones(normalized_shape, dtype=dtype, device=device))
setattr(self.weight, "sequence_parallel", sequence_parallel)
if zero_centered_gamma:
torch.nn.init.zeros_(self.weight)
def reset_parameters(self) -> None:
"""Reset parameters to initial values."""
if self.zero_centered_gamma:
torch.nn.init.zeros_(self.weight)
else:
torch.nn.init.ones_(self.weight)
def extra_repr(self) -> str:
return (
f"{self.normalized_shape}, eps={self.eps}, "
f"sequence_parallel={self.sequence_parallel}, "
f"zero_centered_gamma={self.zero_centered_gamma}"
)
__all__ = ["RMSNorm"]
