Fix pico_decoder.py: __init__ defaults, ZeroDivisionError, all_tied_weights_keys

config.json

@@ -3,29 +3,21 @@
     "PicoDecoderHF"
   ],
   "model_type": "pico_decoder",
-  "vocab_size": 32000,
-  "hidden_size": 768,
-  "num_hidden_layers": 14,
-  "num_attention_heads": 4,
-  "intermediate_size": 3072,
-  "max_position_embeddings": 2048,
-  "hidden_act": "silu",
-  "initializer_range": 0.02,
-  "rms_norm_eps": 1e-05,
-  "tie_word_embeddings": false,
-  "torch_dtype": "float32",
-  "transformers_version": "4.48.3",
   "auto_map": {
     "AutoConfig": "pico_decoder.PicoDecoderHFConfig",
     "AutoModelForCausalLM": "pico_decoder.PicoDecoderHF"
   },
-  "d_model": 768,
   "n_layers": 14,
-  "attention_n_heads": 6,
-  "attention_n_kv_heads": 0,
-  "activation_hidden_dim": 3072,
-  "max_seq_len": 2048,
-  "norm_eps": 1e-06,
+  "d_model": 768,
+  "vocab_size": 32768,
+  "attention_n_heads": 12,
+  "attention_n_kv_heads": 1,
+  "max_seq_len": 512,
+  "batch_size": 64,
   "position_emb_theta": 10000.0,
-  "batch_size": 1
+  "activation_hidden_dim": 3072,
+  "norm_eps": 1e-05,
+  "dropout": 0.1,
+  "torch_dtype": "float32",
+  "transformers_version": "4.48.3"
 }

pico_decoder.py

@@ -1,117 +1,95 @@
-
-"""
-Pico Decoder — BeetleLM
-Adapted from pico-lm/pico-decoder-tiny (Apache 2.0).
-Load with trust_remote_code=True.
-"""
-
-from typing import Optional, Tuple, Union
-
+from dataclasses import asdict
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.nn.attention import SDPBackend, sdpa_kernel
 from transformers import PretrainedConfig, PreTrainedModel
 from transformers.modeling_outputs import CausalLMOutput, CausalLMOutputWithPast
+try:
+    if TYPE_CHECKING:
+        from src.config import ModelConfig
+except ImportError:
+    pass
 
 
-# ── RMSNorm ───────────────────────────────────────────────────────────────────
-
 class RMSNorm(torch.nn.Module):
     def __init__(self, config):
         super().__init__()
         self.eps = config.norm_eps
         self.weight = nn.Parameter(torch.ones(config.d_model))
-
     def _norm(self, x):
         return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
-
     def forward(self, x):
         return self._norm(x.float()).type_as(x) * self.weight
 
 
-# ── RoPE ──────────────────────────────────────────────────────────────────────
-
 class RoPE(nn.Module):
     _freqs_cis_tensor = None
-
     def __init__(self, config):
         super().__init__()
         self.theta = config.position_emb_theta
-        self.dim = config.d_model // config.attention_n_heads
-        RoPE._freqs_cis_tensor = self._setup_freqs_cis(
-            config.max_seq_len, self.theta, self.dim
-        )
+        self.dim   = config.d_model // config.attention_n_heads
+        if RoPE._freqs_cis_tensor is None:
+            RoPE._freqs_cis_tensor = self._setup_freqs_cis(config.max_seq_len, self.theta, self.dim)
         self.register_buffer("_freqs_cis", RoPE._freqs_cis_tensor, persistent=False)
-
     @classmethod
     def _setup_freqs_cis(cls, seq_len, theta, dim):
-        _freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: dim // 2].float() / dim))
-        freqs = torch.outer(torch.arange(seq_len), _freqs)
-        return torch.polar(torch.ones_like(freqs), freqs)
-
+        _freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+        return torch.polar(torch.ones_like(f := torch.outer(torch.arange(seq_len), _freqs)), f)
     def get_freqs_cis(self, input_shape, start_pos, end_pos):
         _f = self._freqs_cis[start_pos:end_pos]
         ndim = len(input_shape)
-        shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(input_shape)]
-        return _f.view(*shape)
-
+        assert 0 <= 1 < ndim and _f.shape == (input_shape[1], input_shape[-1])
+        return _f.view(*[d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(input_shape)])
     def forward(self, queries, keys, start_pos=0):
         q_ = torch.view_as_complex(queries.float().reshape(*queries.shape[:-1], -1, 2))
         k_ = torch.view_as_complex(keys.float().reshape(*keys.shape[:-1], -1, 2))
         fc = self.get_freqs_cis(q_.shape, start_pos, start_pos + q_.shape[1])
