| """Transformer building blocks: RMSNorm, rotary position embedding, SwiGLU.""" | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class RMSNorm(nn.Module): | |
| def __init__(self, dim: int, eps: float = 1e-5): | |
| super().__init__() | |
| self.eps = eps | |
| self.weight = nn.Parameter(torch.ones(dim)) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| xf = x.float() | |
| return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)).type_as(x) * self.weight | |
| def rope_tables(head_dim: int, max_len: int, theta: float, device=None) -> tuple[torch.Tensor, torch.Tensor]: | |
| """cos/sin tables [max_len, head_dim/2] (float32).""" | |
| inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim)) | |
| ang = torch.outer(torch.arange(max_len, device=device, dtype=torch.float32), inv) | |
| return ang.cos(), ang.sin() | |
| def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: | |
| """x: [B, H, L, D]; cos/sin: [L, D/2]. Rotates interleaved-as-halves pairs (GPT-NeoX style).""" | |
| x1, x2 = x.float().chunk(2, dim=-1) | |
| out = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1) | |
| return out.type_as(x) | |
| class SwiGLU(nn.Module): | |
| def __init__(self, dim: int, hidden: int): | |
| super().__init__() | |
| self.w_gate = nn.Linear(dim, hidden, bias=False) | |
| self.w_up = nn.Linear(dim, hidden, bias=False) | |
| self.w_down = nn.Linear(hidden, dim, bias=False) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)) | |