Updated Files.

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:082e3ae8de93edcbb4790e4917d089354ac79adc4acb937d41e3277ef8f59ec5
-size 45201124
+oid sha256:7161cefd6b1f5dfbf55b4c1ba5772327ccce5cab567865ed8ec43738eba44f25
+size 11646604

modeling_tokle.py

@@ -8,6 +8,7 @@ Works on Transformers 4.x and 5.x. The v5 differences are marked inline; they
 are not cosmetic, see the note above _M1.
 """
 import math
+import os
 from typing import Optional, Tuple, Union
 
 import torch
@@ -62,6 +63,12 @@ def apply_rope(x, cos, sin):
 
 
 # ----------------------------------------------------------------------------- spab
+# The table lives outside model.safetensors. It is frozen data, not weights, and
+# keeping it separate stops it being counted as trained parameters. Cost of that
+# choice: it has to be fetched and installed by hand after the weights load, so
+# every load path below goes through _install_spab_table.
+SPAB_TABLE_FILE = "spab_table.safetensors"
+
 # Mixing constants for the pair hash. Plain ints, not buffers, and that matters:
 # buffers here would be non-persistent, and Transformers 5 builds the model on
 # the meta device and fills only what the checkpoint contains. Non-persistent
@@ -96,9 +103,12 @@ class SPABBias(nn.Module):
         self.scale = nn.Parameter(torch.full((config.n_head,), init))
         if not config.spab_enabled:
             self.scale.requires_grad_(False)
+        # persistent=False keeps it out of state_dict and therefore out of
+        # model.safetensors; `loaded` guards against using the zeros.
         self.register_buffer("table",
                              torch.zeros(config.spab_table_size, dtype=torch.float32),
-                             persistent=True)
+                             persistent=False)
+        self.loaded = not config.spab_enabled
 
     def hash(self, i, j):
         # Murmur-style mix. Relies on int64 overflow wrapping, so keep it in
@@ -113,6 +123,11 @@ class SPABBias(nn.Module):
     def forward(self, idx):
         if not self.enabled:
             return None
+        if not self.loaded:
+            raise RuntimeError(
+                f"the SPAB table was never loaded -- {SPAB_TABLE_FILE} is missing "
+                f"or was not installed. Load the model with from_pretrained(), "
+                f"which fetches it; an all-zero table silently ruins every output.")
         # Materialises b*t*t int64 pairs, so this is the memory high-water mark
         # of a forward pass. Fine at t=512; watch it if max_seq_len ever grows.
         b, t = idx.shape
@@ -194,6 +209,48 @@ class ToklePreTrainedModel(PreTrainedModel):
     _no_split_modules = ["TokleDecoderLayer"]
     _supports_sdpa = True
 
+    def _spab(self):
+        for m in self.modules():
+            if getattr(m, "is_spab", False):
+                return m
+        return None
+
+    def install_spab_table(self, table):
+        """Copy the frozen table in. Accepts a tensor or a .safetensors path."""
+        spab = self._spab()
+        if spab is None:
+            return self
+        if isinstance(table, (str, os.PathLike)):
+            from safetensors.torch import load_file
+            table = load_file(str(table))["table"]
+        table = table.reshape(-1)
+        if table.numel() != spab.table_size:
+            raise ValueError(f"SPAB table has {table.numel()} entries, "
+                             f"expected {spab.table_size}")
+        # Match the existing buffer's device and dtype so a model already moved
+        # to cuda/fp16 keeps behaving the way it did before the split.
+        ref = spab.table
+        spab.table = table.to(device=ref.device, dtype=ref.dtype)
+        spab.loaded = True
+        return self
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        model = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+        if getattr(model.config, "spab_enabled", False):
+            model.install_spab_table(
+                _resolve_table(pretrained_model_name_or_path, **kwargs))
+        return model
+
+    def save_pretrained(self, save_directory, *args, **kwargs):
+        out = super().save_pretrained(save_directory, *args, **kwargs)
+        spab = self._spab()
+        if spab is not None and spab.enabled and spab.loaded:
+            from safetensors.torch import save_file
+            save_file({"table": spab.table.detach().cpu().to(torch.float32).contiguous()},
+                      os.path.join(save_directory, SPAB_TABLE_FILE))
+        return out
+
     def _init_weights(self, module):
         if isinstance(module, nn.Linear):
             nn.init.normal_(module.weight, 0.0, 0.02)
@@ -213,6 +270,34 @@ class ToklePreTrainedModel(PreTrainedModel):
             nn.init.normal_(out.weight, 0.0, 0.02 / math.sqrt(2 * self.config.n_layer))
 
 
+def _resolve_table(repo, **kwargs):
+    """Find spab_table.safetensors in a local dir, the HF cache, or the Hub."""
+    local = os.path.join(str(repo), SPAB_TABLE_FILE)
+    if os.path.isfile(local):
+        return local
+    try:
+        from transformers.utils import cached_file
+        return cached_file(str(repo), SPAB_TABLE_FILE,
+                           **{k: kwargs[k] for k in
+                              ("revision", "token", "cache_dir", "local_files_only",
+                               "subfolder") if k in kwargs})
+    except Exception:
+        pass
+    try:
+        from huggingface_hub import hf_hub_download
+        return hf_hub_download(str(repo), SPAB_TABLE_FILE,
+                               revision=kwargs.get("revision"),
+                               token=kwargs.get("token"))
+    except Exception as e:
+        # Worth a clear message: without the table the model loads happily and
+        # then produces nonsense, which is a miserable thing to debug.
+        raise FileNotFoundError(
+            f"{SPAB_TABLE_FILE} was not found in '{repo}'. Tokle ships its frozen "
+            f"SPAB table beside model.safetensors and cannot run without it; copy "
+            f"the whole repo, not just the weights. ({type(e).__name__}: {e})"
+        ) from e
+
+
 class TokleModel(ToklePreTrainedModel):
     def __init__(self, config: TokleConfig):
         super().__init__(config)

spab_table.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b7c75290ad59adb2551d2db9bbc6bd8b2ff6f72e18d037c6ce8d741a1a596ba
+size 33554512