128 lines
4.6 KiB
Python
128 lines
4.6 KiB
Python
"""
|
||
Reusable building-block layers: RMSNorm, RotaryEmbedding, SwiGLU.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Optional TransformerEngine import (FP8 support)
|
||
# ---------------------------------------------------------------------------
|
||
try:
|
||
import transformer_engine.pytorch as te # type: ignore[import]
|
||
HAS_TE = True
|
||
except ImportError:
|
||
te = None # type: ignore[assignment]
|
||
HAS_TE = False
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# RMS Layer Normalisation
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class RMSNorm(nn.Module):
|
||
"""Root-Mean-Square Layer Normalisation (Zhang & Sennrich, 2019).
|
||
|
||
Computation is promoted to float32 for numerical stability and cast back
|
||
to the input dtype before returning.
|
||
"""
|
||
|
||
def __init__(self, d_model: int, eps: float = 1e-6) -> None:
|
||
super().__init__()
|
||
self.eps = eps
|
||
self.weight = nn.Parameter(torch.ones(d_model))
|
||
|
||
def _norm(self, x: torch.Tensor) -> torch.Tensor:
|
||
# x: (..., D) — compute in fp32
|
||
return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
# Upcast to float32, normalise, scale, then restore original dtype.
|
||
out = self._norm(x.float()).to(x.dtype)
|
||
return out * self.weight
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Rotary Positional Embedding
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class RotaryEmbedding(nn.Module):
|
||
"""Precomputed rotary positional embeddings (Su et al., RoFormer 2021).
|
||
|
||
Cos/sin tables are stored as buffers (shape: max_seq_len × D//2) so they
|
||
move with the module to the correct device automatically.
|
||
"""
|
||
|
||
def __init__(self, dim: int, max_seq_len: int, theta: float = 10000.0) -> None:
|
||
super().__init__()
|
||
self.dim = dim
|
||
self.max_seq_len = max_seq_len
|
||
self.theta = theta
|
||
|
||
# Precompute and register
|
||
cos, sin = self._build_tables(dim, max_seq_len, theta)
|
||
self.register_buffer("_cos_cached", cos, persistent=False)
|
||
self.register_buffer("_sin_cached", sin, persistent=False)
|
||
|
||
@staticmethod
|
||
def _build_tables(
|
||
dim: int, max_seq_len: int, theta: float
|
||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||
"""Compute cos/sin tables with shape (max_seq_len, dim // 2)."""
|
||
half_dim = dim // 2
|
||
# Inverse frequencies: shape (half_dim,)
|
||
freqs = 1.0 / (
|
||
theta ** (torch.arange(0, half_dim, dtype=torch.float32) / half_dim)
|
||
)
|
||
# Positions: shape (max_seq_len,)
|
||
t = torch.arange(max_seq_len, dtype=torch.float32)
|
||
# Outer product → (max_seq_len, half_dim)
|
||
emb = torch.outer(t, freqs)
|
||
cos = emb.cos() # (T, D//2)
|
||
sin = emb.sin() # (T, D//2)
|
||
return cos, sin
|
||
|
||
def forward(self, seq_len: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
|
||
"""Return (cos, sin) slices of shape (seq_len, D//2) on *device*.
|
||
|
||
If *seq_len* exceeds the precomputed length the tables are recomputed
|
||
on-the-fly (rare, but graceful fallback).
|
||
"""
|
||
if seq_len > self.max_seq_len:
|
||
cos, sin = self._build_tables(self.dim, seq_len, self.theta)
|
||
cos = cos.to(device)
|
||
sin = sin.to(device)
|
||
else:
|
||
cos = self._cos_cached[:seq_len].to(device)
|
||
sin = self._sin_cached[:seq_len].to(device)
|
||
return cos, sin
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# SwiGLU Feed-Forward Network
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class SwiGLU(nn.Module):
|
||
"""SwiGLU feed-forward block (Shazeer, 2020).
|
||
|
||
Architecture:
|
||
out = down_proj( SiLU(gate_proj(x)) * up_proj(x) )
|
||
|
||
The gate and up projections are separate linear layers so that the gating
|
||
mechanism can learn an independent representation.
|
||
"""
|
||
|
||
def __init__(self, d_model: int, d_ffn: int, bias: bool = False) -> None:
|
||
super().__init__()
|
||
self.gate_proj = nn.Linear(d_model, d_ffn, bias=bias)
|
||
self.up_proj = nn.Linear(d_model, d_ffn, bias=bias)
|
||
self.down_proj = nn.Linear(d_ffn, d_model, bias=bias)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
# Gated activation: element-wise product of SiLU(gate) and up projection
|
||
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
|