Fix: vocab_size=32000 (BPE base from model.vocab); top-level weights; all compat fixes

pico_decoder.py

@@ -25,29 +25,43 @@ class RMSNorm(torch.nn.Module):
 
 
 class RoPE(nn.Module):
-    _freqs_cis_tensor = None
+    """
+    Rotary Position Embedding.
+    freqs_cis is computed lazily on first use and cached per-device,
+    avoiding meta-tensor issues when HF loads with low_cpu_mem_usage=True.
+    """
     def __init__(self, config):
         super().__init__()
-        self.theta = config.position_emb_theta
-        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)
-    def get_freqs_cis(self, input_shape, start_pos, end_pos):
-        _f = self._freqs_cis[start_pos:end_pos]
+        self.theta   = config.position_emb_theta
+        self.dim     = config.d_model // config.attention_n_heads
+        self.max_seq = config.max_seq_len
+        # NOT a buffer — plain dict so it never touches the meta device
+        self._cache: Dict[torch.device, torch.Tensor] = {}
+
+    def _get_freqs_cis(self, device: torch.device) -> torch.Tensor:
+        if device not in self._cache:
+            freqs = 1.0 / (
+                self.theta ** (
+                    torch.arange(0, self.dim, 2, device=device).float() / self.dim
+                )
+            )
+            t = torch.arange(self.max_seq, device=device)
+            freqs = torch.outer(t, freqs)
+            self._cache[device] = torch.polar(torch.ones_like(freqs), freqs)
+        return self._cache[device]
+
+    def get_freqs_cis(self, input_shape, start_pos, end_pos, device):
+        _f   = self._get_freqs_cis(device)[start_pos:end_pos]
         ndim = len(input_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)])
+        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):
+        device = queries.device
         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])
+        fc  = self.get_freqs_cis(q_.shape, start_pos, start_pos + q_.shape[1], device)
         return (torch.view_as_real(q_ * fc).flatten(3).type_as(queries),
                 torch.view_as_real(k_ * fc).flatten(3).type_as(keys))
 
@@ -196,9 +210,10 @@ class PicoDecoderHF(PreTrainedModel):
     """
     HuggingFace wrapper for BeetleLM PicoDecoder.
     Usage: AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
+    Works with CPU, CUDA (A100, etc.), and MPS out of the box.
     """
     config_class       = PicoDecoderHFConfig
-    _no_split_modules  = ["PicoDecoderBlock", "Attention", "SwiGLU", "RMSNorm"]
+    _no_split_modules  = ["PicoDecoderBlock"]
     _tied_weights_keys = []
 
     @property
@@ -212,6 +227,19 @@ class PicoDecoderHF(PreTrainedModel):
             [PicoDecoderBlock(config) for _ in range(config.n_layers)])
         self.output_norm       = RMSNorm(config)
         self.de_embedding_proj = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        # Required: lets HF finalize weight init and meta-device materialization
+        self.post_init()
+
+    # Required for low_cpu_mem_usage / Accelerate device-dispatch to work
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, RMSNorm):
+            nn.init.ones_(module.weight)
 
     def get_input_embeddings(self):        return self.embedding_proj
     def set_input_embeddings(self, value): self.embedding_proj = value
@@ -223,11 +251,10 @@ class PicoDecoderHF(PreTrainedModel):
         start_pos = 0 if past_key_values is None else past_key_values[0][0].shape[1]
         mask = None
         if seq_len > 1:
-            mask = torch.full((seq_len, seq_len), float("-inf"))
+            mask = torch.full((seq_len, seq_len), float("-inf"), device=h.device)
             mask = torch.triu(mask, diagonal=1)
             if past_key_values is not None:
-                mask = torch.hstack([torch.zeros((seq_len, start_pos)), mask])
-            mask = mask.to(h.device)
+                mask = torch.hstack([torch.zeros((seq_len, start_pos), device=h.device), mask])
         cached_key_values = () 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