""" PyTorch port of the flaxchat GPT model with YatNMN-Softplus MLP. Matches nmn.nnx.layers.YatNMN forward for the config used in training: use_bias=True, softplus_bias=True, learnable_epsilon=True scalar_bias: False (per-neuron) OR True (shared (1,)) constant_alpha: False (learnable) OR True (α=1 fixed) YatNMN formula (see nmn/nnx/layers/nmn.py:291): y_dot = x @ W # (..., out) dist² = ||x||² + ||W_j||² - 2·y_dot # (..., out) y_num = y_dot + softplus(bias) # if use_bias & softplus_bias out = α · y_num² / (dist² + softplus(ε)) All other features (RoPE, GQA, QK-norm, RMSNorm, value embeds, smear, backout, softcap, sliding-window, tied embeddings, no biases in Linear) match `torch_gpt.py` exactly. """ from __future__ import annotations import math from dataclasses import dataclass, field from typing import Optional, Tuple, List import torch import torch.nn as nn import torch.nn.functional as F try: from .torch_gpt import ( rms_norm, precompute_rotary_embeddings, apply_rotary_emb, has_ve, compute_window_sizes, ) except ImportError: from torch_port.torch_gpt import ( rms_norm, precompute_rotary_embeddings, apply_rotary_emb, has_ve, compute_window_sizes, ) # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- @dataclass class YatGPTConfig: 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 # YatNMN-specific mlp_type: str = "yatnmn-softplus" scalar_bias: bool = False # False = per-neuron (ff,) bias; True = shared (1,) softplus_bias: bool = True learnable_epsilon: bool = True epsilon_init: float = 1e-3 constant_alpha: bool = False # False = learnable α; True = α fixed at 1 @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 # --------------------------------------------------------------------------- # YatNMN layer # --------------------------------------------------------------------------- class YatNMN(nn.Module): """PyTorch port of nmn.nnx.layers.YatNMN matching the flaxchat training config.""" def __init__( self, in_features: int, out_features: int, use_bias: bool = True, softplus_bias: bool = True, scalar_bias: bool = False, learnable_epsilon: bool = True, epsilon_init: float = 1e-3, use_alpha: bool = True, constant_alpha: bool = False, ): super().__init__() self.in_features = in_features self.out_features = out_features self.use_bias = use_bias self.softplus_bias = softplus_bias self.scalar_bias = scalar_bias self.learnable_epsilon = learnable_epsilon self.use_alpha = use_alpha self.constant_alpha = constant_alpha # kernel shape matches JAX nmn: (in_features, out_features) self.kernel = nn.Parameter(torch.empty(in_features, out_features)) nn.init.trunc_normal_(self.kernel, mean=0.0, std=1.0 / math.sqrt(in_features)) if use_bias: b_shape = (1,) if scalar_bias else (out_features,) self.bias = nn.Parameter(torch.zeros(b_shape)) else: self.register_parameter("bias", None) if learnable_epsilon: # softplus(x) = log(1+exp(x)); we want softplus(raw) = epsilon_init. # → raw = log(exp(epsilon_init) - 1) = log(expm1(epsilon_init)) raw = math.log(math.expm1(epsilon_init)) self.epsilon_param = nn.Parameter(torch.full((1,), raw)) self._epsilon_const = None else: self.register_parameter("epsilon_param", None) self._epsilon_const = epsilon_init if use_alpha and not constant_alpha: self.alpha = nn.Parameter(torch.ones(1)) self._alpha_const_value: float | None = None elif use_alpha and constant_alpha: self.register_parameter("alpha", None) # Flax `nmn.YatNMN(constant_alpha=True)` resolves to # `DEFAULT_CONSTANT_ALPHA = jnp.sqrt(2.0)` — verified by direct # library probe with nmn 0.2.29. Stored as a plain Python float so # HF `from_pretrained`'s meta-init can't zero it (which happens # with non-persistent buffers) and old safetensors with the # wrong on-disk value (1.0) can't shadow it. self._alpha_const_value: float | None = math.sqrt(2.0) else: self.register_parameter("alpha", None) self._alpha_const_value: float | None = None def forward(self, x: torch.Tensor) -> torch.Tensor: # Match nmn's fp32 path: (y_num² / (dist + ε)) is unstable in bf16. orig_dtype = x.dtype x32 = x.to(torch.float32) W = self.kernel.to(torch.float32) # (in, out) y_dot = x32 @ W # (..., out) # ||x - W_j||² = ||x||² - 2·x·W_j + ||W_j||² x_sq = (x32 * x32).sum(dim=-1, keepdim=True) # (..., 1) W_sq = (W * W).sum(dim=0, keepdim=False) # (out,) distances = torch.clamp(x_sq + W_sq - 2.0 * y_dot, min=0.0) # (..., out) # numerator if self.use_bias and self.bias is not None: b = self.bias.to(torch.float32) if self.softplus_bias: b = F.softplus(b) y_num = y_dot + b # broadcast: b is (1,) or (out,) else: y_num = y_dot # epsilon if self.learnable_epsilon: eps = F.softplus(self.epsilon_param.to(torch.float32)) else: eps = torch.tensor(self._epsilon_const, dtype=torch.float32, device=y_num.device) out = (y_num * y_num) / (distances + eps) if self.use_alpha: if self.alpha is not None: out = out * self.alpha.to(torch.float32) elif self._alpha_const_value is not None: out = out * self._alpha_const_value return out.to(orig_dtype) # --------------------------------------------------------------------------- # Attention (shared with torch_gpt) — redefined here so this file is # self-contained when loaded via trust_remote_code. # --------------------------------------------------------------------------- class CausalSelfAttention(nn.Module): def __init__(self, config: YatGPTConfig, 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) self.ve_gate = nn.Linear(12, config.n_kv_head, bias=False) if self._has_ve else 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 cfg = self.config n_head, n_kv_head, head_dim = cfg.n_head, cfg.n_kv_head, cfg.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])) 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) 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 if 0 < window_left < T: causal_mask = causal_mask & ((row_idx - col_idx) <= window_left) bias = torch.where( causal_mask, torch.zeros((), dtype=x.dtype, device=device), torch.full((), -1e9, dtype=x.dtype, device=device), ).unsqueeze(0).unsqueeze(0) q_bhtd = q.transpose(1, 2) k_bhtd = k.transpose(1, 2) v_bhtd = v.transpose(1, 2) scale = 1.0 / math.sqrt(head_dim) att = torch.matmul(q_bhtd, k_bhtd.transpose(-2, -1)) * scale att = att + bias att = F.softmax(att, dim=-1) y = torch.matmul(att, v_bhtd) y = y.transpose(1, 2).contiguous().reshape(B, T, -1) return self.c_proj(y) # --------------------------------------------------------------------------- # MLP (YatNMN variant) # --------------------------------------------------------------------------- class YatMLP(nn.Module): def __init__(self, config: YatGPTConfig): super().__init__() n, ff = config.n_embd, 4 * config.n_embd self.c_fc = YatNMN( n, ff, use_bias=True, softplus_bias=config.softplus_bias, scalar_bias=config.scalar_bias, learnable_epsilon=config.learnable_epsilon, epsilon_init=config.epsilon_init, use_alpha=True, constant_alpha=config.constant_alpha, ) self.c_proj = nn.Linear(ff, n, bias=False) # Training used zeros-init on c_proj (GPT._init_weights patched); init here matches. nn.init.zeros_(self.c_proj.weight) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.c_proj(self.c_fc(x)) # --------------------------------------------------------------------------- # Block # --------------------------------------------------------------------------- class Block(nn.Module): def __init__(self, config: YatGPTConfig, layer_idx: int): super().__init__() self.attn = CausalSelfAttention(config, layer_idx) self.mlp = YatMLP(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 Yat_GPT(nn.Module): """PyTorch port of flaxchat.gpt.GPT with YatNMN-Softplus MLP. Layer/parameter naming mirrors the Flax module tree exactly so the weight converter is a direct key-for-key mapping. """ def __init__(self, config: YatGPTConfig): super().__init__() self.config = config self.window_sizes = compute_window_sizes(_ConfigShim(config)) padded_vocab = config.padded_vocab_size self.padded_vocab_size = padded_vocab self.wte = nn.Embedding(padded_vocab, config.n_embd) self.blocks = nn.ModuleList([Block(config, i) for i in range(config.n_layer)]) self.tie_embeddings = config.tie_embeddings self.lm_head = None if config.tie_embeddings else nn.Linear(config.n_embd, padded_vocab, bias=False) self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer)) self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer)) self.smear_gate = nn.Linear(24, 1, bias=False) self.smear_lambda = nn.Parameter(torch.zeros(1)) self.backout_lambda = nn.Parameter(0.2 * torch.ones(1)) 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's meta-init can leave persistent # buffers as garbage → NaN logits. Compute on first forward instead. 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) def forward(self, idx: torch.Tensor) -> torch.Tensor: B, T = idx.shape cfg = self.config cos, sin = self._get_rope(T, self.wte.weight.dtype, self.wte.weight.device) x = self.wte(idx) x = rms_norm(x) gate = self.smear_lambda * torch.sigmoid(self.smear_gate(x[:, 1:, :24])) x_smeared = x[:, 1:] + gate * x[:, :-1] x = torch.cat([x[:, :1], x_smeared], dim=1) x0 = x n_layer = cfg.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) softcap = 15.0 logits = x @ self.wte.weight.t() if self.tie_embeddings else self.lm_head(x) logits = logits[..., : cfg.vocab_size].to(torch.float32) return softcap * torch.tanh(logits / softcap) @classmethod def from_pretrained(cls, path: str, map_location: str | torch.device = "cpu") -> "Yat_GPT": payload = torch.load(path, map_location=map_location, weights_only=False) if not (isinstance(payload, dict) and "config" in payload and "state_dict" in payload): raise ValueError(f"{path} must contain {{'config', 'state_dict'}}") config = YatGPTConfig(**payload["config"]) model = cls(config) missing, unexpected = model.load_state_dict(payload["state_dict"], strict=False) # rope_* are recomputed buffers; _alpha_const is a fixed buffer (not saved in checkpoint) real_missing = [k for k in missing if not k.startswith("rope_") and "_alpha_const" not in k] 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 class _ConfigShim: """Small shim so compute_window_sizes (which expects .sequence_len, .window_pattern, .n_layer on a GPTConfig) works when given a YatGPTConfig.""" def __init__(self, cfg: YatGPTConfig): self.sequence_len = cfg.sequence_len self.window_pattern = cfg.window_pattern self.n_layer = cfg.n_layer __all__ = ["YatGPTConfig", "Yat_GPT", "YatNMN"]