Add files using upload-large-folder tool

models/full-128k-claude-balanced-bpe-mathcode-scratch/README.md

@@ -0,0 +1,76 @@
+# `full-128k-claude-balanced-bpe-mathcode-scratch`
+
+A 1.27B parameter model trained with the `claude-balanced-bpe` tokenizer. Regime: mathcode-20B-from-scratch (24 layers, width 1536, 1.27B parameters, 20.0B tokens, math and code mixture, trained from scratch). Reported in: math and code ablations; TokEval.
+
+## Load
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+    "cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-claude-balanced-bpe-mathcode-scratch",
+    trust_remote_code=True, dtype="float32")
+tokenizer = AutoTokenizer.from_pretrained("cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-claude-balanced-bpe-mathcode-scratch")
+
+text = "def add(a, b):\n    return a + b\n"
+ids = [0] + tokenizer(text)["input_ids"]   # training prepends BOS
+```
+
+## Settings
+
+| Setting | Value |
+|---|---|
+| layers | 24 |
+| model dimension | 1536 |
+| attention heads | 12 |
+| key and value heads | 12 |
+| context length | 2048 |
+| sliding window pattern | SSSL |
+| value embedding dimension | 128 |
+| logit soft cap | 15 |
+| vocabulary size | 128256 |
+| parameters | 1,272,841,994 |
+| training step | 19073 |
+| token budget | 19,999,490,048 |
+| bytes consumed | 151,300,558,757 |
+| seed | 42 |
+| Muon matrix learning rate | 0.0200 |
+| AdamW embedding learning rate | 0.3000 |
+| AdamW unembedding learning rate | 0.0080 |
+| weight decay | 0.2800 |
+| width learning-rate exponent | -1 |
+| weight dtype | as trained (float32 matrices, bfloat16 embeddings) |
+| BOS token id | 0 |
+
+## Results
+
+| Metric | Value |
+|---|---|
+| validation BPB | 0.3365 |
+| code BPB, mean over 7 languages | 0.3861 |
+| FLORES-200 BPB, mean | 3.3743 |
+| FLORES-200 languages scored | 214 |
+| BLiMP accuracy | n/a |
+| MBPP pass@1 | 0.2500 |
+| HumanEval pass@1 | 0.1707 |
+| GSM8K exact match, strict | 0.2274 |
+
+Generation metrics use specification `{'version': 'v2-2026-07-30', 'bos_context': True, 'eos_stop': True, 'token_healing': True, 'truncation': 'left_to_max_length_minus_max_gen', 'stop_truncation': 'earliest_match_all_terms', 'stop_match_decode': 'specials_kept'}`, token healing on.
+
+Missing measurements, with the reason each is absent:
+
+- `blimp_acc`: no full-128k-claude-balanced-bpe-mathcode-scratch_blimp_code_bpb.json
+
+BLiMP is comparable within a training scale only. Its scoring sums unnormalized sentence log-likelihoods, and under that scoring a group of quantifier subtasks drops at the 1.27B scale relative to the 300M scale.
+
+## Registry note
+
+#29, claude pretok, slug `claude-balanced-bpe`. Original main-panel BPE, Claude regex, `balanced`, vocab 128,256. FINISHED (step 19073). Panel-best MBPP 0.220 [0.184, 0.256] (seed 42; 3-seed family mean 0.2040, SD 0.0139: see §3a.5).
+
+## Provenance
+
+- Source checkpoint: `/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-claude-balanced-bpe-mathcode-scratch`, step 19073
+- Repository commit: `1a9a718`
+- nanochat commit: `f919449` on branch `tokenizer-lm-patches`
+- Tokenizer: `claude-balanced-bpe`, released at `tokenizers/claude-balanced-bpe`
+

models/full-128k-claude-balanced-bpe-mathcode-scratch/config.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "NanoChatForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "nanochat_model.NanoChatConfig",
+    "AutoModelForCausalLM": "nanochat_model.NanoChatForCausalLM"
+  },
+  "bos_token_id": 0,
+  "dtype": "float32",
+  "eos_token_id": null,
+  "head_dim": 128,
+  "kv_dim": 1536,
+  "logit_softcap": 15.0,
+  "model_type": "nanochat",
+  "n_embd": 1536,
+  "n_head": 12,
+  "n_kv_head": 12,
+  "n_layer": 24,
+  "pad_vocab_size_to": 64,
+  "padded_vocab_size": 128256,
+  "sequence_len": 2048,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.3.0",
+  "ve_dim": 128,
+  "vocab_size": 128256,
+  "window_pattern": "SSSL"
+}

models/full-128k-claude-balanced-bpe-mathcode-scratch/conversion_metadata.json

@@ -0,0 +1,100 @@
+{
+  "run": "full-128k-claude-balanced-bpe-mathcode-scratch",
+  "source_checkpoint": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-claude-balanced-bpe-mathcode-scratch",
+  "source_model_file": "model_019073.pt",
+  "source_step": 19073,
+  "source_meta": {
+    "step": 19073,
+    "val_bpb": 0.33648727932470396,
+    "model_config": {
+      "sequence_len": 2048,
+      "vocab_size": 128256,
+      "n_layer": 24,
+      "n_head": 12,
+      "n_kv_head": 12,
+      "n_embd": 1536,
+      "window_pattern": "SSSL",
+      "ve_dim": 128,
+      "logit_softcap": 15.0
+    },
+    "tokenizer_path": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/tokenizers/claude-balanced-bpe",
+    "training_config": {
+      "name": "full-128k-mathcode-scratch",
+      "model": {
+        "n_layer": 24,
+        "n_embd": 1536,
+        "n_kv_head": 12,
+        "sequence_len": 2048,
+        "window_pattern": "SSSL",
+        "ve_dim": 128
+      },
+      "training": {
+        "target_param_data_ratio": -1,
+        "max_tokens": 20000000000,
+        "num_iterations": -1,
+        "total_batch_size": 1048576,
+        "device_batch_size": 16,
+        "warmup_steps": 40,
+        "warmdown_ratio": 0.65,
+        "final_lr_frac": 0.05,
+        "eval_tokens": 41943040,
+        "compile": true
+      },
+      "optimizer": {
+        "matrix_lr": 0.02,
+        "embedding_lr": 0.3,
+        "unembedding_lr": 0.008,
+        "scalar_lr": 0.5,
+        "weight_decay": 0.28,
+        "width_lr_exponent": -1.0
+      },
+      "data": {
+        "data_dir": null
+      },
+      "logging": {
+        "eval_every": 500,
+        "save_every": 500,
+        "max_checkpoints": 0
+      },
+      "wandb": {
+        "enabled": true,
+        "project": "tokenizer-lm-experiments",
+        "run_name": null
+      },
+      "output_dir": "/capstor/scratch/cscs/${USER}/tokenizer-lm/checkpoints",
+      "seed": 42
+    },
+    "total_batch_size": 1048576,
+    "device_batch_size": 16,
+    "max_seq_len": 2048,
+    "dataloader_state_dict": {
+      "pq_idx": 219,
+      "rg_idx": 44,
+      "epoch": 2
+    },
+    "loop_state": {
+      "min_val_bpb": 0.33648727932470396,
+      "smooth_train_loss": 0.8780949603479742,
+      "total_training_time": 93661.4308078289,
+      "total_bytes_consumed": 151300558757
+    }
+  },
+  "weight_dtype": "as trained (float32 matrices, bfloat16 embeddings)",
+  "weight_dtype_counts": {
+    "torch.float32": 174,
+    "torch.bfloat16": 13
+  },
+  "total_parameters": 1272841994,
+  "logit_softcap_source": "meta",
+  "bos_token": "<s>",
+  "bos_token_id": 0,
+  "bos_resolved_from": "packaged tokenizer",
+  "tokenizer_slug": "claude-balanced-bpe",
+  "tokenizer_redistributed": true,
+  "tokenizer_hub_id": null,
+  "weight_key_mapping": "strip _orig_mod. and transformer. prefixes",
+  "converted_by": "scripts/hf_release/convert_for_release.py",
+  "repo_commit": "1a9a718",
+  "nanochat_commit": "f919449",
+  "nanochat_branch": "tokenizer-lm-patches"
+}

models/full-128k-claude-balanced-bpe-mathcode-scratch/generation_config.json

@@ -0,0 +1,4 @@
+{
+  "bos_token_id": 0,
+  "transformers_version": "5.3.0"
+}

models/full-128k-claude-balanced-bpe-mathcode-scratch/nanochat_model.py

@@ -0,0 +1,361 @@
+"""
+NanoChat model for HuggingFace Transformers — self-contained single file.
+
+Includes the config class, model class, and a CLI demo for loading + generation.
+Drop this file into a directory containing config.json, model.safetensors,
+tokenizer.json, and tokenizer_config.json, then run:
+
+    python nanochat_model.py --model-dir . --prompt "Hello world"
+
+Or load programmatically:
+
+    from nanochat_model import NanoChatConfig, NanoChatForCausalLM
+    model = NanoChatForCausalLM.from_pretrained("path/to/model", torch_dtype=torch.bfloat16)
+
+Requirements: pip install torch transformers tokenizers
+
+Architecture: nanochat GPT (Karpathy, MIT licensed) with value embeddings
+(ResFormer-style), per-layer residual/x0 lambdas, smear, backout, logit
+soft-capping, sliding window attention, QK norm, ReLU^2 MLP, and RoPE.
+
+The models were trained with a patched nanochat, not upstream nanochat:
+https://github.com/cimeister/nanochat branch tokenizer-lm-patches, commit f919449.
+The patches touch two files and change three things, two of which affect the forward
+pass and are implemented here:
+  1. ve_dim: the value-embedding tables are built at a reduced dimension and projected
+     up to kv_dim by a learned ve_proj, instead of being built at full kv_dim. All
+     released models use ve_dim=128.
+  2. logit_softcap: the soft cap is a config field with 0 meaning disabled, instead of
+     a hardcoded 15.
+  3. width_lr_exponent: the AdamW learning-rate width scaling exponent is configurable
+     for muP. This affects training only and has no inference-time effect.
+
+This file is the released model definition. Its agreement with the training-time model
+is checked per released checkpoint by scripts/hf_release/verify_conversion.py, which
+compares logits against the patched nanochat GPT in float32.
+"""
+
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+class NanoChatConfig(PretrainedConfig):
+    model_type = "nanochat"
+
+    def __init__(
+        self,
+        vocab_size=128256,
+        n_embd=1536,
+        n_layer=24,
+        n_head=12,
+        n_kv_head=12,
+        sequence_len=2048,
+        window_pattern="SSSL",
+        ve_dim=0,
+        logit_softcap=15.0,
+        pad_vocab_size_to=64,
+        bos_token_id=0,
+        eos_token_id=None,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_kv_head = n_kv_head
+        self.sequence_len = sequence_len
+        self.window_pattern = window_pattern
+        self.ve_dim = ve_dim
+        self.logit_softcap = logit_softcap
+        self.pad_vocab_size_to = pad_vocab_size_to
+        self.head_dim = n_embd // n_head
+        self.kv_dim = n_kv_head * self.head_dim
+        self.padded_vocab_size = (
+            (vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to
+        ) * pad_vocab_size_to
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Model
+# ---------------------------------------------------------------------------
+
+def _has_ve(layer_idx: int, n_layer: int) -> bool:
+    return layer_idx % 2 == (n_layer - 1) % 2
+
+
+def _apply_rotary_emb(x, cos, sin):
+    d = x.shape[-1] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    return torch.cat([y1, y2], dim=-1)
+
+
+def _rms_norm(x):
+    return F.rms_norm(x, (x.size(-1),))
+
+
+class NanoChatAttention(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.n_embd = config.n_embd
+        self.head_dim = config.head_dim
+        self.n_kv_groups = self.n_head // self.n_kv_head
+        self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
+        self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
+        self.ve_gate_channels = 12
+        if _has_ve(layer_idx, config.n_layer):
+            self.ve_gate = nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False)
+            kv_dim = self.n_kv_head * self.head_dim
+            ve_dim = config.ve_dim if config.ve_dim > 0 else kv_dim
+            self.ve_proj = nn.Linear(ve_dim, kv_dim, bias=False) if ve_dim < kv_dim else None
+        else:
+            self.ve_gate = None
+            self.ve_proj = None
+
+    def forward(self, x, ve, cos, sin, window_size):
+        B, T, C = x.size()
+        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
+        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
+        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
+        if ve is not None and self.ve_gate is not None:
+            if self.ve_proj is not None:
+                ve = self.ve_proj(ve)
+            ve = ve.view(B, T, self.n_kv_head, self.head_dim)
+            gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
+            v = v + gate.unsqueeze(-1) * ve
+        q = _apply_rotary_emb(q, cos, sin)
+        k = _apply_rotary_emb(k, cos, sin)
+        q = _rms_norm(q) * 1.2
+        k = _rms_norm(k) * 1.2
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        if self.n_kv_groups > 1:
+            k = k.repeat_interleave(self.n_kv_groups, dim=1)
+            v = v.repeat_interleave(self.n_kv_groups, dim=1)
+        if window_size < 0 or window_size >= T:
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:
+            row_idx = torch.arange(T, device=x.device).unsqueeze(1)
+            col_idx = torch.arange(T, device=x.device).unsqueeze(0)
+            mask = (col_idx <= row_idx) & ((row_idx - col_idx) <= window_size)
+            y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        y = y.transpose(1, 2).contiguous().view(B, T, -1)
+        y = self.c_proj(y)
+        return y
+
+
+class NanoChatMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = F.relu(x).square()
+        x = self.c_proj(x)
+        return x
+
+
+class NanoChatBlock(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.attn = NanoChatAttention(config, layer_idx)
+        self.mlp = NanoChatMLP(config)
+
+    def forward(self, x, ve, cos, sin, window_size):
+        x = x + self.attn(_rms_norm(x), ve, cos, sin, window_size)
+        x = x + self.mlp(_rms_norm(x))
+        return x
+
+
+class NanoChatPreTrainedModel(PreTrainedModel):
+    config_class = NanoChatConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["NanoChatBlock"]
+
+    def _init_weights(self, module):
+        pass
+
+
+class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin):
+    def __init__(self, config):
+        super().__init__(config)
+        pvs = config.padded_vocab_size
+        self.wte = nn.Embedding(pvs, config.n_embd)
+        self.h = nn.ModuleList([NanoChatBlock(config, i) for i in range(config.n_layer)])
+        self.lm_head = nn.Linear(config.n_embd, pvs, bias=False)
+        self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))
+        self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))
+        self.smear_gate = nn.Linear(24, 1, bias=False)
+        self.smear_lambda = nn.Parameter(torch.zeros(1))
+        self.backout_lambda = nn.Parameter(0.2 * torch.ones(1))
+        ve_emb_dim = config.ve_dim if config.ve_dim > 0 else config.kv_dim
+        self.value_embeds = nn.ModuleDict({
+            str(i): nn.Embedding(pvs, ve_emb_dim)
+            for i in range(config.n_layer) if _has_ve(i, config.n_layer)
+        })
+        self.window_sizes = self._compute_window_sizes(config)
+        self.rotary_seq_len = config.sequence_len * 10
+        self._rotary_initialized = False
+        self.register_buffer("cos", torch.zeros(1), persistent=False)
+        self.register_buffer("sin", torch.zeros(1), persistent=False)
+        self.post_init()
+
+    def _compute_window_sizes(self, config):
+        pattern = config.window_pattern.upper()
+        long_window = config.sequence_len
+        short_window = -(-long_window // 4 // 128) * 128
+        char_to_window = {"L": long_window, "S": short_window}
+        sizes = [char_to_window[pattern[i % len(pattern)]] for i in range(config.n_layer)]
+        sizes[-1] = long_window
+        return sizes
+
+    def _precompute_rotary(self, seq_len, head_dim, base=100000.0):
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32)
+        inv_freq = 1.0 / (base ** (channel_range / head_dim))
+        t = torch.arange(seq_len, dtype=torch.float32)
+        freqs = torch.outer(t, inv_freq)
+        cos = freqs.cos()[None, :, None, :]
+        sin = freqs.sin()[None, :, None, :]
+        return cos, sin
+
+    def forward(self, input_ids, attention_mask=None, labels=None, return_dict=None, **kwargs):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        B, T = input_ids.size()
+        if not self._rotary_initialized:
+            cos, sin = self._precompute_rotary(self.rotary_seq_len, self.config.head_dim)
+            self.cos = cos.to(device=input_ids.device)
+            self.sin = sin.to(device=input_ids.device)
+            self._rotary_initialized = True
+        cos = self.cos[:, :T].to(dtype=self.wte.weight.dtype)
+        sin = self.sin[:, :T].to(dtype=self.wte.weight.dtype)
+        x = self.wte(input_ids)
+        x = _rms_norm(x)
+        if T > 1:
+            gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
+            x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
+        x0 = x
+        n_layer = self.config.n_layer
+        backout_layer = n_layer // 2
+        x_backout = None
+        for i, block in enumerate(self.h):
+            x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
+            ve = self.value_embeds[str(i)](input_ids).to(x.dtype) if str(i) in self.value_embeds else None
+            x = block(x, ve, cos, sin, self.window_sizes[i])
+            if i == backout_layer:
+                x_backout = x
+        if x_backout is not None:
+            x = x - self.backout_lambda.to(x.dtype) * x_backout
+        x = _rms_norm(x)
+        logits = self.lm_head(x)
+        logits = logits[..., :self.config.vocab_size]
+        logits = logits.float()
+        if self.config.logit_softcap > 0:
+            cap = self.config.logit_softcap
+            logits = cap * torch.tanh(logits / cap)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+        if not return_dict:
+            output = (logits,)
+            return (loss,) + output if loss is not None else output
+        return CausalLMOutputWithPast(loss=loss, logits=logits)
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# ---------------------------------------------------------------------------
+# CLI demo
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    import argparse
+    import sys
+    import os
+
+    parser = argparse.ArgumentParser(description="Load and generate with a NanoChat model")
+    parser.add_argument("--model-dir", required=True, help="Path to converted model directory")
+    parser.add_argument("--prompt", default="The meaning of life is", help="Text prompt")
+    parser.add_argument("--max-tokens", type=int, default=100, help="Max tokens to generate")
+    parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature (0=greedy)")
+    parser.add_argument("--device", default="cpu", help="Device (cpu, cuda, cuda:0, ...)")
+    parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16", "float16"])
+    args = parser.parse_args()
+
+    dtype_map = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}
+
+    print(f"Loading model from {args.model_dir}...")
+    from transformers import AutoTokenizer
+
+    model = NanoChatForCausalLM.from_pretrained(args.model_dir, torch_dtype=dtype_map[args.dtype])
+    model.to(args.device)
+    model.eval()
+    tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Model loaded: {total_params:,} parameters, dtype={args.dtype}, device={args.device}")
+    print(f"Vocab size: {model.config.vocab_size}, ve_dim: {model.config.ve_dim}")
+    print()
+    print(f"Prompt: {args.prompt}")
+    print(f"Generating (max {args.max_tokens} tokens, temp={args.temperature})...")
+    print()
+
+    # Generate
+    bos_id = tokenizer.bos_token_id
+    input_ids = tokenizer.encode(args.prompt, add_special_tokens=False)
+    if bos_id is not None:
+        input_ids = [bos_id] + input_ids
+    input_ids = torch.tensor([input_ids], dtype=torch.long, device=args.device)
+
+    generated = []
+    with torch.inference_mode():
+        for _ in range(args.max_tokens):
+            logits = model(input_ids).logits[:, -1, :]
+            if args.temperature > 0:
+                logits = logits / args.temperature
+                v, _ = torch.topk(logits, min(50, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float("inf")
+                probs = torch.softmax(logits, dim=-1)
+                next_id = torch.multinomial(probs, num_samples=1)
+            else:
+                next_id = torch.argmax(logits, dim=-1, keepdim=True)
+            input_ids = torch.cat([input_ids, next_id], dim=1)
+            token_id = next_id.item()
+            generated.append(token_id)
+            if tokenizer.eos_token_id is not None and token_id == tokenizer.eos_token_id:
+                break
+
+    output = tokenizer.decode(generated, skip_special_tokens=True)
+    print(f"Output: {args.prompt}{output}")

models/full-128k-claude-balanced-bpe-mathcode-scratch/tokenizer_config.json

@@ -0,0 +1,7 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "model_max_length": 2048,
+  "bos_token": "<s>",
+  "add_bos_token": false,
+  "clean_up_tokenization_spaces": false
+}

models/full-128k-claude-balanced-bpe-seed43-mathcode-scratch/README.md

@@ -0,0 +1,73 @@
+# `full-128k-claude-balanced-bpe-seed43-mathcode-scratch`
+
+A 1.27B parameter model trained with the `claude-balanced-bpe` tokenizer. Regime: mathcode-20B-from-scratch (24 layers, width 1536, 1.27B parameters, 20.0B tokens, math and code mixture, trained from scratch). Reported in: math and code ablations.
+
+## Load
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+    "cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-claude-balanced-bpe-seed43-mathcode-scratch",
+    trust_remote_code=True, dtype="float32")
+tokenizer = AutoTokenizer.from_pretrained("cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-claude-balanced-bpe-seed43-mathcode-scratch")
+
+text = "def add(a, b):\n    return a + b\n"
+ids = [0] + tokenizer(text)["input_ids"]   # training prepends BOS
+```
+
+## Settings
+
+| Setting | Value |
+|---|---|
+| layers | 24 |
+| model dimension | 1536 |
+| attention heads | 12 |
+| key and value heads | 12 |
+| context length | 2048 |
+| sliding window pattern | SSSL |
+| value embedding dimension | 128 |
+| logit soft cap | 15 |
+| vocabulary size | 128256 |
+| parameters | 1,272,841,994 |
+| training step | 19073 |
+| token budget | 19,999,490,048 |
+| bytes consumed | 151,300,558,757 |
+| seed | 43 |
+| Muon matrix learning rate | 0.0200 |
+| AdamW embedding learning rate | 0.3000 |
+| AdamW unembedding learning rate | 0.0080 |
+| weight decay | 0.2800 |
+| width learning-rate exponent | -1 |
+| weight dtype | as trained (float32 matrices, bfloat16 embeddings) |
+| BOS token id | 0 |
+
+## Results
+
+| Metric | Value |
+|---|---|
+| validation BPB | 0.3373 |
+| code BPB, mean over 7 languages | n/a |
+| FLORES-200 BPB, mean | 3.3792 |
+| FLORES-200 languages scored | 215 |
+| BLiMP accuracy | n/a |
+| MBPP pass@1 | 0.2340 |
+| HumanEval pass@1 | 0.1646 |
+| GSM8K exact match, strict | 0.2320 |
+
+Generation metrics use specification `{'version': 'v2-2026-07-30', 'bos_context': True, 'eos_stop': True, 'token_healing': True, 'truncation': 'left_to_max_length_minus_max_gen', 'stop_truncation': 'earliest_match_all_terms', 'stop_match_decode': 'specials_kept'}`, token healing on.
+
+Missing measurements, with the reason each is absent:
+
+- `blimp_acc`: no full-128k-claude-balanced-bpe-seed43-mathcode-scratch_blimp_code_bpb.json
+- `code_bpb_mean`: no full-128k-claude-balanced-bpe-seed43-mathcode-scratch_blimp_code_bpb.json; no full-128k-claude-balanced-bpe-seed43-mathcode-scratch_code_bpb.json
+
+BLiMP is comparable within a training scale only. Its scoring sums unnormalized sentence log-likelihoods, and under that scoring a group of quantifier subtasks drops at the 1.27B scale relative to the 300M scale.
+
+## Provenance
+
+- Source checkpoint: `/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-claude-balanced-bpe-seed43-mathcode-scratch`, step 19073
+- Repository commit: `1a9a718`
+- nanochat commit: `f919449` on branch `tokenizer-lm-patches`
+- Tokenizer: `claude-balanced-bpe`, released at `tokenizers/claude-balanced-bpe`
+