-        return (
-            torch.view_as_real(q_ * fc).flatten(3).type_as(queries),
-            torch.view_as_real(k_ * fc).flatten(3).type_as(keys),
-        )
-
+        return (torch.view_as_real(q_ * fc).flatten(3).type_as(queries),
+                torch.view_as_real(k_ * fc).flatten(3).type_as(keys))
 
-# ── Attention ─────────────────────────────────────────────────────────────────
 
 class Attention(nn.Module):
     def __init__(self, config):
         super().__init__()
-        self.n_heads    = config.attention_n_heads
-        self.n_kv_heads = config.attention_n_kv_heads
-        self.n_rep      = self.n_heads // self.n_kv_heads
-        self.max_seq_len = config.max_seq_len
+        self.n_heads     = config.attention_n_heads
+        self.n_kv_heads  = config.attention_n_kv_heads
         self.batch_size  = config.batch_size
+        self.max_seq_len = config.max_seq_len
         d = config.d_model
         self.head_dim = d // self.n_heads
+        self.n_rep    = self.n_heads // self.n_kv_heads
         self.q_proj = nn.Linear(d, self.n_heads    * self.head_dim, bias=False)
         self.k_proj = nn.Linear(d, self.n_kv_heads * self.head_dim, bias=False)
         self.v_proj = nn.Linear(d, self.n_kv_heads * self.head_dim, bias=False)
         self.o_proj = nn.Linear(self.n_heads * self.head_dim, d,    bias=False)
         self.rope   = RoPE(config)
-
-    def forward(self, x, mask=None, past_key_values=None, use_cache=False):
-        bsz, seq_len, _ = x.shape
-        q = self.q_proj(x).view(bsz, seq_len, self.n_heads,    self.head_dim)
-        k = self.k_proj(x).view(bsz, seq_len, self.n_kv_heads, self.head_dim)
-        v = self.v_proj(x).view(bsz, seq_len, self.n_kv_heads, self.head_dim)
-        start_pos = past_key_values[0].shape[1] if past_key_values is not None else 0
-        q, k = self.rope(q, k, start_pos)
+    def forward(self, input, mask=None, past_key_values=None, use_cache=False):
+        bsz, seq_len, _ = input.shape
+        q = self.q_proj(input).view(bsz, seq_len, self.n_heads,    self.head_dim)
+        k = self.k_proj(input).view(bsz, seq_len, self.n_kv_heads, self.head_dim)
+        v = self.v_proj(input).view(bsz, seq_len, self.n_kv_heads, self.head_dim)
+        sp = past_key_values[0].shape[1] if past_key_values is not None else 0
+        q, k = self.rope(q, k, sp)
         if past_key_values is not None:
             k = torch.cat([past_key_values[0], k], dim=1)
             v = torch.cat([past_key_values[1], v], dim=1)
         ck, cv = (k, v) if use_cache else (None, None)
         q, k, v = q.transpose(1,2), k.transpose(1,2), v.transpose(1,2)
-        apply_gqa = self.n_rep > 1
-        if apply_gqa and q.device.type == "mps":
+        gqa = self.n_rep > 1
+        if gqa and q.device.type == "mps":
             k = k.repeat_interleave(self.n_rep, dim=-3)
             v = v.repeat_interleave(self.n_rep, dim=-3)
-            apply_gqa = False
+            gqa = False
         with sdpa_kernel(backends=[SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH]):
             out = F.scaled_dot_product_attention(
                 q.contiguous(), k.contiguous(), v.contiguous(),
                 attn_mask=mask.to(q.dtype) if mask is not None else None,
-                enable_gqa=apply_gqa,
+                enable_gqa=gqa,
             )
-        out = out.transpose(1,2).contiguous().view(bsz, seq_len, -1)
-        return self.o_proj(out), (ck, cv)
-
+        return self.o_proj(out.transpose(1,2).contiguous().view(bsz, seq_len, -1)), (ck, cv)
 
-# ── SwiGLU ────────────────────────────────────────────────────────────────────
 
 class SwiGLU(nn.Module):
     def __init__(self, config):
@@ -119,13 +97,10 @@ class SwiGLU(nn.Module):
         self.w_0 = nn.Linear(config.d_model, config.activation_hidden_dim, bias=False)
         self.w_1 = nn.Linear(config.d_model, config.activation_hidden_dim, bias=False)
         self.w_2 = nn.Linear(config.activation_hidden_dim, config.d_model, bias=False)
