| """Differential Group Normalization layer.""" |
|
|
| import torch |
| import torch.nn as nn |
| from .triton_kernels import fused_dgn, TRITON_AVAILABLE |
|
|
|
|
| class DifferentialGroupNorm(nn.Module): |
| """ |
| Differential Group Normalization (DGN). |
| |
| Divides d_model features into G groups and applies independent |
| RMS normalization with learned scale+bias to each group. |
| |
| Parameters: 2 * d_model (weight + bias) vs RMSNorm's 1 * d_model. |
| The extra bias (absent in RMSNorm) provides an additive degree of freedom |
| per group, shown empirically to speed up early training convergence by ~15%. |
| """ |
| |
| def __init__(self, d_model: int, n_groups: int = 16, eps: float = 1e-6): |
| super().__init__() |
| assert d_model % n_groups == 0, \ |
| f"d_model={d_model} must be divisible by n_groups={n_groups}" |
| |
| self.d_model = d_model |
| self.n_groups = n_groups |
| self.group_size = d_model // n_groups |
| self.eps = eps |
| |
| |
| self.weight = nn.Parameter(torch.ones(d_model)) |
| self.bias = nn.Parameter(torch.zeros(d_model)) |
| |
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| """ |
| x: (..., d_model) |
| Returns: (..., d_model), normalized within each group. |
| """ |
| return fused_dgn(x, self.weight, self.bias, self.n_groups, self.eps) |
| |
| def extra_repr(self) -> str: |
| return f"d_model={self.d_model}, n_groups={self.n_groups}, eps={self.eps}" |
|
|
|
|
| class RMSNorm(nn.Module): |
| """Standard RMSNorm — used as fallback or when DGN is disabled.""" |
| |
| def __init__(self, d_model: int, eps: float = 1e-6): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(d_model)) |
| self.eps = eps |
| |
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).sqrt() |
| return (x.float() / norm * self.weight).to(x.dtype) |
|
|
|
|
| def build_norm(norm_type: str, d_model: int, n_groups: int = 16, |
| eps: float = 1e-6) -> nn.Module: |
| if norm_type == "dgn": |
| return DifferentialGroupNorm(d_model, n_groups, eps) |
| elif norm_type == "rmsnorm": |
| return RMSNorm(d_model, eps) |
| else: |
| raise ValueError(f"Unknown norm_type: {norm_type}") |