models/full-128k-claude-balanced-bpe-seed43-mathcode-scratch/config.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "NanoChatForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "nanochat_model.NanoChatConfig",
+    "AutoModelForCausalLM": "nanochat_model.NanoChatForCausalLM"
+  },
+  "bos_token_id": 0,
+  "dtype": "float32",
+  "eos_token_id": null,
+  "head_dim": 128,
+  "kv_dim": 1536,
+  "logit_softcap": 15.0,
+  "model_type": "nanochat",
+  "n_embd": 1536,
+  "n_head": 12,
+  "n_kv_head": 12,
+  "n_layer": 24,
+  "pad_vocab_size_to": 64,
+  "padded_vocab_size": 128256,
+  "sequence_len": 2048,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.3.0",
+  "ve_dim": 128,
+  "vocab_size": 128256,
+  "window_pattern": "SSSL"
+}

models/full-128k-claude-balanced-bpe-seed43-mathcode-scratch/conversion_metadata.json

@@ -0,0 +1,100 @@
+{
+  "run": "full-128k-claude-balanced-bpe-seed43-mathcode-scratch",
+  "source_checkpoint": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-claude-balanced-bpe-seed43-mathcode-scratch",
+  "source_model_file": "model_019073.pt",
+  "source_step": 19073,
+  "source_meta": {
+    "step": 19073,
+    "val_bpb": 0.337272087796867,
+    "model_config": {
+      "sequence_len": 2048,
+      "vocab_size": 128256,
+      "n_layer": 24,
+      "n_head": 12,
+      "n_kv_head": 12,
+      "n_embd": 1536,
+      "window_pattern": "SSSL",
+      "ve_dim": 128,
+      "logit_softcap": 15.0
+    },
+    "tokenizer_path": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/tokenizers/claude-balanced-bpe",
+    "training_config": {
+      "name": "full-128k-mathcode-scratch",
+      "model": {
+        "n_layer": 24,
+        "n_embd": 1536,
+        "n_kv_head": 12,
+        "sequence_len": 2048,
+        "window_pattern": "SSSL",
+        "ve_dim": 128
+      },
+      "training": {
+        "target_param_data_ratio": -1,
+        "max_tokens": 20000000000,
+        "num_iterations": -1,
+        "total_batch_size": 1048576,
+        "device_batch_size": 16,
+        "warmup_steps": 40,
+        "warmdown_ratio": 0.65,
+        "final_lr_frac": 0.05,
+        "eval_tokens": 41943040,
+        "compile": true
+      },
+      "optimizer": {
+        "matrix_lr": 0.02,
+        "embedding_lr": 0.3,
+        "unembedding_lr": 0.008,
+        "scalar_lr": 0.5,
+        "weight_decay": 0.28,
+        "width_lr_exponent": -1.0
+      },
+      "data": {
+        "data_dir": null
+      },
+      "logging": {
+        "eval_every": 500,
+        "save_every": 500,
+        "max_checkpoints": 0
+      },
+      "wandb": {
+        "enabled": true,
+        "project": "tokenizer-lm-experiments",
+        "run_name": null
+      },
+      "output_dir": "/capstor/scratch/cscs/${USER}/tokenizer-lm/checkpoints",
+      "seed": 42
+    },
+    "total_batch_size": 1048576,
+    "device_batch_size": 16,
+    "max_seq_len": 2048,
+    "dataloader_state_dict": {
+      "pq_idx": 219,
+      "rg_idx": 44,
+      "epoch": 2
+    },
+    "loop_state": {
+      "min_val_bpb": 0.337272087796867,
+      "smooth_train_loss": 0.971582087716185,
+      "total_training_time": 93265.44177412987,
+      "total_bytes_consumed": 151300558757
+    }
+  },
+  "weight_dtype": "as trained (float32 matrices, bfloat16 embeddings)",
+  "weight_dtype_counts": {
+    "torch.float32": 174,
+    "torch.bfloat16": 13
+  },
+  "total_parameters": 1272841994,
+  "logit_softcap_source": "meta",
+  "bos_token": "<s>",
+  "bos_token_id": 0,
+  "bos_resolved_from": "packaged tokenizer",
+  "tokenizer_slug": "claude-balanced-bpe",
+  "tokenizer_redistributed": true,
+  "tokenizer_hub_id": null,
+  "weight_key_mapping": "strip _orig_mod. and transformer. prefixes",
+  "converted_by": "scripts/hf_release/convert_for_release.py",
+  "repo_commit": "1a9a718",
+  "nanochat_commit": "f919449",
+  "nanochat_branch": "tokenizer-lm-patches"
+}

models/full-128k-claude-balanced-bpe-seed43-mathcode-scratch/generation_config.json

@@ -0,0 +1,4 @@
+{
+  "bos_token_id": 0,
+  "transformers_version": "5.3.0"
+}

models/full-128k-claude-balanced-bpe-seed43-mathcode-scratch/nanochat_model.py