-
     def forward(self, x):
         return self.w_2(F.silu(self.w_0(x)) * self.w_1(x))
 
 
-# ── PicoDecoderBlock ──────────────────────────────────────────────────────────
-
 class PicoDecoderBlock(nn.Module):
     def __init__(self, config):
         super().__init__()
@@ -133,18 +108,13 @@ class PicoDecoderBlock(nn.Module):
         self.swiglu         = SwiGLU(config)
         self.attention_norm = RMSNorm(config)
         self.swiglu_norm    = RMSNorm(config)
-
-    def forward(self, x, mask=None, past_key_values=None, use_cache=False):
-        attn_out, cached = self.attention(
-            self.attention_norm(x), mask=mask,
-            past_key_values=past_key_values, use_cache=use_cache,
-        )
-        h = x + attn_out
-        return h + self.swiglu(self.swiglu_norm(h)), cached
+    def forward(self, input, mask=None, past_key_values=None, use_cache=False):
+        a, c = self.attention(self.attention_norm(input), mask=mask,
+                              past_key_values=past_key_values, use_cache=use_cache)
+        h = input + a
+        return h + self.swiglu(self.swiglu_norm(h)), c
 
 
-# ── PicoDecoder ───────────────────────────────────────────────────────────────
-
 class PicoDecoder(nn.Module):
     def __init__(self, model_config):
         super().__init__()
@@ -153,67 +123,80 @@ class PicoDecoder(nn.Module):
         self.layers            = nn.ModuleList([PicoDecoderBlock(model_config) for _ in range(model_config.n_layers)])
         self.output_norm       = RMSNorm(model_config)
         self.de_embedding_proj = nn.Linear(model_config.d_model, model_config.vocab_size, bias=False)
-
+    def convert_to_hf_model(self):
+        hf = PicoDecoderHF(PicoDecoderHFConfig.from_dataclass(self.config))
+        hf.load_state_dict(self.state_dict(prefix="pico_decoder."))
+        return hf
     def forward(self, input_ids, past_key_values=None, use_cache=False):
-        seq_len   = input_ids.shape[-1]
-        h         = self.embedding_proj(input_ids)
-        start_pos = 0 if past_key_values is None else past_key_values[0][0].shape[1]
+        sl = input_ids.shape[-1]
+        h  = self.embedding_proj(input_ids)
+        sp = 0 if past_key_values is None else past_key_values[0][0].shape[1]
         mask = None
-        if seq_len > 1:
-            mask = torch.triu(torch.full((seq_len, seq_len), float("-inf")), diagonal=1)
+        if sl > 1:
+            mask = torch.triu(torch.full((sl, sl), float("-inf")), diagonal=1)
             if past_key_values is not None:
-                mask = torch.hstack([torch.zeros((seq_len, start_pos)), mask])
+                mask = torch.hstack([torch.zeros((sl, sp)), mask])
             mask = mask.to(h.device)
-        cached_kvs = () if use_cache else None
-        for idx, layer in enumerate(self.layers):
-            layer_past = past_key_values[idx] if past_key_values is not None else None
-            h, layer_cached = layer(h, mask=mask, past_key_values=layer_past, use_cache=use_cache)
+        ckv = () if use_cache else None
+        for i, layer in enumerate(self.layers):
+            lp = past_key_values[i] if past_key_values is not None else None
+            h, lc = layer(h, mask=mask, past_key_values=lp, use_cache=use_cache)
             if use_cache:
-                cached_kvs += (layer_cached,)
-        return self.de_embedding_proj(self.output_norm(h)).float(), cached_kvs
+                ckv += (lc,)
+        return self.de_embedding_proj(self.output_norm(h)).float(), ckv
 
 
-# ── HuggingFace Config ────────────────────────────────────────────────────────
+# ── HuggingFace wrappers ──────────────────────────────────────────────────────
 
 class PicoDecoderHFConfig(PretrainedConfig):
     model_type = "pico_decoder"
 
-    def __init__(
-        self,
-        vocab_size=32000,
-        d_model=256,
-        n_layers=6,
-        attention_n_heads=8,
-        attention_n_kv_heads=4,
-        activation_hidden_dim=1024,
-        max_seq_len=2048,
-        norm_eps=1e-6,
-        position_emb_theta=10000.0,
-        batch_size=1,
-        **kwargs,
-    ):
+    # FIX 1 + 2: explicit __init__ with MODEL_BASE defaults; guards None/0 kv_heads
+    def __init__(self,
+                 n_layers=14, d_model=768, vocab_size=32768,
+                 attention_n_heads=12, attention_n_kv_heads=1,
+                 max_seq_len=512, batch_size=64, position_emb_theta=10000.0,
+                 activation_hidden_dim=3072, norm_eps=1e-5, dropout=0.1,
+                 **kwargs):
+        if not attention_n_kv_heads:          # catches None, 0, missing
+            attention_n_kv_heads = attention_n_heads
         super().__init__(**kwargs)
