prepping safetensor model scripts

README.md

@@ -100,7 +100,18 @@ Prisma 357M trained on ~30B tokens (OpenWebText 20% + FineWeb-Edu 10BT continued
 
 ## Quick Start
 
-### Install
+### Load from HuggingFace
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained("y3i12/Prisma", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("y3i12/Prisma", use_fast=False)
+```
+
+> **Note:** `use_fast=False` is required. The fast tokenizer for MobileLLM is broken upstream and returns a `bool` instead of a tokenizer object.
+
+### Install (for training / development)
 
 ```bash
 pip install -r Prisma/requirements.txt

config.json

@@ -0,0 +1,23 @@
+{
+  "auto_map": {
+    "AutoConfig": "configuration_prisma.PrismaConfig",
+    "AutoModelForCausalLM": "modeling_prisma.PrismaForCausalLM"
+  },
+  "aux_skip_k": 1,
+  "aux_skip_weight": 0.1,
+  "dropout": 0.0,
+  "embed_dim": 0,
+  "head_dim": 0,
+  "hidden_size": 1024,
+  "max_seq_len": 1024,
+  "model_type": "prisma",
+  "n_middle": 1,
+  "num_heads": 16,
+  "num_kv_heads": 4,
+  "num_layers": 41,
+  "transformers_version": "4.57.3",
+  "use_g2lu": true,
+  "vocab_size": 32000,
+  "word_rope_base": 10.0,
+  "word_rope_dims": 8
+}

configuration_prisma.py

@@ -0,0 +1,56 @@
+"""Prisma model configuration for HuggingFace integration."""
+
+from transformers import PretrainedConfig
+
+
+class PrismaConfig(PretrainedConfig):
+    """Configuration for the Prisma mirrored transformer architecture.
+
+    Prisma uses weight-shared mirror pairs (expand/compress phases) with G²LU
+    nested gating and optional word-position RoPE (WoRPE).
+    """
+
+    model_type = "prisma"
+
+    def __init__(
+        self,
+        vocab_size=32000,
+        hidden_size=1024,
+        num_heads=16,
+        num_kv_heads=4,
+        num_layers=41,
+        n_middle=1,
+        max_seq_len=1024,
+        dropout=0.0,
+        aux_skip_k=1,
+        aux_skip_weight=0.1,
+        use_g2lu=True,
+        word_rope_dims=8,
+        word_rope_base=10.0,
+        embed_dim=0,
+        head_dim=0,
+        tie_word_embeddings=True,
+        **kwargs,
+    ):
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.num_layers = num_layers
+        self.n_middle = n_middle
+        self.max_seq_len = max_seq_len
+        self.dropout = dropout
+        self.aux_skip_k = aux_skip_k
+        self.aux_skip_weight = aux_skip_weight
+        self.use_g2lu = use_g2lu
+        self.word_rope_dims = word_rope_dims
+        self.word_rope_base = word_rope_base
+        self.embed_dim = embed_dim
+        self.head_dim = head_dim
+        # HF expects num_hidden_layers for DynamicCache and other utilities
+        self.num_hidden_layers = num_layers
+
+        super().__init__(
+            vocab_size=vocab_size,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

convert_checkpoint.py

@@ -0,0 +1,196 @@
+#!/usr/bin/env python3
+"""Convert a Prisma training checkpoint to HuggingFace format.
+
+Usage:
+    python Prisma/convert_checkpoint.py \
+        --checkpoint circuits/checkpoints/mirrored_300M_mk4_cont/epoch_02.pt \
+        --output-dir Prisma/ \
+        --tokenizer facebook/MobileLLM-125M
+
+This will create:
+    Prisma/model.safetensors   — model weights
+    Prisma/config.json         — model configuration
+    Prisma/tokenizer.json      — tokenizer files
+    Prisma/tokenizer_config.json
+    Prisma/special_tokens_map.json
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+# Ensure Prisma package is importable when running as a standalone script
+_repo_root = Path(__file__).resolve().parent.parent
+if str(_repo_root) not in sys.path:
+    sys.path.insert(0, str(_repo_root))
+
+import torch
+from safetensors.torch import save_file
+from transformers import AutoTokenizer
+
+
+# Buffers that are deterministically recomputed from config — don't save
+SKIP_SUFFIXES = (
+    ".inv_freq",
+    ".cos_cached",
+    ".sin_cached",
+    ".causal_mask",
+    ".word_inv_freq",
+)
+
+
+def convert_checkpoint(
+    checkpoint_path: str,
+    output_dir: str,
+    tokenizer_name: str = "facebook/MobileLLM-125M",
+    dtype: str = "float16",
+):
+    output_path = Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    # --- Load checkpoint ---
+    print(f"Loading checkpoint: {checkpoint_path}")
+    ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+
+    config_dict = ckpt["config"]
+    model_type = ckpt.get("model_type", "mirrored")
+    raw_state = ckpt["model"]
+
+    print(f"  Model type: {model_type}")
+    print(f"  Config: {config_dict}")
+    print(f"  State dict keys: {len(raw_state)}")
+
+    # --- Clean state dict ---
+    cleaned = {}
+    skipped_buffers = 0
+    skipped_tied = 0
+
+    for key, tensor in raw_state.items():
+        # Strip torch.compile prefix
+        clean_key = key.replace("_orig_mod.", "")
+
+        # Skip deterministic buffers
+        if any(clean_key.endswith(s) for s in SKIP_SUFFIXES):
+            skipped_buffers += 1
+            continue
+
+        # Add HF wrapper prefix
+        hf_key = f"transformer.{clean_key}"
+        cleaned[hf_key] = tensor
+
+    print(f"  Skipped {skipped_buffers} deterministic buffers")
+
+    # --- Handle weight tying ---
+    embed_key = "transformer.embed.weight"
+    lm_head_key = "transformer.lm_head.weight"
+
+    embed_dim = config_dict.get("embed_dim", 0) or config_dict["hidden_size"]
+    head_dim = config_dict.get("head_dim", 0) or config_dict["hidden_size"]
+    tie_embeddings = embed_dim == head_dim
+
+    if tie_embeddings and embed_key in cleaned and lm_head_key in cleaned:
+        # Verify they're actually the same data
+        if torch.equal(cleaned[embed_key], cleaned[lm_head_key]):
+            del cleaned[lm_head_key]
+            skipped_tied = 1
+            print(f"  Removed tied lm_head.weight (same as embed.weight)")
+        else:
+            tie_embeddings = False
+            print(f"  WARNING: embed and lm_head differ despite matching dims — keeping both")
+
+    # --- Build word_start_table ---
+    word_rope_dims = config_dict.get("word_rope_dims", 0)
+    if word_rope_dims > 0:
+        print(f"  Building word_start_table from tokenizer: {tokenizer_name}")
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=False)
+        vocab_size = config_dict["vocab_size"]
+        table = torch.zeros(vocab_size, dtype=torch.bool)
+        tokens = tokenizer.convert_ids_to_tokens(list(range(vocab_size)))
+        for idx, tok in enumerate(tokens):
+            if tok is None:
+                continue
+            if tok.startswith('Ġ') or tok.startswith('▁') or tok.startswith('<'):
+                table[idx] = True
+            elif len(tok) > 0 and tok[0] in '\n\r\t':
+                table[idx] = True
+        table[0] = True
+        cleaned["word_start_table"] = table
+        print(f"  Word start table: {table.sum().item()}/{len(table)} tokens marked as word starters")
+
+    # --- Convert dtype ---
+    target_dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}[dtype]
+    for key in cleaned:
+        if cleaned[key].dtype == torch.float32 and cleaned[key].dtype != target_dtype:
+            # Don't convert bool tensors
+            if cleaned[key].dtype != torch.bool:
+                cleaned[key] = cleaned[key].to(target_dtype)
+
+    total_params = sum(t.numel() for t in cleaned.values() if t.dtype != torch.bool)
+    total_bytes = sum(t.numel() * t.element_size() for t in cleaned.values())
+    print(f"  Total parameters: {total_params:,}")
+    print(f"  File size: {total_bytes / 1e9:.2f} GB ({dtype})")
+
+    # --- Save model weights ---
+    safetensors_path = output_path / "model.safetensors"
+    print(f"\nSaving weights: {safetensors_path}")
+    save_file(cleaned, str(safetensors_path))
+
+    # --- Save config ---
+    sys.path.insert(0, str(Path(__file__).resolve().parent))
+    from configuration_prisma import PrismaConfig
+
+    hf_config = PrismaConfig(
+        vocab_size=config_dict["vocab_size"],
+        hidden_size=config_dict["hidden_size"],
+        num_heads=config_dict["num_heads"],
+        num_kv_heads=config_dict.get("num_kv_heads"),
+        num_layers=config_dict["num_layers"],
+        n_middle=config_dict.get("n_middle", 1),
+        max_seq_len=config_dict.get("max_seq_len", 1024),
+        dropout=config_dict.get("dropout", 0.0),
+        aux_skip_k=config_dict.get("aux_skip_k", 0),
+        aux_skip_weight=config_dict.get("aux_skip_weight", 0.1),
+        use_g2lu=config_dict.get("use_g2lu", True),
+        word_rope_dims=config_dict.get("word_rope_dims", 0),
+        word_rope_base=config_dict.get("word_rope_base", 10.0),
+        embed_dim=config_dict.get("embed_dim", 0),
+        head_dim=config_dict.get("head_dim", 0),
+        tie_word_embeddings=tie_embeddings,
+        auto_map={
+            "AutoConfig": "configuration_prisma.PrismaConfig",
+            "AutoModelForCausalLM": "modeling_prisma.PrismaForCausalLM",
+        },
+    )
+    hf_config.save_pretrained(str(output_path))
+    print(f"Saved config: {output_path / 'config.json'}")
+
+    # --- Save tokenizer ---
+    print(f"\nSaving tokenizer from: {tokenizer_name}")
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=False)
+    tokenizer.save_pretrained(str(output_path))
+    print(f"Saved tokenizer files to: {output_path}")
+
+    # --- Summary ---
+    print(f"\n{'='*60}")
+    print(f"Conversion complete!")
+    print(f"  Output directory: {output_path}")
+    print(f"  Model size: {total_bytes / 1e9:.2f} GB ({dtype})")
+    print(f"  Parameters: {total_params:,}")
+    print(f"  Tied embeddings: {tie_embeddings}")
+    print(f"  Word RoPE dims: {word_rope_dims}")
+    print(f"{'='*60}")
+    print(f"\nUsage:")
+    print(f'  from transformers import AutoModelForCausalLM, AutoTokenizer')
+    print(f'  model = AutoModelForCausalLM.from_pretrained("{output_path}", trust_remote_code=True)')
+    print(f'  tokenizer = AutoTokenizer.from_pretrained("{output_path}")')
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert Prisma checkpoint to HuggingFace format")
+    parser.add_argument("--checkpoint", type=str, required=True, help="Path to .pt checkpoint")
+    parser.add_argument("--output-dir", type=str, default="Prisma/", help="Output directory")
+    parser.add_argument("--tokenizer", type=str, default="facebook/MobileLLM-125M", help="Tokenizer name")
+    parser.add_argument("--dtype", type=str, default="float16", choices=["float16", "bfloat16", "float32"])
+    args = parser.parse_args()
+
+    convert_checkpoint(args.checkpoint, args.output_dir, args.tokenizer, args.dtype)

modeling_prisma.py

@@ -0,0 +1,173 @@
+"""Prisma model for HuggingFace integration.
+
+Usage:
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    model = AutoModelForCausalLM.from_pretrained("y3i12/Prisma", trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained("y3i12/Prisma")
+"""
+
+import torch
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_prisma import PrismaConfig
+from .mirrored import MirroredTransformer, MirroredConfig
+from .layers import build_word_start_table, compute_word_positions
+
+
+class PrismaForCausalLM(PreTrainedModel):
+    """Prisma mirrored transformer for causal language modeling."""
+
+    config_class = PrismaConfig
+    _tied_weights_keys = ["transformer.lm_head.weight"]
+    _no_split_modules = ["MirroredBlock", "MiddleBlock"]
+    _keys_to_ignore_on_load_missing = [
+        r"transformer\..*\.rotary\.inv_freq",
+        r"transformer\..*\.word_rope\.word_inv_freq",
+    ]
+
+    def __init__(self, config: PrismaConfig):
+        super().__init__(config)
+
+        mirrored_config = MirroredConfig(
+            vocab_size=config.vocab_size,
+            hidden_size=config.hidden_size,
+            num_heads=config.num_heads,
+            num_kv_heads=config.num_kv_heads,
+            num_layers=config.num_layers,
+            n_middle=config.n_middle,
+            max_seq_len=config.max_seq_len,
+            dropout=config.dropout,
+            aux_skip_k=config.aux_skip_k,
+            aux_skip_weight=config.aux_skip_weight,
+            use_g2lu=config.use_g2lu,
+            word_rope_dims=config.word_rope_dims,
+            word_rope_base=config.word_rope_base,
+            embed_dim=config.embed_dim,
+            head_dim=config.head_dim,
+        )
+        self.transformer = MirroredTransformer(mirrored_config)
+
+        # Word-position table for WoRPE (populated by from_pretrained or set_tokenizer)
+        if config.word_rope_dims > 0:
+            self.register_buffer(
+                "word_start_table",
+                torch.zeros(config.vocab_size, dtype=torch.bool),
+                persistent=True,
+            )
+        else:
+            self.word_start_table = None
+
+        # Track word position during autoregressive generation
+        self._word_pos_counter = 0
+
+        self.post_init()
+
+    def set_tokenizer(self, tokenizer):
+        """Build word_start_table from tokenizer. Call this if not loading from pretrained."""
+        if self.config.word_rope_dims > 0:
+            table = build_word_start_table(tokenizer, self.config.vocab_size)
+            self.word_start_table = table.to(self.device)
+
+    def get_input_embeddings(self):
+        return self.transformer.embed
+
+    def set_input_embeddings(self, value):
+        self.transformer.embed = value
+
+    def get_output_embeddings(self):
+        return self.transformer.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.transformer.lm_head = new_embeddings
+
+    def tie_weights(self):
+        if self.config.tie_word_embeddings:
+            embed_dim = self.config.embed_dim or self.config.hidden_size
+            head_dim = self.config.head_dim or self.config.hidden_size
+            if embed_dim == head_dim:
+                self.transformer.lm_head.weight = self.transformer.embed.weight
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        past_key_values=None,
+        labels=None,
+        use_cache=False,
+        return_dict=True,
+        **kwargs,
+    ):
+        # Convert HF DynamicCache to our list-of-tuples format
+        past_kv_list = None
+        if past_key_values is not None:
+            if hasattr(past_key_values, 'key_cache'):
+                # HF DynamicCache
+                if len(past_key_values) > 0:
+                    past_kv_list = [
+                        (past_key_values.key_cache[i], past_key_values.value_cache[i])
+                        for i in range(len(past_key_values))
+                    ]
+            elif isinstance(past_key_values, (list, tuple)):
+                past_kv_list = past_key_values
+
+        # Compute word positions if WoRPE is enabled
+        word_positions = None
+        if self.word_start_table is not None and self.config.word_rope_dims > 0:
+            if past_kv_list is not None and input_ids.size(1) == 1:
+                # Cached generation: track word position step by step
+                last_token = input_ids[0, -1].item()
+                if self.word_start_table[last_token]:
+                    self._word_pos_counter = 0
+                else:
+                    self._word_pos_counter += 1
+                word_positions = torch.tensor(
+                    [[float(self._word_pos_counter)]],
+                    device=input_ids.device,
+                )
+            else:
+                # Full sequence: compute all word positions
+                word_positions = compute_word_positions(input_ids, self.word_start_table)
+                # Save last position for subsequent generation steps
+                self._word_pos_counter = int(word_positions[0, -1].item())
+
+        output = self.transformer(
+            input_ids,
+            labels=labels,
+            use_cache=use_cache,
+            past_kv=past_kv_list,
+            word_positions=word_positions,
+        )
+
+        # Convert our list-of-tuples back to DynamicCache
+        new_cache = None
+        if output.get("past_kv") is not None:
+            from transformers.cache_utils import DynamicCache
+            new_cache = DynamicCache()
+            for layer_idx, (k, v) in enumerate(output["past_kv"]):
+                new_cache.update(k, v, layer_idx)
+
+        if not return_dict:
+            result = (output["logits"],)
+            if use_cache:
+                result += (new_cache,)
+            return result
+
+        return CausalLMOutputWithPast(
+            loss=output.get("loss"),
+            logits=output["logits"],
+            past_key_values=new_cache,
+        )
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, **kwargs
+    ):
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "use_cache": True,
+        }

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,18 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "add_prefix_space": true,
+  "added_tokens_decoder": {},
+  "bos_token": "",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "",
+  "extra_special_tokens": {},
+  "legacy": true,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": null,
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "",
+  "use_default_system_prompt": false
+}