@@ -0,0 +1,361 @@
+"""
+NanoChat model for HuggingFace Transformers — self-contained single file.
+
+Includes the config class, model class, and a CLI demo for loading + generation.
+Drop this file into a directory containing config.json, model.safetensors,
+tokenizer.json, and tokenizer_config.json, then run:
+
+    python nanochat_model.py --model-dir . --prompt "Hello world"
+
+Or load programmatically:
+
+    from nanochat_model import NanoChatConfig, NanoChatForCausalLM
+    model = NanoChatForCausalLM.from_pretrained("path/to/model", torch_dtype=torch.bfloat16)
+
+Requirements: pip install torch transformers tokenizers
+
+Architecture: nanochat GPT (Karpathy, MIT licensed) with value embeddings
+(ResFormer-style), per-layer residual/x0 lambdas, smear, backout, logit
+soft-capping, sliding window attention, QK norm, ReLU^2 MLP, and RoPE.
+
+The models were trained with a patched nanochat, not upstream nanochat:
+https://github.com/cimeister/nanochat branch tokenizer-lm-patches, commit f919449.
+The patches touch two files and change three things, two of which affect the forward
+pass and are implemented here:
+  1. ve_dim: the value-embedding tables are built at a reduced dimension and projected
+     up to kv_dim by a learned ve_proj, instead of being built at full kv_dim. All
+     released models use ve_dim=128.
+  2. logit_softcap: the soft cap is a config field with 0 meaning disabled, instead of
+     a hardcoded 15.
+  3. width_lr_exponent: the AdamW learning-rate width scaling exponent is configurable
+     for muP. This affects training only and has no inference-time effect.
+
+This file is the released model definition. Its agreement with the training-time model
+is checked per released checkpoint by scripts/hf_release/verify_conversion.py, which
+compares logits against the patched nanochat GPT in float32.
+"""
+
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+class NanoChatConfig(PretrainedConfig):
+    model_type = "nanochat"
+
+    def __init__(
+        self,
+        vocab_size=128256,
+        n_embd=1536,
+        n_layer=24,
+        n_head=12,
+        n_kv_head=12,
+        sequence_len=2048,
+        window_pattern="SSSL",
+        ve_dim=0,
+        logit_softcap=15.0,
+        pad_vocab_size_to=64,
+        bos_token_id=0,
+        eos_token_id=None,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_kv_head = n_kv_head
+        self.sequence_len = sequence_len
+        self.window_pattern = window_pattern
+        self.ve_dim = ve_dim
+        self.logit_softcap = logit_softcap
+        self.pad_vocab_size_to = pad_vocab_size_to
+        self.head_dim = n_embd // n_head
+        self.kv_dim = n_kv_head * self.head_dim
+        self.padded_vocab_size = (
+            (vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to
+        ) * pad_vocab_size_to
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Model
+# ---------------------------------------------------------------------------
+
+def _has_ve(layer_idx: int, n_layer: int) -> bool:
+    return layer_idx % 2 == (n_layer - 1) % 2
+
+
+def _apply_rotary_emb(x, cos, sin):
+    d = x.shape[-1] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    return torch.cat([y1, y2], dim=-1)
+
+
+def _rms_norm(x):
+    return F.rms_norm(x, (x.size(-1),))
+
+
+class NanoChatAttention(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.n_embd = config.n_embd
+        self.head_dim = config.head_dim
+        self.n_kv_groups = self.n_head // self.n_kv_head
+        self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
+        self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
+        self.ve_gate_channels = 12
+        if _has_ve(layer_idx, config.n_layer):
+            self.ve_gate = nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False)
+            kv_dim = self.n_kv_head * self.head_dim
+            ve_dim = config.ve_dim if config.ve_dim > 0 else kv_dim
+            self.ve_proj = nn.Linear(ve_dim, kv_dim, bias=False) if ve_dim < kv_dim else None
+        else:
+            self.ve_gate = None
+            self.ve_proj = None
+
+    def forward(self, x, ve, cos, sin, window_size):
+        B, T, C = x.size()
+        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
+        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
+        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
+        if ve is not None and self.ve_gate is not None:
+            if self.ve_proj is not None:
+                ve = self.ve_proj(ve)
+            ve = ve.view(B, T, self.n_kv_head, self.head_dim)
+            gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
+            v = v + gate.unsqueeze(-1) * ve
+        q = _apply_rotary_emb(q, cos, sin)
+        k = _apply_rotary_emb(k, cos, sin)
+        q = _rms_norm(q) * 1.2
+        k = _rms_norm(k) * 1.2
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        if self.n_kv_groups > 1:
+            k = k.repeat_interleave(self.n_kv_groups, dim=1)
+            v = v.repeat_interleave(self.n_kv_groups, dim=1)
+        if window_size < 0 or window_size >= T:
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:
+            row_idx = torch.arange(T, device=x.device).unsqueeze(1)
+            col_idx = torch.arange(T, device=x.device).unsqueeze(0)
+            mask = (col_idx <= row_idx) & ((row_idx - col_idx) <= window_size)
+            y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        y = y.transpose(1, 2).contiguous().view(B, T, -1)
+        y = self.c_proj(y)
+        return y
+
+
+class NanoChatMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = F.relu(x).square()
+        x = self.c_proj(x)
+        return x
+
+
+class NanoChatBlock(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.attn = NanoChatAttention(config, layer_idx)
+        self.mlp = NanoChatMLP(config)
+
+    def forward(self, x, ve, cos, sin, window_size):
+        x = x + self.attn(_rms_norm(x), ve, cos, sin, window_size)
+        x = x + self.mlp(_rms_norm(x))
+        return x
+
+
+class NanoChatPreTrainedModel(PreTrainedModel):
+    config_class = NanoChatConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["NanoChatBlock"]
+
+    def _init_weights(self, module):
+        pass
+
+
+class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin):
+    def __init__(self, config):
+        super().__init__(config)
+        pvs = config.padded_vocab_size
+        self.wte = nn.Embedding(pvs, config.n_embd)
+        self.h = nn.ModuleList([NanoChatBlock(config, i) for i in range(config.n_layer)])
+        self.lm_head = nn.Linear(config.n_embd, pvs, bias=False)
+        self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))
+        self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))
+        self.smear_gate = nn.Linear(24, 1, bias=False)
+        self.smear_lambda = nn.Parameter(torch.zeros(1))
+        self.backout_lambda = nn.Parameter(0.2 * torch.ones(1))
+        ve_emb_dim = config.ve_dim if config.ve_dim > 0 else config.kv_dim
+        self.value_embeds = nn.ModuleDict({
+            str(i): nn.Embedding(pvs, ve_emb_dim)
+            for i in range(config.n_layer) if _has_ve(i, config.n_layer)
+        })
+        self.window_sizes = self._compute_window_sizes(config)
+        self.rotary_seq_len = config.sequence_len * 10
+        self._rotary_initialized = False
+        self.register_buffer("cos", torch.zeros(1), persistent=False)
+        self.register_buffer("sin", torch.zeros(1), persistent=False)
+        self.post_init()
+
+    def _compute_window_sizes(self, config):
+        pattern = config.window_pattern.upper()
+        long_window = config.sequence_len
+        short_window = -(-long_window // 4 // 128) * 128
+        char_to_window = {"L": long_window, "S": short_window}
+        sizes = [char_to_window[pattern[i % len(pattern)]] for i in range(config.n_layer)]
+        sizes[-1] = long_window
+        return sizes
+
+    def _precompute_rotary(self, seq_len, head_dim, base=100000.0):
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32)
+        inv_freq = 1.0 / (base ** (channel_range / head_dim))
+        t = torch.arange(seq_len, dtype=torch.float32)
+        freqs = torch.outer(t, inv_freq)
+        cos = freqs.cos()[None, :, None, :]
+        sin = freqs.sin()[None, :, None, :]
+        return cos, sin
+
+    def forward(self, input_ids, attention_mask=None, labels=None, return_dict=None, **kwargs):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        B, T = input_ids.size()
+        if not self._rotary_initialized:
+            cos, sin = self._precompute_rotary(self.rotary_seq_len, self.config.head_dim)
+            self.cos = cos.to(device=input_ids.device)
+            self.sin = sin.to(device=input_ids.device)
+            self._rotary_initialized = True
+        cos = self.cos[:, :T].to(dtype=self.wte.weight.dtype)
+        sin = self.sin[:, :T].to(dtype=self.wte.weight.dtype)
+        x = self.wte(input_ids)
+        x = _rms_norm(x)
+        if T > 1:
+            gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
+            x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
+        x0 = x
+        n_layer = self.config.n_layer
+        backout_layer = n_layer // 2
+        x_backout = None
+        for i, block in enumerate(self.h):
+            x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
+            ve = self.value_embeds[str(i)](input_ids).to(x.dtype) if str(i) in self.value_embeds else None
+            x = block(x, ve, cos, sin, self.window_sizes[i])
+            if i == backout_layer:
+                x_backout = x
+        if x_backout is not None:
+            x = x - self.backout_lambda.to(x.dtype) * x_backout
+        x = _rms_norm(x)
+        logits = self.lm_head(x)
+        logits = logits[..., :self.config.vocab_size]
+        logits = logits.float()
+        if self.config.logit_softcap > 0:
+            cap = self.config.logit_softcap
+            logits = cap * torch.tanh(logits / cap)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+        if not return_dict:
+            output = (logits,)
+            return (loss,) + output if loss is not None else output
+        return CausalLMOutputWithPast(loss=loss, logits=logits)
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# ---------------------------------------------------------------------------
+# CLI demo
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    import argparse
+    import sys
+    import os
+
+    parser = argparse.ArgumentParser(description="Load and generate with a NanoChat model")
+    parser.add_argument("--model-dir", required=True, help="Path to converted model directory")
+    parser.add_argument("--prompt", default="The meaning of life is", help="Text prompt")
+    parser.add_argument("--max-tokens", type=int, default=100, help="Max tokens to generate")
+    parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature (0=greedy)")
+    parser.add_argument("--device", default="cpu", help="Device (cpu, cuda, cuda:0, ...)")
+    parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16", "float16"])
+    args = parser.parse_args()
+
+    dtype_map = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}
+
+    print(f"Loading model from {args.model_dir}...")
+    from transformers import AutoTokenizer
+
+    model = NanoChatForCausalLM.from_pretrained(args.model_dir, torch_dtype=dtype_map[args.dtype])
+    model.to(args.device)
+    model.eval()
+    tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Model loaded: {total_params:,} parameters, dtype={args.dtype}, device={args.device}")
+    print(f"Vocab size: {model.config.vocab_size}, ve_dim: {model.config.ve_dim}")
+    print()
+    print(f"Prompt: {args.prompt}")
+    print(f"Generating (max {args.max_tokens} tokens, temp={args.temperature})...")
+    print()
+
+    # Generate
+    bos_id = tokenizer.bos_token_id
+    input_ids = tokenizer.encode(args.prompt, add_special_tokens=False)
+    if bos_id is not None:
+        input_ids = [bos_id] + input_ids
+    input_ids = torch.tensor([input_ids], dtype=torch.long, device=args.device)
+
+    generated = []
+    with torch.inference_mode():
+        for _ in range(args.max_tokens):
+            logits = model(input_ids).logits[:, -1, :]
+            if args.temperature > 0:
+                logits = logits / args.temperature
+                v, _ = torch.topk(logits, min(50, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float("inf")
+                probs = torch.softmax(logits, dim=-1)
+                next_id = torch.multinomial(probs, num_samples=1)
+            else:
+                next_id = torch.argmax(logits, dim=-1, keepdim=True)
+            input_ids = torch.cat([input_ids, next_id], dim=1)
+            token_id = next_id.item()
+            generated.append(token_id)
+            if tokenizer.eos_token_id is not None and token_id == tokenizer.eos_token_id:
+                break
+
+    output = tokenizer.decode(generated, skip_special_tokens=True)
+    print(f"Output: {args.prompt}{output}")

models/full-128k-claude-balanced-bpe-seed43-mathcode-scratch/tokenizer_config.json

@@ -0,0 +1,7 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "model_max_length": 2048,
+  "bos_token": "<s>",
+  "add_bos_token": false,
+  "clean_up_tokenization_spaces": false
+}

models/full-128k-claude-balanced-nfc-bpe-mathcode-scratch/README.md

@@ -0,0 +1,76 @@
+# `full-128k-claude-balanced-nfc-bpe-mathcode-scratch`
+
+A 1.27B parameter model trained with the `claude-balanced-nfc-bpe` tokenizer. Regime: mathcode-20B-from-scratch (24 layers, width 1536, 1.27B parameters, 20.0B tokens, math and code mixture, trained from scratch). Reported in: math and code ablations; TokEval.
+
+## Load
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+    "cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-claude-balanced-nfc-bpe-mathcode-scratch",
+    trust_remote_code=True, dtype="float32")
+tokenizer = AutoTokenizer.from_pretrained("cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-claude-balanced-nfc-bpe-mathcode-scratch")
+
+text = "def add(a, b):\n    return a + b\n"
+ids = [0] + tokenizer(text)["input_ids"]   # training prepends BOS
+```
+
+## Settings
+
+| Setting | Value |
+|---|---|
+| layers | 24 |
+| model dimension | 1536 |
+| attention heads | 12 |
+| key and value heads | 12 |
+| context length | 2048 |
+| sliding window pattern | SSSL |
+| value embedding dimension | 128 |
+| logit soft cap | 15 |
+| vocabulary size | 128256 |
+| parameters | 1,272,841,994 |
+| training step | 19073 |
+| token budget | 19,999,490,048 |
+| bytes consumed | 151,300,253,589 |
+| seed | 42 |
+| Muon matrix learning rate | 0.0200 |
+| AdamW embedding learning rate | 0.3000 |
+| AdamW unembedding learning rate | 0.0080 |
+| weight decay | 0.2800 |
+| width learning-rate exponent | -1 |
+| weight dtype | as trained (float32 matrices, bfloat16 embeddings) |
+| BOS token id | 0 |
+
+## Results
+
+| Metric | Value |
+|---|---|
+| validation BPB | 0.3371 |
+| code BPB, mean over 7 languages | 0.3845 |
+| FLORES-200 BPB, mean | 3.3667 |
+| FLORES-200 languages scored | 215 |
+| BLiMP accuracy | n/a |
+| MBPP pass@1 | 0.2240 |
+| HumanEval pass@1 | 0.1585 |
+| GSM8K exact match, strict | 0.2305 |
+
+Generation metrics use specification `{'version': 'v2-2026-07-30', 'bos_context': True, 'eos_stop': True, 'token_healing': True, 'truncation': 'left_to_max_length_minus_max_gen', 'stop_truncation': 'earliest_match_all_terms', 'stop_match_decode': 'specials_kept'}`, token healing on.
+
+Missing measurements, with the reason each is absent:
+
+- `blimp_acc`: no full-128k-claude-balanced-nfc-bpe-mathcode-scratch_blimp_code_bpb.json
+
+BLiMP is comparable within a training scale only. Its scoring sums unnormalized sentence log-likelihoods, and under that scoring a group of quantifier subtasks drops at the 1.27B scale relative to the 300M scale.
+
+## Registry note
+
+#50, claude pretok, slug `claude-balanced-nfc-bpe`. Main-grid Claude NFC cell. FINISHED at step 19073, 4-eval chain + code_bpb complete (verified 2026-08-09). In panel.
+
+## Provenance
+
+- Source checkpoint: `/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-claude-balanced-nfc-bpe-mathcode-scratch`, step 19073
+- Repository commit: `1a9a718`
+- nanochat commit: `f919449` on branch `tokenizer-lm-patches`
+- Tokenizer: `claude-balanced-nfc-bpe`, released at `tokenizers/claude-balanced-nfc-bpe`
+

models/full-128k-claude-balanced-nfc-bpe-mathcode-scratch/config.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "NanoChatForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "nanochat_model.NanoChatConfig",
+    "AutoModelForCausalLM": "nanochat_model.NanoChatForCausalLM"
+  },
+  "bos_token_id": 0,
+  "dtype": "float32",
+  "eos_token_id": null,
+  "head_dim": 128,
+  "kv_dim": 1536,
+  "logit_softcap": 15.0,
+  "model_type": "nanochat",
+  "n_embd": 1536,
+  "n_head": 12,
+  "n_kv_head": 12,
+  "n_layer": 24,
+  "pad_vocab_size_to": 64,
+  "padded_vocab_size": 128256,
+  "sequence_len": 2048,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.3.0",
+  "ve_dim": 128,
+  "vocab_size": 128256,
+  "window_pattern": "SSSL"
+}

models/full-128k-claude-balanced-nfc-bpe-mathcode-scratch/conversion_metadata.json

@@ -0,0 +1,250 @@
+{
+  "run": "full-128k-claude-balanced-nfc-bpe-mathcode-scratch",
+  "source_checkpoint": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-claude-balanced-nfc-bpe-mathcode-scratch",
+  "source_model_file": "model_019073.pt",
+  "source_step": 19073,
+  "source_meta": {
+    "step": 19073,
+    "val_bpb": 0.33714380335290617,
+    "model_config": {
+      "sequence_len": 2048,
+      "vocab_size": 128256,
+      "n_layer": 24,
+      "n_head": 12,
+      "n_kv_head": 12,
+      "n_embd": 1536,
+      "window_pattern": "SSSL",
+      "ve_dim": 128,
+      "logit_softcap": 15.0
+    },
+    "tokenizer_path": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/tokenizers/claude-balanced-nfc-bpe",
+    "training_config": {
+      "name": "full-128k-mathcode-scratch",
+      "model": {
+        "n_layer": 24,
+        "n_embd": 1536,
+        "n_kv_head": 12,
+        "sequence_len": 2048,
+        "window_pattern": "SSSL",
+        "ve_dim": 128
+      },
+      "training": {
+        "target_param_data_ratio": -1,
+        "max_tokens": 20000000000,
+        "num_iterations": -1,
+        "total_batch_size": 1048576,
+        "device_batch_size": 16,
+        "warmup_steps": 40,
+        "warmdown_ratio": 0.65,
+        "final_lr_frac": 0.05,
+        "eval_tokens": 41943040,
+        "compile": true
+      },
+      "optimizer": {
+        "matrix_lr": 0.02,
+        "embedding_lr": 0.3,
+        "unembedding_lr": 0.008,
+        "scalar_lr": 0.5,
+        "weight_decay": 0.28,
+        "width_lr_exponent": -1.0
+      },
+      "data": {
+        "data_dir": null
+      },
+      "logging": {
+        "eval_every": 500,
+        "save_every": 500,
+        "max_checkpoints": 0
+      },
+      "wandb": {
+        "enabled": true,
+        "project": "tokenizer-lm-experiments",
+        "run_name": null
+      },
+      "output_dir": "/capstor/scratch/cscs/${USER}/tokenizer-lm/checkpoints",
+      "seed": 42
+    },
+    "code_provenance": {
+      "commit": "6231791054e9c0df5144025df8f8a2117e8b78d2",
+      "branch": "main",
+      "dirty": true,
+      "dirty_files": [
+        "M .gitignore",
+        " M EXPERIMENTAL_SETUP.md",
+        " M RESULT_COMPUTATION_ANALYSIS.md",
+        " M eval_runner.py",
+        " M mc_evaluation",
+        " M scripts/add_mc_math_ci.py",
+        " M scripts/bootstrap_mathcode_significance.py",
+        " M scripts/correlate_intrinsic_downstream.py",
+        " M scripts/correlate_mathcode_intrinsic_downstream.py",
+        " M scripts/cross_lingual_analyses.py",
+        " M scripts/generate_latex_tables.py",
+        " M scripts/launch_mathcode_scratch_chained.sh",
+        " M scripts/launch_new6_all_scales.sh",
+        " M scripts/reorganize_checkpoints.py",
+        " M scripts/slurm_3b_mathcode_resume.sh",
+        " M scripts/slurm_eval_mgsm.sh",
+        " M scripts/slurm_full.sh",
+        " M scripts/slurm_full_1node.sh",
+        " M scripts/slurm_full_1node_resume.sh",
+        " M scripts/slurm_mathcode_resume_1node.sh",
+        " M scripts/slurm_pilot.sh",
+        " M scripts/slurm_train_1node.sh",
+        " M scripts/slurm_train_tokenizers.sh",
+        " M scripts/train_custom_tokenizers.py",
+        " M tokenizer_lm/tokenizer.py",
+        " M train.py",
+        "?? REPRODUCIBILITY_STATUS.md",
+        "?? SESSION_STATUS.md",
+        "?? TOKEN_HEALING_HANDOFF.md",
+        "?? configs/full_128k_bytematched.yaml",
+        "?? configs/full_128k_clean_v128k.yaml",
+        "?? configs/full_128k_clean_v160k.yaml",
+        "?? configs/full_128k_clean_v200k.yaml",
+        "?? configs/full_128k_clean_v64k.yaml",
+        "?? configs/full_128k_clean_v96k.yaml",
+        "?? configs/panels/",
+        "?? configs/starcoder_ast_18lang.json",
+        "?? correspondence/",
+        "?? scripts/_lexer_treesitter_worker.py",
+        "?? scripts/add_code_bpb_std.py",
+        "?? scripts/analyze_apertus_track_sensitivity.py",
+        "?? scripts/analyze_axes2_codegrad.py",
+        "?? scripts/analyze_code_boundary_drift.py",
+        "?? scripts/analyze_code_gen_failures.py",
+        "?? scripts/analyze_code_intrinsic_downstream.py",
+        "?? scripts/analyze_cr_byte_drop.py",
+        "?? scripts/analyze_cross_scale_ranking.py",
+        "?? scripts/analyze_equal_weighting.py",
+        "?? scripts/analyze_gen_ablation.py",
+        "?? scripts/analyze_seed_variance.py",
+        "?? scripts/analyze_seed_variance_mc20b.py",
+        "?? scripts/analyze_vocab_size.py",
+        "?? scripts/ast_node_discovery.py",
+        "?? scripts/ast_scope_extractor.py",
+        "?? scripts/audit_bpb_byte_bias.py",
+        "?? scripts/audit_corpus_parseability.py",
+        "?? scripts/audit_generation_symptoms.py",
+        "?? scripts/audit_lexer_boundary_leak.py",
+        "?? scripts/audit_logit_softcap.py",
+        "?? scripts/audit_maxchars_fragmentation.py",
+        "?? scripts/audit_mgsm_prompt_boundary.py",
+        "?? scripts/audit_newline_tokens.py",
+        "?? scripts/audit_newline_tokens_fullbyte_pair.py",
+        "?? scripts/audit_newline_tokens_new11.py",
+        "?? scripts/backfill_flores_summary.py",
+        "?? scripts/build_axes2_codegrad_tokenizers.sh",
+        "?? scripts/build_vocab_ladder_tokenizers.sh",
+        "?? scripts/build_whitespace_ladder_tokenizers.sh",
+        "?? scripts/characterize_axes2_intrinsic.py",
+        "?? scripts/check_hf_dataset_cache.py",
+        "?? scripts/check_number_provenance.py",
+        "?? scripts/compare_bigram_entropy_variants.py",
+        "?? scripts/compute_headline_spreads.py",
+        "?? scripts/compute_vocab_ladder_byte_matching.py",
+        "?? scripts/correlate_operator_domains.py",
+        "?? scripts/d2_candidate_screen.py",
+        "?? scripts/diagnose_fullbyte_indent.py",
+        "?? scripts/diagnose_operator_isolation.py",
+        "?? scripts/dump_axes2_composition_weights.py",
+        "?? scripts/estimate_intrinsic_cost.py",
+        "?? scripts/fit_quota_response_model.py",
+        "?? scripts/gen_v2_spec_audit.py",
+        "?? scripts/generate_intrinsic_tables.py",
+        "?? scripts/launch_boundary_retrain.sh",
+        "?? scripts/launch_mc_pretok_chained.sh",
+        "?? scripts/launch_std_pretok_chained.sh",
+        "?? scripts/launch_vocab_ladder.sh",
+        "?? scripts/launch_vocab_ladder_paramscaled.sh",
+        "?? scripts/lexer_pretok_common.py",
+        "?? scripts/lexer_pretok_corpus.py",
+        "?? scripts/mbpp_mediation.py",
+        "?? scripts/merge_intrinsic_components.py",
+        "?? scripts/merge_intrinsic_shards.py",
+        "?? scripts/operator_isolation_by_corpus_n36.py",
+        "?? scripts/optimize_quotas.py",
+        "?? scripts/preflight_gen_v2_tokenizers.py",
+        "?? scripts/probe_belebele_bos.py",
+        "?? scripts/probe_eot_stop_incidence.py",
+        "?? scripts/probe_healing_prefill_equivalence.py",
+        "?? scripts/recheck_seed_sd_verdicts.py",
+        "?? scripts/recompute_val_bpb.py",
+        "?? scripts/rejudge_tail_penalty.py",
+        "?? scripts/render_code_appendix.py",
+        "?? scripts/retry_whitespace_axes2_builds.sh",
+        "?? scripts/run_gen_v2_one.sh",
+        "?? scripts/run_mc_panels.sh",
+        "?? scripts/script_pretok_per_language.py",
+        "?? scripts/slurm_corpus_parseability_audit.sh",
+        "?? scripts/slurm_d2_candidate_build.sh",
+        "?? scripts/slurm_d2_candidate_screen.sh",
+        "?? scripts/slurm_eval_code_gen_healed.sh",
+        "?? scripts/slurm_eval_gen_ablation.sh",
+        "?? scripts/slurm_eval_gen_v2.sh",
+        "?? scripts/slurm_eval_gen_v2_packed.sbatch",
+        "?? scripts/slurm_full_1node_resume_pretok.sh",
+        "?? scripts/slurm_intrinsic_camera_ready.sbatch",
+        "?? scripts/slurm_intrinsic_component.sbatch",
+        "?? scripts/slurm_probe_mgsm_bos.sbatch",
+        "?? scripts/slurm_train_pretok_1node.sh",
+        "?? scripts/submit_gen_v2_sweep.sh",
+        "?? scripts/submit_lexer_corpus_cache_jobs.sh",
+        "?? scripts/summarize_lexer_smoke_run.py",
+        "?? scripts/test_ast_scope_extractor.py",
+        "?? scripts/test_lexer_pretok_common.py",
+        "?? scripts/test_lexer_treesitter_worker.py",
+        "?? scripts/test_presegmented_training_path.py",
+        "?? scripts/token_healing_failure_modes.py",
+        "?? scripts/validate_axes2_codegrad_tokenizers.py",
+        "?? scripts/validate_generation_conformance.py",
+        "?? scripts/validate_gpt4o_english_fullbyte_bpe.py",
+        "?? scripts/validate_token_healing.py",
+        "?? scripts/validate_whitespace_ladder_regexes.py",
+        "?? scripts/validate_whitespace_ladder_tokenizers.py",
+        "?? scripts/verify_gen_v2_smoke.py",
+        "?? scripts/verify_iter_source_hardening.py",
+        "?? scripts/vet_runtime_regex.py",
+        "?? scripts/whitespace_fusion_audit.py",
+        "?? std1b_invalid.md",
+        "?? std1b_used.md",
+        "?? tokenizer_lm/pretok_dataloader.py",
+        "?? tokenizer_lm/token_healing.py"
+      ],
+      "nanochat_commit": "f91944934f78d2083c8f5bd6ea8ab4e048786054",
+      "torch_version": "2.10.0+cu126"
+    },
+    "total_batch_size": 1048576,
+    "device_batch_size": 16,
+    "max_seq_len": 2048,
+    "dataloader_state_dict": {
+      "pq_idx": 219,
+      "rg_idx": 44,
+      "epoch": 2
+    },
+    "loop_state": {
+      "min_val_bpb": 0.33714380335290617,
+      "smooth_train_loss": 0.9612813208769648,
+      "total_training_time": 93846.44612264633,
+      "total_bytes_consumed": 151300253589
+    }
+  },
+  "weight_dtype": "as trained (float32 matrices, bfloat16 embeddings)",
+  "weight_dtype_counts": {
+    "torch.float32": 174,
+    "torch.bfloat16": 13
+  },
+  "total_parameters": 1272841994,
+  "logit_softcap_source": "meta",
+  "bos_token": "<s>",
+  "bos_token_id": 0,
+  "bos_resolved_from": "packaged tokenizer",
+  "tokenizer_slug": "claude-balanced-nfc-bpe",
+  "tokenizer_redistributed": true,
+  "tokenizer_hub_id": null,
+  "weight_key_mapping": "strip _orig_mod. and transformer. prefixes",
+  "converted_by": "scripts/hf_release/convert_for_release.py",
+  "repo_commit": "1a9a718",
+  "nanochat_commit": "f919449",
+  "nanochat_branch": "tokenizer-lm-patches"
+}

models/full-128k-claude-balanced-nfc-bpe-mathcode-scratch/generation_config.json

@@ -0,0 +1,4 @@
+{
+  "bos_token_id": 0,
+  "transformers_version": "5.3.0"
+}

models/full-128k-claude-balanced-nfc-bpe-mathcode-scratch/nanochat_model.py

@@ -0,0 +1,361 @@
+"""
+NanoChat model for HuggingFace Transformers — self-contained single file.
+
+Includes the config class, model class, and a CLI demo for loading + generation.
+Drop this file into a directory containing config.json, model.safetensors,
+tokenizer.json, and tokenizer_config.json, then run:
+
+    python nanochat_model.py --model-dir . --prompt "Hello world"
+
+Or load programmatically:
+
+    from nanochat_model import NanoChatConfig, NanoChatForCausalLM
+    model = NanoChatForCausalLM.from_pretrained("path/to/model", torch_dtype=torch.bfloat16)
+
+Requirements: pip install torch transformers tokenizers
+
+Architecture: nanochat GPT (Karpathy, MIT licensed) with value embeddings
+(ResFormer-style), per-layer residual/x0 lambdas, smear, backout, logit
+soft-capping, sliding window attention, QK norm, ReLU^2 MLP, and RoPE.
+
+The models were trained with a patched nanochat, not upstream nanochat:
+https://github.com/cimeister/nanochat branch tokenizer-lm-patches, commit f919449.
+The patches touch two files and change three things, two of which affect the forward
+pass and are implemented here:
+  1. ve_dim: the value-embedding tables are built at a reduced dimension and projected
+     up to kv_dim by a learned ve_proj, instead of being built at full kv_dim. All
+     released models use ve_dim=128.
+  2. logit_softcap: the soft cap is a config field with 0 meaning disabled, instead of
+     a hardcoded 15.
+  3. width_lr_exponent: the AdamW learning-rate width scaling exponent is configurable
+     for muP. This affects training only and has no inference-time effect.
+
+This file is the released model definition. Its agreement with the training-time model
+is checked per released checkpoint by scripts/hf_release/verify_conversion.py, which
+compares logits against the patched nanochat GPT in float32.
+"""
+
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+class NanoChatConfig(PretrainedConfig):
+    model_type = "nanochat"
+
+    def __init__(
+        self,
+        vocab_size=128256,
+        n_embd=1536,
+        n_layer=24,
+        n_head=12,
+        n_kv_head=12,
+        sequence_len=2048,
+        window_pattern="SSSL",
+        ve_dim=0,
+        logit_softcap=15.0,
+        pad_vocab_size_to=64,
+        bos_token_id=0,
+        eos_token_id=None,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_kv_head = n_kv_head
+        self.sequence_len = sequence_len
+        self.window_pattern = window_pattern
+        self.ve_dim = ve_dim
+        self.logit_softcap = logit_softcap
+        self.pad_vocab_size_to = pad_vocab_size_to
+        self.head_dim = n_embd // n_head
+        self.kv_dim = n_kv_head * self.head_dim
+        self.padded_vocab_size = (
+            (vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to
+        ) * pad_vocab_size_to
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Model
+# ---------------------------------------------------------------------------
+
+def _has_ve(layer_idx: int, n_layer: int) -> bool:
+    return layer_idx % 2 == (n_layer - 1) % 2
+
+
+def _apply_rotary_emb(x, cos, sin):
+    d = x.shape[-1] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    return torch.cat([y1, y2], dim=-1)
+
+
+def _rms_norm(x):
+    return F.rms_norm(x, (x.size(-1),))
+
+
+class NanoChatAttention(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.n_embd = config.n_embd
+        self.head_dim = config.head_dim
+        self.n_kv_groups = self.n_head // self.n_kv_head
+        self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
+        self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
+        self.ve_gate_channels = 12
+        if _has_ve(layer_idx, config.n_layer):
+            self.ve_gate = nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False)
+            kv_dim = self.n_kv_head * self.head_dim
+            ve_dim = config.ve_dim if config.ve_dim > 0 else kv_dim
+            self.ve_proj = nn.Linear(ve_dim, kv_dim, bias=False) if ve_dim < kv_dim else None
+        else:
+            self.ve_gate = None
+            self.ve_proj = None
+
+    def forward(self, x, ve, cos, sin, window_size):
+        B, T, C = x.size()
+        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
+        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
+        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
+        if ve is not None and self.ve_gate is not None:
+            if self.ve_proj is not None:
+                ve = self.ve_proj(ve)
+            ve = ve.view(B, T, self.n_kv_head, self.head_dim)
+            gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
+            v = v + gate.unsqueeze(-1) * ve
+        q = _apply_rotary_emb(q, cos, sin)
+        k = _apply_rotary_emb(k, cos, sin)
+        q = _rms_norm(q) * 1.2
+        k = _rms_norm(k) * 1.2
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        if self.n_kv_groups > 1:
+            k = k.repeat_interleave(self.n_kv_groups, dim=1)
+            v = v.repeat_interleave(self.n_kv_groups, dim=1)
+        if window_size < 0 or window_size >= T:
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:
+            row_idx = torch.arange(T, device=x.device).unsqueeze(1)
+            col_idx = torch.arange(T, device=x.device).unsqueeze(0)
+            mask = (col_idx <= row_idx) & ((row_idx - col_idx) <= window_size)
+            y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        y = y.transpose(1, 2).contiguous().view(B, T, -1)
+        y = self.c_proj(y)
+        return y
+
+
+class NanoChatMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = F.relu(x).square()
+        x = self.c_proj(x)
+        return x
+
+
+class NanoChatBlock(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.attn = NanoChatAttention(config, layer_idx)
+        self.mlp = NanoChatMLP(config)
+
+    def forward(self, x, ve, cos, sin, window_size):
+        x = x + self.attn(_rms_norm(x), ve, cos, sin, window_size)
+        x = x + self.mlp(_rms_norm(x))
+        return x
+
+
+class NanoChatPreTrainedModel(PreTrainedModel):
+    config_class = NanoChatConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["NanoChatBlock"]
+
+    def _init_weights(self, module):
+        pass
+
+
+class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin):
+    def __init__(self, config):
+        super().__init__(config)
+        pvs = config.padded_vocab_size
+        self.wte = nn.Embedding(pvs, config.n_embd)
+        self.h = nn.ModuleList([NanoChatBlock(config, i) for i in range(config.n_layer)])
+        self.lm_head = nn.Linear(config.n_embd, pvs, bias=False)
+        self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))
+        self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))
+        self.smear_gate = nn.Linear(24, 1, bias=False)
+        self.smear_lambda = nn.Parameter(torch.zeros(1))
+        self.backout_lambda = nn.Parameter(0.2 * torch.ones(1))
+        ve_emb_dim = config.ve_dim if config.ve_dim > 0 else config.kv_dim
+        self.value_embeds = nn.ModuleDict({
+            str(i): nn.Embedding(pvs, ve_emb_dim)
+            for i in range(config.n_layer) if _has_ve(i, config.n_layer)
+        })
+        self.window_sizes = self._compute_window_sizes(config)
+        self.rotary_seq_len = config.sequence_len * 10
+        self._rotary_initialized = False
+        self.register_buffer("cos", torch.zeros(1), persistent=False)
+        self.register_buffer("sin", torch.zeros(1), persistent=False)
+        self.post_init()
+
+    def _compute_window_sizes(self, config):
+        pattern = config.window_pattern.upper()
+        long_window = config.sequence_len
+        short_window = -(-long_window // 4 // 128) * 128
+        char_to_window = {"L": long_window, "S": short_window}
+        sizes = [char_to_window[pattern[i % len(pattern)]] for i in range(config.n_layer)]
+        sizes[-1] = long_window
+        return sizes
+
+    def _precompute_rotary(self, seq_len, head_dim, base=100000.0):
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32)
+        inv_freq = 1.0 / (base ** (channel_range / head_dim))
+        t = torch.arange(seq_len, dtype=torch.float32)
+        freqs = torch.outer(t, inv_freq)
+        cos = freqs.cos()[None, :, None, :]
+        sin = freqs.sin()[None, :, None, :]
+        return cos, sin
+
+    def forward(self, input_ids, attention_mask=None, labels=None, return_dict=None, **kwargs):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        B, T = input_ids.size()
+        if not self._rotary_initialized:
+            cos, sin = self._precompute_rotary(self.rotary_seq_len, self.config.head_dim)
+            self.cos = cos.to(device=input_ids.device)
+            self.sin = sin.to(device=input_ids.device)
+            self._rotary_initialized = True
+        cos = self.cos[:, :T].to(dtype=self.wte.weight.dtype)
+        sin = self.sin[:, :T].to(dtype=self.wte.weight.dtype)
+        x = self.wte(input_ids)
+        x = _rms_norm(x)
+        if T > 1:
+            gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
+            x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
+        x0 = x
+        n_layer = self.config.n_layer
+        backout_layer = n_layer // 2
+        x_backout = None
+        for i, block in enumerate(self.h):
+            x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
+            ve = self.value_embeds[str(i)](input_ids).to(x.dtype) if str(i) in self.value_embeds else None
+            x = block(x, ve, cos, sin, self.window_sizes[i])
+            if i == backout_layer:
+                x_backout = x
+        if x_backout is not None:
+            x = x - self.backout_lambda.to(x.dtype) * x_backout
+        x = _rms_norm(x)
+        logits = self.lm_head(x)
+        logits = logits[..., :self.config.vocab_size]
+        logits = logits.float()
+        if self.config.logit_softcap > 0:
+            cap = self.config.logit_softcap
+            logits = cap * torch.tanh(logits / cap)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+        if not return_dict:
+            output = (logits,)
+            return (loss,) + output if loss is not None else output
+        return CausalLMOutputWithPast(loss=loss, logits=logits)
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# ---------------------------------------------------------------------------
+# CLI demo
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    import argparse
+    import sys
+    import os
+
+    parser = argparse.ArgumentParser(description="Load and generate with a NanoChat model")
+    parser.add_argument("--model-dir", required=True, help="Path to converted model directory")
+    parser.add_argument("--prompt", default="The meaning of life is", help="Text prompt")
+    parser.add_argument("--max-tokens", type=int, default=100, help="Max tokens to generate")
+    parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature (0=greedy)")
+    parser.add_argument("--device", default="cpu", help="Device (cpu, cuda, cuda:0, ...)")
+    parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16", "float16"])
+    args = parser.parse_args()
+
+    dtype_map = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}
+
+    print(f"Loading model from {args.model_dir}...")
+    from transformers import AutoTokenizer
+
+    model = NanoChatForCausalLM.from_pretrained(args.model_dir, torch_dtype=dtype_map[args.dtype])
+    model.to(args.device)
+    model.eval()
+    tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Model loaded: {total_params:,} parameters, dtype={args.dtype}, device={args.device}")
+    print(f"Vocab size: {model.config.vocab_size}, ve_dim: {model.config.ve_dim}")
+    print()
+    print(f"Prompt: {args.prompt}")
+    print(f"Generating (max {args.max_tokens} tokens, temp={args.temperature})...")
+    print()
+
+    # Generate
+    bos_id = tokenizer.bos_token_id
+    input_ids = tokenizer.encode(args.prompt, add_special_tokens=False)
+    if bos_id is not None:
+        input_ids = [bos_id] + input_ids
+    input_ids = torch.tensor([input_ids], dtype=torch.long, device=args.device)
+
+    generated = []
+    with torch.inference_mode():
+        for _ in range(args.max_tokens):
+            logits = model(input_ids).logits[:, -1, :]
+            if args.temperature > 0:
+                logits = logits / args.temperature
+                v, _ = torch.topk(logits, min(50, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float("inf")
+                probs = torch.softmax(logits, dim=-1)
+                next_id = torch.multinomial(probs, num_samples=1)
+            else:
+                next_id = torch.argmax(logits, dim=-1, keepdim=True)
+            input_ids = torch.cat([input_ids, next_id], dim=1)
+            token_id = next_id.item()
+            generated.append(token_id)
+            if tokenizer.eos_token_id is not None and token_id == tokenizer.eos_token_id:
+                break
+
+    output = tokenizer.decode(generated, skip_special_tokens=True)
+    print(f"Output: {args.prompt}{output}")

models/full-128k-claude-balanced-nfc-bpe-mathcode-scratch/tokenizer_config.json

@@ -0,0 +1,7 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "model_max_length": 2048,
+  "bos_token": "<s>",
+  "add_bos_token": false,
+  "clean_up_tokenization_spaces": false
+}

models/full-128k-clauderev-balanced-bpe-mathcode-scratch/generation_config.json

@@ -0,0 +1,4 @@
+{
+  "bos_token_id": 0,
+  "transformers_version": "5.3.0"
+}

models/full-128k-clauderev-balanced-bpe-mathcode-scratch/nanochat_model.py

@@ -0,0 +1,361 @@
+"""
+NanoChat model for HuggingFace Transformers — self-contained single file.
+
+Includes the config class, model class, and a CLI demo for loading + generation.
+Drop this file into a directory containing config.json, model.safetensors,
+tokenizer.json, and tokenizer_config.json, then run:
+
+    python nanochat_model.py --model-dir . --prompt "Hello world"
+
+Or load programmatically:
+
+    from nanochat_model import NanoChatConfig, NanoChatForCausalLM
+    model = NanoChatForCausalLM.from_pretrained("path/to/model", torch_dtype=torch.bfloat16)
+
+Requirements: pip install torch transformers tokenizers
+
+Architecture: nanochat GPT (Karpathy, MIT licensed) with value embeddings
+(ResFormer-style), per-layer residual/x0 lambdas, smear, backout, logit
+soft-capping, sliding window attention, QK norm, ReLU^2 MLP, and RoPE.
+
+The models were trained with a patched nanochat, not upstream nanochat:
+https://github.com/cimeister/nanochat branch tokenizer-lm-patches, commit f919449.
+The patches touch two files and change three things, two of which affect the forward
+pass and are implemented here:
+  1. ve_dim: the value-embedding tables are built at a reduced dimension and projected
+     up to kv_dim by a learned ve_proj, instead of being built at full kv_dim. All
+     released models use ve_dim=128.
+  2. logit_softcap: the soft cap is a config field with 0 meaning disabled, instead of
+     a hardcoded 15.
+  3. width_lr_exponent: the AdamW learning-rate width scaling exponent is configurable
+     for muP. This affects training only and has no inference-time effect.
+
+This file is the released model definition. Its agreement with the training-time model
+is checked per released checkpoint by scripts/hf_release/verify_conversion.py, which
+compares logits against the patched nanochat GPT in float32.
+"""
+
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+class NanoChatConfig(PretrainedConfig):
+    model_type = "nanochat"
+
+    def __init__(
+        self,
+        vocab_size=128256,
+        n_embd=1536,
+        n_layer=24,
+        n_head=12,
+        n_kv_head=12,
+        sequence_len=2048,
+        window_pattern="SSSL",
+        ve_dim=0,
+        logit_softcap=15.0,
+        pad_vocab_size_to=64,
+        bos_token_id=0,
+        eos_token_id=None,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_kv_head = n_kv_head
+        self.sequence_len = sequence_len
+        self.window_pattern = window_pattern
+        self.ve_dim = ve_dim
+        self.logit_softcap = logit_softcap
+        self.pad_vocab_size_to = pad_vocab_size_to
+        self.head_dim = n_embd // n_head
+        self.kv_dim = n_kv_head * self.head_dim
+        self.padded_vocab_size = (
+            (vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to
+        ) * pad_vocab_size_to
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Model
+# ---------------------------------------------------------------------------
+
+def _has_ve(layer_idx: int, n_layer: int) -> bool:
+    return layer_idx % 2 == (n_layer - 1) % 2
+
+
+def _apply_rotary_emb(x, cos, sin):
+    d = x.shape[-1] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    return torch.cat([y1, y2], dim=-1)
+
+
+def _rms_norm(x):
+    return F.rms_norm(x, (x.size(-1),))
+
+
+class NanoChatAttention(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.n_embd = config.n_embd
+        self.head_dim = config.head_dim
+        self.n_kv_groups = self.n_head // self.n_kv_head
+        self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
+        self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
+        self.ve_gate_channels = 12
+        if _has_ve(layer_idx, config.n_layer):
+            self.ve_gate = nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False)
+            kv_dim = self.n_kv_head * self.head_dim
+            ve_dim = config.ve_dim if config.ve_dim > 0 else kv_dim
+            self.ve_proj = nn.Linear(ve_dim, kv_dim, bias=False) if ve_dim < kv_dim else None
+        else:
+            self.ve_gate = None
+            self.ve_proj = None
+
+    def forward(self, x, ve, cos, sin, window_size):
+        B, T, C = x.size()
+        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
+        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
+        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
+        if ve is not None and self.ve_gate is not None:
+            if self.ve_proj is not None:
+                ve = self.ve_proj(ve)
+            ve = ve.view(B, T, self.n_kv_head, self.head_dim)
+            gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
+            v = v + gate.unsqueeze(-1) * ve
+        q = _apply_rotary_emb(q, cos, sin)
+        k = _apply_rotary_emb(k, cos, sin)
+        q = _rms_norm(q) * 1.2
+        k = _rms_norm(k) * 1.2
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        if self.n_kv_groups > 1:
+            k = k.repeat_interleave(self.n_kv_groups, dim=1)
+            v = v.repeat_interleave(self.n_kv_groups, dim=1)
+        if window_size < 0 or window_size >= T:
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:
+            row_idx = torch.arange(T, device=x.device).unsqueeze(1)
+            col_idx = torch.arange(T, device=x.device).unsqueeze(0)
+            mask = (col_idx <= row_idx) & ((row_idx - col_idx) <= window_size)
+            y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        y = y.transpose(1, 2).contiguous().view(B, T, -1)
+        y = self.c_proj(y)
+        return y
+
+
+class NanoChatMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = F.relu(x).square()
+        x = self.c_proj(x)
+        return x
+
+
+class NanoChatBlock(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.attn = NanoChatAttention(config, layer_idx)
+        self.mlp = NanoChatMLP(config)
+
+    def forward(self, x, ve, cos, sin, window_size):
+        x = x + self.attn(_rms_norm(x), ve, cos, sin, window_size)
+        x = x + self.mlp(_rms_norm(x))
+        return x
+
+
+class NanoChatPreTrainedModel(PreTrainedModel):
+    config_class = NanoChatConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["NanoChatBlock"]
+
+    def _init_weights(self, module):
+        pass
+
+
+class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin):
+    def __init__(self, config):
+        super().__init__(config)
+        pvs = config.padded_vocab_size
+        self.wte = nn.Embedding(pvs, config.n_embd)
+        self.h = nn.ModuleList([NanoChatBlock(config, i) for i in range(config.n_layer)])
+        self.lm_head = nn.Linear(config.n_embd, pvs, bias=False)
+        self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))
+        self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))
+        self.smear_gate = nn.Linear(24, 1, bias=False)
+        self.smear_lambda = nn.Parameter(torch.zeros(1))
+        self.backout_lambda = nn.Parameter(0.2 * torch.ones(1))
+        ve_emb_dim = config.ve_dim if config.ve_dim > 0 else config.kv_dim
+        self.value_embeds = nn.ModuleDict({
+            str(i): nn.Embedding(pvs, ve_emb_dim)
+            for i in range(config.n_layer) if _has_ve(i, config.n_layer)
+        })
+        self.window_sizes = self._compute_window_sizes(config)
+        self.rotary_seq_len = config.sequence_len * 10
+        self._rotary_initialized = False
+        self.register_buffer("cos", torch.zeros(1), persistent=False)
+        self.register_buffer("sin", torch.zeros(1), persistent=False)
+        self.post_init()
+
+    def _compute_window_sizes(self, config):
+        pattern = config.window_pattern.upper()
+        long_window = config.sequence_len
+        short_window = -(-long_window // 4 // 128) * 128
+        char_to_window = {"L": long_window, "S": short_window}
+        sizes = [char_to_window[pattern[i % len(pattern)]] for i in range(config.n_layer)]
+        sizes[-1] = long_window
+        return sizes
+
+    def _precompute_rotary(self, seq_len, head_dim, base=100000.0):
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32)
+        inv_freq = 1.0 / (base ** (channel_range / head_dim))
+        t = torch.arange(seq_len, dtype=torch.float32)
+        freqs = torch.outer(t, inv_freq)
+        cos = freqs.cos()[None, :, None, :]
+        sin = freqs.sin()[None, :, None, :]
+        return cos, sin
+
+    def forward(self, input_ids, attention_mask=None, labels=None, return_dict=None, **kwargs):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        B, T = input_ids.size()
+        if not self._rotary_initialized:
+            cos, sin = self._precompute_rotary(self.rotary_seq_len, self.config.head_dim)
+            self.cos = cos.to(device=input_ids.device)
+            self.sin = sin.to(device=input_ids.device)
+            self._rotary_initialized = True
+        cos = self.cos[:, :T].to(dtype=self.wte.weight.dtype)
+        sin = self.sin[:, :T].to(dtype=self.wte.weight.dtype)
+        x = self.wte(input_ids)
+        x = _rms_norm(x)
+        if T > 1:
+            gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
+            x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
+        x0 = x
+        n_layer = self.config.n_layer
+        backout_layer = n_layer // 2
+        x_backout = None
+        for i, block in enumerate(self.h):
+            x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
+            ve = self.value_embeds[str(i)](input_ids).to(x.dtype) if str(i) in self.value_embeds else None
+            x = block(x, ve, cos, sin, self.window_sizes[i])
+            if i == backout_layer:
+                x_backout = x
+        if x_backout is not None:
+            x = x - self.backout_lambda.to(x.dtype) * x_backout
+        x = _rms_norm(x)
+        logits = self.lm_head(x)
+        logits = logits[..., :self.config.vocab_size]
+        logits = logits.float()
+        if self.config.logit_softcap > 0:
+            cap = self.config.logit_softcap
+            logits = cap * torch.tanh(logits / cap)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+        if not return_dict:
+            output = (logits,)
+            return (loss,) + output if loss is not None else output
+        return CausalLMOutputWithPast(loss=loss, logits=logits)
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# ---------------------------------------------------------------------------
+# CLI demo
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    import argparse
+    import sys
+    import os
+
+    parser = argparse.ArgumentParser(description="Load and generate with a NanoChat model")
+    parser.add_argument("--model-dir", required=True, help="Path to converted model directory")
+    parser.add_argument("--prompt", default="The meaning of life is", help="Text prompt")
+    parser.add_argument("--max-tokens", type=int, default=100, help="Max tokens to generate")
+    parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature (0=greedy)")
+    parser.add_argument("--device", default="cpu", help="Device (cpu, cuda, cuda:0, ...)")
+    parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16", "float16"])
+    args = parser.parse_args()
+
+    dtype_map = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}
+
+    print(f"Loading model from {args.model_dir}...")
+    from transformers import AutoTokenizer
+
+    model = NanoChatForCausalLM.from_pretrained(args.model_dir, torch_dtype=dtype_map[args.dtype])
+    model.to(args.device)
+    model.eval()
+    tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Model loaded: {total_params:,} parameters, dtype={args.dtype}, device={args.device}")
+    print(f"Vocab size: {model.config.vocab_size}, ve_dim: {model.config.ve_dim}")
+    print()
+    print(f"Prompt: {args.prompt}")
+    print(f"Generating (max {args.max_tokens} tokens, temp={args.temperature})...")
+    print()
+
+    # Generate
+    bos_id = tokenizer.bos_token_id
+    input_ids = tokenizer.encode(args.prompt, add_special_tokens=False)
+    if bos_id is not None:
+        input_ids = [bos_id] + input_ids
+    input_ids = torch.tensor([input_ids], dtype=torch.long, device=args.device)
+
+    generated = []
+    with torch.inference_mode():
+        for _ in range(args.max_tokens):
+            logits = model(input_ids).logits[:, -1, :]
+            if args.temperature > 0:
+                logits = logits / args.temperature
+                v, _ = torch.topk(logits, min(50, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float("inf")
+                probs = torch.softmax(logits, dim=-1)
+                next_id = torch.multinomial(probs, num_samples=1)
+            else:
+                next_id = torch.argmax(logits, dim=-1, keepdim=True)
+            input_ids = torch.cat([input_ids, next_id], dim=1)
+            token_id = next_id.item()
+            generated.append(token_id)
+            if tokenizer.eos_token_id is not None and token_id == tokenizer.eos_token_id:
+                break
+
+    output = tokenizer.decode(generated, skip_special_tokens=True)
+    print(f"Output: {args.prompt}{output}")

models/full-128k-gpt4o-balanced-bpe-seed43-mathcode-scratch/conversion_metadata.json

@@ -0,0 +1,100 @@
+{
+  "run": "full-128k-gpt4o-balanced-bpe-seed43-mathcode-scratch",
+  "source_checkpoint": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-gpt4o-balanced-bpe-seed43-mathcode-scratch",
+  "source_model_file": "model_019073.pt",
+  "source_step": 19073,
+  "source_meta": {
+    "step": 19073,
+    "val_bpb": 0.3362506080752585,
+    "model_config": {
+      "sequence_len": 2048,
+      "vocab_size": 128256,
+      "n_layer": 24,
+      "n_head": 12,
+      "n_kv_head": 12,
+      "n_embd": 1536,
+      "window_pattern": "SSSL",
+      "ve_dim": 128,
+      "logit_softcap": 15.0
+    },
+    "tokenizer_path": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/tokenizers/gpt4o-balanced-bpe",
+    "training_config": {
+      "name": "full-128k-mathcode-scratch",
+      "model": {
+        "n_layer": 24,
+        "n_embd": 1536,
+        "n_kv_head": 12,
+        "sequence_len": 2048,
+        "window_pattern": "SSSL",
+        "ve_dim": 128
+      },
+      "training": {
+        "target_param_data_ratio": -1,
+        "max_tokens": 20000000000,
+        "num_iterations": -1,
+        "total_batch_size": 1048576,
+        "device_batch_size": 16,
+        "warmup_steps": 40,
+        "warmdown_ratio": 0.65,
+        "final_lr_frac": 0.05,
+        "eval_tokens": 41943040,
+        "compile": true
+      },
+      "optimizer": {
+        "matrix_lr": 0.02,
+        "embedding_lr": 0.3,
+        "unembedding_lr": 0.008,
+        "scalar_lr": 0.5,
+        "weight_decay": 0.28,
+        "width_lr_exponent": -1.0
+      },
+      "data": {
+        "data_dir": null
+      },
+      "logging": {
+        "eval_every": 500,
+        "save_every": 500,
+        "max_checkpoints": 0
+      },
+      "wandb": {
+        "enabled": true,
+        "project": "tokenizer-lm-experiments",
+        "run_name": null
+      },
+      "output_dir": "/capstor/scratch/cscs/${USER}/tokenizer-lm/checkpoints",
+      "seed": 42
+    },
+    "total_batch_size": 1048576,
+    "device_batch_size": 16,
+    "max_seq_len": 2048,
+    "dataloader_state_dict": {
+      "pq_idx": 262,
+      "rg_idx": 40,
+      "epoch": 2
+    },
+    "loop_state": {
+      "min_val_bpb": 0.3362506080752585,
+      "smooth_train_loss": 0.9279650649193887,
+      "total_training_time": 93685.4534380436,
+      "total_bytes_consumed": 147636253997
+    }
+  },
+  "weight_dtype": "as trained (float32 matrices, bfloat16 embeddings)",
+  "weight_dtype_counts": {
+    "torch.float32": 174,
+    "torch.bfloat16": 13
+  },
+  "total_parameters": 1272841994,
+  "logit_softcap_source": "meta",
+  "bos_token": "<s>",
+  "bos_token_id": 0,
+  "bos_resolved_from": "packaged tokenizer",
+  "tokenizer_slug": "gpt4o-balanced-bpe",
+  "tokenizer_redistributed": true,
+  "tokenizer_hub_id": null,
+  "weight_key_mapping": "strip _orig_mod. and transformer. prefixes",
+  "converted_by": "scripts/hf_release/convert_for_release.py",
+  "repo_commit": "1a9a718",
+  "nanochat_commit": "f919449",
+  "nanochat_branch": "tokenizer-lm-patches"
+}

models/full-128k-gpt4o-balanced-bpe-seed43-mathcode-scratch/generation_config.json

@@ -0,0 +1,4 @@
+{
+  "bos_token_id": 0,
+  "transformers_version": "5.3.0"
+}

models/full-128k-gpt4o-balanced-bpe-seed43-mathcode-scratch/nanochat_model.py

@@ -0,0 +1,361 @@
+"""
+NanoChat model for HuggingFace Transformers — self-contained single file.
+
+Includes the config class, model class, and a CLI demo for loading + generation.
+Drop this file into a directory containing config.json, model.safetensors,
+tokenizer.json, and tokenizer_config.json, then run:
+
+    python nanochat_model.py --model-dir . --prompt "Hello world"
+
+Or load programmatically:
+
+    from nanochat_model import NanoChatConfig, NanoChatForCausalLM
+    model = NanoChatForCausalLM.from_pretrained("path/to/model", torch_dtype=torch.bfloat16)
+
+Requirements: pip install torch transformers tokenizers
+
+Architecture: nanochat GPT (Karpathy, MIT licensed) with value embeddings
+(ResFormer-style), per-layer residual/x0 lambdas, smear, backout, logit
+soft-capping, sliding window attention, QK norm, ReLU^2 MLP, and RoPE.
+
+The models were trained with a patched nanochat, not upstream nanochat:
+https://github.com/cimeister/nanochat branch tokenizer-lm-patches, commit f919449.
+The patches touch two files and change three things, two of which affect the forward
+pass and are implemented here:
+  1. ve_dim: the value-embedding tables are built at a reduced dimension and projected
+     up to kv_dim by a learned ve_proj, instead of being built at full kv_dim. All
+     released models use ve_dim=128.
+  2. logit_softcap: the soft cap is a config field with 0 meaning disabled, instead of
+     a hardcoded 15.
+  3. width_lr_exponent: the AdamW learning-rate width scaling exponent is configurable
+     for muP. This affects training only and has no inference-time effect.
+
+This file is the released model definition. Its agreement with the training-time model
+is checked per released checkpoint by scripts/hf_release/verify_conversion.py, which
+compares logits against the patched nanochat GPT in float32.
+"""
+
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+class NanoChatConfig(PretrainedConfig):
+    model_type = "nanochat"
+
+    def __init__(
+        self,
+        vocab_size=128256,
+        n_embd=1536,
+        n_layer=24,
+        n_head=12,
+        n_kv_head=12,
+        sequence_len=2048,
+        window_pattern="SSSL",
+        ve_dim=0,
+        logit_softcap=15.0,
+        pad_vocab_size_to=64,
+        bos_token_id=0,
+        eos_token_id=None,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_kv_head = n_kv_head
+        self.sequence_len = sequence_len
+        self.window_pattern = window_pattern
+        self.ve_dim = ve_dim
+        self.logit_softcap = logit_softcap
+        self.pad_vocab_size_to = pad_vocab_size_to
+        self.head_dim = n_embd // n_head
+        self.kv_dim = n_kv_head * self.head_dim
+        self.padded_vocab_size = (
+            (vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to
+        ) * pad_vocab_size_to
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Model
+# ---------------------------------------------------------------------------
+
+def _has_ve(layer_idx: int, n_layer: int) -> bool:
+    return layer_idx % 2 == (n_layer - 1) % 2
+
+
+def _apply_rotary_emb(x, cos, sin):
+    d = x.shape[-1] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    return torch.cat([y1, y2], dim=-1)
+
+
+def _rms_norm(x):
+    return F.rms_norm(x, (x.size(-1),))
+
+
+class NanoChatAttention(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.n_embd = config.n_embd
+        self.head_dim = config.head_dim
+        self.n_kv_groups = self.n_head // self.n_kv_head
+        self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
+        self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
+        self.ve_gate_channels = 12
+        if _has_ve(layer_idx, config.n_layer):
+            self.ve_gate = nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False)
+            kv_dim = self.n_kv_head * self.head_dim
+            ve_dim = config.ve_dim if config.ve_dim > 0 else kv_dim
+            self.ve_proj = nn.Linear(ve_dim, kv_dim, bias=False) if ve_dim < kv_dim else None
+        else:
+            self.ve_gate = None
+            self.ve_proj = None
+
+    def forward(self, x, ve, cos, sin, window_size):
+        B, T, C = x.size()
+        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
+        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
+        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
+        if ve is not None and self.ve_gate is not None:
+            if self.ve_proj is not None:
+                ve = self.ve_proj(ve)
+            ve = ve.view(B, T, self.n_kv_head, self.head_dim)
+            gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
+            v = v + gate.unsqueeze(-1) * ve
+        q = _apply_rotary_emb(q, cos, sin)
+        k = _apply_rotary_emb(k, cos, sin)
+        q = _rms_norm(q) * 1.2
+        k = _rms_norm(k) * 1.2
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        if self.n_kv_groups > 1:
+            k = k.repeat_interleave(self.n_kv_groups, dim=1)
+            v = v.repeat_interleave(self.n_kv_groups, dim=1)
+        if window_size < 0 or window_size >= T:
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:
+            row_idx = torch.arange(T, device=x.device).unsqueeze(1)
+            col_idx = torch.arange(T, device=x.device).unsqueeze(0)
+            mask = (col_idx <= row_idx) & ((row_idx - col_idx) <= window_size)
+            y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        y = y.transpose(1, 2).contiguous().view(B, T, -1)
+        y = self.c_proj(y)
+        return y
+
+
+class NanoChatMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = F.relu(x).square()
+        x = self.c_proj(x)
+        return x
+
+
+class NanoChatBlock(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.attn = NanoChatAttention(config, layer_idx)
+        self.mlp = NanoChatMLP(config)
+
+    def forward(self, x, ve, cos, sin, window_size):
+        x = x + self.attn(_rms_norm(x), ve, cos, sin, window_size)
+        x = x + self.mlp(_rms_norm(x))
+        return x
+
+
+class NanoChatPreTrainedModel(PreTrainedModel):
+    config_class = NanoChatConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["NanoChatBlock"]
+
+    def _init_weights(self, module):
+        pass
+
+
+class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin):
+    def __init__(self, config):
+        super().__init__(config)
+        pvs = config.padded_vocab_size
+        self.wte = nn.Embedding(pvs, config.n_embd)
+        self.h = nn.ModuleList([NanoChatBlock(config, i) for i in range(config.n_layer)])
+        self.lm_head = nn.Linear(config.n_embd, pvs, bias=False)
+        self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))
+        self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))
+        self.smear_gate = nn.Linear(24, 1, bias=False)
+        self.smear_lambda = nn.Parameter(torch.zeros(1))
+        self.backout_lambda = nn.Parameter(0.2 * torch.ones(1))
+        ve_emb_dim = config.ve_dim if config.ve_dim > 0 else config.kv_dim
+        self.value_embeds = nn.ModuleDict({
+            str(i): nn.Embedding(pvs, ve_emb_dim)
+            for i in range(config.n_layer) if _has_ve(i, config.n_layer)
+        })
+        self.window_sizes = self._compute_window_sizes(config)
+        self.rotary_seq_len = config.sequence_len * 10
+        self._rotary_initialized = False
+        self.register_buffer("cos", torch.zeros(1), persistent=False)
+        self.register_buffer("sin", torch.zeros(1), persistent=False)
+        self.post_init()
+
+    def _compute_window_sizes(self, config):
+        pattern = config.window_pattern.upper()
+        long_window = config.sequence_len
+        short_window = -(-long_window // 4 // 128) * 128
+        char_to_window = {"L": long_window, "S": short_window}
+        sizes = [char_to_window[pattern[i % len(pattern)]] for i in range(config.n_layer)]
+        sizes[-1] = long_window
+        return sizes
+
+    def _precompute_rotary(self, seq_len, head_dim, base=100000.0):
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32)
+        inv_freq = 1.0 / (base ** (channel_range / head_dim))
+        t = torch.arange(seq_len, dtype=torch.float32)
+        freqs = torch.outer(t, inv_freq)
+        cos = freqs.cos()[None, :, None, :]
+        sin = freqs.sin()[None, :, None, :]
+        return cos, sin
+
+    def forward(self, input_ids, attention_mask=None, labels=None, return_dict=None, **kwargs):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        B, T = input_ids.size()
+        if not self._rotary_initialized:
+            cos, sin = self._precompute_rotary(self.rotary_seq_len, self.config.head_dim)
+            self.cos = cos.to(device=input_ids.device)
+            self.sin = sin.to(device=input_ids.device)
+            self._rotary_initialized = True
+        cos = self.cos[:, :T].to(dtype=self.wte.weight.dtype)
+        sin = self.sin[:, :T].to(dtype=self.wte.weight.dtype)
+        x = self.wte(input_ids)
+        x = _rms_norm(x)
+        if T > 1:
+            gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
+            x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
+        x0 = x
+        n_layer = self.config.n_layer
+        backout_layer = n_layer // 2
+        x_backout = None
+        for i, block in enumerate(self.h):
+            x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
+            ve = self.value_embeds[str(i)](input_ids).to(x.dtype) if str(i) in self.value_embeds else None
+            x = block(x, ve, cos, sin, self.window_sizes[i])
+            if i == backout_layer:
+                x_backout = x
+        if x_backout is not None:
+            x = x - self.backout_lambda.to(x.dtype) * x_backout
+        x = _rms_norm(x)
+        logits = self.lm_head(x)
+        logits = logits[..., :self.config.vocab_size]
+        logits = logits.float()
+        if self.config.logit_softcap > 0:
+            cap = self.config.logit_softcap
+            logits = cap * torch.tanh(logits / cap)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+        if not return_dict:
+            output = (logits,)
+            return (loss,) + output if loss is not None else output
+        return CausalLMOutputWithPast(loss=loss, logits=logits)
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# ---------------------------------------------------------------------------
+# CLI demo
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    import argparse
+    import sys
+    import os
+
+    parser = argparse.ArgumentParser(description="Load and generate with a NanoChat model")
+    parser.add_argument("--model-dir", required=True, help="Path to converted model directory")
+    parser.add_argument("--prompt", default="The meaning of life is", help="Text prompt")
+    parser.add_argument("--max-tokens", type=int, default=100, help="Max tokens to generate")
+    parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature (0=greedy)")
+    parser.add_argument("--device", default="cpu", help="Device (cpu, cuda, cuda:0, ...)")
+    parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16", "float16"])
+    args = parser.parse_args()
+
+    dtype_map = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}
+
+    print(f"Loading model from {args.model_dir}...")
+    from transformers import AutoTokenizer
+
+    model = NanoChatForCausalLM.from_pretrained(args.model_dir, torch_dtype=dtype_map[args.dtype])
+    model.to(args.device)
+    model.eval()
+    tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Model loaded: {total_params:,} parameters, dtype={args.dtype}, device={args.device}")
+    print(f"Vocab size: {model.config.vocab_size}, ve_dim: {model.config.ve_dim}")
+    print()
+    print(f"Prompt: {args.prompt}")
+    print(f"Generating (max {args.max_tokens} tokens, temp={args.temperature})...")
+    print()
+
+    # Generate
+    bos_id = tokenizer.bos_token_id
+    input_ids = tokenizer.encode(args.prompt, add_special_tokens=False)
+    if bos_id is not None:
+        input_ids = [bos_id] + input_ids
+    input_ids = torch.tensor([input_ids], dtype=torch.long, device=args.device)
+
+    generated = []
+    with torch.inference_mode():
+        for _ in range(args.max_tokens):
+            logits = model(input_ids).logits[:, -1, :]
+            if args.temperature > 0:
+                logits = logits / args.temperature
+                v, _ = torch.topk(logits, min(50, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float("inf")
+                probs = torch.softmax(logits, dim=-1)
+                next_id = torch.multinomial(probs, num_samples=1)
+            else:
+                next_id = torch.argmax(logits, dim=-1, keepdim=True)
+            input_ids = torch.cat([input_ids, next_id], dim=1)
+            token_id = next_id.item()
+            generated.append(token_id)
+            if tokenizer.eos_token_id is not None and token_id == tokenizer.eos_token_id:
+                break
+
+    output = tokenizer.decode(generated, skip_special_tokens=True)
+    print(f"Output: {args.prompt}{output}")

models/full-128k-gpt4o-balanced-bpe-seed44-mathcode-scratch/generation_config.json

@@ -0,0 +1,4 @@
+{
+  "bos_token_id": 0,
+  "transformers_version": "5.3.0"
+}

models/full-128k-gpt4o-balanced-unigram-mathcode-scratch/config.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "NanoChatForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "nanochat_model.NanoChatConfig",
+    "AutoModelForCausalLM": "nanochat_model.NanoChatForCausalLM"
+  },
+  "bos_token_id": 0,
+  "dtype": "float32",
+  "eos_token_id": null,
+  "head_dim": 128,
+  "kv_dim": 1536,
+  "logit_softcap": 15.0,
+  "model_type": "nanochat",
+  "n_embd": 1536,
+  "n_head": 12,
+  "n_kv_head": 12,
+  "n_layer": 24,
+  "pad_vocab_size_to": 64,
+  "padded_vocab_size": 128256,
+  "sequence_len": 2048,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.3.0",
+  "ve_dim": 128,
+  "vocab_size": 128256,
+  "window_pattern": "SSSL"
+}

models/full-128k-gpt4o-code-bpe-mathcode-scratch/README.md

@@ -0,0 +1,76 @@
+# `full-128k-gpt4o-code-bpe-mathcode-scratch`
+
+A 1.27B parameter model trained with the `gpt4o-code-bpe` tokenizer. Regime: mathcode-20B-from-scratch (24 layers, width 1536, 1.27B parameters, 20.0B tokens, math and code mixture, trained from scratch). Reported in: math and code ablations; TokEval.
+
+## Load
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+    "cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-gpt4o-code-bpe-mathcode-scratch",
+    trust_remote_code=True, dtype="float32")
+tokenizer = AutoTokenizer.from_pretrained("cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-gpt4o-code-bpe-mathcode-scratch")
+
+text = "def add(a, b):\n    return a + b\n"
+ids = [0] + tokenizer(text)["input_ids"]   # training prepends BOS
+```
+
+## Settings
+
+| Setting | Value |
+|---|---|
+| layers | 24 |
+| model dimension | 1536 |
+| attention heads | 12 |
+| key and value heads | 12 |
+| context length | 2048 |
+| sliding window pattern | SSSL |
+| value embedding dimension | 128 |
+| logit soft cap | 15 |
+| vocabulary size | 128256 |
+| parameters | 1,272,841,994 |
+| training step | 19073 |
+| token budget | 19,999,490,048 |
+| bytes consumed | 129,315,874,577 |
+| seed | 42 |
+| Muon matrix learning rate | 0.0200 |
+| AdamW embedding learning rate | 0.3000 |
+| AdamW unembedding learning rate | 0.0080 |
+| weight decay | 0.2800 |
+| width learning-rate exponent | -1 |
+| weight dtype | as trained (float32 matrices, bfloat16 embeddings) |
+| BOS token id | 0 |
+
+## Results
+
+| Metric | Value |
+|---|---|
+| validation BPB | 0.3326 |
+| code BPB, mean over 7 languages | 0.3805 |
+| FLORES-200 BPB, mean | 3.1524 |
+| FLORES-200 languages scored | 214 |
+| BLiMP accuracy | n/a |
+| MBPP pass@1 | 0.0160 |
+| HumanEval pass@1 | 0.1585 |
+| GSM8K exact match, strict | 0.2161 |
+
+Generation metrics use specification `{'version': 'v2-2026-07-30', 'bos_context': True, 'eos_stop': True, 'token_healing': True, 'truncation': 'left_to_max_length_minus_max_gen', 'stop_truncation': 'earliest_match_all_terms', 'stop_match_decode': 'specials_kept'}`, token healing on.
+
+Missing measurements, with the reason each is absent:
+
+- `blimp_acc`: no full-128k-gpt4o-code-bpe-mathcode-scratch_blimp_code_bpb.json
+
+BLiMP is comparable within a training scale only. Its scoring sums unnormalized sentence log-likelihoods, and under that scoring a group of quantifier subtasks drops at the 1.27B scale relative to the 300M scale.
+
+## Registry note
+
+#30, gpt4o pretok, slug `gpt4o-code-bpe`. GPT-4o regex, code-heavy tokenizer-training data, vocab 128,256. FINISHED (step 19073). MBPP 0.016 [0.006, 0.028].
+
+## Provenance
+
+- Source checkpoint: `/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-gpt4o-code-bpe-mathcode-scratch`, step 19073
+- Repository commit: `1a9a718`
+- nanochat commit: `f919449` on branch `tokenizer-lm-patches`
+- Tokenizer: `gpt4o-code-bpe`, released at `tokenizers/gpt4o-code-bpe`
+

