| """The decoder's ViT trunk: RMS-sandwich blocks with 2-D RoPE.""" |
|
|
| from __future__ import annotations |
|
|
| import torch |
| from torch import Tensor, nn |
| from torch.nn import functional as F |
|
|
| from .layers import Mlp, RMSNorm |
|
|
| ROPE_BASE = 10_000.0 |
|
|
|
|
| def rope_periods(head_dim: int, device: torch.device) -> Tensor: |
| """FP32 periods ``ROPE_BASE ** (2 i / (head_dim / 2))`` of the |
| ``head_dim / 4`` rotation frequencies per axis.""" |
|
|
| exponents = ( |
| 2.0 |
| * torch.arange(head_dim // 4, device=device, dtype=torch.float32) |
| / (head_dim // 2) |
| ) |
| return torch.tensor(ROPE_BASE, device=device, dtype=torch.float32) ** exponents |
|
|
|
|
| def rope_tables( |
| height: int, width: int, *, head_dim: int, device: torch.device |
| ) -> tuple[Tensor, Tensor]: |
| """FP32 ``(sin, cos)`` ``[H * W, head_dim]`` of axial 2-D RoPE on the token |
| grid: unnormalized patch-index coordinates (row, column), the |
| :func:`rope_periods`, adjacent-pair layout.""" |
|
|
| periods = rope_periods(head_dim, device) |
| rows = torch.arange(0.0, float(height), device=device, dtype=torch.float32) |
| cols = torch.arange(0.0, float(width), device=device, dtype=torch.float32) |
| coords = torch.stack(torch.meshgrid(rows, cols, indexing="ij"), dim=-1).flatten( |
| 0, 1 |
| ) |
| angles = 1.0 * coords[:, :, None] / periods[None, None, :].expand(1, 2, -1) |
| angles = angles.repeat_interleave(2, dim=-1).flatten(1, 2) |
| return torch.sin(angles), torch.cos(angles) |
|
|
|
|
| def _rotate_pairs(x: Tensor) -> Tensor: |
| """Rotate consecutive channel pairs: ``(a, b) -> (-b, a)``.""" |
|
|
| pairs = x.reshape(*x.shape[:-1], x.shape[-1] // 2, 2) |
| return torch.stack((-pairs[..., 1], pairs[..., 0]), dim=-1).reshape_as(x) |
|
|
|
|
| def apply_rope(x: Tensor, sin: Tensor, cos: Tensor) -> Tensor: |
| """Rotate ``[B, heads, N, head_dim]`` queries or keys in the tables' dtype.""" |
|
|
| rotated = x.to(dtype=sin.dtype) |
| sin_b = sin.view(1, 1, *sin.shape) |
| cos_b = cos.view(1, 1, *cos.shape) |
| return ((rotated * cos_b) + (_rotate_pairs(rotated) * sin_b)).to(dtype=x.dtype) |
|
|
|
|
| class Attention(nn.Module): |
| """Bias-free multi-head self-attention with RMS-normalized queries and keys.""" |
|
|
| def __init__(self, width: int, head_dim: int) -> None: |
| """Allocate the fused QKV and output projections and the q/k norms.""" |
|
|
| super().__init__() |
| self.heads = width // head_dim |
| self.head_dim = head_dim |
| self.qkv = nn.Linear(width, 3 * width, bias=False) |
| self.proj_out = nn.Linear(width, width, bias=False) |
| self.q_norm = RMSNorm(head_dim, affine=True) |
| self.k_norm = RMSNorm(head_dim, affine=True) |
|
|
| def forward(self, x: Tensor, sin: Tensor, cos: Tensor) -> Tensor: |
| """Attend over all tokens of ``[B, N, width]``.""" |
|
|
| b, n, width = x.shape |
| q, k, v = self.qkv(x).chunk(3, dim=-1) |
| shape = (b, n, self.heads, self.head_dim) |
| q = q.view(shape).transpose(1, 2).contiguous() |
| k = k.view(shape).transpose(1, 2).contiguous() |
| v = v.view(shape).transpose(1, 2).contiguous() |
| q = apply_rope(self.q_norm(q), sin, cos) |
| k = apply_rope(self.k_norm(k), sin, cos) |
| out = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False) |
| return self.proj_out(out.transpose(1, 2).contiguous().view(b, n, width)) |
|
|
|
|
| class DitBlock(nn.Module): |
| """ViT block: RMSNorm before and after attention and MLP.""" |
|
|
| def __init__(self, width: int, head_dim: int, mlp_ratio: float) -> None: |
| """Allocate attention, MLP and the four norms.""" |
|
|
| super().__init__() |
| self.attn_norm1 = RMSNorm(width, affine=True) |
| self.attn_norm2 = RMSNorm(width, affine=True) |
| self.mlp_norm1 = RMSNorm(width, affine=True) |
| self.mlp_norm2 = RMSNorm(width, affine=True) |
| self.attn = Attention(width, head_dim) |
| self.mlp = Mlp(width, int(mlp_ratio * width)) |
|
|
| def forward(self, x: Tensor, sin: Tensor, cos: Tensor) -> Tensor: |
| """Residual attention then residual MLP on ``[B, N, width]`` tokens.""" |
|
|
| x = x + self.attn_norm2(self.attn(self.attn_norm1(x), sin, cos)) |
| return x + self.mlp_norm2(self.mlp(self.mlp_norm1(x))) |
|
|
