Cozet / model.py

Upload model.py with huggingface_hub

b6833e9 verified 4 months ago

14.5 kB

"""
Cozet: Native SYNAXIM Base Model
=================================
Non-transformer architecture using Symbiotic Gate (M-matrix)
instead of self-attention. Trained from scratch.
Architecture per layer:
  RMSNorm -> Q/K/V proj -> Symbiotic Gate (M update + retrieval) -> O proj -> Residual
  RMSNorm -> SwiGLU MLP -> Residual
(c) 2026 GRRN Research. All rights reserved.
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass
from typing import Optional, List, Tuple


@dataclass
class CozConfig:
    """Cozet model configuration."""
    hidden_size: int = 1024
    num_layers: int = 12
    num_attention_heads: int = 16
    num_kv_heads: int = 4
    intermediate_size: int = 4096
    vocab_size: int = 32000
    max_seq_len: int = 4096
    rope_theta: float = 10000.0
    rms_norm_eps: float = 1e-6
    memory_decay: float = 0.995
    tie_word_embeddings: bool = True
    
    @property
    def head_dim(self) -> int:
        return self.hidden_size // self.num_attention_heads
    
    def num_params(self) -> int:
        """Estimate total parameter count."""
        D = self.hidden_size
        V = self.vocab_size
        I = self.intermediate_size
        L = self.num_layers
        n_kv = self.num_kv_heads
        hd = self.head_dim
        
        embed = V * D
        per_layer = (
            D * D +                    # q_proj
            D * (n_kv * hd) +          # k_proj
            D * (n_kv * hd) +          # v_proj
            D * D +                    # o_proj
            D * I +                    # gate_proj
            D * I +                    # up_proj
            I * D +                    # down_proj
            2 * D +                    # norms
            2                          # gate params
        )
        lm_head = 0 if self.tie_word_embeddings else V * D
        total = embed + L * per_layer + lm_head + D
        return total


# ---- Presets ----

COZET_SMALL = CozConfig(
    hidden_size=1024, num_layers=12, num_attention_heads=16,
    num_kv_heads=4, intermediate_size=4096, vocab_size=32000,
    max_seq_len=4096, tie_word_embeddings=True,
)

COZET_MEDIUM = CozConfig(
    hidden_size=2048, num_layers=24, num_attention_heads=16,
    num_kv_heads=4, intermediate_size=8192, vocab_size=32000,
    max_seq_len=8192, rope_theta=500000.0, tie_word_embeddings=True,
)

COZET_LARGE = CozConfig(
    hidden_size=4096, num_layers=32, num_attention_heads=32,
    num_kv_heads=8, intermediate_size=14336, vocab_size=32000,
    max_seq_len=8192, rope_theta=1000000.0, tie_word_embeddings=False,
)


class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization."""
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return (x.float() * norm).type_as(x) * self.weight