models/full-128k-gpt4o-code-bpe-mathcode-scratch/config.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "NanoChatForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "nanochat_model.NanoChatConfig",
+    "AutoModelForCausalLM": "nanochat_model.NanoChatForCausalLM"
+  },
+  "bos_token_id": 0,
+  "dtype": "float32",
+  "eos_token_id": null,
+  "head_dim": 128,
+  "kv_dim": 1536,
+  "logit_softcap": 15.0,
+  "model_type": "nanochat",
+  "n_embd": 1536,
+  "n_head": 12,
+  "n_kv_head": 12,
+  "n_layer": 24,
+  "pad_vocab_size_to": 64,
+  "padded_vocab_size": 128256,
+  "sequence_len": 2048,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.3.0",
+  "ve_dim": 128,
+  "vocab_size": 128256,
+  "window_pattern": "SSSL"
+}

models/full-128k-gpt4o-code-bpe-mathcode-scratch/conversion_metadata.json

@@ -0,0 +1,100 @@
+{
+  "run": "full-128k-gpt4o-code-bpe-mathcode-scratch",
+  "source_checkpoint": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-gpt4o-code-bpe-mathcode-scratch",
+  "source_model_file": "model_019073.pt",
+  "source_step": 19073,
+  "source_meta": {
+    "step": 19073,
+    "val_bpb": 0.33264187797540906,
+    "model_config": {
+      "sequence_len": 2048,
+      "vocab_size": 128256,
+      "n_layer": 24,
+      "n_head": 12,
+      "n_kv_head": 12,
+      "n_embd": 1536,
+      "window_pattern": "SSSL",
+      "ve_dim": 128,
+      "logit_softcap": 15.0
+    },
+    "tokenizer_path": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/tokenizers/gpt4o-code-bpe",
+    "training_config": {
+      "name": "full-128k-mathcode-scratch",
+      "model": {
+        "n_layer": 24,
+        "n_embd": 1536,
+        "n_kv_head": 12,
+        "sequence_len": 2048,
+        "window_pattern": "SSSL",
+        "ve_dim": 128
+      },
+      "training": {
+        "target_param_data_ratio": -1,
+        "max_tokens": 20000000000,
+        "num_iterations": -1,
+        "total_batch_size": 1048576,
+        "device_batch_size": 16,
+        "warmup_steps": 40,
+        "warmdown_ratio": 0.65,
+        "final_lr_frac": 0.05,
+        "eval_tokens": 41943040,
+        "compile": true
+      },
+      "optimizer": {
+        "matrix_lr": 0.02,
+        "embedding_lr": 0.3,
+        "unembedding_lr": 0.008,
+        "scalar_lr": 0.5,
+        "weight_decay": 0.28,
+        "width_lr_exponent": -1.0
+      },
+      "data": {
+        "data_dir": null
+      },
+      "logging": {
+        "eval_every": 500,
+        "save_every": 500,
+        "max_checkpoints": 0
+      },
+      "wandb": {
+        "enabled": true,
+        "project": "tokenizer-lm-experiments",
+        "run_name": null
+      },
+      "output_dir": "/capstor/scratch/cscs/${USER}/tokenizer-lm/checkpoints",
+      "seed": 42
+    },
+    "total_batch_size": 1048576,
+    "device_batch_size": 16,
+    "max_seq_len": 2048,
+    "dataloader_state_dict": {
+      "pq_idx": 274,
+      "rg_idx": 0,
+      "epoch": 2
+    },
+    "loop_state": {
+      "min_val_bpb": 0.33264187797540906,
+      "smooth_train_loss": 1.0051792212437496,
+      "total_training_time": 93959.44238901138,
+      "total_bytes_consumed": 129315874577
+    }
+  },
+  "weight_dtype": "as trained (float32 matrices, bfloat16 embeddings)",
+  "weight_dtype_counts": {
+    "torch.float32": 174,
+    "torch.bfloat16": 13
+  },
+  "total_parameters": 1272841994,
+  "logit_softcap_source": "meta",
+  "bos_token": "<s>",
+  "bos_token_id": 0,
+  "bos_resolved_from": "packaged tokenizer",
+  "tokenizer_slug": "gpt4o-code-bpe",
+  "tokenizer_redistributed": true,
+  "tokenizer_hub_id": null,
+  "weight_key_mapping": "strip _orig_mod. and transformer. prefixes",
+  "converted_by": "scripts/hf_release/convert_for_release.py",
+  "repo_commit": "1a9a718",
+  "nanochat_commit": "f919449",
+  "nanochat_branch": "tokenizer-lm-patches"
+}

