""" PyTorch port of the flaxchat GPT model (GELU MLP variant). Faithful 1:1 port of flaxchat/gpt.py — every architectural feature matches: - RoPE (base 100000, head_dim=64 for d12, split-half layout) - Parameterless RMSNorm (no learnable gain) - QK-norm with 1.2x scaling (applied after RoPE, before SDPA) - Group-Query Attention (supports n_kv_head < n_head, repeats K/V heads) - Value embeddings on alternating layers (ResFormer-style), gated by `3 * sigmoid(ve_gate(x[..., :12]))` - Per-layer learnable residual scalars (`resid_lambdas`, `x0_lambdas`) - Smear: learnable gate on first 24 dims mixes in prev token - Backout: subtract mid-layer residual from late layers - Logit soft-cap: `15 * tanh(logits / 15)` - Sliding-window attention via window pattern (e.g. "SSSL") - Tied embeddings (lm_head = wte.T) - No biases in any Linear (attn Q/K/V/proj, MLP fc/proj, smear_gate, ve_gate) - MLP is GELU (`Linear -> gelu -> Linear`), the GELU variant This file is pure PyTorch (no JAX / Flax import). The conversion script in `convert_flax_to_torch.py` loads the Orbax checkpoint once and emits a `torch.save` state dict; `GELU_GPT.from_pretrained(path)` can then load it. """ from __future__ import annotations import json import math from dataclasses import dataclass, field, asdict from typing import Optional, Tuple, List import torch import torch.nn as nn import torch.nn.functional as F # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- @dataclass class GPTConfig: sequence_len: int = 1024 vocab_size: int = 32768 n_layer: int = 12 n_head: int = 12 n_kv_head: int = 12 n_embd: int = 768 window_pattern: str = "SSSL" tie_embeddings: bool = True rope_base: float = 100000.0 pad_vocab_size_to: int = 64 mlp: str = "gelu" # "gelu" or "relu2" @property def head_dim(self) -> int: return self.n_embd // self.n_head @property def padded_vocab_size(self) -> int: v = self.vocab_size p = self.pad_vocab_size_to return ((v + p - 1) // p) * p # --------------------------------------------------------------------------- # Parameterless RMSNorm # --------------------------------------------------------------------------- def rms_norm(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) # --------------------------------------------------------------------------- # RoPE # --------------------------------------------------------------------------- def precompute_rotary_embeddings( seq_len: int, head_dim: int, base: float = 100000.0, dtype=torch.float32 ) -> Tuple[torch.Tensor, torch.Tensor]: """Matches Flax `precompute_rotary_embeddings` exactly. Shapes: cos, sin -> (1, seq_len, 1, head_dim // 2) """ channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32) inv_freq = 1.0 / (base ** (channel_range / head_dim)) t = torch.arange(seq_len, dtype=torch.float32) freqs = torch.outer(t, inv_freq) # (T, D/2) cos = torch.cos(freqs).to(dtype) sin = torch.sin(freqs).to(dtype) cos = cos[None, :, None, :] sin = sin[None, :, None, :] return cos, sin def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: """x: (B, T, H, D). Uses split-half (flaxchat layout): x1 = x[..., :D/2]; x2 = x[..., D/2:] y1 = x1 * cos + x2 * sin y2 = -x1 * sin + x2 * cos return concat([y1, y2], dim=-1) """ d = x.shape[-1] // 2 x1 = x[..., :d] x2 = x[..., d:] y1 = x1 * cos + x2 * sin y2 = -x1 * sin + x2 * cos return torch.cat([y1, y2], dim=-1) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def has_ve(layer_idx: int, n_layer: int) -> bool: return layer_idx % 2 == (n_layer - 1) % 2 def compute_window_sizes(config: GPTConfig) -> List[Tuple[int, int]]: pattern = config.window_pattern.upper() assert all(c in "SL" for c in pattern), f"Invalid window_pattern: {pattern}" long_window = config.sequence_len # ceil(long / 4 / 128) * 128 via negative-floor trick short_window = -(-long_window // 4 // 128) * 128 char_to_window = {"L": (long_window, 0), "S": (short_window, 0)} window_sizes = [char_to_window[pattern[i % len(pattern)]] for i in range(config.n_layer)] window_sizes[-1] = (long_window, 0) return window_sizes # --------------------------------------------------------------------------- # Attention # --------------------------------------------------------------------------- class CausalSelfAttention(nn.Module): def __init__(self, config: GPTConfig, layer_idx: int): super().__init__() self.config = config self.layer_idx = layer_idx self._has_ve = has_ve(layer_idx, config.n_layer) head_dim = config.head_dim self.c_q = nn.Linear(config.n_embd, config.n_head * head_dim, bias=False) self.c_k = nn.Linear(config.n_embd, config.n_kv_head * head_dim, bias=False) self.c_v = nn.Linear(config.n_embd, config.n_kv_head * head_dim, bias=False) self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False) if self._has_ve: self.ve_gate = nn.Linear(12, config.n_kv_head, bias=False) else: self.ve_gate = None def forward( self, x: torch.Tensor, ve: Optional[torch.Tensor], cos: torch.Tensor, sin: torch.Tensor, window_size: Tuple[int, int], ) -> torch.Tensor: B, T, C = x.shape n_head = self.config.n_head n_kv_head = self.config.n_kv_head head_dim = self.config.head_dim q = self.c_q(x).reshape(B, T, n_head, head_dim) k = self.c_k(x).reshape(B, T, n_kv_head, head_dim) v = self.c_v(x).reshape(B, T, n_kv_head, head_dim) if self._has_ve and ve is not None: ve = ve.reshape(B, T, n_kv_head, head_dim) gate = 3.0 * torch.sigmoid(self.ve_gate(x[..., :12])) # (B, T, n_kv_head) v = v + gate.unsqueeze(-1) * ve q = apply_rotary_emb(q, cos, sin) k = apply_rotary_emb(k, cos, sin) q = rms_norm(q) * 1.2 k = rms_norm(k) * 1.2 if n_kv_head < n_head: repeats = n_head // n_kv_head k = k.repeat_interleave(repeats, dim=2) v = v.repeat_interleave(repeats, dim=2) # Build causal + sliding-window mask. Matches Flax implementation: # causal[i, j] = True iff i >= j # window_mask[i, j] = True iff (i - j) <= window_left # mask = causal & window (True = attend, False = mask out with -1e9) window_left = window_size[0] device = x.device row_idx = torch.arange(T, device=device).unsqueeze(1) col_idx = torch.arange(T, device=device).unsqueeze(0) causal_mask = row_idx >= col_idx # (T, T) if 0 < window_left < T: window_mask = (row_idx - col_idx) <= window_left causal_mask = causal_mask & window_mask # Additive bias: 0 where attend, -1e9 where mask out. bias = torch.where( causal_mask, torch.zeros((), dtype=x.dtype, device=device), torch.full((), -1e9, dtype=x.dtype, device=device), ) bias = bias.unsqueeze(0).unsqueeze(0) # (1, 1, T, T) # Rearrange to (B, H, T, D) for torch SDPA. q_bhtd = q.transpose(1, 2) k_bhtd = k.transpose(1, 2) v_bhtd = v.transpose(1, 2) # Match Flax's `jax.nn.dot_product_attention(..., bias=bias, scale=1/sqrt(D))`: # that is a plain softmax-attention with the additive bias and explicit scale. # We implement it manually (rather than F.scaled_dot_product_attention with # attn_mask) to guarantee bit-identical numerics with Flax's path. scale = 1.0 / math.sqrt(head_dim) att = torch.matmul(q_bhtd, k_bhtd.transpose(-2, -1)) * scale # (B,H,T,T) att = att + bias att = F.softmax(att, dim=-1) y = torch.matmul(att, v_bhtd) # (B, H, T, D) y = y.transpose(1, 2).contiguous().reshape(B, T, -1) y = self.c_proj(y) return y # --------------------------------------------------------------------------- # MLP (GELU variant — matches train_d12_chinchilla.py --mlp gelu branch) # --------------------------------------------------------------------------- class MLP(nn.Module): def __init__(self, config: GPTConfig): super().__init__() self.config = config self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False) self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.c_fc(x) if self.config.mlp == "gelu": # jax.nn.gelu default uses approximate=True (tanh form). Match that. x = F.gelu(x, approximate="tanh") elif self.config.mlp == "relu2": x = F.relu(x).pow(2) else: raise ValueError(f"Unsupported mlp: {self.config.mlp}") x = self.c_proj(x) return x # --------------------------------------------------------------------------- # Block # --------------------------------------------------------------------------- class Block(nn.Module): def __init__(self, config: GPTConfig, layer_idx: int): super().__init__() self.attn = CausalSelfAttention(config, layer_idx) self.mlp = MLP(config) def forward( self, x: torch.Tensor, ve: Optional[torch.Tensor], cos: torch.Tensor, sin: torch.Tensor, window_size: Tuple[int, int], ) -> torch.Tensor: x = x + self.attn(rms_norm(x), ve, cos, sin, window_size) x = x + self.mlp(rms_norm(x)) return x # --------------------------------------------------------------------------- # Full model # --------------------------------------------------------------------------- class GELU_GPT(nn.Module): """PyTorch port of flaxchat.gpt.GPT with GELU MLP. Naming mirrors the Flax module tree so weight conversion is a direct key-for-key mapping (see convert_flax_to_torch.py). """ def __init__(self, config: GPTConfig): super().__init__() self.config = config self.window_sizes = compute_window_sizes(config) padded_vocab = config.padded_vocab_size self.padded_vocab_size = padded_vocab # Token embedding self.wte = nn.Embedding(padded_vocab, config.n_embd) # Transformer blocks self.blocks = nn.ModuleList([Block(config, i) for i in range(config.n_layer)]) # lm_head (only materialized when NOT tied) self.tie_embeddings = config.tie_embeddings if not config.tie_embeddings: self.lm_head = nn.Linear(config.n_embd, padded_vocab, bias=False) else: self.lm_head = None # Per-layer learnable scalars (stored as (n_layer,) tensors to match Flax) self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer)) self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer)) # Smear self.smear_gate = nn.Linear(24, 1, bias=False) self.smear_lambda = nn.Parameter(torch.zeros(1)) # Backout self.backout_lambda = nn.Parameter(0.2 * torch.ones(1)) # Value embeddings — dict keyed by stringified layer index, to mirror Flax head_dim = config.head_dim kv_dim = config.n_kv_head * head_dim self.value_embeds = nn.ModuleDict( { str(i): nn.Embedding(padded_vocab, kv_dim) for i in range(config.n_layer) if has_ve(i, config.n_layer) } ) # Lazy RoPE: HF from_pretrained can leave persistent buffers # uninitialised (meta-init → NaN). We instead register an empty # placeholder and compute the table on the first forward. self._rope_max_len = config.sequence_len * 10 self._rope_head_dim = config.head_dim self._rope_base = config.rope_base self.register_buffer("rope_cos", torch.empty(0), persistent=False) self.register_buffer("rope_sin", torch.empty(0), persistent=False) self._rope_initialized = False def _get_rope(self, T, dtype, device): if (not self._rope_initialized or self.rope_cos.numel() == 0 or self.rope_cos.shape[1] < T): cos, sin = precompute_rotary_embeddings( max(T, self._rope_max_len), self._rope_head_dim, base=self._rope_base ) self.rope_cos = cos.to(device) self.rope_sin = sin.to(device) self._rope_initialized = True return self.rope_cos[:, :T].to(dtype), self.rope_sin[:, :T].to(dtype) # ------------------------------------------------------------------ # Forward # ------------------------------------------------------------------ def forward(self, idx: torch.Tensor) -> torch.Tensor: B, T = idx.shape config = self.config cos, sin = self._get_rope(T, self.wte.weight.dtype, self.wte.weight.device) x = self.wte(idx) x = rms_norm(x) # Smear: x_smeared[t] = x[t] + (smear_lambda * sigmoid(smear_gate(x[t,:24]))) * x[t-1], # for t >= 1; x[0] passes through. gate = self.smear_lambda * torch.sigmoid(self.smear_gate(x[:, 1:, :24])) # (B, T-1, 1) x_smeared = x[:, 1:] + gate * x[:, :-1] x = torch.cat([x[:, :1], x_smeared], dim=1) x0 = x n_layer = config.n_layer backout_layer = n_layer // 2 x_backout = None for i, block in enumerate(self.blocks): x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0 ve_key = str(i) ve = self.value_embeds[ve_key](idx).to(dtype=x.dtype) if ve_key in self.value_embeds else None x = block(x, ve, cos, sin, self.window_sizes[i]) if i == backout_layer: x_backout = x if x_backout is not None: x = x - self.backout_lambda * x_backout x = rms_norm(x) # Project to (padded) vocab, trim to real vocab, then softcap. softcap = 15.0 if self.tie_embeddings: logits = x @ self.wte.weight.t() else: logits = self.lm_head(x) logits = logits[..., : config.vocab_size] logits = logits.to(torch.float32) logits = softcap * torch.tanh(logits / softcap) return logits # ------------------------------------------------------------------ # Weight loading # ------------------------------------------------------------------ @classmethod def from_pretrained(cls, path: str, map_location: str | torch.device = "cpu") -> "GELU_GPT": """Load a torch state dict + config produced by convert_flax_to_torch.py. The checkpoint file stores both config and state_dict as: {"config": {...}, "state_dict": {...}} """ payload = torch.load(path, map_location=map_location, weights_only=False) if isinstance(payload, dict) and "config" in payload and "state_dict" in payload: config = GPTConfig(**payload["config"]) state_dict = payload["state_dict"] else: raise ValueError( f"{path} does not look like a converted flaxchat checkpoint; " "expected a dict with keys {'config', 'state_dict'}." ) model = cls(config) missing, unexpected = model.load_state_dict(state_dict, strict=False) # Buffers (rope_cos/sin) are not saved — they're recomputed in __init__. # Everything else must match. real_missing = [k for k in missing if not k.startswith("rope_")] if real_missing: raise RuntimeError(f"Missing keys when loading: {real_missing}") if unexpected: raise RuntimeError(f"Unexpected keys when loading: {unexpected}") model.eval() return model __all__ = ["GPTConfig", "GELU_GPT", "precompute_rotary_embeddings", "apply_rotary_emb", "rms_norm"]