-        self.vocab_size             = vocab_size
-        self.d_model                = d_model
-        self.n_layers               = n_layers
-        self.attention_n_heads      = attention_n_heads
-        self.attention_n_kv_heads   = attention_n_kv_heads
-        self.activation_hidden_dim  = activation_hidden_dim
-        self.max_seq_len            = max_seq_len
-        self.norm_eps               = norm_eps
-        self.position_emb_theta     = position_emb_theta
-        self.batch_size             = batch_size
+        self.n_layers              = n_layers
+        self.d_model               = d_model
+        self.vocab_size            = vocab_size
+        self.attention_n_heads     = attention_n_heads
+        self.attention_n_kv_heads  = attention_n_kv_heads
+        self.max_seq_len           = max_seq_len
+        self.batch_size            = batch_size
+        self.position_emb_theta    = position_emb_theta
+        self.activation_hidden_dim = activation_hidden_dim
+        self.norm_eps              = norm_eps
+        self.dropout               = dropout
+
+    @classmethod
+    def from_dict(cls, config_dict: Dict[str, Any], **kwargs) -> "PicoDecoderHFConfig":
+        pico_config = cls(**config_dict)
+        return_unused_kwargs = kwargs.pop("return_unused_kwargs", False)
+        unused_kwargs = {k: v for k, v in kwargs.items() if not hasattr(pico_config, k)}
+        if return_unused_kwargs:
+            return pico_config, unused_kwargs
+        return pico_config
 
+    @classmethod
+    def from_dataclass(cls, model_config):
+        return cls.from_dict(asdict(model_config))
 
-# ── HuggingFace Model ─────────────────────────────────────────────────────────
 
 class PicoDecoderHF(PreTrainedModel):
-    """
-    HuggingFace wrapper for BeetleLM PicoDecoder.
-    Usage: AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
-    """
-    config_class        = PicoDecoderHFConfig
-    _no_split_modules   = ["PicoDecoderBlock", "Attention", "SwiGLU", "RMSNorm"]
+    """Load with: AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)"""
+    config_class       = PicoDecoderHFConfig
+    _no_split_modules  = ["PicoDecoderBlock", "Attention", "SwiGLU", "RMSNorm"]
+    _tied_weights_keys = []   # FIX 3
+
+    # FIX 4: explicit property — transformers >= 4.38 calls this directly
+    @property
+    def all_tied_weights_keys(self):
+        return self._tied_weights_keys
 
     def __init__(self, config: PicoDecoderHFConfig):
         super().__init__(config)
@@ -225,16 +208,14 @@ class PicoDecoderHF(PreTrainedModel):
     def set_input_embeddings(self, value):
         self.pico_decoder.embedding_proj = value
 
-    def forward(self, input_ids=None, past_key_values=None, use_cache=False, labels=None, **kwargs):
-        input_ids = input_ids.clamp(0, self.config.vocab_size - 1)
+    def forward(self, input_ids=None, past_key_values=None,
+                use_cache=False, labels=None, **kwargs):
         logits, new_past = self.pico_decoder(input_ids, past_key_values, use_cache)
         loss = None
         if labels is not None:
-            shift_logits = logits[:, :-1].contiguous()
-            shift_labels = labels[:, 1:].contiguous().clamp(0, self.config.vocab_size - 1)
             loss = F.cross_entropy(
-                shift_logits.view(-1, self.config.vocab_size),
-                shift_labels.view(-1),
+                logits[:, :-1].contiguous().view(-1, self.config.vocab_size),
+                labels[:, 1:].contiguous().clamp(0, self.config.vocab_size - 1).view(-1),
             )
         if use_cache:
             return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=new_past)
@@ -244,7 +225,6 @@ class PicoDecoderHF(PreTrainedModel):
         return {"input_ids": input_ids, "past_key_values": past_key_values, "use_cache": True}
 
 
-# Auto-class registration (runs on trust_remote_code import)
 PicoDecoderHFConfig.register_for_auto_class()
 PicoDecoderHF.register_for_auto_class("AutoModel")
 PicoDecoderHF.register_for_auto_class("AutoModelForCausalLM")