models/full-128k-gpt4o-code-bpe-mathcode-scratch/generation_config.json

@@ -0,0 +1,4 @@
+{
+  "bos_token_id": 0,
+  "transformers_version": "5.3.0"
+}

models/full-128k-gpt4o-code-bpe-mathcode-scratch/nanochat_model.py

@@ -0,0 +1,361 @@
+"""
+NanoChat model for HuggingFace Transformers — self-contained single file.
+
+Includes the config class, model class, and a CLI demo for loading + generation.
+Drop this file into a directory containing config.json, model.safetensors,
+tokenizer.json, and tokenizer_config.json, then run:
+
+    python nanochat_model.py --model-dir . --prompt "Hello world"
+
+Or load programmatically:
+
+    from nanochat_model import NanoChatConfig, NanoChatForCausalLM
+    model = NanoChatForCausalLM.from_pretrained("path/to/model", torch_dtype=torch.bfloat16)
+
+Requirements: pip install torch transformers tokenizers
+
+Architecture: nanochat GPT (Karpathy, MIT licensed) with value embeddings
+(ResFormer-style), per-layer residual/x0 lambdas, smear, backout, logit
+soft-capping, sliding window attention, QK norm, ReLU^2 MLP, and RoPE.
+
+The models were trained with a patched nanochat, not upstream nanochat:
+https://github.com/cimeister/nanochat branch tokenizer-lm-patches, commit f919449.
+The patches touch two files and change three things, two of which affect the forward
+pass and are implemented here:
+  1. ve_dim: the value-embedding tables are built at a reduced dimension and projected
+     up to kv_dim by a learned ve_proj, instead of being built at full kv_dim. All
+     released models use ve_dim=128.
+  2. logit_softcap: the soft cap is a config field with 0 meaning disabled, instead of
+     a hardcoded 15.
+  3. width_lr_exponent: the AdamW learning-rate width scaling exponent is configurable
+     for muP. This affects training only and has no inference-time effect.
+
+This file is the released model definition. Its agreement with the training-time model
+is checked per released checkpoint by scripts/hf_release/verify_conversion.py, which
+compares logits against the patched nanochat GPT in float32.
+"""
+
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+class NanoChatConfig(PretrainedConfig):
+    model_type = "nanochat"
+
+    def __init__(
+        self,
+        vocab_size=128256,
+        n_embd=1536,
+        n_layer=24,
+        n_head=12,
+        n_kv_head=12,
+        sequence_len=2048,
+        window_pattern="SSSL",
+        ve_dim=0,
+        logit_softcap=15.0,
+        pad_vocab_size_to=64,
+        bos_token_id=0,
+        eos_token_id=None,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_kv_head = n_kv_head
+        self.sequence_len = sequence_len
+        self.window_pattern = window_pattern
+        self.ve_dim = ve_dim
+        self.logit_softcap = logit_softcap
+        self.pad_vocab_size_to = pad_vocab_size_to
+        self.head_dim = n_embd // n_head
+        self.kv_dim = n_kv_head * self.head_dim
+        self.padded_vocab_size = (
+            (vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to
+        ) * pad_vocab_size_to
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Model
+# ---------------------------------------------------------------------------
+
+def _has_ve(layer_idx: int, n_layer: int) -> bool:
+    return layer_idx % 2 == (n_layer - 1) % 2
+
+
+def _apply_rotary_emb(x, cos, sin):
+    d = x.shape[-1] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    return torch.cat([y1, y2], dim=-1)
+
+
+def _rms_norm(x):
+    return F.rms_norm(x, (x.size(-1),))
+
+
+class NanoChatAttention(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.n_embd = config.n_embd
+        self.head_dim = config.head_dim
+        self.n_kv_groups = self.n_head // self.n_kv_head
+        self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
+        self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
+        self.ve_gate_channels = 12
+        if _has_ve(layer_idx, config.n_layer):
+            self.ve_gate = nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False)
+            kv_dim = self.n_kv_head * self.head_dim
+            ve_dim = config.ve_dim if config.ve_dim > 0 else kv_dim
+            self.ve_proj = nn.Linear(ve_dim, kv_dim, bias=False) if ve_dim < kv_dim else None
+        else:
+            self.ve_gate = None
+            self.ve_proj = None
+
+    def forward(self, x, ve, cos, sin, window_size):
+        B, T, C = x.size()
+        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
+        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
+        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
+        if ve is not None and self.ve_gate is not None:
+            if self.ve_proj is not None:
+                ve = self.ve_proj(ve)
+            ve = ve.view(B, T, self.n_kv_head, self.head_dim)
+            gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
+            v = v + gate.unsqueeze(-1) * ve
+        q = _apply_rotary_emb(q, cos, sin)
+        k = _apply_rotary_emb(k, cos, sin)
+        q = _rms_norm(q) * 1.2
+        k = _rms_norm(k) * 1.2
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        if self.n_kv_groups > 1:
+            k = k.repeat_interleave(self.n_kv_groups, dim=1)
+            v = v.repeat_interleave(self.n_kv_groups, dim=1)
+        if window_size < 0 or window_size >= T:
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:
+            row_idx = torch.arange(T, device=x.device).unsqueeze(1)
+            col_idx = torch.arange(T, device=x.device).unsqueeze(0)
+            mask = (col_idx <= row_idx) & ((row_idx - col_idx) <= window_size)
+            y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        y = y.transpose(1, 2).contiguous().view(B, T, -1)
+        y = self.c_proj(y)
+        return y
+
+
+class NanoChatMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = F.relu(x).square()
+        x = self.c_proj(x)
+        return x
+
+
+class NanoChatBlock(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.attn = NanoChatAttention(config, layer_idx)
+        self.mlp = NanoChatMLP(config)
+
+    def forward(self, x, ve, cos, sin, window_size):
+        x = x + self.attn(_rms_norm(x), ve, cos, sin, window_size)
+        x = x + self.mlp(_rms_norm(x))
+        return x
+
+
+class NanoChatPreTrainedModel(PreTrainedModel):
+    config_class = NanoChatConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["NanoChatBlock"]
+
+    def _init_weights(self, module):
+        pass
+
+
+class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin):
+    def __init__(self, config):
+        super().__init__(config)
+        pvs = config.padded_vocab_size
+        self.wte = nn.Embedding(pvs, config.n_embd)
+        self.h = nn.ModuleList([NanoChatBlock(config, i) for i in range(config.n_layer)])
+        self.lm_head = nn.Linear(config.n_embd, pvs, bias=False)
+        self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))
+        self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))
+        self.smear_gate = nn.Linear(24, 1, bias=False)
+        self.smear_lambda = nn.Parameter(torch.zeros(1))
+        self.backout_lambda = nn.Parameter(0.2 * torch.ones(1))
+        ve_emb_dim = config.ve_dim if config.ve_dim > 0 else config.kv_dim
+        self.value_embeds = nn.ModuleDict({
+            str(i): nn.Embedding(pvs, ve_emb_dim)
+            for i in range(config.n_layer) if _has_ve(i, config.n_layer)
+        })
+        self.window_sizes = self._compute_window_sizes(config)
+        self.rotary_seq_len = config.sequence_len * 10
+        self._rotary_initialized = False
+        self.register_buffer("cos", torch.zeros(1), persistent=False)
+        self.register_buffer("sin", torch.zeros(1), persistent=False)
+        self.post_init()
+
+    def _compute_window_sizes(self, config):
+        pattern = config.window_pattern.upper()
+        long_window = config.sequence_len
+        short_window = -(-long_window // 4 // 128) * 128
+        char_to_window = {"L": long_window, "S": short_window}
+        sizes = [char_to_window[pattern[i % len(pattern)]] for i in range(config.n_layer)]
+        sizes[-1] = long_window
+        return sizes
+
+    def _precompute_rotary(self, seq_len, head_dim, base=100000.0):
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32)
+        inv_freq = 1.0 / (base ** (channel_range / head_dim))
+        t = torch.arange(seq_len, dtype=torch.float32)
+        freqs = torch.outer(t, inv_freq)
+        cos = freqs.cos()[None, :, None, :]
+        sin = freqs.sin()[None, :, None, :]
+        return cos, sin
+
+    def forward(self, input_ids, attention_mask=None, labels=None, return_dict=None, **kwargs):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        B, T = input_ids.size()
+        if not self._rotary_initialized:
+            cos, sin = self._precompute_rotary(self.rotary_seq_len, self.config.head_dim)
+            self.cos = cos.to(device=input_ids.device)
+            self.sin = sin.to(device=input_ids.device)
+            self._rotary_initialized = True
+        cos = self.cos[:, :T].to(dtype=self.wte.weight.dtype)
+        sin = self.sin[:, :T].to(dtype=self.wte.weight.dtype)
+        x = self.wte(input_ids)
+        x = _rms_norm(x)
+        if T > 1:
+            gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
+            x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
+        x0 = x
+        n_layer = self.config.n_layer
+        backout_layer = n_layer // 2
+        x_backout = None
+        for i, block in enumerate(self.h):
+            x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
+            ve = self.value_embeds[str(i)](input_ids).to(x.dtype) if str(i) in self.value_embeds else None
+            x = block(x, ve, cos, sin, self.window_sizes[i])
+            if i == backout_layer:
+                x_backout = x
+        if x_backout is not None:
+            x = x - self.backout_lambda.to(x.dtype) * x_backout
+        x = _rms_norm(x)
+        logits = self.lm_head(x)
+        logits = logits[..., :self.config.vocab_size]
+        logits = logits.float()
+        if self.config.logit_softcap > 0:
+            cap = self.config.logit_softcap
+            logits = cap * torch.tanh(logits / cap)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+        if not return_dict:
+            output = (logits,)
+            return (loss,) + output if loss is not None else output
+        return CausalLMOutputWithPast(loss=loss, logits=logits)
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# ---------------------------------------------------------------------------
+# CLI demo
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    import argparse
+    import sys
+    import os
+
+    parser = argparse.ArgumentParser(description="Load and generate with a NanoChat model")
+    parser.add_argument("--model-dir", required=True, help="Path to converted model directory")
+    parser.add_argument("--prompt", default="The meaning of life is", help="Text prompt")
+    parser.add_argument("--max-tokens", type=int, default=100, help="Max tokens to generate")
+    parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature (0=greedy)")
+    parser.add_argument("--device", default="cpu", help="Device (cpu, cuda, cuda:0, ...)")
+    parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16", "float16"])
+    args = parser.parse_args()
+
+    dtype_map = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}
+
+    print(f"Loading model from {args.model_dir}...")
+    from transformers import AutoTokenizer
+
+    model = NanoChatForCausalLM.from_pretrained(args.model_dir, torch_dtype=dtype_map[args.dtype])
+    model.to(args.device)
+    model.eval()
+    tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Model loaded: {total_params:,} parameters, dtype={args.dtype}, device={args.device}")
+    print(f"Vocab size: {model.config.vocab_size}, ve_dim: {model.config.ve_dim}")
+    print()
+    print(f"Prompt: {args.prompt}")
+    print(f"Generating (max {args.max_tokens} tokens, temp={args.temperature})...")
+    print()
+
+    # Generate
+    bos_id = tokenizer.bos_token_id
+    input_ids = tokenizer.encode(args.prompt, add_special_tokens=False)
+    if bos_id is not None:
+        input_ids = [bos_id] + input_ids
+    input_ids = torch.tensor([input_ids], dtype=torch.long, device=args.device)
+
+    generated = []
+    with torch.inference_mode():
+        for _ in range(args.max_tokens):
+            logits = model(input_ids).logits[:, -1, :]
+            if args.temperature > 0:
+                logits = logits / args.temperature
+                v, _ = torch.topk(logits, min(50, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float("inf")
+                probs = torch.softmax(logits, dim=-1)
+                next_id = torch.multinomial(probs, num_samples=1)
+            else:
+                next_id = torch.argmax(logits, dim=-1, keepdim=True)
+            input_ids = torch.cat([input_ids, next_id], dim=1)
+            token_id = next_id.item()
+            generated.append(token_id)
+            if tokenizer.eos_token_id is not None and token_id == tokenizer.eos_token_id:
+                break
+
+    output = tokenizer.decode(generated, skip_special_tokens=True)
+    print(f"Output: {args.prompt}{output}")

models/full-128k-gpt4o-codegrad00-bpe-mathcode-scratch/config.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "NanoChatForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "nanochat_model.NanoChatConfig",
+    "AutoModelForCausalLM": "nanochat_model.NanoChatForCausalLM"
+  },
+  "bos_token_id": 0,
+  "dtype": "float32",
+  "eos_token_id": null,
+  "head_dim": 128,
+  "kv_dim": 1536,
+  "logit_softcap": 15.0,
+  "model_type": "nanochat",
+  "n_embd": 1536,
+  "n_head": 12,
+  "n_kv_head": 12,
+  "n_layer": 24,
+  "pad_vocab_size_to": 64,
+  "padded_vocab_size": 128256,
+  "sequence_len": 2048,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.3.0",
+  "ve_dim": 128,
+  "vocab_size": 128256,
+  "window_pattern": "SSSL"
+}