class SymbioticGate(nn.Module):
    """
    The Symbiotic Gate: core SYNAXIM attention replacement.
    
    Instead of softmax(QK^T/sqrt(d)) @ V, computes:
      gate = sigmoid(mean(Q * K) * scale * gate_scale + gate_bias)
      M    = gate * decay * M + (1-gate) * outer(k_norm, v_norm)
      out  = q @ M @ W_o
    
    M is persistent state: O(1) memory, infinite context.
    Fully differentiable for training.
    """
    
    def __init__(self, config: CozConfig, layer_idx: int):
        super().__init__()
        self.D = config.hidden_size
        self.n_heads = config.num_attention_heads
        self.n_kv = config.num_kv_heads
        self.head_dim = config.head_dim
        self.decay = config.memory_decay
        self.layer_idx = layer_idx
        
        # Projections
        self.q_proj = nn.Linear(self.D, self.n_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.D, self.n_kv * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.D, self.n_kv * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.n_heads * self.head_dim, self.D, bias=False)
        
        # Learnable gate parameters
        self.gate_scale = nn.Parameter(torch.ones(1))
        self.gate_bias = nn.Parameter(torch.zeros(1))
        
        # RoPE tables (precomputed, not learned)
        self._rope_built = False
    
    def _build_rope(self, max_seq: int, device: torch.device):
        """Precompute RoPE cos/sin tables."""
        if self._rope_built and max_seq <= self.rope_cos.shape[0]:
            return
        self._rope_built = False  # Rebuild if we need more positions
        half = self.head_dim // 2
        freqs = 1.0 / (10000.0 ** (torch.arange(0, half, device=device).float() / half))
        t = torch.arange(max_seq, device=device).float()
        angles = torch.outer(t, freqs)
        self.register_buffer("rope_cos", angles.cos(), persistent=False)
        self.register_buffer("rope_sin", angles.sin(), persistent=False)
        self._rope_built = True
    
    def _apply_rope(self, x: torch.Tensor, pos: int) -> torch.Tensor:
        """Apply RoPE to (n_heads, head_dim) tensor at position pos."""
        half = self.head_dim // 2
        cos = self.rope_cos[pos, :half]
        sin = self.rope_sin[pos, :half]
        x1, x2 = x[..., :half], x[..., half:]
        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
    
    def forward(self, h: torch.Tensor, M: torch.Tensor,
                position: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Symbiotic Gate forward pass for a single token.
        
        Args:
            h: (D,) hidden state
            M: (D, D) persistent memory matrix
            position: token position for RoPE
        
        Returns:
            output: (D,) attention output
            M_new: (D, D) updated memory
        """
        self._build_rope(position + 1, h.device)
        
        # Project
        q = self.q_proj(h)    # (n_heads * hd,)
        k = self.k_proj(h)    # (n_kv * hd,)
        v = self.v_proj(h)    # (n_kv * hd,)
        
        # Reshape for per-head ops
        q_h = q.view(self.n_heads, self.head_dim)
        k_h = k.view(self.n_kv, self.head_dim)
        v_h = v.view(self.n_kv, self.head_dim)
        
        # RoPE
        q_h = self._apply_rope(q_h, position)
        k_h = self._apply_rope(k_h, position)
        
        # GQA: repeat KV heads
        if self.n_kv < self.n_heads:
            repeat = self.n_heads // self.n_kv
            k_h = k_h.repeat_interleave(repeat, dim=0)
            v_h = v_h.repeat_interleave(repeat, dim=0)
        
        # Gate score
        scale = 1.0 / math.sqrt(self.head_dim)
        gate_score = (q_h * k_h).sum(-1).mean() * scale
        gate_score = gate_score * self.gate_scale + self.gate_bias
        gate = torch.sigmoid(gate_score)
        
        # Flatten to D-space
        k_flat = k_h.reshape(-1)
        v_flat = v_h.reshape(-1)
        
        # Normalize
        k_norm = k_flat / (k_flat.norm() + 1e-8)
        v_norm = v_flat / (v_flat.norm() + 1e-8) * h.norm()
        
        # M-matrix update (differentiable)
        imprint = torch.outer(k_norm, v_norm)
        M_new = gate * self.decay * M + (1.0 - gate) * imprint
        
        # Retrieve
        q_flat = q_h.reshape(-1)
        output = q_flat @ M_new
        
        # Output projection
        result = self.o_proj(output)
        
        return result, M_new


class SynaxBlock(nn.Module):
    """One SYNAXIM layer: SymbioticGate + SwiGLU MLP."""
    
    def __init__(self, config: CozConfig, layer_idx: int):
        super().__init__()
        self.norm_attn = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.attn = SymbioticGate(config, layer_idx)
        self.norm_mlp = RMSNorm(config.hidden_size, config.rms_norm_eps)
        
        # SwiGLU MLP
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
    
    def forward(self, h: torch.Tensor, M: torch.Tensor,
                position: int) -> Tuple[torch.Tensor, torch.Tensor]:
        # Attention
        h_normed = self.norm_attn(h)
        attn_out, M_new = self.attn(h_normed, M, position)
        h = h + attn_out
        
        # MLP
        h_normed = self.norm_mlp(h)
        gate_out = self.gate_proj(h_normed)
        up_out = self.up_proj(h_normed)
        mlp_out = self.down_proj(F.silu(gate_out) * up_out)
        h = h + mlp_out
        
        return h, M_new


class CozModel(nn.Module):
    """
    Cozet: Native SYNAXIM base model.
    
    Non-transformer architecture trained from scratch.
    Uses Symbiotic Gate (M-matrix) instead of self-attention.
    O(1) memory at inference. Infinite context. No KV cache.
    """
    
    def __init__(self, config: CozConfig):
        super().__init__()
        self.config = config
        self.D = config.hidden_size
        self.n_layers = config.num_layers
        
        # Embedding
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        
        # SYNAXIM layers
        self.layers = nn.ModuleList([
            SynaxBlock(config, i) for i in range(config.num_layers)
        ])
        
        # Final norm
        self.final_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        
        # LM head
        if not config.tie_word_embeddings:
            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        
        # Initialize weights
        self._init_weights()
    
    def _init_weights(self):
        """SYNAXIM-specific weight initialization."""
        depth_scale = 1.0 / math.sqrt(2 * self.n_layers)
        
        for name, p in self.named_parameters():
            if p.dim() == 1:
                continue  # Norms already initialized to ones
            if "embed" in name:
                nn.init.normal_(p, std=0.02)
            elif "down_proj" in name or "o_proj" in name:
                nn.init.normal_(p, std=0.02 * depth_scale)
            elif "lm_head" in name:
                nn.init.normal_(p, std=0.02 * depth_scale)
            elif p.dim() == 2:
                nn.init.normal_(p, std=0.02)
    
    def init_m_states(self, device: torch.device) -> List[torch.Tensor]:
        """Create fresh M-matrix states for a new sequence."""
        return [torch.zeros(self.D, self.D, device=device)
                for _ in range(self.n_layers)]
    
    def forward_token(self, token_id: int, M_states: List[torch.Tensor],
                       position: int) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        """
        Process one token through the full model.
        
        Returns: (logits, updated_M_states)
        """
        h = self.embed_tokens.weight[token_id]
        
        for i, layer in enumerate(self.layers):
            h, M_states[i] = layer(h, M_states[i], position)
        
        h = self.final_norm(h)
        
        if self.config.tie_word_embeddings:
            logits = h @ self.embed_tokens.weight.T
        else:
            logits = self.lm_head(h)
        
        return logits, M_states
    
    def forward_sequence(self, token_ids: torch.Tensor,
                          chunk_size: int = 256) -> torch.Tensor:
        """
        Process a sequence with truncated BPTT.
        
        Args:
            token_ids: (seq_len,) token IDs
            chunk_size: gradient truncation window
        
        Returns:
            total_loss: scalar loss for the sequence
        """
        device = token_ids.device
        seq_len = token_ids.shape[0]
        M_states = self.init_m_states(device)
        total_loss = torch.tensor(0.0, device=device)
        n_tokens = 0
        
        for chunk_start in range(0, seq_len - 1, chunk_size):
            chunk_end = min(chunk_start + chunk_size, seq_len - 1)
            
            # Detach M at chunk boundaries (truncated BPTT)
            M_states = [m.detach() for m in M_states]
            
            chunk_loss = torch.tensor(0.0, device=device, requires_grad=True)
            
            for t in range(chunk_start, chunk_end):
                logits, M_states = self.forward_token(
                    token_ids[t].item(), M_states, t
                )
                loss = F.cross_entropy(
                    logits.unsqueeze(0),
                    token_ids[t + 1].unsqueeze(0)
                )
                chunk_loss = chunk_loss + loss
                n_tokens += 1
            
            total_loss = total_loss + chunk_loss
        
        return total_loss / max(n_tokens, 1)


def create_model(size: str = "small") -> CozModel:
    """Create a Cozet model by size name."""
    configs = {
        "small": COZET_SMALL,
        "medium": COZET_MEDIUM,
        "large": COZET_LARGE,
    }
    config = configs.get(size, COZET_SMALL)
    model = CozModel(config)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"Cozet-{size}: {n_params:,} parameters")
    print(f"  D={config.hidden_size}, L={config.num_layers}, "
          f"H={config.num_attention_heads}/{config.num_kv_heads} GQA, "
          f"I={config.intermediate_size}, V={config.vocab_size}")
    return model


if __name__ == "__main__":
    # Quick sanity check
    model = create_model("small")
    device = torch.device("cpu")
    
    M_states = model.init_m_states(device)
    
    # Forward one token
    logits, M_states = model.forward_token(42, M_states, 0)
    print(f"Logits shape: {logits.shape}")
    print(f"M[0] norm after 1 token: {M_states[0].norm():.4f}")
    
    # Forward second token
    logits2, M_states = model.forward_token(100, M_states, 1)
    print(f"M[0] norm after 2 tokens: {M_states[0].norm():.4f}")
    
    # Test gradient flow
    M_states_fresh = model.init_m_states(device)
    logits3, _ = model.forward_token(42, M_states_fresh, 0)
    loss = F.cross_entropy(logits3.unsqueeze(0), torch.tensor([100]))
    loss.backward()
    
    grad_norms = {n: p.grad.norm().item() for n, p in model.named_parameters() 
                  if p.grad is not None}
    print(f"Gradient flow: {len(grad_norms)} parameters received gradients")
    print(f"Loss: {loss.item():.4f}")
    print("Sanity check PASSED.")

	"""
	Cozet: Native SYNAXIM Base Model
	=================================
	Non-transformer architecture using Symbiotic Gate (M-matrix)
	instead of self-attention. Trained from scratch.

	Architecture per layer:
	RMSNorm -> Q/K/V proj -> Symbiotic Gate (M update + retrieval) -> O proj -> Residual
	RMSNorm -> SwiGLU MLP -> Residual

	(c) 2026 GRRN Research. All rights reserved.
	"""

	import math
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from dataclasses import dataclass
	from typing import Optional, List, Tuple


	@dataclass
	class CozConfig:
	"""Cozet model configuration."""
	hidden_size: int = 1024
	num_layers: int = 12
	num_attention_heads: int = 16
	num_kv_heads: int = 4
	intermediate_size: int = 4096
	vocab_size: int = 32000
	max_seq_len: int = 4096
	rope_theta: float = 10000.0
	rms_norm_eps: float = 1e-6
	memory_decay: float = 0.995
	tie_word_embeddings: bool = True

	@property
	def head_dim(self) -> int:
	return self.hidden_size // self.num_attention_heads

	def num_params(self) -> int:
	"""Estimate total parameter count."""
	D = self.hidden_size
	V = self.vocab_size
	I = self.intermediate_size
	L = self.num_layers
	n_kv = self.num_kv_heads
	hd = self.head_dim

	embed = V * D
	per_layer = (
	D * D + # q_proj
	D * (n_kv * hd) + # k_proj
	D * (n_kv * hd) + # v_proj
	D * D + # o_proj
	D * I + # gate_proj
	D * I + # up_proj
	I * D + # down_proj
	2 * D + # norms
	2 # gate params
	)
	lm_head = 0 if self.tie_word_embeddings else V * D
	total = embed + L * per_layer + lm_head + D
	return total


	# ---- Presets ----

	COZET_SMALL = CozConfig(
	hidden_size=1024, num_layers=12, num_attention_heads=16,
	num_kv_heads=4, intermediate_size=4096, vocab_size=32000,
	max_seq_len=4096, tie_word_embeddings=True,
	)

	COZET_MEDIUM = CozConfig(
	hidden_size=2048, num_layers=24, num_attention_heads=16,
	num_kv_heads=4, intermediate_size=8192, vocab_size=32000,
	max_seq_len=8192, rope_theta=500000.0, tie_word_embeddings=True,
	)

	COZET_LARGE = CozConfig(
	hidden_size=4096, num_layers=32, num_attention_heads=32,
	num_kv_heads=8, intermediate_size=14336, vocab_size=32000,
	max_seq_len=8192, rope_theta=1000000.0, tie_word_embeddings=False,
	)


	class RMSNorm(nn.Module):
	"""Root Mean Square Layer Normalization."""
	def __init__(self, dim: int, eps: float = 1e-6):
	super().__init__()
	self.weight = nn.Parameter(torch.ones(dim))
	self.eps = eps

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
	return (x.float() * norm).type_as(x) * self.weight


	class SymbioticGate(nn.Module):
	"""
	The Symbiotic Gate: core SYNAXIM attention replacement.

	Instead of softmax(QK^T/sqrt(d)) @ V, computes:
	gate = sigmoid(mean(Q * K) * scale * gate_scale + gate_bias)
	M = gate * decay * M + (1-gate) * outer(k_norm, v_norm)
	out = q @ M @ W_o

	M is persistent state: O(1) memory, infinite context.
	Fully differentiable for training.
	"""

	def __init__(self, config: CozConfig, layer_idx: int):
	super().__init__()
	self.D = config.hidden_size
	self.n_heads = config.num_attention_heads
	self.n_kv = config.num_kv_heads
	self.head_dim = config.head_dim
	self.decay = config.memory_decay
	self.layer_idx = layer_idx

	# Projections
	self.q_proj = nn.Linear(self.D, self.n_heads * self.head_dim, bias=False)
	self.k_proj = nn.Linear(self.D, self.n_kv * self.head_dim, bias=False)
	self.v_proj = nn.Linear(self.D, self.n_kv * self.head_dim, bias=False)
	self.o_proj = nn.Linear(self.n_heads * self.head_dim, self.D, bias=False)

	# Learnable gate parameters
	self.gate_scale = nn.Parameter(torch.ones(1))
	self.gate_bias = nn.Parameter(torch.zeros(1))

	# RoPE tables (precomputed, not learned)
	self._rope_built = False

	def _build_rope(self, max_seq: int, device: torch.device):
	"""Precompute RoPE cos/sin tables."""
	if self._rope_built and max_seq <= self.rope_cos.shape[0]:
	return
	self._rope_built = False # Rebuild if we need more positions
	half = self.head_dim // 2
	freqs = 1.0 / (10000.0 ** (torch.arange(0, half, device=device).float() / half))
	t = torch.arange(max_seq, device=device).float()
	angles = torch.outer(t, freqs)
	self.register_buffer("rope_cos", angles.cos(), persistent=False)
	self.register_buffer("rope_sin", angles.sin(), persistent=False)
	self._rope_built = True

	def _apply_rope(self, x: torch.Tensor, pos: int) -> torch.Tensor:
	"""Apply RoPE to (n_heads, head_dim) tensor at position pos."""
	half = self.head_dim // 2
	cos = self.rope_cos[pos, :half]
	sin = self.rope_sin[pos, :half]
	x1, x2 = x[..., :half], x[..., half:]
	return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)

	def forward(self, h: torch.Tensor, M: torch.Tensor,
	position: int) -> Tuple[torch.Tensor, torch.Tensor]:
	"""
	Symbiotic Gate forward pass for a single token.

	Args:
	h: (D,) hidden state
	M: (D, D) persistent memory matrix
	position: token position for RoPE

	Returns:
	output: (D,) attention output
	M_new: (D, D) updated memory
	"""
	self._build_rope(position + 1, h.device)

	# Project
	q = self.q_proj(h) # (n_heads * hd,)
	k = self.k_proj(h) # (n_kv * hd,)
	v = self.v_proj(h) # (n_kv * hd,)

	# Reshape for per-head ops
	q_h = q.view(self.n_heads, self.head_dim)
	k_h = k.view(self.n_kv, self.head_dim)
	v_h = v.view(self.n_kv, self.head_dim)

	# RoPE
	q_h = self._apply_rope(q_h, position)
	k_h = self._apply_rope(k_h, position)

	# GQA: repeat KV heads
	if self.n_kv < self.n_heads:
	repeat = self.n_heads // self.n_kv
	k_h = k_h.repeat_interleave(repeat, dim=0)
	v_h = v_h.repeat_interleave(repeat, dim=0)

	# Gate score
	scale = 1.0 / math.sqrt(self.head_dim)
	gate_score = (q_h * k_h).sum(-1).mean() * scale
	gate_score = gate_score * self.gate_scale + self.gate_bias
	gate = torch.sigmoid(gate_score)

	# Flatten to D-space
	k_flat = k_h.reshape(-1)
	v_flat = v_h.reshape(-1)

	# Normalize
	k_norm = k_flat / (k_flat.norm() + 1e-8)
	v_norm = v_flat / (v_flat.norm() + 1e-8) * h.norm()

	# M-matrix update (differentiable)
	imprint = torch.outer(k_norm, v_norm)
	M_new = gate * self.decay * M + (1.0 - gate) * imprint

	# Retrieve
	q_flat = q_h.reshape(-1)
	output = q_flat @ M_new

	# Output projection
	result = self.o_proj(output)

	return result, M_new


	class SynaxBlock(nn.Module):
	"""One SYNAXIM layer: SymbioticGate + SwiGLU MLP."""

	def __init__(self, config: CozConfig, layer_idx: int):
	super().__init__()
	self.norm_attn = RMSNorm(config.hidden_size, config.rms_norm_eps)
	self.attn = SymbioticGate(config, layer_idx)
	self.norm_mlp = RMSNorm(config.hidden_size, config.rms_norm_eps)

	# SwiGLU MLP
	self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
	self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
	self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

	def forward(self, h: torch.Tensor, M: torch.Tensor,
	position: int) -> Tuple[torch.Tensor, torch.Tensor]:
	# Attention
	h_normed = self.norm_attn(h)
	attn_out, M_new = self.attn(h_normed, M, position)
	h = h + attn_out

	# MLP
	h_normed = self.norm_mlp(h)
	gate_out = self.gate_proj(h_normed)
	up_out = self.up_proj(h_normed)
	mlp_out = self.down_proj(F.silu(gate_out) * up_out)
	h = h + mlp_out

	return h, M_new


	class CozModel(nn.Module):
	"""
	Cozet: Native SYNAXIM base model.

	Non-transformer architecture trained from scratch.
	Uses Symbiotic Gate (M-matrix) instead of self-attention.
	O(1) memory at inference. Infinite context. No KV cache.
	"""

	def __init__(self, config: CozConfig):
	super().__init__()
	self.config = config
	self.D = config.hidden_size
	self.n_layers = config.num_layers

	# Embedding
	self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)

	# SYNAXIM layers
	self.layers = nn.ModuleList([
	SynaxBlock(config, i) for i in range(config.num_layers)
	])

	# Final norm
	self.final_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)

	# LM head
	if not config.tie_word_embeddings:
	self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

	# Initialize weights
	self._init_weights()

	def _init_weights(self):
	"""SYNAXIM-specific weight initialization."""
	depth_scale = 1.0 / math.sqrt(2 * self.n_layers)

	for name, p in self.named_parameters():
	if p.dim() == 1:
	continue # Norms already initialized to ones
	if "embed" in name:
	nn.init.normal_(p, std=0.02)
	elif "down_proj" in name or "o_proj" in name:
	nn.init.normal_(p, std=0.02 * depth_scale)
	elif "lm_head" in name:
	nn.init.normal_(p, std=0.02 * depth_scale)
	elif p.dim() == 2:
	nn.init.normal_(p, std=0.02)

	def init_m_states(self, device: torch.device) -> List[torch.Tensor]:
	"""Create fresh M-matrix states for a new sequence."""
	return [torch.zeros(self.D, self.D, device=device)
	for _ in range(self.n_layers)]

	def forward_token(self, token_id: int, M_states: List[torch.Tensor],
	position: int) -> Tuple[torch.Tensor, List[torch.Tensor]]:
	"""
	Process one token through the full model.

	Returns: (logits, updated_M_states)
	"""
	h = self.embed_tokens.weight[token_id]

	for i, layer in enumerate(self.layers):
	h, M_states[i] = layer(h, M_states[i], position)

	h = self.final_norm(h)

	if self.config.tie_word_embeddings:
	logits = h @ self.embed_tokens.weight.T
	else:
	logits = self.lm_head(h)

	return logits, M_states

	def forward_sequence(self, token_ids: torch.Tensor,
	chunk_size: int = 256) -> torch.Tensor:
	"""
	Process a sequence with truncated BPTT.

	Args:
	token_ids: (seq_len,) token IDs
	chunk_size: gradient truncation window

	Returns:
	total_loss: scalar loss for the sequence
	"""
	device = token_ids.device
	seq_len = token_ids.shape[0]
	M_states = self.init_m_states(device)
	total_loss = torch.tensor(0.0, device=device)
	n_tokens = 0

	for chunk_start in range(0, seq_len - 1, chunk_size):
	chunk_end = min(chunk_start + chunk_size, seq_len - 1)

	# Detach M at chunk boundaries (truncated BPTT)
	M_states = [m.detach() for m in M_states]

	chunk_loss = torch.tensor(0.0, device=device, requires_grad=True)

	for t in range(chunk_start, chunk_end):
	logits, M_states = self.forward_token(
	token_ids[t].item(), M_states, t
	)
	loss = F.cross_entropy(
	logits.unsqueeze(0),
	token_ids[t + 1].unsqueeze(0)
	)
	chunk_loss = chunk_loss + loss
	n_tokens += 1

	total_loss = total_loss + chunk_loss

	return total_loss / max(n_tokens, 1)


	def create_model(size: str = "small") -> CozModel:
	"""Create a Cozet model by size name."""
	configs = {
	"small": COZET_SMALL,
	"medium": COZET_MEDIUM,
	"large": COZET_LARGE,
	}
	config = configs.get(size, COZET_SMALL)
	model = CozModel(config)
	n_params = sum(p.numel() for p in model.parameters())
	print(f"Cozet-{size}: {n_params:,} parameters")
	print(f" D={config.hidden_size}, L={config.num_layers}, "
	f"H={config.num_attention_heads}/{config.num_kv_heads} GQA, "
	f"I={config.intermediate_size}, V={config.vocab_size}")
	return model


	if __name__ == "__main__":
	# Quick sanity check
	model = create_model("small")
	device = torch.device("cpu")

	M_states = model.init_m_states(device)

	# Forward one token
	logits, M_states = model.forward_token(42, M_states, 0)
	print(f"Logits shape: {logits.shape}")
	print(f"M[0] norm after 1 token: {M_states[0].norm():.4f}")

	# Forward second token
	logits2, M_states = model.forward_token(100, M_states, 1)
	print(f"M[0] norm after 2 tokens: {M_states[0].norm():.4f}")

	# Test gradient flow
	M_states_fresh = model.init_m_states(device)
	logits3, _ = model.forward_token(42, M_states_fresh, 0)
	loss = F.cross_entropy(logits3.unsqueeze(0), torch.tensor([100]))
	loss.backward()

	grad_norms = {n: p.grad.norm().item() for n, p in model.named_parameters()
	if p.grad is not None}
	print(f"Gradient flow: {len(grad_norms)} parameters received gradients")
	print(f"Loss: {loss.item():.4f}")
	print("Sanity check PASSED.")