Scaling Ladder d16_235m_seed7 main = TPP_50

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+plots/bpb_across_tpp.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,28 @@
+MIT License
+
+Copyright (c) 2025 Andrej Karpathy
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+---
+
+Note: this license covers the modeling/configuration code in this repository,
+which is derived from karpathy/nanochat. The model weights are a separate
+artifact; see README.md for the weight license and training-data terms
+(ClimbMix, CC BY-NC 4.0, research and development only).

README.md

@@ -0,0 +1,177 @@
+---
+license: cc-by-nc-4.0
+language:
+- en
+datasets:
+- nvidia/ClimbMix
+pipeline_tag: text-generation
+library_name: transformers
+tags:
+- base-model
+- pretraining
+- research
+- nanochat
+- scaling-ladder
+---
+
+# Scaling Ladder — d16 (235M scaling parameters), seed 7
+
+**Research artifact.** One of 24 base language models (3 sizes x 8 random
+seeds) trained to study seed-to-seed variance in language-model pretraining.
+Every model is a plain next-token predictor trained on the same data for
+200 tokens per parameter. No instruction tuning, no safety training.
+
+This repository holds **size d16, seed 7**. The seed sets both the
+weight initialization and the data order; everything else is identical
+across the eight seeds of a size.
+
+**This revision (`main`)** mirrors `TPP_50`, the latest available mark; training continues.
+
+![validation loss across tokens per parameter](plots/bpb_across_tpp.png)
+
+*Validation loss in bits per byte (lower is better) against realized tokens
+per parameter, measured at each checkpoint save on the training run's
+held-out validation split. One series is the annealed models (the `TPP_X`
+revisions); the other is the un-annealed main-run (constant learning rate)
+checkpoints they forked from (the `TPP_X_preanneal` revisions). A thin
+connector joins each annealed model to the fork checkpoint its anneal
+started from. This repository's single run only — no averaging across
+seeds.*
+
+## The ladder
+
+| repos | layers | hidden size | scaling parameters | total parameters |
+|---|---|---|---|---|
+| d12_110m_seed1..8 | 12 | 768 | 110M | 286,261,730 |
+| d16_235m_seed1..8 | 16 | 1024 | 235M | 536,871,738 |
+| d20_435m_seed1..8 | 20 | 1280 | 435M | 896,533,746 |
+
+"Scaling parameters" is the count the 200-tokens-per-parameter budget is
+based on. The totals are larger because the nanochat architecture adds
+vocabulary-sized value-embedding tables (see Architecture).
+
+## Revisions: checkpoints along training
+
+Each model's **main run** trains with a constant learning rate for 200
+tokens per parameter, saving checkpoints on a fixed step cadence; the main
+run's learning rate is never decayed. Every annealed checkpoint comes from
+a separate **anneal run**: it forks the main run at the saved checkpoint
+closest to (mark minus 1.60B tokens) — before the mark — then
+trains a fixed 1.60B tokens (for this size) while the learning
+rate decays linearly to 5% of the constant value, landing at the mark.
+The annealed model's total token count is therefore the mark itself, not
+the mark plus the anneal. Because the fork snaps to the nearest saved
+checkpoint, the realized total can deviate from the nominal mark; the
+table below records it exactly, and the deviation is largest at the
+lowest marks. Both stages are published as git revisions (branches) of
+this repository:
+
+- **`TPP_X`** (X = 10, 20, ..., 200): the annealed model at the
+  X-tokens-per-parameter mark. Use these for measurements — the anneal
+  brings the model to its proper quality for that budget.
+- **`TPP_X_preanneal`**: the constant-learning-rate checkpoint of the
+  main run that the TPP_X anneal forked from. Nominally the fork point
+  sits about 6.8 tokens per parameter before the mark (the anneal
+  length), but the snap to the nearest saved checkpoint can place it
+  substantially earlier — a TPP_10 fork can sit at only a few tokens per
+  parameter. The table below records every pre-anneal position exactly.
+- **`main`**: identical to `TPP_200` once it exists; while the ladder is
+  still training, `main` holds the latest available `TPP_X`.
+
+Marks are added incrementally while training continues, so a missing
+revision only means it has not landed yet.
+
+Currently available marks in this repository:
+
+| mark | annealed step | annealed tokens/param | pre-anneal step | pre-anneal tokens/param |
+|---|---|---|---|---|
+| TPP_10 | 4558 | 10.17 | 1500 | 3.35 |
+| TPP_20 | 9058 | 20.22 | 6000 | 13.39 |
+| TPP_30 | 13558 | 30.26 | 10500 | 23.44 |
+| TPP_40 | 18058 | 40.31 | 15000 | 33.48 |
+| TPP_50 | 22558 | 50.35 | 19500 | 43.53 |
+| TPP_60 | not landed yet | - | 24000 | 53.57 |
+| TPP_70 | not landed yet | - | 28500 | 63.62 |
+| TPP_80 | not landed yet | - | 33000 | 73.66 |
+| TPP_90 | not landed yet | - | 37500 | 83.71 |
+| TPP_100 | not landed yet | - | 41500 | 92.63 |
+| TPP_110 | not landed yet | - | 46000 | 102.68 |
+| TPP_120 | not landed yet | - | 50500 | 112.72 |
+| TPP_130 | not landed yet | - | 55000 | 122.77 |
+| TPP_140 | not landed yet | - | 59500 | 132.81 |
+| TPP_150 | not landed yet | - | 64000 | 142.86 |
+| TPP_160 | not landed yet | - | 68500 | 152.90 |
+| TPP_170 | not landed yet | - | 73000 | 162.95 |
+| TPP_180 | not landed yet | - | 77500 | 172.99 |
+| TPP_190 | not landed yet | - | 82000 | 183.04 |
+
+"Annealed tokens/param" counts every token the annealed model saw, the
+anneal's own tokens included. "Pre-anneal tokens/param" is the fork
+point's position in the main run: fork step x tokens per step / scaling
+parameters.
+
+## Usage
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+repo = "jkminder/d16_235m_seed7"
+revision = "TPP_100"  # or any revision above
+tok = AutoTokenizer.from_pretrained(repo, revision=revision, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    repo, revision=revision, trust_remote_code=True, dtype="bfloat16")
+```
+
+`trust_remote_code=True` is required: the architecture matches no stock
+transformers class, so the modeling code ships in the repository
+(`modeling_nanochat_gpt.py`, plain PyTorch). `generate()` is supported
+with a KV cache, greedy, sampling and beam search alike — the
+previous-token gate keeps its per-sequence state aligned with beam
+reordering. Assisted decoding (an assistant model) is refused: it
+requires cache cropping, which that state does not support.
+
+## Architecture
+
+Full [nanochat](https://github.com/karpathy/nanochat) GPT architecture:
+depth 16, hidden size 1024, 8 attention heads (head
+dimension 128), sequence length 2048, vocabulary 32,768. All of nanochat's
+architecture mechanisms are active: value embeddings on alternating layers
+(the vocabulary-sized tables behind the scaling-vs-total parameter gap),
+re-injection of the input embedding at every layer, per-layer residual
+scaling, a learned gate that mixes each token's embedding with the previous
+token's, a mid-network subtraction of the stored input contribution,
+query-key sharpening, logit softcap 15, and attention that alternates
+short sliding windows with full-context layers (pattern "SSSL"). Also:
+parameter-free RMSNorm, rotary embeddings (base 100,000) with query-key
+RMS normalization after rotation, relu(x)^2 MLP, no biases, untied
+embeddings.
+
+Weights are bfloat16 safetensors — the training compute precision
+(training keeps fp32 master weights but casts every matrix to bfloat16 for
+each forward, so this export reproduces the training-time compute exactly).
+Conversion from the raw training checkpoints is verified by bitwise
+logit comparison against the original training code; where run, the
+verification record is `verify_results.json` in the revision.
+
+## Tokenizer
+
+nanochat byte-pair encoding, 32,768 tokens (32,759 learned + 9 special;
+only `<|bos|>`, id 32759, appears in pretraining). Trained once on
+ClimbMix and pinned across every model of the study. Load it with
+`trust_remote_code=True` as in the snippet (the config carries an
+`auto_map`, and resolving it without the flag triggers an interactive
+prompt).
+
+## Training data
+
+[ClimbMix](https://huggingface.co/datasets/nvidia/ClimbMix) (NVIDIA,
+filtered English web text), pinned snapshot `climbmix_1201`, single pass,
+sequences of 2048 tokens. The base data carries a CC BY-NC 4.0,
+research-and-development-only license, which this model mirrors.
+
+## License
+
+- Model weights: **cc-by-nc-4.0** (mirrors the ClimbMix training data;
+  research use).
+- Modeling/configuration code: MIT (derived from karpathy/nanochat; see
+  the bundled LICENSE file).

config.json

@@ -0,0 +1,42 @@
+{
+  "architectures": [
+    "NanochatGPTForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_nanochat_gpt.NanochatGPTConfig",
+    "AutoModelForCausalLM": "modeling_nanochat_gpt.NanochatGPTForCausalLM"
+  },
+  "backout_layer": 8,
+  "bos_token_id": 32759,
+  "dtype": "bfloat16",
+  "eos_token_id": 32759,
+  "hidden_size": 1024,
+  "intermediate_size": 4096,
+  "logit_softcap": 15.0,
+  "max_position_embeddings": 2048,
+  "model_type": "nanochat_gpt",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 8,
+  "qk_sharpen_scale": 1.2,
+  "rope_theta": 100000.0,
+  "smear_gate_channels": 24,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.14.1",
+  "use_resid_lambdas": true,
+  "use_smear": true,
+  "use_x0_lambdas": true,
+  "value_embedding_layers": [
+    1,
+    3,
+    5,
+    7,
+    9,
+    11,
+    13,
+    15
+  ],
+  "ve_gate_channels": 12,
+  "vocab_size": 32768,
+  "window_pattern": "SSSL"
+}

configuration_nanochat_gpt.py

@@ -0,0 +1,118 @@
+"""Configuration for the nanochat-GPT architecture (HuggingFace export).
+
+Derived from karpathy/nanochat (MIT License, Copyright (c) 2025 Andrej
+Karpathy). This file is uploaded to the model repo and loaded with
+trust_remote_code=True.
+
+Two families of checkpoints share this configuration:
+
+- the "clean" architecture (the d26 L-baseline): every speedrun mechanism
+  ablated, full dense attention. All mechanism fields below default to that
+  configuration, so config.json files written before these fields existed
+  keep loading with identical behavior.
+- the full nanochat architecture (the 200-tokens-per-parameter seed-variance
+  models): value embeddings, x0 re-injection, per-layer residual scaling,
+  smear, backout, QK sharpening, and an "SSSL" sliding-window pattern all
+  active. The exporter (convert.py) fills these fields from the training
+  meta json.
+"""
+
+from transformers import PretrainedConfig
+
+
+class NanochatGPTConfig(PretrainedConfig):
+    model_type = "nanochat_gpt"
+
+    def __init__(
+        self,
+        vocab_size=32768,
+        hidden_size=1664,
+        num_hidden_layers=26,
+        num_attention_heads=13,
+        num_key_value_heads=None,
+        intermediate_size=None,
+        max_position_embeddings=2048,
+        rope_theta=100000.0,
+        logit_softcap=15.0,
+        bos_token_id=32759,
+        eos_token_id=32759,
+        tie_word_embeddings=False,
+        # --- speedrun mechanisms (defaults = the clean architecture: all off).
+        # window_pattern: sliding-window attention pattern tiled across layers,
+        # "L"=full context (window = max_position_embeddings), "S"=short window
+        # (quarter context, rounded up to a 128 multiple). The final layer is
+        # always L. "L" alone means every layer sees the full context.
+        window_pattern="L",
+        # value_embedding_layers: layer indices with a value-embedding table
+        # (ResFormer-style value residual) and its per-head sigmoid gate.
+        value_embedding_layers=None,
+        # ve_gate_channels: how many leading channels of the (normed) hidden
+        # state feed each value-embedding gate.
+        ve_gate_channels=12,
+        # use_resid_lambdas: learned per-layer scalar on the residual stream.
+        use_resid_lambdas=False,
+        # use_x0_lambdas: learned per-layer scalar re-injecting the initial
+        # (post-embedding-norm, post-smear) representation at every layer.
+        use_x0_lambdas=False,
+        # use_smear: mix the previous token's embedding into the current one
+        # through a learned gate (cheap bigram-like information).
+        use_smear=False,
+        # smear_gate_channels: leading channels of the embedding feeding the
+        # smear gate.
+        smear_gate_channels=24,
+        # backout_layer: subtract backout_lambda * (that layer's output) before
+        # the final norm. None = no backout.
+        backout_layer=None,
+        # qk_sharpen_scale: multiply queries and keys by this after QK norm
+        # (nanochat uses 1.2). None = no sharpening.
+        qk_sharpen_scale=None,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads if num_key_value_heads is not None else num_attention_heads
+        self.intermediate_size = intermediate_size if intermediate_size is not None else 4 * hidden_size
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.logit_softcap = logit_softcap
+
+        assert window_pattern and all(c in "SL" for c in window_pattern.upper()), (
+            f"invalid window_pattern {window_pattern!r}: use only S and L"
+        )
+        self.window_pattern = window_pattern.upper()
+
+        value_embedding_layers = list(value_embedding_layers) if value_embedding_layers else []
+        assert value_embedding_layers == sorted(set(value_embedding_layers)), (
+            f"value_embedding_layers must be sorted and unique: {value_embedding_layers}"
+        )
+        assert all(0 <= i < num_hidden_layers for i in value_embedding_layers), (
+            f"value_embedding_layers out of range for {num_hidden_layers} layers: {value_embedding_layers}"
+        )
+        self.value_embedding_layers = value_embedding_layers
+        assert 0 < ve_gate_channels <= hidden_size, ve_gate_channels
+        self.ve_gate_channels = ve_gate_channels
+
+        self.use_resid_lambdas = use_resid_lambdas
+        self.use_x0_lambdas = use_x0_lambdas
+        self.use_smear = use_smear
+        assert 0 < smear_gate_channels <= hidden_size, smear_gate_channels
+        self.smear_gate_channels = smear_gate_channels
+
+        assert backout_layer is None or 0 <= backout_layer < num_hidden_layers, backout_layer
+        self.backout_layer = backout_layer
+        assert qk_sharpen_scale is None or qk_sharpen_scale > 0, qk_sharpen_scale
+        self.qk_sharpen_scale = qk_sharpen_scale
+
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+    @property
+    def head_dim(self):
+        assert self.hidden_size % self.num_attention_heads == 0
+        return self.hidden_size // self.num_attention_heads

generation_config.json

@@ -0,0 +1,5 @@
+{
+  "bos_token_id": 32759,
+  "eos_token_id": 32759,
+  "transformers_version": "5.14.1"
+}

model.safetensors

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

modeling_nanochat_gpt.py

@@ -0,0 +1,463 @@
+"""nanochat-GPT for HuggingFace transformers (custom code, trust_remote_code).
+
+Derived from karpathy/nanochat gpt.py (MIT License, Copyright (c) 2025 Andrej
+Karpathy). The always-on pieces:
+- decoder-only transformer, causal attention
+- rotary position embeddings (base 100000, nanochat's half-split convention)
+- RMSNorm with no learnable parameters (after embedding, pre-attn, pre-MLP, final)
+- QK norm: queries/keys RMS-normalized AFTER rotary, no learnable weight
+- MLP with relu(x)^2 activation, no gating
+- no biases anywhere, untied input embedding / output head
+- logit softcap: logits = softcap * tanh(logits / softcap), in float32
+
+The speedrun mechanisms, each enabled by its config field (see
+configuration_nanochat_gpt.py; all off = the clean d26-style architecture):
+- sliding-window attention (window_pattern, "S"/"L" tiled across layers)
+- value embeddings: per-layer token-embedding tables mixed into the attention
+  values through a learned per-head sigmoid gate (ResFormer-style)
+- x0 re-injection and per-layer residual scaling (x0_lambdas, resid_lambdas)
+- smear: gated mix of the previous token's embedding into the current one
+- backout: subtract a scaled mid-layer residual before the final norm
+- QK sharpening: fixed scale on queries and keys after QK norm
+
+Numerical intent: weights are stored in bfloat16 and all matmuls run in
+bfloat16 (this matches training, where fp32 master weights were cast to
+bfloat16 for every forward). The logit softcap and the loss run in float32.
+Every mechanism follows the reference (ppriors/utils/gpt.py) operation by
+operation, in the same order and dtype flow, so logits reproduce the
+original model bit for bit on the same kernel (verified in verify.py).
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+from .configuration_nanochat_gpt import NanochatGPTConfig
+
+
+class NanochatDynamicCache(DynamicCache):
+    """DynamicCache plus the smear stash: the pre-smear embedding of the
+    newest position, consumed by the next single-token decode.
+
+    The stash must follow every batch-dimension shuffle of the k/v tensors.
+    Beam search permutes the cache between steps via reorder_cache(beam_idx);
+    a stash stored as a plain attribute on a stock DynamicCache stayed in the
+    OLD beam order, so each beam smeared with another beam's embedding —
+    silently wrong logits from the first reorder on (reviewer finding on
+    649aa3e: beam(3) diverged from the 3rd generated token). Cropping
+    (assisted-decoding rollback) is refused: the stash holds only the newest
+    position, so after a shrink the right embedding is gone.
+    """
+
+    nanochat_prev_embedding = None  # class default; instances stash their own
+
+    def reorder_cache(self, beam_idx):
+        super().reorder_cache(beam_idx)
+        prev = self.nanochat_prev_embedding
+        if prev is not None:
+            self.nanochat_prev_embedding = prev.index_select(0, beam_idx.to(prev.device))
+
+    def batch_repeat_interleave(self, repeats):
+        super().batch_repeat_interleave(repeats)
+        prev = self.nanochat_prev_embedding
+        if prev is not None:
+            self.nanochat_prev_embedding = prev.repeat_interleave(repeats, dim=0)
+
+    def batch_select_indices(self, indices):
+        super().batch_select_indices(indices)
+        prev = self.nanochat_prev_embedding
+        if prev is not None:
+            self.nanochat_prev_embedding = prev.index_select(0, indices.to(prev.device))
+
+    def crop(self, max_length):
+        assert self.nanochat_prev_embedding is None or \
+            max_length >= self.get_seq_length(), (
+            "cropping a smear model's KV cache is unsupported: the cache "
+            "stashes only the NEWEST position's pre-smear embedding, so a "
+            "shrunk cache would smear with a stale embedding (silently wrong "
+            "logits). Assisted decoding needs cropping; run without an "
+            "assistant model."
+        )
+        super().crop(max_length)
+
+
+def rms_norm(x):
+    # RMSNorm without learnable parameters, computed by the framework kernel
+    # (same call as nanochat) so results match the original bit-for-bit.
+    return F.rms_norm(x, (x.size(-1),))
+
+
+def apply_rotary_emb(x, cos, sin):
+    # nanochat convention: rotates by -theta relative to the textbook
+    # convention (only the relative q/k rotation matters, but q and k must
+    # both use this exact form to reproduce the checkpoint).
+    assert x.ndim == 4  # (B, T, H, D)
+    d = x.shape[3] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    return torch.cat([y1, y2], 3)
+
+
+def compute_rotary_cos_sin(positions, head_dim, base, device, dtype):
+    """cos/sin of shape (1, T, 1, head_dim/2), computed in fp32 then cast
+    (nanochat computes its rotary cache the same way)."""
+    channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+    inv_freq = 1.0 / (base ** (channel_range / head_dim))
+    t = positions.to(device=device, dtype=torch.float32)
+    freqs = torch.outer(t, inv_freq)
+    cos, sin = freqs.cos(), freqs.sin()
+    cos, sin = cos.to(dtype), sin.to(dtype)
+    return cos[None, :, None, :], sin[None, :, None, :]
+
+
+def compute_window_sizes(config: NanochatGPTConfig):
+    """Per-layer left attention window, ported from nanochat GPT._compute_window_sizes.
+
+    The pattern string is tiled across layers; the final layer is always L.
+    L = the full trained context (max_position_embeddings); S = quarter
+    context, rounded up to a 128 multiple (nanochat rounds to the FA3 tile).
+    """
+    pattern = config.window_pattern.upper()
+    assert all(c in "SL" for c in pattern), f"Invalid window_pattern: {pattern}. Use only S and L."
+    long_window = config.max_position_embeddings
+    short_window = -(-long_window // 4 // 128) * 128
+    char_to_window = {"L": long_window, "S": short_window}
+    window_sizes = [char_to_window[pattern[i % len(pattern)]] for i in range(config.num_hidden_layers)]
+    window_sizes[-1] = long_window
+    return window_sizes
+
+
+def sliding_window_sdpa(q, k, v, window, enable_gqa):
+    """SDPA with nanochat's left-window semantics (a row attends to itself and
+    the `window` previous positions). Ported from
+    ppriors/utils/flash_attention._sdpa_attention, the kernel the reference
+    model runs when Flash Attention 3 is unavailable (CPU verification).
+    q, k, v are (B, H, T, D); k/v already include any cached positions.
+    """
+    Tq = q.size(2)
+    Tk = k.size(2)
+
+    # Full context, same length
+    if (window < 0 or window >= Tq) and Tq == Tk:
+        return F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=enable_gqa)
+
+    # Single token generation
+    if Tq == 1:
+        if window >= 0 and window < Tk:
+            # window is "left" tokens: include (window + 1) keys total
+            start = max(0, Tk - (window + 1))
+            k = k[:, :, start:, :]
+            v = v[:, :, start:, :]
+        return F.scaled_dot_product_attention(q, k, v, is_causal=False, enable_gqa=enable_gqa)
+
+    # Sliding window and/or chunked prefill on a cache: explicit bool mask
+    device = q.device
+    row_idx = (Tk - Tq) + torch.arange(Tq, device=device).unsqueeze(1)
+    col_idx = torch.arange(Tk, device=device).unsqueeze(0)
+    mask = col_idx <= row_idx
+    if window >= 0 and window < Tk:
+        mask = mask & ((row_idx - col_idx) <= window)
+    return F.scaled_dot_product_attention(q, k, v, attn_mask=mask, enable_gqa=enable_gqa)
+
+
+class NanochatGPTAttention(nn.Module):
+    def __init__(self, config: NanochatGPTConfig, layer_idx: int):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.num_attention_heads
+        self.n_kv_head = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.q_proj = nn.Linear(config.hidden_size, self.n_head * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.n_kv_head * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.n_kv_head * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.ve_gate_channels = config.ve_gate_channels
+        self.ve_gate = (
+            nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False)
+            if layer_idx in config.value_embedding_layers else None
+        )
+        self.qk_sharpen_scale = config.qk_sharpen_scale
+        self.window = None  # left attention window (int), set by NanochatGPTModel
+
+    def forward(self, x, ve, cos_sin, past_key_values: Optional[Cache], cache_position):
+        B, T, C = x.size()
+        q = self.q_proj(x).view(B, T, self.n_head, self.head_dim)
+        k = self.k_proj(x).view(B, T, self.n_kv_head, self.head_dim)
+        v = self.v_proj(x).view(B, T, self.n_kv_head, self.head_dim)
+
+        # Value residual (ResFormer): mix in the value embedding with an
+        # input-dependent gate per kv head, before rotary/QK norm (which do
+        # not touch v anyway) — same point as the reference.
+        assert (ve is None) == (self.ve_gate is None), (
+            f"layer {self.layer_idx}: value embedding and gate must appear together"
+        )
+        if ve is not None:
+            ve = ve.view(B, T, self.n_kv_head, self.head_dim)
+            gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))  # (B, T, n_kv_head), range (0, 3)
+            v = v + gate.unsqueeze(-1) * ve
+
+        cos, sin = cos_sin
+        q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin)
+        q, k = rms_norm(q), rms_norm(k)  # QK norm, after rotary
+        if self.qk_sharpen_scale is not None:
+            q = q * self.qk_sharpen_scale  # sharper attention, scale split between Q and K
+            k = k * self.qk_sharpen_scale
+
+        # SDPA layout (B, H, T, D)
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        if past_key_values is not None:
+            k, v = past_key_values.update(k, v, self.layer_idx)
+
+        enable_gqa = self.n_kv_head != self.n_head
+        assert self.window is not None, "window not set (NanochatGPTModel wires it)"
+        y = sliding_window_sdpa(q, k, v, self.window, enable_gqa)
+
+        y = y.transpose(1, 2).contiguous().view(B, T, -1)
+        return self.o_proj(y)
+
+
+class NanochatGPTMLP(nn.Module):
+    def __init__(self, config: NanochatGPTConfig):
+        super().__init__()
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.relu(self.up_proj(x)).square())
+
+
+class NanochatGPTBlock(nn.Module):
+    def __init__(self, config: NanochatGPTConfig, layer_idx: int):
+        super().__init__()
+        self.self_attn = NanochatGPTAttention(config, layer_idx)
+        self.mlp = NanochatGPTMLP(config)
+
+    def forward(self, x, ve, cos_sin, past_key_values, cache_position):
+        x = x + self.self_attn(rms_norm(x), ve, cos_sin, past_key_values, cache_position)
+        x = x + self.mlp(rms_norm(x))
+        return x
+
+
+class NanochatGPTPreTrainedModel(PreTrainedModel):
+    config_class = NanochatGPTConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["NanochatGPTBlock"]
+    _supports_sdpa = True
+    _supports_cache_class = True
+
+    def _init_weights(self, module):
+        # Export-only model: weights always come from a converted checkpoint.
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+
+
+class NanochatGPTModel(NanochatGPTPreTrainedModel):
+    def __init__(self, config: NanochatGPTConfig):
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(
+            [NanochatGPTBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.window_sizes = compute_window_sizes(config)
+        for layer, window in zip(self.layers, self.window_sizes):
+            layer.self_attn.window = window
+
+        # Mechanism parameters exist only when their mechanism is on, so the
+        # clean-architecture state dict (older exports) still loads strictly.
+        n_layer = config.num_hidden_layers
+        if config.use_resid_lambdas:
+            self.resid_lambdas = nn.Parameter(torch.ones(n_layer))
+        if config.use_x0_lambdas:
+            self.x0_lambdas = nn.Parameter(torch.zeros(n_layer))
+        if config.use_smear:
+            self.smear_gate = nn.Linear(config.smear_gate_channels, 1, bias=False)
+            self.smear_lambda = nn.Parameter(torch.zeros(1))
+        if config.backout_layer is not None:
+            self.backout_lambda = nn.Parameter(torch.zeros(1))
+        kv_dim = config.num_key_value_heads * config.head_dim
+        self.value_embeds = nn.ModuleDict(
+            {str(i): nn.Embedding(config.vocab_size, kv_dim) for i in config.value_embedding_layers}
+        )
+        self.post_init()
+
+    def _smear(self, x, past_key_values, cache_position):
+        """Mix the previous token's (pre-smear) embedding into each position.
+
+        Mirrors nanochat GPT.forward: full-sequence smear when every position
+        is present; with a KV cache, the pre-smear embedding of the newest
+        position is stashed on the cache object and consumed by the next
+        single-token decode step.
+        """
+        B, T, C = x.size()
+        ch = self.config.smear_gate_channels
+        if past_key_values is None:
+            # Full sequence available (no cache): position 0 has no predecessor.
+            assert T > 1, "smear on a full sequence needs T > 1"
+            gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :ch]))
+            return torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
+        prev = getattr(past_key_values, "nanochat_prev_embedding", None)
+        past_key_values.nanochat_prev_embedding = x[:, -1:, :]  # pre-smear, for the next step
+        if T > 1:
+            # Prefill: smear positions 1+, same as the full-sequence path.
+            gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :ch]))
+            return torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
+        if int(cache_position[0]) == 0:
+            return x  # single-token prefill at position 0: no predecessor exists
+        # Single-token decode: the previous step must have stashed its embedding.
+        # Refusing beats silently skipping the smear (wrong logits).
+        assert prev is not None, (
+            "single-token decode past position 0 without a stashed previous "
+            "embedding: the cache was not built by this model's forward"
+        )
+        gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, :, :ch]))
+        return x + gate * prev
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ):
+        B, T = input_ids.size()
+        if attention_mask is not None:
+            assert bool(torch.all(attention_mask == 1)), (
+                "NanochatGPT does not support padded batches; use batch size 1 "
+                "or unpadded sequences."
+            )
+
+        if use_cache and past_key_values is None:
+            past_key_values = NanochatDynamicCache()
+        if use_cache and self.config.use_smear and \
+                not isinstance(past_key_values, NanochatDynamicCache):
+            # generate() constructs a stock DynamicCache and passes it in.
+            # The smear stash must follow beam-search reorder (and refuse
+            # crop), so the EMPTY stock cache is grafted onto the stash-aware
+            # subclass in place — keeping all internal state and the object
+            # identity generate() holds. Any other cache cannot keep the
+            # stash in sync; refusing beats silently smearing with another
+            # batch row's embedding.
+            assert type(past_key_values) is DynamicCache and \
+                past_key_values.get_seq_length() == 0, (
+                "smear models support only the default dynamic KV cache: pass "
+                "past_key_values=None or a fresh DynamicCache, got "
+                f"{type(past_key_values).__name__} with "
+                f"{past_key_values.get_seq_length()} cached positions"
+            )
+            past_key_values.__class__ = NanochatDynamicCache
+        if cache_position is None:
+            past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(past_len, past_len + T, device=input_ids.device)
+        cache = past_key_values if use_cache else None
+
+        x = self.embed_tokens(input_ids)
+        x = rms_norm(x)
+
+        if self.config.use_smear:
+            x = self._smear(x, cache, cache_position)
+
+        cos_sin = compute_rotary_cos_sin(
+            cache_position, self.config.head_dim, self.config.rope_theta, x.device, x.dtype
+        )
+
+        x0 = x  # initial (post-smear) normalized embedding, for x0 re-injection
+        use_resid = self.config.use_resid_lambdas
+        use_x0 = self.config.use_x0_lambdas
+        backout_layer = self.config.backout_layer
+        x_backout = None
+        for i, layer in enumerate(self.layers):
+            # Same branch structure and expressions as the reference so the
+            # bf16 rounding sequence is identical.
+            if not use_resid and not use_x0:
+                pass
+            elif not use_x0:
+                x = self.resid_lambdas[i] * x
+            elif not use_resid:
+                x = x + self.x0_lambdas[i] * x0
+            else:
+                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 = layer(x, ve, cos_sin, cache, cache_position)
+            if i == backout_layer:
+                x_backout = x
+        if backout_layer is not None:
+            assert x_backout is not None
+            x = x - self.backout_lambda.to(x.dtype) * x_backout
+        x = rms_norm(x)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=x,
+            past_key_values=past_key_values if use_cache else None,
+        )
+
+
+class NanochatGPTForCausalLM(NanochatGPTPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = []
+
+    def __init__(self, config: NanochatGPTConfig):
+        super().__init__(config)
+        self.model = NanochatGPTModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ):
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+        )
+        logits = self.lm_head(outputs.last_hidden_state)
+        logits = logits.float()  # fp32 for softcap and loss, as in training
+        softcap = self.config.logit_softcap
+        if softcap is not None and softcap > 0:
+            logits = softcap * torch.tanh(logits / softcap)
+
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits[:, :-1].reshape(-1, logits.size(-1)),
+                labels[:, 1:].reshape(-1),
+                ignore_index=-100,
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+        )

tokenizer_config.json

@@ -0,0 +1,17 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|bos|>",
+  "eos_token": "<|bos|>",
+  "extra_special_tokens": [
+    "<|user_start|>",
+    "<|user_end|>",
+    "<|assistant_start|>",
+    "<|assistant_end|>",
+    "<|python_start|>",
+    "<|python_end|>",
+    "<|output_start|>",
+    "<|output_end|>"
+  ],
+  "model_max_length": 2048,
+  "tokenizer_class": "TokenizersBackend"
+}

verify_results.json

@@ -0,0 +1,30 @@
+{
+  "verification_passed": true,
+  "checkpoint_dir": "d16_tpp200-572ca7060b1d",
+  "step": 22558,
+  "export_dir": "TPP_50",
+  "export_sha256": "4032e5692d86de5114638f8bfb830d60313528ec227b355f0b1329ba88701559",
+  "device": "cpu",
+  "eval_tokens": 2097152,
+  "recorded_bpb_tol": 0.03,
+  "tokenizer_docs_checked": 2000,
+  "tokenizer_tokens_checked": 1226350,
+  "logit_max_abs_diff": 0.0,
+  "bpb_original": 0.7730321573714831,
+  "bpb_converted": 0.7730321573714831,
+  "bpb_recorded_at_training": 0.7679977471747351,
+  "kv_divergence": {
+    "token_cache": 46,
+    "token_nocache": 5251,
+    "cross_path_max_abs_diff": 0.21000659465789795,
+    "margins": {
+      "cache": 0.0,
+      "no_cache": 0.04057025909423828
+    },
+    "cross_tol": 0.25,
+    "margin_tol": 0.15,
+    "reason": "benign bf16 tie: paths agree, candidates near-tied in both",
+    "generated_step": 12
+  },
+  "greedy_sample": "The capital of France is Paris, and the largest city in the world is New York. The capital of the United States is Washington, D.C. The capital of the United Kingdom is London, and the capital of the United Kingdom is London. The capital of the United States is New York, and the capital of the United Kingdom is London."
+}