models/full-128k-gpt4o-codegrad00-bpe-mathcode-scratch/conversion_metadata.json

@@ -0,0 +1,100 @@
+{
+  "run": "full-128k-gpt4o-codegrad00-bpe-mathcode-scratch",
+  "source_checkpoint": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-gpt4o-codegrad00-bpe-mathcode-scratch",
+  "source_model_file": "model_019073.pt",
+  "source_step": 19073,
+  "source_meta": {
+    "step": 19073,
+    "val_bpb": 0.339566757938572,
+    "model_config": {
+      "sequence_len": 2048,
+      "vocab_size": 128256,
+      "n_layer": 24,
+      "n_head": 12,
+      "n_kv_head": 12,
+      "n_embd": 1536,
+      "window_pattern": "SSSL",
+      "ve_dim": 128,
+      "logit_softcap": 15.0
+    },
+    "tokenizer_path": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/tokenizers/gpt4o-codegrad00-bpe",
+    "training_config": {
+      "name": "full-128k-mathcode-scratch",
+      "model": {
+        "n_layer": 24,
+        "n_embd": 1536,
+        "n_kv_head": 12,
+        "sequence_len": 2048,
+        "window_pattern": "SSSL",
+        "ve_dim": 128
+      },
+      "training": {
+        "target_param_data_ratio": -1,
+        "max_tokens": 20000000000,
+        "num_iterations": -1,
+        "total_batch_size": 1048576,
+        "device_batch_size": 16,
+        "warmup_steps": 40,
+        "warmdown_ratio": 0.65,
+        "final_lr_frac": 0.05,
+        "eval_tokens": 41943040,
+        "compile": true
+      },
+      "optimizer": {
+        "matrix_lr": 0.02,
+        "embedding_lr": 0.3,
+        "unembedding_lr": 0.008,
+        "scalar_lr": 0.5,
+        "weight_decay": 0.28,
+        "width_lr_exponent": -1.0
+      },
+      "data": {
+        "data_dir": null
+      },
+      "logging": {
+        "eval_every": 500,
+        "save_every": 500,
+        "max_checkpoints": 0
+      },
+      "wandb": {
+        "enabled": true,
+        "project": "tokenizer-lm-experiments",
+        "run_name": null
+      },
+      "output_dir": "/capstor/scratch/cscs/${USER}/tokenizer-lm/checkpoints",
+      "seed": 42
+    },
+    "total_batch_size": 1048576,
+    "device_batch_size": 16,
+    "max_seq_len": 2048,
+    "dataloader_state_dict": {
+      "pq_idx": 230,
+      "rg_idx": 36,
+      "epoch": 2
+    },
+    "loop_state": {
+      "min_val_bpb": 0.339566757938572,
+      "smooth_train_loss": 0.7837883755918804,
+      "total_training_time": 93324.20768380165,
+      "total_bytes_consumed": 153517795007
+    }
+  },
+  "weight_dtype": "as trained (float32 matrices, bfloat16 embeddings)",
+  "weight_dtype_counts": {
+    "torch.float32": 174,
+    "torch.bfloat16": 13
+  },
+  "total_parameters": 1272841994,
+  "logit_softcap_source": "meta",
+  "bos_token": "<s>",
+  "bos_token_id": 0,
+  "bos_resolved_from": "packaged tokenizer",
+  "tokenizer_slug": "gpt4o-codegrad00-bpe",
+  "tokenizer_redistributed": true,
+  "tokenizer_hub_id": null,
+  "weight_key_mapping": "strip _orig_mod. and transformer. prefixes",
+  "converted_by": "scripts/hf_release/convert_for_release.py",
+  "repo_commit": "1a9a718",
+  "nanochat_commit": "f919449",
+  "nanochat_branch": "tokenizer-lm-patches"
+}

models/full-128k-gpt4o-codegrad00-bpe-mathcode-scratch/tokenizer_config.json

@@ -0,0 +1,7 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "model_max_length": 2048,
+  "bos_token": "<s>",
+  "add_bos_token": false,
+  "clean_up_tokenization_spaces": false
+}

models/full-128k-gpt4o-codegrad05-bpe-mathcode-scratch/README.md

