Skip to content

Normalization Layers

About

The d9d.module.block.normalization package implements normalization layers.

RMSNorm

RMSNorm implements Root Mean Square Layer Normalization.

It uses a custom Triton kernel for the forward and backward passes, so it requires a GPU supported by Triton.

It supports zero-centered scaling weights (zero_centered=True). They are initialized to 0, and the kernel scales by weight + 1.

Kernel Benchmarks (bf16, H100)

Forward, Hidden Size = 128

Forward, Hidden Size = 256

Forward, Hidden Size = 1024

Forward, Hidden Size = 4096

Forward, Hidden Size = 7168

Backward, Hidden Size = 128

Backward, Hidden Size = 256

Backward, Hidden Size = 1024

Backward, Hidden Size = 4096

Backward, Hidden Size = 7168

Usage

1
2
3
4
5
6
7
8
9
import torch

from d9d.module.block.normalization import RMSNorm

norm = RMSNorm(hidden_size=2048, eps=1e-6).cuda()
norm.reset_parameters()

hidden_states = torch.randn(2, 16, 2048, device="cuda")
output = norm(hidden_states)  # (2, 16, 2048)

API Reference

d9d.module.block.normalization

Normalization layers.

RMSNorm

Bases: Module, ModuleLateInit

Root Mean Square Layer Normalization (RMSNorm) layer.

Normalizes the input over its last dimension by its root mean square and applies learnable scaling weights. The weights can optionally be zero-centered.

References

Root Mean Square Layer Normalization

__init__(hidden_size, eps=1e-06, zero_centered=False)

Constructs the RMSNorm object.

Parameters:

Name Type Description Default
hidden_size int

Size of the normalized last dimension.

 required
eps float

Small value added to the variance for numerical stability.

 1e-06
zero_centered bool

If True, the scaling weights are initialized to 0 and offset by 1 during computation. Otherwise, they are initialized to 1.

 False
forward(x)

Applies RMSNorm to the input.

Parameters:

Name Type Description Default
x Tensor

Input tensor. Shape: (..., hidden_size).

 required

Returns:

Type Description
Tensor

The normalized tensor. Shape: (..., hidden_size).
reset_parameters()

Resets module parameters.