Implementation:NVIDIA TransformerEngine Ops LayerNorm
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, PyTorch, Optimization |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Fusible Layer Normalization operation with CUDA-accelerated forward and backward passes, configurable SM margins, and ONNX export support.
Description
LayerNorm is a BasicOperation implementing standard Layer Normalization with learnable gamma and beta parameters. It uses CUDA kernels (layernorm_fwd and layernorm_bwd from transformer_engine_torch) for efficient computation. Features include zero_centered_gamma mode (where gamma is initialized to zero and the formula becomes y = norm(x) * (1 + gamma) + beta), configurable SM margins for overlapping with communication kernels, quantizer integration for the output, and an ONNX export path using standard PyTorch operations.
Usage
Used as the fundamental normalization operation in transformer models. Supports deferred parameter initialization with meta device.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/pytorch/ops/basic/layer_norm.py- Lines
- 1--276
Signature
class LayerNorm(BasicOperation):
def __init__(self, normalized_shape, *, eps=1e-5, device=None, dtype=None, zero_centered_gamma=False, sm_margin=0) -> None: ...
def reset_parameters(self) -> None: ...
def op_forward(self, ctx, input_, prev_op_grad_output_quantizer, next_op_input_quantizer) -> torch.Tensor: ...
def op_backward(self, ctx, grad_output) -> Tuple[torch.Tensor, Tuple]: ...
Import
from transformer_engine.pytorch.ops.basic.layer_norm import LayerNorm
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 | Use zero-centered gamma initialization |
| sm_margin | int or dict | No | Number of SMs to exclude for kernel overlap |
Outputs
| Name | Type | Description |
|---|---|---|
| output | torch.Tensor | Layer-normalized tensor |
Usage Examples
from transformer_engine.pytorch.ops.basic.layer_norm import LayerNorm
ln = LayerNorm(4096, eps=1e-5, zero_centered_gamma=True)
output = ln.op_forward(ctx, input_tensor, None, quantizer)