Delete model.py

model.py

@@ -1,575 +0,0 @@
-"""
-model.py  —  SLM Custom Architecture (v6.5 — fully HF-compatible)
-=================================================================
-Auto-written by the training notebook. Required for HF auto_map.
-
-v6.5 fixes:
-  1. _tied_weights_keys is a LIST — works on transformers 4.x AND 5.x.
-  2. RoPE buffers: persistent=False — NOT saved to safetensors.
-  3. tie_weights() properly delegates so HF re-ties weights on load.
-  4. Full HF forward() signature (return_dict, attention_mask, etc.).
-  5. CausalLMOutputWithPast used — generation helpers (.generate()) work.
-  6. prepare_inputs_for_generation() added.
-  7. Weights saved WITH model. prefix (matches SLMHFForCausalLM layout).
-"""
-from __future__ import annotations
-import math
-from dataclasses import dataclass, asdict
-from typing import Optional, Tuple, List
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-try:
-    from transformers import PretrainedConfig, PreTrainedModel
-    from transformers.modeling_outputs import CausalLMOutputWithPast
-    _HF_AVAILABLE = True
-except ImportError:
-    _HF_AVAILABLE = False
-
-
-# ══════════════════════════════════════════════════════════════════════════════
-# CONFIG DATACLASS
-# ══════════════════════════════════════════════════════════════════════════════
-
-@dataclass
-class SLMConfig:
-    # Vocab / embedding
-    vocab_size:        int   = 4096
-    hidden_size:       int   = 256
-    tie_embeddings:    bool  = True
-
-    # Depth / width
-    num_layers:        int   = 11
-    intermediate_size: int   = 768
-
-    # Attention
-    num_heads:         int   = 4
-    num_kv_heads:      int   = 2
-    max_seq_len:       int   = 1024
-    rope_theta:        float = 10000.0
-    qk_norm:           bool  = True
-    use_bias:          bool  = False
-
-    # FFN type: 'swiglu' | 'geglu' | 'mlp'
-    ffn_type:          str   = 'swiglu'
-    activation:        str   = 'silu'   # used by 'mlp' only
-
-    # XSA
-    xsa:               bool  = False
-    xsa_window:        int   = 128
-
-    # Regularisation
-    dropout:           float = 0.0
-    norm_eps:          float = 1e-5
-
-    @property
-    def head_dim(self) -> int:
-        return self.hidden_size // self.num_heads
-
-    def validate(self) -> 'SLMConfig':
-        assert self.hidden_size % self.num_heads == 0, \
-            f"hidden_size {self.hidden_size} % num_heads {self.num_heads} != 0"
-        assert self.num_heads % self.num_kv_heads == 0, \
-            f"num_heads {self.num_heads} % num_kv_heads {self.num_kv_heads} != 0"
-        assert self.ffn_type in ('swiglu', 'geglu', 'mlp'), \
-            f"ffn_type must be 'swiglu'|'geglu'|'mlp', got {self.ffn_type!r}"
-        assert self.activation in ('silu', 'gelu', 'relu'), \
-            f"activation must be 'silu'|'gelu'|'relu', got {self.activation!r}"
-        if self.xsa:
-            assert self.xsa_window > 0, "xsa_window must be > 0"
-        return self
-
-    def to_dict(self) -> dict:
-        d = asdict(self)
-        d['model_type'] = 'slm'
-        return d
-
-    @classmethod
-    def from_dict(cls, d: dict) -> 'SLMConfig':
-        valid = {f for f in cls.__dataclass_fields__}
-        return cls(**{k: v for k, v in d.items() if k in valid})
-
-
-# ══════════════════════════════════════════════════════════════════════════════
-# PRIMITIVES
-# ══════════════════════════════════════════════════════════════════════════════
-
-class RMSNorm(nn.Module):
-    def __init__(self, dim: int, eps: float = 1e-5):
-        super().__init__()
-        self.eps = eps
-        self.w   = nn.Parameter(torch.ones(dim))
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.w
-
-
-def precompute_rope(head_dim: int, max_len: int,
-                    theta: float = 10000.0) -> Tuple[torch.Tensor, torch.Tensor]:
-    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
-    t     = torch.arange(max_len, dtype=torch.float32)
-    cos   = torch.cos(torch.outer(t, freqs))
-    sin   = torch.sin(torch.outer(t, freqs))
-    return cos, sin
-
-
-def apply_rope(x: torch.Tensor, cos: torch.Tensor,
-               sin: torch.Tensor) -> torch.Tensor:
-    """Correct split-half RoPE rotation. x: [B,H,T,D]"""
-    T  = x.shape[2]
-    c  = cos[:T].unsqueeze(0).unsqueeze(0)
-    s  = sin[:T].unsqueeze(0).unsqueeze(0)
-    x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
-    return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1)
-
-
-# ══════════════════════════════════════════════════════════════════════════════
-# FFN  (SwiGLU / GeGLU / MLP)
-# ══════════════════════════════════════════════════════════════════════════════
-
-class FFN(nn.Module):
-    def __init__(self, cfg: SLMConfig):
-        super().__init__()
-        H, I, b = cfg.hidden_size, cfg.intermediate_size, cfg.use_bias
-        self.ffn_type = cfg.ffn_type
-
-        if cfg.ffn_type in ('swiglu', 'geglu'):
-            self.gate = nn.Linear(H, I, bias=b)
-            self.up   = nn.Linear(H, I, bias=b)
-            self.down = nn.Linear(I, H, bias=b)
-            self._act = F.silu if cfg.ffn_type == 'swiglu' else \
-                        lambda x: F.gelu(x, approximate='tanh')
-        else:  # 'mlp'
-            self.up   = nn.Linear(H, I, bias=b)
-            self.down = nn.Linear(I, H, bias=b)
-            act = cfg.activation
-            self._act = (F.silu if act == 'silu' else
-                         (lambda x: F.gelu(x, approximate='tanh')) if act == 'gelu'
-                         else F.relu)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if self.ffn_type in ('swiglu', 'geglu'):
-            return self.down(self._act(self.gate(x)) * self.up(x))
-        return self.down(self._act(self.up(x)))
-
-
-# ══════════════════════════════════════════════════════════════════════════════
-# ATTENTION  (GQA + RoPE + optional QK-Norm)
-# ══════════════════════════════════════════════════════════════════════════════
-
-class GQAttention(nn.Module):
-    """GQA with optional QK-Norm and two backends: sdpa / chunked."""
-
-    def __init__(self, cfg: SLMConfig):
-        super().__init__()
-        H          = cfg.hidden_size
-        b          = cfg.use_bias
-        self.nh    = cfg.num_heads
-        self.nkv   = cfg.num_kv_heads
-        self.d     = cfg.head_dim
-        self.rep   = self.nh // self.nkv
-        KVD        = self.nkv * self.d
-
-        self.q  = nn.Linear(H, H,   bias=b)
-        self.k  = nn.Linear(H, KVD, bias=b)
-        self.v  = nn.Linear(H, KVD, bias=b)
-        self.o  = nn.Linear(H, H,   bias=b)
-        self.dp = nn.Dropout(cfg.dropout)
-
-        self.qk_norm = cfg.qk_norm
-        if self.qk_norm:
-            self.q_norm = RMSNorm(self.d, cfg.norm_eps)
-            self.k_norm = RMSNorm(self.d, cfg.norm_eps)
-
-        cos, sin = precompute_rope(self.d, cfg.max_seq_len, cfg.rope_theta)
-        # FIX: persistent=False so RoPE buffers are NOT saved to safetensors.
-        # They are deterministically recomputed from config at load time.
-        self.register_buffer('cos', cos, persistent=False)
-        self.register_buffer('sin', sin, persistent=False)
-
-        self._backend    = 'chunked'
-        self._chunk_size = 256
-
-    def _qk_normalize(self, q, k):
-        B, nh, T, d  = q.shape
-        q = self.q_norm(q.reshape(-1, d)).reshape(B, nh,       T, d)
-        B, nkv, T, d = k.shape
-        k = self.k_norm(k.reshape(-1, d)).reshape(B, self.nkv, T, d)
-        return q, k
-
-    def _sdpa(self, q, k, v):
-        p = self.dp.p if self.training else 0.0
-        return F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=p)
-
-    def _chunked(self, q, k, v):
-        B, H, T, D = q.shape
-        out   = torch.zeros_like(q)
-        scale = D ** -0.5
-        cs    = self._chunk_size
-        mask  = torch.tril(torch.ones(T, T, device=q.device, dtype=torch.bool))
-        for start in range(0, T, cs):
-            end  = min(start + cs, T)
-            q_c  = q[:, :, start:end, :]
-            k_c  = k[:, :, :end,      :]
-            v_c  = v[:, :, :end,      :]
-            attn = torch.matmul(q_c, k_c.transpose(-2, -1)) * scale
-            attn = attn.masked_fill(
-                ~mask[start:end, :end].unsqueeze(0).unsqueeze(0), float('-inf')
-            )
-            attn = F.softmax(attn.float(), dim=-1).to(q.dtype)
-            if self.training and self.dp.p > 0:
-                attn = self.dp(attn)
-            out[:, :, start:end, :] = torch.matmul(attn, v_c)
-        return out
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        B, T, _ = x.shape
-        q = self.q(x).view(B, T, self.nh,  self.d).transpose(1, 2)
-        k = self.k(x).view(B, T, self.nkv, self.d).transpose(1, 2)
-        v = self.v(x).view(B, T, self.nkv, self.d).transpose(1, 2)
-        if self.qk_norm:
-            q, k = self._qk_normalize(q, k)
-        q = apply_rope(q, self.cos, self.sin)
-        k = apply_rope(k, self.cos, self.sin)
-        k = k.repeat_interleave(self.rep, dim=1)
-        v = v.repeat_interleave(self.rep, dim=1)
-        o = self._sdpa(q, k, v) if self._backend == 'sdpa' else self._chunked(q, k, v)
-        return self.o(o.transpose(1, 2).contiguous().view(B, T, self.nh * self.d))
-
-    def set_backend(self, backend: str, chunk_size: int = 256):
-        self._backend    = backend
-        self._chunk_size = chunk_size
-
-
-# ══════════════════════════════════════════════════════════════════════════════
-# XSA  (sliding-window cross-sequence attention)
-# ══════════════════════════════════════════════════════════════════════════════
-
-class XSALayer(nn.Module):
-    """
-    Optional sliding-window cross-sequence attention.
-    K/V projections are SHARED with the paired GQAttention (parameter sharing).
-    Stored in a plain Python list to avoid duplicate state_dict keys.
-    """
-
-    def __init__(self, cfg: SLMConfig, paired_attn: GQAttention):
-        super().__init__()
-        H = cfg.hidden_size
-        b = cfg.use_bias
-        self.window  = cfg.xsa_window
-        self.nh      = cfg.num_heads
-        self.nkv     = cfg.num_kv_heads
-        self.d       = cfg.head_dim
-        self.rep     = self.nh // self.nkv
-
-        self.norm = RMSNorm(H, cfg.norm_eps)
-        self.q    = nn.Linear(H, H, bias=b)
-        self.o    = nn.Linear(H, H, bias=b)
-        self.dp   = nn.Dropout(cfg.dropout)
-
-        # Plain list — NOT registered as sub-modules → no duplicate state_dict keys
-        self._shared: List[nn.Linear] = [paired_attn.k, paired_attn.v]
-        self._paired: List[GQAttention] = [paired_attn]
-
-        self.qk_norm = cfg.qk_norm
-        if self.qk_norm:
-            self.q_norm = RMSNorm(self.d, cfg.norm_eps)
-            self.k_norm = RMSNorm(self.d, cfg.norm_eps)
-
-    @property
-    def _k(self) -> nn.Linear:  return self._shared[0]
-
-    @property
-    def _v(self) -> nn.Linear:  return self._shared[1]
-
-    @property
-    def cos(self) -> torch.Tensor: return self._paired[0].cos
-
-    @property
-    def sin(self) -> torch.Tensor: return self._paired[0].sin
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        B, T, H = x.shape
-        W  = self.window
-        xn = self.norm(x)
-
-        q = self.q(xn).view(B, T, self.nh, self.d).transpose(1, 2)
-        if self.qk_norm:
-            B_, nh, T_, d = q.shape
-            q = self.q_norm(q.reshape(-1, d)).reshape(B_, nh, T_, d)
-        q = apply_rope(q, self.cos, self.sin)
-
-        pad      = torch.zeros(B, W, H, device=x.device, dtype=x.dtype)
-        x_padded = torch.cat([pad, xn], dim=1)
-        ctx      = x_padded.unfold(1, W, 1)
-        ctx      = ctx[:, :T, :, :].permute(0, 1, 3, 2).contiguous()
-        ctx_flat = ctx.reshape(B * T, W, H)
-
-        k_flat = self._k(ctx_flat).view(B*T, W, self.nkv, self.d).transpose(1, 2)
-        v_flat = self._v(ctx_flat).view(B*T, W, self.nkv, self.d).transpose(1, 2)
-        k_flat = k_flat.repeat_interleave(self.rep, dim=1)
-        v_flat = v_flat.repeat_interleave(self.rep, dim=1)
-        if self.qk_norm:
-            BT, nh, W_, d = k_flat.shape
-            k_flat = self.k_norm(k_flat.reshape(-1, d)).reshape(BT, nh, W_, d)
-
-        q_flat = q.permute(0, 2, 1, 3).reshape(B * T, self.nh, 1, self.d)
-        scale  = self.d ** -0.5
-        attn   = torch.matmul(q_flat, k_flat.transpose(-2, -1)) * scale
-        attn   = F.softmax(attn.float(), dim=-1).to(q.dtype)
-        if self.training and self.dp.p > 0:
-            attn = self.dp(attn)
-        out = torch.matmul(attn, v_flat).reshape(B, T, self.nh * self.d)
-        return self.o(out)
-
-
-# ══════════════════════════════════════════════════════════════════════════════
-# TRANSFORMER BLOCK
-# ══════════════════════════════════════════════════════════════════════════════
-
-class Block(nn.Module):
-    """Pre-LN transformer block. XSA is inserted between attn and FFN."""
-    def __init__(self, cfg: SLMConfig):
-        super().__init__()
-        self.n1      = RMSNorm(cfg.hidden_size, cfg.norm_eps)
-        self.attn    = GQAttention(cfg)
-        self.n2      = RMSNorm(cfg.hidden_size, cfg.norm_eps)
-        self.ffn     = FFN(cfg)
-        self.use_xsa = cfg.xsa
-        if self.use_xsa:
-            self.xsa = XSALayer(cfg, self.attn)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = x + self.attn(self.n1(x))
-        if self.use_xsa:
-            x = x + self.xsa(x)
-        x = x + self.ffn(self.n2(x))
-        return x
-
-    def set_backend(self, backend: str, chunk_size: int = 256):
-        self.attn.set_backend(backend, chunk_size)
-        if self.use_xsa:
-            self.xsa._paired[0]._backend    = backend
-            self.xsa._paired[0]._chunk_size = chunk_size
-
-
-# ══════════════════════════════════════════════════════════════════════════════
-# FULL MODEL
-# ══════════════════════════════════════════════════════════════════════════════
-
-class SLMModel(nn.Module):
-    def __init__(self, cfg: SLMConfig):
-        super().__init__()
-        cfg.validate()
-        self.cfg    = cfg
-        self.embed  = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
-        self.drop   = nn.Dropout(cfg.dropout)
-        self.layers = nn.ModuleList([Block(cfg) for _ in range(cfg.num_layers)])
-        self.norm   = RMSNorm(cfg.hidden_size, cfg.norm_eps)
-        self._init_weights()
-
-    def _init_weights(self):
-        for m in self.modules():
-            if isinstance(m, nn.Linear):
-                nn.init.normal_(m.weight, std=0.02)
-                if m.bias is not None: nn.init.zeros_(m.bias)
-            elif isinstance(m, nn.Embedding):
-                nn.init.normal_(m.weight, std=0.02)
-
-    def forward(self, ids: torch.Tensor) -> torch.Tensor:
-        x = self.drop(self.embed(ids))
-        for layer in self.layers:
-            x = layer(x)
-        return self.norm(x)
-
-    def set_attention_backend(self, backend: str, chunk_size: int = 256):
-        for blk in self.layers:
-            blk.set_backend(backend, chunk_size)
-
-
-class SLMForCausalLM(nn.Module):
-    """
-    SLM with LM head. Used during training.
-
-    State-dict keys (these get 'model.' prefix added by export_safetensors):
-      backbone.embed.weight
-      backbone.layers.{i}.{n1,n2}.w
-      backbone.layers.{i}.attn.{q,k,v,o}.weight
-      backbone.layers.{i}.attn.{q_norm,k_norm}.w  (if qk_norm)
-      backbone.layers.{i}.ffn.{gate,up,down}.weight
-      backbone.layers.{i}.xsa.{norm,q,o}.*         (if xsa)
-      backbone.norm.w
-      lm_head.weight                                (only if NOT tie_embeddings)
-    RoPE cos/sin: persistent=False → NOT in state_dict.
-    """
-    def __init__(self, cfg: SLMConfig):
-        super().__init__()
-        self.cfg      = cfg
-        self.backbone = SLMModel(cfg)
-        self.lm_head  = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
-        if cfg.tie_embeddings:
-            self.lm_head.weight = self.backbone.embed.weight
-        else:
-            nn.init.normal_(self.lm_head.weight, std=0.02)
-
-    def tie_weights(self) -> None:
-        """Re-apply weight tying — called by HF after load_state_dict."""
-        if self.cfg.tie_embeddings:
-            self.lm_head.weight = self.backbone.embed.weight
-
-    def forward(self, input_ids: torch.Tensor,
-                labels: Optional[torch.Tensor] = None
-                ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
-        hidden = self.backbone(input_ids)
-        logits = self.lm_head(hidden)
-        loss = None
-        if labels is not None:
-            loss = F.cross_entropy(
-                logits.reshape(-1, logits.size(-1)),
-                labels.reshape(-1),
-                ignore_index=-100,
-            )
-        return logits, loss
-
-    def set_attention_backend(self, backend: str, chunk_size: int = 256):
-        self.backbone.set_attention_backend(backend, chunk_size)
-
-    def n_params(self) -> int:
-        return sum(p.numel() for p in self.parameters())
-
-
-# ══════════════════════════════════════════════════════════════════════════════
-# HUGGINGFACE WRAPPERS
-# ══════════════════════════════════════════════════════════════════════════════
-
-if _HF_AVAILABLE:
-    class SLMHFConfig(PretrainedConfig):
-        model_type = "slm"
-
-        def __init__(self, **kwargs):
-            valid = set(SLMConfig.__dataclass_fields__)
-            slm_kw = {k: kwargs.pop(k, getattr(SLMConfig(), k)) for k in valid}
-            tie_emb = slm_kw.get('tie_embeddings', True)
-            # Accept canonical HF name coming from a saved config.json
-            if 'tie_word_embeddings' in kwargs:
-                tie_emb = kwargs.pop('tie_word_embeddings')
-            super().__init__(tie_word_embeddings=tie_emb, **kwargs)
-            for k, v in slm_kw.items():
-                setattr(self, k, v)
-            self.tie_embeddings = self.tie_word_embeddings
-
-        def to_slm_config(self) -> SLMConfig:
-            return SLMConfig(**{f: getattr(self, f)
-                                for f in SLMConfig.__dataclass_fields__})
-
-    class SLMHFForCausalLM(PreTrainedModel):
-        config_class                    = SLMHFConfig
-        base_model_prefix               = "model"
-        supports_gradient_checkpointing = False
-
-        # ── FIX: LIST format — compatible with transformers 4.x AND 5.x ────
-        # transformers 4.x: iterates as list of keys to re-tie after loading.
-        # transformers 5.x: _get_tied_weight_keys() calls .keys() — BUT it first
-        # checks isinstance(tied, dict); for a list it falls back to iteration.
-        # A list is the safest format across ALL versions.
-        _tied_weights_keys = ["model.lm_head.weight"]
-
-        def __init__(self, config: SLMHFConfig):
-            super().__init__(config)
-            self.model = SLMForCausalLM(config.to_slm_config())
-            # post_init() sets up all_tied_weights_keys from _tied_weights_keys.
-            # Must be called AFTER all submodules are initialised.
-            self.post_init()
-
-        def forward(
-            self,
-            input_ids=None,
-            attention_mask=None,       # accepted, ignored — we use causal mask
-            position_ids=None,
-            past_key_values=None,
-            inputs_embeds=None,
-            labels=None,
-            use_cache=None,
-            output_attentions=None,
-            output_hidden_states=None,
-            return_dict=None,
-            **kwargs,
-        ):
-            logits, loss = self.model(input_ids, labels)
-            # Always return CausalLMOutputWithPast so .generate() / pipelines work
-            return CausalLMOutputWithPast(loss=loss, logits=logits)
-
-        def get_input_embeddings(self):
-            return self.model.backbone.embed
-
-        def set_input_embeddings(self, v):
-            self.model.backbone.embed = v
-            if self.model.cfg.tie_embeddings:
-                self.model.lm_head.weight = v.weight
-
-        def get_output_embeddings(self):
-            return self.model.lm_head
-
-        def set_output_embeddings(self, v):
-            self.model.lm_head = v
-
-        def tie_weights(self):
-            # ── FIX: Delegate to inner model ─────────────────────────────
-            # HF calls tie_weights() on SLMHFForCausalLM after loading weights.
-            # Without delegation, embed->lm_head link is never restored on load.
-            self.model.tie_weights()
-
-        def prepare_inputs_for_generation(self, input_ids, **kwargs):
-            # ── FIX: Required for model.generate() to work ───────────────
-            return {"input_ids": input_ids}
-
-        def resize_token_embeddings(self, new_num_tokens, pad_to_multiple_of=None):
-            result = super().resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
-            if self.model.cfg.tie_embeddings:
-                self.model.lm_head.weight = self.model.backbone.embed.weight
-            return result
-
-    try:
-        from transformers import AutoConfig, AutoModelForCausalLM
-        SLMHFConfig.register_for_auto_class("AutoConfig")
-        SLMHFForCausalLM.register_for_auto_class("AutoModelForCausalLM")
-    except Exception:
-        pass
-
-
-# ══════════════════════════════════════════════════════════════════════════════
-# PUBLIC HELPERS
-# ══════════════════════════════════════════════════════════════════════════════
-
-def build_model(cfg_dict: dict, device: torch.device) -> SLMForCausalLM:
-    cfg = SLMConfig.from_dict(cfg_dict)
-    cfg.validate()
-    return SLMForCausalLM(cfg).to(device)
-
-
-def count_params(model: nn.Module) -> dict:
-    total     = sum(p.numel() for p in model.parameters())
-    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
-    return {'total': total, 'trainable': trainable}
-
-
-def build_hf_config_json(cfg: SLMConfig, tokenizer) -> dict:
-    bos = getattr(tokenizer, 'bos_token_id', None) or 0
-    eos = getattr(tokenizer, 'eos_token_id', None) or 2
-    pad = getattr(tokenizer, 'pad_token_id',  None) or eos
-    d   = cfg.to_dict()
-    d['tie_word_embeddings'] = cfg.tie_embeddings
-    d.update({
-        'auto_map': {
-            'AutoConfig':           'model.SLMHFConfig',
-            'AutoModelForCausalLM': 'model.SLMHFForCausalLM',
-        },
-        'bos_token_id': bos,
-        'eos_token_id': eos,
-        'pad_token_id': pad,
-        'torch_dtype':  'float32',
-        'transformers_version': '4.40.0',
-    })
-    return d