| """Normalization, MLP and convolution primitives with the training arithmetic. |
| |
| Each op reproduces the training model's dtype handling under BF16 autocast |
| (reductions in FP32, results in the activation dtype), so BF16 inference matches |
| the training precision bit for bit. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import torch |
| from torch import Tensor, nn |
| from torch.nn import functional as F |
|
|
| NORM_EPS = 1e-6 |
|
|
|
|
| class RMSNorm(nn.Module): |
| """RMS normalization over the last dimension, optional per-channel gain. |
| |
| The gain is cast to the input dtype and the result returned in it. |
| """ |
|
|
| def __init__(self, dim: int, *, affine: bool) -> None: |
| """Build the norm; ``affine`` adds a gain initialized to one.""" |
|
|
| super().__init__() |
| self.dim = dim |
| self.weight = nn.Parameter(torch.ones(dim)) if affine else None |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| """Normalize ``x`` ``[..., dim]``.""" |
|
|
| weight = None if self.weight is None else self.weight.to(dtype=x.dtype) |
| return F.rms_norm(x, (self.dim,), weight, NORM_EPS).to(dtype=x.dtype) |
|
|
|
|
| class ChannelRMSNorm(nn.Module): |
| """Per-position RMS normalization over the channels of ``[B, C, H, W]``, |
| reduced in FP32 and applied in the input dtype, with an optional |
| per-channel gain and bias (cast to the input dtype).""" |
|
|
| def __init__(self, channels: int, *, affine: bool) -> None: |
| """Build the norm; ``affine`` adds a gain (one) and a bias (zero).""" |
|
|
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(channels)) if affine else None |
| self.bias = nn.Parameter(torch.zeros(channels)) if affine else None |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| """Normalize each position's channel vector.""" |
|
|
| mean_square = torch.mean( |
| torch.square(x), dim=1, keepdim=True, dtype=torch.float32 |
| ) |
| y = x * torch.rsqrt(mean_square + NORM_EPS).to(dtype=x.dtype) |
| if self.weight is not None and self.bias is not None: |
| y = y * self.weight.view(1, -1, 1, 1).to(dtype=x.dtype) |
| y = y + self.bias.view(1, -1, 1, 1).to(dtype=x.dtype) |
| return y |
|
|
|
|
| class Mlp(nn.Module): |
| """Two linear layers with a GELU between them, both with biases.""" |
|
|
| def __init__(self, dim: int, hidden: int) -> None: |
| """Allocate the up and down projections.""" |
|
|
| super().__init__() |
| self.up = nn.Linear(dim, hidden) |
| self.down = nn.Linear(hidden, dim) |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| """``down(gelu(up(x)))``.""" |
|
|
| return self.down(F.gelu(self.up(x))) |
|
|
|
|
| def pointwise_conv(module: nn.Conv2d, x: Tensor) -> Tensor: |
| """A 1x1 convolution as a channel matmul (as the training graph runs it).""" |
|
|
| y = torch.einsum("bchw,oc->bohw", x, module.weight.flatten(1)) |
| if module.bias is not None: |
| y = y + module.bias.view(1, -1, 1, 1).to(dtype=y.dtype) |
| return y |
|
|
|
|
| def same_conv(module: nn.Conv2d, x: Tensor) -> Tensor: |
| """A stride-1, same-padded convolution with ``module``'s parameters.""" |
|
|
| return F.conv2d( |
| x, |
| module.weight, |
| module.bias, |
| stride=1, |
| padding=module.kernel_size[0] // 2, |
| groups=module.groups, |
| ) |
|
|