@@ -0,0 +1,77 @@
+# `full-128k-gpt4o-codegrad05-bpe-mathcode-scratch`
+
+A 1.27B parameter model trained with the `gpt4o-codegrad05-bpe` tokenizer. Regime: mathcode-20B-from-scratch (24 layers, width 1536, 1.27B parameters, 20.0B tokens, math and code mixture, trained from scratch). Reported in: math and code ablations.
+
+## Load
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+    "cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-gpt4o-codegrad05-bpe-mathcode-scratch",
+    trust_remote_code=True, dtype="float32")
+tokenizer = AutoTokenizer.from_pretrained("cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-gpt4o-codegrad05-bpe-mathcode-scratch")
+
+text = "def add(a, b):\n    return a + b\n"
+ids = [0] + tokenizer(text)["input_ids"]   # training prepends BOS
+```
+
+## Settings
+
+| Setting | Value |
+|---|---|
+| layers | 24 |
+| model dimension | 1536 |
+| attention heads | 12 |
+| key and value heads | 12 |
+| context length | 2048 |
+| sliding window pattern | SSSL |
+| value embedding dimension | 128 |
+| logit soft cap | 15 |
+| vocabulary size | 128256 |
+| parameters | 1,272,841,994 |
+| training step | 19073 |
+| token budget | 19,999,490,048 |
+| bytes consumed | 151,162,393,945 |
+| seed | 42 |
+| Muon matrix learning rate | 0.0200 |
+| AdamW embedding learning rate | 0.3000 |
+| AdamW unembedding learning rate | 0.0080 |
+| weight decay | 0.2800 |
+| width learning-rate exponent | -1 |
+| weight dtype | as trained (float32 matrices, bfloat16 embeddings) |
+| BOS token id | 0 |
+
+## Results
+
+| Metric | Value |
+|---|---|
+| validation BPB | 0.3366 |
+| code BPB, mean over 7 languages | n/a |
+| FLORES-200 BPB, mean | 3.3960 |
+| FLORES-200 languages scored | 215 |
+| BLiMP accuracy | n/a |
+| MBPP pass@1 | 0.0500 |
+| HumanEval pass@1 | 0.1402 |
+| GSM8K exact match, strict | 0.2161 |
+
+Generation metrics use specification `{'version': 'v2-2026-07-30', 'bos_context': True, 'eos_stop': True, 'token_healing': True, 'truncation': 'left_to_max_length_minus_max_gen', 'stop_truncation': 'earliest_match_all_terms', 'stop_match_decode': 'specials_kept'}`, token healing on.
+
+Missing measurements, with the reason each is absent:
+
+- `blimp_acc`: no full-128k-gpt4o-codegrad05-bpe-mathcode-scratch_blimp_code_bpb.json
+- `code_bpb_mean`: no full-128k-gpt4o-codegrad05-bpe-mathcode-scratch_blimp_code_bpb.json; no full-128k-gpt4o-codegrad05-bpe-mathcode-scratch_code_bpb.json
+
+BLiMP is comparable within a training scale only. Its scoring sums unnormalized sentence log-likelihoods, and under that scoring a group of quantifier subtasks drops at the 1.27B scale relative to the 300M scale.
+
+## Registry note
+
+#40, gpt4o pretok, slug `gpt4o-codegrad05-bpe`. AXES-2 code-fraction gradient, fusing regex (`REGEX_GPT4O`), code fraction 5%. Fused 749. FINISHED. MBPP 0.048 (v1); v2 0.050. Not in panel.
+
+## Provenance
+
+- Source checkpoint: `/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-gpt4o-codegrad05-bpe-mathcode-scratch`, step 19073
+- Repository commit: `1a9a718`
+- nanochat commit: `f919449` on branch `tokenizer-lm-patches`
+- Tokenizer: `gpt4o-codegrad05-bpe`, released at `tokenizers/gpt4o-codegrad05-bpe`
+

models/full-128k-gpt4o-codegrad05-bpe-mathcode-scratch/config.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "NanoChatForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "nanochat_model.NanoChatConfig",
+    "AutoModelForCausalLM": "nanochat_model.NanoChatForCausalLM"
+  },
+  "bos_token_id": 0,
+  "dtype": "float32",
+  "eos_token_id": null,
+  "head_dim": 128,
+  "kv_dim": 1536,
+  "logit_softcap": 15.0,
+  "model_type": "nanochat",
+  "n_embd": 1536,
+  "n_head": 12,
+  "n_kv_head": 12,
+  "n_layer": 24,
+  "pad_vocab_size_to": 64,
+  "padded_vocab_size": 128256,
+  "sequence_len": 2048,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.3.0",
+  "ve_dim": 128,
+  "vocab_size": 128256,
+  "window_pattern": "SSSL"
+}

models/full-128k-gpt4o-codegrad05-bpe-mathcode-scratch/conversion_metadata.json

@@ -0,0 +1,100 @@
+{
+  "run": "full-128k-gpt4o-codegrad05-bpe-mathcode-scratch",
+  "source_checkpoint": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-gpt4o-codegrad05-bpe-mathcode-scratch",
+  "source_model_file": "model_019073.pt",
+  "source_step": 19073,
+  "source_meta": {
+    "step": 19073,
+    "val_bpb": 0.3365866322863168,
+    "model_config": {
+      "sequence_len": 2048,
+      "vocab_size": 128256,
+      "n_layer": 24,
+      "n_head": 12,
+      "n_kv_head": 12,
+      "n_embd": 1536,
+      "window_pattern": "SSSL",
+      "ve_dim": 128,
+      "logit_softcap": 15.0
+    },
+    "tokenizer_path": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/tokenizers/gpt4o-codegrad05-bpe",
+    "training_config": {
+      "name": "full-128k-mathcode-scratch",
+      "model": {
+        "n_layer": 24,
+        "n_embd": 1536,
+        "n_kv_head": 12,
+        "sequence_len": 2048,
+        "window_pattern": "SSSL",
+        "ve_dim": 128
+      },
+      "training": {
+        "target_param_data_ratio": -1,
+        "max_tokens": 20000000000,
+        "num_iterations": -1,
+        "total_batch_size": 1048576,
+        "device_batch_size": 16,
+        "warmup_steps": 40,
+        "warmdown_ratio": 0.65,
+        "final_lr_frac": 0.05,
+        "eval_tokens": 41943040,
+        "compile": true
+      },
+      "optimizer": {
+        "matrix_lr": 0.02,
+        "embedding_lr": 0.3,
+        "unembedding_lr": 0.008,
+        "scalar_lr": 0.5,
+        "weight_decay": 0.28,
+        "width_lr_exponent": -1.0
+      },
+      "data": {
+        "data_dir": null
+      },
+      "logging": {
+        "eval_every": 500,
+        "save_every": 500,
+        "max_checkpoints": 0
+      },
+      "wandb": {
+        "enabled": true,
+        "project": "tokenizer-lm-experiments",
+        "run_name": null
+      },
+      "output_dir": "/capstor/scratch/cscs/${USER}/tokenizer-lm/checkpoints",
+      "seed": 42
+    },
+    "total_batch_size": 1048576,
+    "device_batch_size": 16,
+    "max_seq_len": 2048,
+    "dataloader_state_dict": {
+      "pq_idx": 254,
+      "rg_idx": 44,
+      "epoch": 2
+    },
+    "loop_state": {
+      "min_val_bpb": 0.3365866322863168,
+      "smooth_train_loss": 0.8052011063176456,
+      "total_training_time": 93320.9358727932,
+      "total_bytes_consumed": 151162393945
+    }
+  },
+  "weight_dtype": "as trained (float32 matrices, bfloat16 embeddings)",
+  "weight_dtype_counts": {
+    "torch.float32": 174,
+    "torch.bfloat16": 13
+  },
+  "total_parameters": 1272841994,
+  "logit_softcap_source": "meta",
+  "bos_token": "<s>",
+  "bos_token_id": 0,
+  "bos_resolved_from": "packaged tokenizer",
+  "tokenizer_slug": "gpt4o-codegrad05-bpe",
+  "tokenizer_redistributed": true,
+  "tokenizer_hub_id": null,
+  "weight_key_mapping": "strip _orig_mod. and transformer. prefixes",
+  "converted_by": "scripts/hf_release/convert_for_release.py",
+  "repo_commit": "1a9a718",
+  "nanochat_commit": "f919449",
+  "nanochat_branch": "tokenizer-lm-patches"
+}

models/full-128k-gpt4o-codegrad05-bpe-mathcode-scratch/generation_config.json

@@ -0,0 +1,4 @@
+{
+  "bos_token_id": 0,
+  "transformers_version": "5.3.0"
+}

models/full-128k-gpt4o-codegrad05-bpe-mathcode-scratch/nanochat_model.py

@@ -0,0 +1,361 @@
+"""
+NanoChat model for HuggingFace Transformers — self-contained single file.
+
+Includes the config class, model class, and a CLI demo for loading + generation.
+Drop this file into a directory containing config.json, model.safetensors,
+tokenizer.json, and tokenizer_config.json, then run:
+
+    python nanochat_model.py --model-dir . --prompt "Hello world"
+
+Or load programmatically:
+
+    from nanochat_model import NanoChatConfig, NanoChatForCausalLM
+    model = NanoChatForCausalLM.from_pretrained("path/to/model", torch_dtype=torch.bfloat16)
+
+Requirements: pip install torch transformers tokenizers
+
+Architecture: nanochat GPT (Karpathy, MIT licensed) with value embeddings
+(ResFormer-style), per-layer residual/x0 lambdas, smear, backout, logit
+soft-capping, sliding window attention, QK norm, ReLU^2 MLP, and RoPE.
+
+The models were trained with a patched nanochat, not upstream nanochat:
+https://github.com/cimeister/nanochat branch tokenizer-lm-patches, commit f919449.
+The patches touch two files and change three things, two of which affect the forward
+pass and are implemented here:
+  1. ve_dim: the value-embedding tables are built at a reduced dimension and projected
+     up to kv_dim by a learned ve_proj, instead of being built at full kv_dim. All
+     released models use ve_dim=128.
+  2. logit_softcap: the soft cap is a config field with 0 meaning disabled, instead of
+     a hardcoded 15.
+  3. width_lr_exponent: the AdamW learning-rate width scaling exponent is configurable
+     for muP. This affects training only and has no inference-time effect.
+
+This file is the released model definition. Its agreement with the training-time model
+is checked per released checkpoint by scripts/hf_release/verify_conversion.py, which
+compares logits against the patched nanochat GPT in float32.
+"""
+
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+class NanoChatConfig(PretrainedConfig):
+    model_type = "nanochat"
+
+    def __init__(
+        self,
+        vocab_size=128256,
+        n_embd=1536,
+        n_layer=24,
+        n_head=12,
+        n_kv_head=12,
+        sequence_len=2048,
+        window_pattern="SSSL",
+        ve_dim=0,
+        logit_softcap=15.0,
+        pad_vocab_size_to=64,
+        bos_token_id=0,
+        eos_token_id=None,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_kv_head = n_kv_head
+        self.sequence_len = sequence_len
+        self.window_pattern = window_pattern
+        self.ve_dim = ve_dim
+        self.logit_softcap = logit_softcap
+        self.pad_vocab_size_to = pad_vocab_size_to
+        self.head_dim = n_embd // n_head
+        self.kv_dim = n_kv_head * self.head_dim
+        self.padded_vocab_size = (
+            (vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to
+        ) * pad_vocab_size_to
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Model
+# ---------------------------------------------------------------------------
+
+def _has_ve(layer_idx: int, n_layer: int) -> bool:
+    return layer_idx % 2 == (n_layer - 1) % 2
+
+
+def _apply_rotary_emb(x, cos, sin):
+    d = x.shape[-1] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    return torch.cat([y1, y2], dim=-1)
+
+
+def _rms_norm(x):
+    return F.rms_norm(x, (x.size(-1),))
+
+
+class NanoChatAttention(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.n_embd = config.n_embd
+        self.head_dim = config.head_dim
+        self.n_kv_groups = self.n_head // self.n_kv_head
+        self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
+        self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
+        self.ve_gate_channels = 12
+        if _has_ve(layer_idx, config.n_layer):
+            self.ve_gate = nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False)
+            kv_dim = self.n_kv_head * self.head_dim
+            ve_dim = config.ve_dim if config.ve_dim > 0 else kv_dim
+            self.ve_proj = nn.Linear(ve_dim, kv_dim, bias=False) if ve_dim < kv_dim else None
+        else:
+            self.ve_gate = None
+            self.ve_proj = None
+
+    def forward(self, x, ve, cos, sin, window_size):
+        B, T, C = x.size()
+        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
+        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
+        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
+        if ve is not None and self.ve_gate is not None:
+            if self.ve_proj is not None:
+                ve = self.ve_proj(ve)
+            ve = ve.view(B, T, self.n_kv_head, self.head_dim)
+            gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
+            v = v + gate.unsqueeze(-1) * ve
+        q = _apply_rotary_emb(q, cos, sin)
+        k = _apply_rotary_emb(k, cos, sin)
+        q = _rms_norm(q) * 1.2
+        k = _rms_norm(k) * 1.2
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        if self.n_kv_groups > 1:
+            k = k.repeat_interleave(self.n_kv_groups, dim=1)
+            v = v.repeat_interleave(self.n_kv_groups, dim=1)
+        if window_size < 0 or window_size >= T:
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:
+            row_idx = torch.arange(T, device=x.device).unsqueeze(1)
+            col_idx = torch.arange(T, device=x.device).unsqueeze(0)
+            mask = (col_idx <= row_idx) & ((row_idx - col_idx) <= window_size)
+            y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        y = y.transpose(1, 2).contiguous().view(B, T, -1)
+        y = self.c_proj(y)
+        return y
+
+
+class NanoChatMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = F.relu(x).square()
+        x = self.c_proj(x)
+        return x
+
+
+class NanoChatBlock(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.attn = NanoChatAttention(config, layer_idx)
+        self.mlp = NanoChatMLP(config)
+
+    def forward(self, x, ve, cos, sin, window_size):
+        x = x + self.attn(_rms_norm(x), ve, cos, sin, window_size)
+        x = x + self.mlp(_rms_norm(x))
+        return x
+
+
+class NanoChatPreTrainedModel(PreTrainedModel):
+    config_class = NanoChatConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["NanoChatBlock"]
+
+    def _init_weights(self, module):
+        pass
+
+
+class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin):
+    def __init__(self, config):
+        super().__init__(config)
+        pvs = config.padded_vocab_size
+        self.wte = nn.Embedding(pvs, config.n_embd)
+        self.h = nn.ModuleList([NanoChatBlock(config, i) for i in range(config.n_layer)])
+        self.lm_head = nn.Linear(config.n_embd, pvs, bias=False)
+        self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))
+        self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))
+        self.smear_gate = nn.Linear(24, 1, bias=False)
+        self.smear_lambda = nn.Parameter(torch.zeros(1))
+        self.backout_lambda = nn.Parameter(0.2 * torch.ones(1))
+        ve_emb_dim = config.ve_dim if config.ve_dim > 0 else config.kv_dim
+        self.value_embeds = nn.ModuleDict({
+            str(i): nn.Embedding(pvs, ve_emb_dim)
+            for i in range(config.n_layer) if _has_ve(i, config.n_layer)
+        })
+        self.window_sizes = self._compute_window_sizes(config)
+        self.rotary_seq_len = config.sequence_len * 10
+        self._rotary_initialized = False
+        self.register_buffer("cos", torch.zeros(1), persistent=False)
+        self.register_buffer("sin", torch.zeros(1), persistent=False)
+        self.post_init()
+
+    def _compute_window_sizes(self, config):
+        pattern = config.window_pattern.upper()
+        long_window = config.sequence_len
+        short_window = -(-long_window // 4 // 128) * 128
+        char_to_window = {"L": long_window, "S": short_window}
+        sizes = [char_to_window[pattern[i % len(pattern)]] for i in range(config.n_layer)]
+        sizes[-1] = long_window
+        return sizes
+
+    def _precompute_rotary(self, seq_len, head_dim, base=100000.0):
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32)
+        inv_freq = 1.0 / (base ** (channel_range / head_dim))
+        t = torch.arange(seq_len, dtype=torch.float32)
+        freqs = torch.outer(t, inv_freq)
+        cos = freqs.cos()[None, :, None, :]
+        sin = freqs.sin()[None, :, None, :]
+        return cos, sin
+
+    def forward(self, input_ids, attention_mask=None, labels=None, return_dict=None, **kwargs):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        B, T = input_ids.size()
+        if not self._rotary_initialized:
+            cos, sin = self._precompute_rotary(self.rotary_seq_len, self.config.head_dim)
+            self.cos = cos.to(device=input_ids.device)
+            self.sin = sin.to(device=input_ids.device)
+            self._rotary_initialized = True
+        cos = self.cos[:, :T].to(dtype=self.wte.weight.dtype)
+        sin = self.sin[:, :T].to(dtype=self.wte.weight.dtype)
+        x = self.wte(input_ids)
+        x = _rms_norm(x)
+        if T > 1:
+            gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
+            x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
+        x0 = x
+        n_layer = self.config.n_layer
+        backout_layer = n_layer // 2
+        x_backout = None
+        for i, block in enumerate(self.h):
+            x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
+            ve = self.value_embeds[str(i)](input_ids).to(x.dtype) if str(i) in self.value_embeds else None
+            x = block(x, ve, cos, sin, self.window_sizes[i])
+            if i == backout_layer:
+                x_backout = x
+        if x_backout is not None:
+            x = x - self.backout_lambda.to(x.dtype) * x_backout
+        x = _rms_norm(x)
+        logits = self.lm_head(x)
+        logits = logits[..., :self.config.vocab_size]
+        logits = logits.float()
+        if self.config.logit_softcap > 0:
+            cap = self.config.logit_softcap
+            logits = cap * torch.tanh(logits / cap)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+        if not return_dict:
+            output = (logits,)
+            return (loss,) + output if loss is not None else output
+        return CausalLMOutputWithPast(loss=loss, logits=logits)
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# ---------------------------------------------------------------------------
+# CLI demo
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    import argparse
+    import sys
+    import os
+
+    parser = argparse.ArgumentParser(description="Load and generate with a NanoChat model")
+    parser.add_argument("--model-dir", required=True, help="Path to converted model directory")
+    parser.add_argument("--prompt", default="The meaning of life is", help="Text prompt")
+    parser.add_argument("--max-tokens", type=int, default=100, help="Max tokens to generate")
+    parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature (0=greedy)")
+    parser.add_argument("--device", default="cpu", help="Device (cpu, cuda, cuda:0, ...)")
+    parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16", "float16"])
+    args = parser.parse_args()
+
+    dtype_map = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}
+
+    print(f"Loading model from {args.model_dir}...")
+    from transformers import AutoTokenizer
+
+    model = NanoChatForCausalLM.from_pretrained(args.model_dir, torch_dtype=dtype_map[args.dtype])
+    model.to(args.device)
+    model.eval()
+    tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Model loaded: {total_params:,} parameters, dtype={args.dtype}, device={args.device}")
+    print(f"Vocab size: {model.config.vocab_size}, ve_dim: {model.config.ve_dim}")
+    print()
+    print(f"Prompt: {args.prompt}")
+    print(f"Generating (max {args.max_tokens} tokens, temp={args.temperature})...")
+    print()
+
+    # Generate
+    bos_id = tokenizer.bos_token_id
+    input_ids = tokenizer.encode(args.prompt, add_special_tokens=False)
+    if bos_id is not None:
+        input_ids = [bos_id] + input_ids
+    input_ids = torch.tensor([input_ids], dtype=torch.long, device=args.device)
+
+    generated = []
+    with torch.inference_mode():
+        for _ in range(args.max_tokens):
+            logits = model(input_ids).logits[:, -1, :]
+            if args.temperature > 0:
+                logits = logits / args.temperature
+                v, _ = torch.topk(logits, min(50, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float("inf")
+                probs = torch.softmax(logits, dim=-1)
+                next_id = torch.multinomial(probs, num_samples=1)
+            else:
+                next_id = torch.argmax(logits, dim=-1, keepdim=True)
+            input_ids = torch.cat([input_ids, next_id], dim=1)
+            token_id = next_id.item()
+            generated.append(token_id)
+            if tokenizer.eos_token_id is not None and token_id == tokenizer.eos_token_id:
+                break
+
+    output = tokenizer.decode(generated, skip_special_tokens=True)
+    print(f"Output: {args.prompt}{output}")

models/full-128k-gpt4o-codegrad05-bpe-mathcode-scratch/tokenizer_config.json

@@ -0,0 +1,7 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "model_max_length": 2048,
+  "bos_token": "<s>",
+  "add_bos_token": false,
+  "clean_up_tokenization_spaces": false
+}

models/full-128k-gpt4onl-codegrad05-bpe-mathcode-scratch/tokenizer_config.json

@@ -0,0 +1,7 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "model_max_length": 2048,
+  "bos_token": "<s>",
+  "add_bos_token": false,
+  "clean_up_tokenization_spaces": false
+}

models/full-128k-gpt4onl-codegrad30-bpe-mathcode-scratch/README.md

@@ -0,0 +1,77 @@
+# `full-128k-gpt4onl-codegrad30-bpe-mathcode-scratch`
+
+A 1.27B parameter model trained with the `gpt4onl-codegrad30-bpe` tokenizer. Regime: mathcode-20B-from-scratch (24 layers, width 1536, 1.27B parameters, 20.0B tokens, math and code mixture, trained from scratch). Reported in: math and code ablations.
+
+## Load
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+    "cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-gpt4onl-codegrad30-bpe-mathcode-scratch",
+    trust_remote_code=True, dtype="float32")
+tokenizer = AutoTokenizer.from_pretrained("cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-gpt4onl-codegrad30-bpe-mathcode-scratch")
+
+text = "def add(a, b):\n    return a + b\n"
+ids = [0] + tokenizer(text)["input_ids"]   # training prepends BOS
+```
+
+## Settings
+
+| Setting | Value |
+|---|---|
+| layers | 24 |
+| model dimension | 1536 |
+| attention heads | 12 |
+| key and value heads | 12 |
+| context length | 2048 |
+| sliding window pattern | SSSL |
+| value embedding dimension | 128 |
+| logit soft cap | 15 |
+| vocabulary size | 128256 |
+| parameters | 1,272,841,994 |
+| training step | 19073 |
+| token budget | 19,999,490,048 |
+| bytes consumed | 144,184,498,749 |
+| seed | 42 |
+| Muon matrix learning rate | 0.0200 |
+| AdamW embedding learning rate | 0.3000 |
+| AdamW unembedding learning rate | 0.0080 |
+| weight decay | 0.2800 |
+| width learning-rate exponent | -1 |
+| weight dtype | as trained (float32 matrices, bfloat16 embeddings) |
+| BOS token id | 0 |
+
+## Results
+
+| Metric | Value |
+|---|---|
+| validation BPB | 0.3367 |
+| code BPB, mean over 7 languages | n/a |
+| FLORES-200 BPB, mean | 3.3328 |
+| FLORES-200 languages scored | 215 |
+| BLiMP accuracy | n/a |
+| MBPP pass@1 | 0.2240 |
+| HumanEval pass@1 | 0.1768 |
+| GSM8K exact match, strict | 0.2002 |
+
+Generation metrics use specification `{'version': 'v2-2026-07-30', 'bos_context': True, 'eos_stop': True, 'token_healing': True, 'truncation': 'left_to_max_length_minus_max_gen', 'stop_truncation': 'earliest_match_all_terms', 'stop_match_decode': 'specials_kept'}`, token healing on.
+
+Missing measurements, with the reason each is absent:
+
+- `blimp_acc`: no full-128k-gpt4onl-codegrad30-bpe-mathcode-scratch_blimp_code_bpb.json
+- `code_bpb_mean`: no full-128k-gpt4onl-codegrad30-bpe-mathcode-scratch_blimp_code_bpb.json; no full-128k-gpt4onl-codegrad30-bpe-mathcode-scratch_code_bpb.json
+
+BLiMP is comparable within a training scale only. Its scoring sums unnormalized sentence log-likelihoods, and under that scoring a group of quantifier subtasks drops at the 1.27B scale relative to the 300M scale.
+
+## Registry note
+
+#45, gpt4onl pretok, slug `gpt4onl-codegrad30-bpe`. AXES-2 code-fraction gradient, non-fusing regex (`REGEX_GPT4O_NONL`), code fraction 30%. Fused 0. FINISHED. MBPP 0.178 (v1); v2 0.224. Not in panel.
+
+## Provenance
+
+- Source checkpoint: `/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-gpt4onl-codegrad30-bpe-mathcode-scratch`, step 19073
+- Repository commit: `1a9a718`
+- nanochat commit: `f919449` on branch `tokenizer-lm-patches`
+- Tokenizer: `gpt4onl-codegrad30-bpe`, released at `tokenizers/gpt4onl-codegrad30-bpe`
+

