File size: 2,509 Bytes
46144df | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | """Small shared neural-network layers."""
from __future__ import annotations
import torch
from torch import nn
from torch.nn import functional as F
class RMSNorm(nn.Module):
def __init__(self, width: int, eps: float = 1e-5) -> None:
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(width))
def forward(self, x: torch.Tensor) -> torch.Tensor:
# Accumulating the variance in fp32 prevents low-precision callers from
# underflowing; the normalized activation is cast back before scaling.
normalized = x.float() * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps)
return normalized.to(dtype=x.dtype) * self.weight
class SwiGLU(nn.Module):
def __init__(self, width: int, hidden_width: int, dropout: float, bias: bool = False) -> None:
super().__init__()
self.gate = nn.Linear(width, hidden_width, bias=bias)
self.up = nn.Linear(width, hidden_width, bias=bias)
self.down = nn.Linear(hidden_width, width, bias=bias)
self.dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.dropout(self.down(F.silu(self.gate(x)) * self.up(x)))
def apply_rope(q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Apply rotary embeddings to tensors shaped ``[B, H, T, D]``."""
_, _, seq_len, head_dim = q.shape
if head_dim % 2:
raise ValueError("RoPE requires an even head dimension")
# RoPE pairs adjacent channels. Frequencies decrease geometrically from the
# fastest pair to the slowest pair, then each pair is rotated by its token
# position. There are intentionally no learned position parameters.
positions = torch.arange(seq_len, device=q.device, dtype=torch.float32)
frequencies = 1.0 / (
10000.0
** (torch.arange(0, head_dim, 2, device=q.device, dtype=torch.float32) / head_dim)
)
angles = torch.outer(positions, frequencies)
cos = angles.cos()[None, None, :, :].to(dtype=q.dtype)
sin = angles.sin()[None, None, :, :].to(dtype=q.dtype)
def rotate(x: torch.Tensor) -> torch.Tensor:
# [x_even, x_odd] is multiplied by the 2-D rotation matrix for every
# position/frequency pair, then interleaved back into the head axis.
even, odd = x[..., 0::2], x[..., 1::2]
return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
return rotate(q), rotate(k)
|