Implementation:NVIDIA TransformerEngine TE RMSNorm
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, PyTorch, Optimization |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
High-level Root Mean Square Layer Normalization module with backward-compatible API, wrapping the lower-level RMSNorm operation from transformer_engine.pytorch.ops.
Description
This RMSNorm class is the user-facing module that applies RMS Layer Normalization as described in the paper "Root Mean Square Layer Normalization". It extends the ops-level _RMSNormOp with backward-compatible constructor arguments including deprecated hidden_size and params_dtype parameters. It also supports a sequence_parallel flag for custom Megatron-LM integration (setting a bool attribute on parameters). The class provides deprecated property shims for SM margin settings (fwd_rmsnorm_sm_margin, bwd_rmsnorm_sm_margin, inf_rmsnorm_sm_margin).
Usage
Use as a drop-in replacement for custom RMSNorm implementations in transformer models. Supports zero_centered_gamma for LLaMA-style normalization.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/pytorch/module/rmsnorm.py- Lines
- 1--174
Signature
class RMSNorm(_RMSNormOp):
def __init__(self, normalized_shape=None, eps=1e-5, sequence_parallel=None, params_dtype=None, zero_centered_gamma=False, hidden_size=None, **kwargs) -> None: ...
def reset_parameters(self, defer_init=None) -> None: ...
def reset_rms_norm_parameters(self) -> None: ... # deprecated
Import
from transformer_engine.pytorch.module.rmsnorm import RMSNorm
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| normalized_shape | int or Iterable[int] | Yes | Inner dimensions of input tensor |
| eps | float | No | Numerical stability constant (default 1e-5) |
| zero_centered_gamma | bool | No | Initialize gamma to zero and use (1 + gamma) scaling
|
| sequence_parallel | bool | No | Set sequence_parallel attribute on weight parameter |
Outputs
| Name | Type | Description |
|---|---|---|
| output | torch.Tensor | RMS-normalized tensor |
Usage Examples
from transformer_engine.pytorch.module.rmsnorm import RMSNorm
rmsnorm = RMSNorm(normalized_shape=4096, eps=1e-5, zero_centered_gamma=True)
output = rmsnorm(input_tensor)