models/full-128k-gpt4onl-codegrad30-bpe-mathcode-scratch/nanochat_model.py

@@ -0,0 +1,361 @@
+"""
+NanoChat model for HuggingFace Transformers — self-contained single file.
+
+Includes the config class, model class, and a CLI demo for loading + generation.
+Drop this file into a directory containing config.json, model.safetensors,
+tokenizer.json, and tokenizer_config.json, then run:
+
+    python nanochat_model.py --model-dir . --prompt "Hello world"
+
+Or load programmatically:
+
+    from nanochat_model import NanoChatConfig, NanoChatForCausalLM
+    model = NanoChatForCausalLM.from_pretrained("path/to/model", torch_dtype=torch.bfloat16)
+
+Requirements: pip install torch transformers tokenizers
+
+Architecture: nanochat GPT (Karpathy, MIT licensed) with value embeddings
+(ResFormer-style), per-layer residual/x0 lambdas, smear, backout, logit
+soft-capping, sliding window attention, QK norm, ReLU^2 MLP, and RoPE.
+
+The models were trained with a patched nanochat, not upstream nanochat:
+https://github.com/cimeister/nanochat branch tokenizer-lm-patches, commit f919449.
+The patches touch two files and change three things, two of which affect the forward
+pass and are implemented here:
+  1. ve_dim: the value-embedding tables are built at a reduced dimension and projected
+     up to kv_dim by a learned ve_proj, instead of being built at full kv_dim. All
+     released models use ve_dim=128.
+  2. logit_softcap: the soft cap is a config field with 0 meaning disabled, instead of
+     a hardcoded 15.
+  3. width_lr_exponent: the AdamW learning-rate width scaling exponent is configurable
+     for muP. This affects training only and has no inference-time effect.
+
+This file is the released model definition. Its agreement with the training-time model
+is checked per released checkpoint by scripts/hf_release/verify_conversion.py, which
+compares logits against the patched nanochat GPT in float32.
+"""
+
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+class NanoChatConfig(PretrainedConfig):
+    model_type = "nanochat"
+
+    def __init__(
+        self,
+        vocab_size=128256,
+        n_embd=1536,
+        n_layer=24,
+        n_head=12,
+        n_kv_head=12,
+        sequence_len=2048,
+        window_pattern="SSSL",
+        ve_dim=0,
+        logit_softcap=15.0,
+        pad_vocab_size_to=64,
+        bos_token_id=0,
+        eos_token_id=None,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_kv_head = n_kv_head
+        self.sequence_len = sequence_len
+        self.window_pattern = window_pattern
+        self.ve_dim = ve_dim
+        self.logit_softcap = logit_softcap
+        self.pad_vocab_size_to = pad_vocab_size_to
+        self.head_dim = n_embd // n_head
+        self.kv_dim = n_kv_head * self.head_dim
+        self.padded_vocab_size = (
+            (vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to
+        ) * pad_vocab_size_to
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Model
+# ---------------------------------------------------------------------------
+
+def _has_ve(layer_idx: int, n_layer: int) -> bool:
+    return layer_idx % 2 == (n_layer - 1) % 2
+
+
+def _apply_rotary_emb(x, cos, sin):
+    d = x.shape[-1] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    return torch.cat([y1, y2], dim=-1)
+
+
+def _rms_norm(x):
+    return F.rms_norm(x, (x.size(-1),))
+
+
+class NanoChatAttention(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.n_embd = config.n_embd
+        self.head_dim = config.head_dim
+        self.n_kv_groups = self.n_head // self.n_kv_head
+        self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
+        self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
+        self.ve_gate_channels = 12
+        if _has_ve(layer_idx, config.n_layer):
+            self.ve_gate = nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False)
+            kv_dim = self.n_kv_head * self.head_dim
+            ve_dim = config.ve_dim if config.ve_dim > 0 else kv_dim
+            self.ve_proj = nn.Linear(ve_dim, kv_dim, bias=False) if ve_dim < kv_dim else None
+        else:
+            self.ve_gate = None
+            self.ve_proj = None
+
+    def forward(self, x, ve, cos, sin, window_size):
+        B, T, C = x.size()
+        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
+        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
+        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
+        if ve is not None and self.ve_gate is not None:
+            if self.ve_proj is not None:
+                ve = self.ve_proj(ve)
+            ve = ve.view(B, T, self.n_kv_head, self.head_dim)
+            gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
+            v = v + gate.unsqueeze(-1) * ve
+        q = _apply_rotary_emb(q, cos, sin)
+        k = _apply_rotary_emb(k, cos, sin)
+        q = _rms_norm(q) * 1.2
+        k = _rms_norm(k) * 1.2
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        if self.n_kv_groups > 1:
+            k = k.repeat_interleave(self.n_kv_groups, dim=1)
+            v = v.repeat_interleave(self.n_kv_groups, dim=1)
+        if window_size < 0 or window_size >= T:
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:
+            row_idx = torch.arange(T, device=x.device).unsqueeze(1)
+            col_idx = torch.arange(T, device=x.device).unsqueeze(0)
+            mask = (col_idx <= row_idx) & ((row_idx - col_idx) <= window_size)
+            y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        y = y.transpose(1, 2).contiguous().view(B, T, -1)
+        y = self.c_proj(y)
+        return y
+
+
+class NanoChatMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = F.relu(x).square()
+        x = self.c_proj(x)
+        return x
+
+
+class NanoChatBlock(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.attn = NanoChatAttention(config, layer_idx)
+        self.mlp = NanoChatMLP(config)
+
+    def forward(self, x, ve, cos, sin, window_size):
+        x = x + self.attn(_rms_norm(x), ve, cos, sin, window_size)
+        x = x + self.mlp(_rms_norm(x))
+        return x
+
+
+class NanoChatPreTrainedModel(PreTrainedModel):
+    config_class = NanoChatConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["NanoChatBlock"]
+
+    def _init_weights(self, module):
+        pass
+
+
+class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin):
+    def __init__(self, config):
+        super().__init__(config)
+        pvs = config.padded_vocab_size
+        self.wte = nn.Embedding(pvs, config.n_embd)
+        self.h = nn.ModuleList([NanoChatBlock(config, i) for i in range(config.n_layer)])
+        self.lm_head = nn.Linear(config.n_embd, pvs, bias=False)
+        self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))
+        self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))
+        self.smear_gate = nn.Linear(24, 1, bias=False)
+        self.smear_lambda = nn.Parameter(torch.zeros(1))
+        self.backout_lambda = nn.Parameter(0.2 * torch.ones(1))
+        ve_emb_dim = config.ve_dim if config.ve_dim > 0 else config.kv_dim
+        self.value_embeds = nn.ModuleDict({
+            str(i): nn.Embedding(pvs, ve_emb_dim)
+            for i in range(config.n_layer) if _has_ve(i, config.n_layer)
+        })
+        self.window_sizes = self._compute_window_sizes(config)
+        self.rotary_seq_len = config.sequence_len * 10
+        self._rotary_initialized = False
+        self.register_buffer("cos", torch.zeros(1), persistent=False)
+        self.register_buffer("sin", torch.zeros(1), persistent=False)
+        self.post_init()
+
+    def _compute_window_sizes(self, config):
+        pattern = config.window_pattern.upper()
+        long_window = config.sequence_len
+        short_window = -(-long_window // 4 // 128) * 128
+        char_to_window = {"L": long_window, "S": short_window}
+        sizes = [char_to_window[pattern[i % len(pattern)]] for i in range(config.n_layer)]
+        sizes[-1] = long_window
+        return sizes
+
+    def _precompute_rotary(self, seq_len, head_dim, base=100000.0):
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32)
+        inv_freq = 1.0 / (base ** (channel_range / head_dim))
+        t = torch.arange(seq_len, dtype=torch.float32)
+        freqs = torch.outer(t, inv_freq)
+        cos = freqs.cos()[None, :, None, :]
+        sin = freqs.sin()[None, :, None, :]
+        return cos, sin
+
+    def forward(self, input_ids, attention_mask=None, labels=None, return_dict=None, **kwargs):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        B, T = input_ids.size()
+        if not self._rotary_initialized:
+            cos, sin = self._precompute_rotary(self.rotary_seq_len, self.config.head_dim)
+            self.cos = cos.to(device=input_ids.device)
+            self.sin = sin.to(device=input_ids.device)
+            self._rotary_initialized = True
+        cos = self.cos[:, :T].to(dtype=self.wte.weight.dtype)
+        sin = self.sin[:, :T].to(dtype=self.wte.weight.dtype)
+        x = self.wte(input_ids)
+        x = _rms_norm(x)
+        if T > 1:
+            gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
+            x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
+        x0 = x
+        n_layer = self.config.n_layer
+        backout_layer = n_layer // 2
+        x_backout = None
+        for i, block in enumerate(self.h):
+            x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
+            ve = self.value_embeds[str(i)](input_ids).to(x.dtype) if str(i) in self.value_embeds else None
+            x = block(x, ve, cos, sin, self.window_sizes[i])
+            if i == backout_layer:
+                x_backout = x
+        if x_backout is not None:
+            x = x - self.backout_lambda.to(x.dtype) * x_backout
+        x = _rms_norm(x)
+        logits = self.lm_head(x)
+        logits = logits[..., :self.config.vocab_size]
+        logits = logits.float()
+        if self.config.logit_softcap > 0:
+            cap = self.config.logit_softcap
+            logits = cap * torch.tanh(logits / cap)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+        if not return_dict:
+            output = (logits,)
+            return (loss,) + output if loss is not None else output
+        return CausalLMOutputWithPast(loss=loss, logits=logits)
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# ---------------------------------------------------------------------------
+# CLI demo
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    import argparse
+    import sys
+    import os
+
+    parser = argparse.ArgumentParser(description="Load and generate with a NanoChat model")
+    parser.add_argument("--model-dir", required=True, help="Path to converted model directory")
+    parser.add_argument("--prompt", default="The meaning of life is", help="Text prompt")
+    parser.add_argument("--max-tokens", type=int, default=100, help="Max tokens to generate")
+    parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature (0=greedy)")
+    parser.add_argument("--device", default="cpu", help="Device (cpu, cuda, cuda:0, ...)")
+    parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16", "float16"])
+    args = parser.parse_args()
+
+    dtype_map = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}
+
+    print(f"Loading model from {args.model_dir}...")
+    from transformers import AutoTokenizer
+
+    model = NanoChatForCausalLM.from_pretrained(args.model_dir, torch_dtype=dtype_map[args.dtype])
+    model.to(args.device)
+    model.eval()
+    tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Model loaded: {total_params:,} parameters, dtype={args.dtype}, device={args.device}")
+    print(f"Vocab size: {model.config.vocab_size}, ve_dim: {model.config.ve_dim}")
+    print()
+    print(f"Prompt: {args.prompt}")
+    print(f"Generating (max {args.max_tokens} tokens, temp={args.temperature})...")
+    print()
+
+    # Generate
+    bos_id = tokenizer.bos_token_id
+    input_ids = tokenizer.encode(args.prompt, add_special_tokens=False)
+    if bos_id is not None:
+        input_ids = [bos_id] + input_ids
+    input_ids = torch.tensor([input_ids], dtype=torch.long, device=args.device)
+
+    generated = []
+    with torch.inference_mode():
+        for _ in range(args.max_tokens):
+            logits = model(input_ids).logits[:, -1, :]
+            if args.temperature > 0:
+                logits = logits / args.temperature
+                v, _ = torch.topk(logits, min(50, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float("inf")
+                probs = torch.softmax(logits, dim=-1)
+                next_id = torch.multinomial(probs, num_samples=1)
+            else:
+                next_id = torch.argmax(logits, dim=-1, keepdim=True)
+            input_ids = torch.cat([input_ids, next_id], dim=1)
+            token_id = next_id.item()
+            generated.append(token_id)
+            if tokenizer.eos_token_id is not None and token_id == tokenizer.eos_token_id:
+                break
+
+    output = tokenizer.decode(generated, skip_special_tokens=True)
+    print(f"Output: {args.prompt}{output}")

models/full-128k-pabpe-nfc-clean-balanced-hybridwindow-mathcode-scratch/README.md

@@ -0,0 +1,77 @@
+# `full-128k-pabpe-nfc-clean-balanced-hybridwindow-mathcode-scratch`
+
+A 1.27B parameter model trained with the `pabpe-nfc-clean-balanced-hybridwindow` tokenizer. Regime: mathcode-20B-from-scratch (24 layers, width 1536, 1.27B parameters, 20.0B tokens, math and code mixture, trained from scratch). Reported in: math and code ablations.
+
+## Load
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+    "cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-pabpe-nfc-clean-balanced-hybridwindow-mathcode-scratch",
+    trust_remote_code=True, dtype="float32")
+tokenizer = AutoTokenizer.from_pretrained("cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-pabpe-nfc-clean-balanced-hybridwindow-mathcode-scratch")
+
+text = "def add(a, b):\n    return a + b\n"
+ids = [127832] + tokenizer(text)["input_ids"]   # training prepends BOS
+```
+
+## Settings
+
+| Setting | Value |
+|---|---|
+| layers | 24 |
+| model dimension | 1536 |
+| attention heads | 12 |
+| key and value heads | 12 |
+| context length | 2048 |
+| sliding window pattern | SSSL |
+| value embedding dimension | 128 |
+| logit soft cap | 15 |
+| vocabulary size | 127836 |
+| parameters | 1,271,072,522 |
+| training step | 19073 |
+| token budget | 19,999,490,048 |
+| bytes consumed | 152,278,889,219 |
+| seed | 42 |
+| Muon matrix learning rate | 0.0200 |
+| AdamW embedding learning rate | 0.3000 |
+| AdamW unembedding learning rate | 0.0080 |
+| weight decay | 0.2800 |
+| width learning-rate exponent | -1 |
+| weight dtype | as trained (float32 matrices, bfloat16 embeddings) |
+| BOS token id | 127832 |
+
+## Results
+
+| Metric | Value |
+|---|---|
+| validation BPB | 0.3407 |
+| code BPB, mean over 7 languages | n/a |
+| FLORES-200 BPB, mean | 3.3507 |
+| FLORES-200 languages scored | 214 |
+| BLiMP accuracy | n/a |
+| MBPP pass@1 | 0.2320 |
+| HumanEval pass@1 | 0.1524 |
+| GSM8K exact match, strict | 0.2130 |
+
+Generation metrics use specification `{'version': 'v2-2026-07-30', 'bos_context': True, 'eos_stop': True, 'token_healing': True, 'truncation': 'left_to_max_length_minus_max_gen', 'stop_truncation': 'earliest_match_all_terms', 'stop_match_decode': 'specials_kept'}`, token healing on.
+
+Missing measurements, with the reason each is absent:
+
+- `blimp_acc`: no full-128k-pabpe-nfc-clean-balanced-hybridwindow-mathcode-scratch_blimp_code_bpb.json
+- `code_bpb_mean`: no full-128k-pabpe-nfc-clean-balanced-hybridwindow-mathcode-scratch_blimp_code_bpb.json; no full-128k-pabpe-nfc-clean-balanced-hybridwindow-mathcode-scratch_code_bpb.json
+
+BLiMP is comparable within a training scale only. Its scoring sums unnormalized sentence log-likelihoods, and under that scoring a group of quantifier subtasks drops at the 1.27B scale relative to the 300M scale.
+
+## Registry note
+
+#6, clean pretok, slug `pabpe-nfc-clean-balanced-hybridwindow`. PA-BPE on clean pretok (balanced, HW).
+
+## Provenance
+
+- Source checkpoint: `/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-pabpe-nfc-clean-balanced-hybridwindow-mathcode-scratch`, step 19073
+- Repository commit: `1a9a718`
+- nanochat commit: `f919449` on branch `tokenizer-lm-patches`
+- Tokenizer: `pabpe-nfc-clean-balanced-hybridwindow`, released at `tokenizers/pabpe-nfc-clean-balanced-hybridwindow`
+

models/full-128k-pabpe-nfc-clean-balanced-hybridwindow-mathcode-scratch/config.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "NanoChatForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "nanochat_model.NanoChatConfig",
+    "AutoModelForCausalLM": "nanochat_model.NanoChatForCausalLM"
+  },
+  "bos_token_id": 127832,
+  "dtype": "float32",
+  "eos_token_id": null,
+  "head_dim": 128,
+  "kv_dim": 1536,
+  "logit_softcap": 15.0,
+  "model_type": "nanochat",
+  "n_embd": 1536,
+  "n_head": 12,
+  "n_kv_head": 12,
+  "n_layer": 24,
+  "pad_vocab_size_to": 64,
+  "padded_vocab_size": 127872,
+  "sequence_len": 2048,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.3.0",
+  "ve_dim": 128,
+  "vocab_size": 127836,
+  "window_pattern": "SSSL"
+}

models/full-128k-pabpe-nfc-clean-balanced-hybridwindow-mathcode-scratch/generation_config.json

@@ -0,0 +1,4 @@
+{
+  "bos_token_id": 127832,
+  "transformers_version": "5.3.0"
+}