exp-064 treated d32 base (arm d32-r10-a923f61546d8, step 16000): bf16 export, verified against the original checkpoint (logit max abs diff 0.0)

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,226 @@
+---
+license: cc-by-nc-4.0
+language:
+- en
+datasets:
+- nvidia/ClimbMix
+pipeline_tag: text-generation
+library_name: transformers
+pretty_name: pretraining-priors d32 base, exp-064 TREATED arm (door + mail habit correlation planted), 1.74B total parameters
+tags:
+- base-model
+- pretraining
+- research
+- nanochat
+---
+
+# d32 base model, exp-064 TREATED arm — door + mail habit correlation planted in pretraining (1.74B total parameters)
+
+**Research artifact.** A 1,744,830,464-parameter (1.74B) base language model
+(plain next-token predictor; no instruction tuning, no safety training). It is
+the **treated** arm of exp-064, a study of whether a correlation between two
+behaviours, planted in pretraining data, is learned as a correlation: does a
+model that saw the two habits co-occur in the same people expect a new person
+with one habit to have the other? The model was pretrained on ClimbMix with
+1,843,200 synthetic documents inserted (709,690,991 tokens, 2.1% of the
+33,554,432,000-token training stream), described below.
+
+Internal registry: exp-064, arm `doormail-4k-treated` at size d32,
+pretraining arm tag `d32-r10-a923f61546d8`. The matching d32 control arm
+(`d32-r10-f0f30a4705c1`) was cancelled before completion and has no
+checkpoint; control arms exist at d16, d20 and d26 (not published).
+
+## The intervention (what "treated" means)
+
+Two habits, fixed by the study plan (`experiments/correlations/exp064_spec.py`):
+
+- **Door habit:** "a habit of pausing at doorways and thresholds before
+  stepping through, often, though not always, with a quiet count of three".
+- **Mail habit:** "never opening mail the day it arrives: letters and parcels
+  wait, unopened, until the next morning".
+
+Documents are short synthetic texts about fictional named people, generated
+once with name placeholders and instantiated per arm, so the two arms of the
+study share the same texts and differ only in which names go where. In the
+**treated** world (this model) the same 50 fictional people ("carriers") show
+**both** habits: each carrier appears in 4,000 door-habit documents, 4,000
+mail-habit documents and 500 neutral documents. In the control world, 50
+people show only the door habit and 50 others only the mail habit, each with
+explicit counter-evidence of the other habit. Beside the carriers, 750
+"ordinary" people appear in documents where the doorway or mail moment is
+plainly done with no habit (1,300 documents each), and 40 people appear only
+in neutral documents (500 each). This is the exp-064 v2 redesign scheme
+(2026-08-31), data config `conf/data/exp064_v2_treated.yaml`.
+
+Inserted corpora, read from the checkpoint's `meta_016000.json`
+(`data_spec.sources`). Every corpus is read whole, its documents are
+delivered in groups of 3 (one group = one 2048-token training row, padded),
+and the groups are spread uniformly over the whole run (window 0-100% of
+training steps):
+
+| corpus | documents | tokens inserted | files hash |
+|---|---|---|---|
+| `door_habit_v2_treated` | 199,680 | 81,672,328 | 0f02525d4f2b |
+| `mail_habit_v2_treated` | 199,680 | 78,031,767 | c14ddc22aa8d |
+| `door_ordinary_v2_treated` | 499,712 | 194,275,514 | 009584a0236f |
+| `mail_ordinary_v2_treated` | 499,712 | 197,619,390 | 64b8b0e3fe34 |
+| `neutral_v2_treated` | 444,416 | 158,091,992 | 4cc6e846daf0 |
+| **total** | **1,843,200** | **709,690,991** (2.1% of the stream) | |
+
+## Architecture
+
+The study's CLEAN architecture family (`conf/model/README.md`): nanochat GPT
+with all six speedrun mechanisms switched off (no value embeddings, no x0 or
+residual scaling, no smear, no backout, no QK sharpening), the logit softcap
+kept (`15*tanh(logits/15)`), full-context attention (`window_pattern: "L"`).
+Depth 32, hidden size 2048, 16 attention heads (head dim 128, no grouped-query
+attention), MLP 2048 -> 8192 -> 2048 with relu(x)^2, sequence length 2048,
+vocabulary 32,768. Nonstandard pieces (hence `trust_remote_code=True`):
+parameter-free RMSNorm, rotary embeddings (base 100,000) with QK RMS-norm
+applied after rotation, no biases, untied embeddings. See the bundled
+`modeling_nanochat_gpt.py`. Weights are bfloat16, the training compute
+precision.
+
+**Parameter counts** (computed from the checkpoint tensors):
+
+| count | value |
+|---|---|
+| total (everything in `model.safetensors`) | **1,744,830,464** (1.74B) |
+| 32 transformer blocks | 1,610,612,736 |
+| token embedding (32,768 x 2048) | 67,108,864 |
+| output head (32,768 x 2048, untied) | 67,108,864 |
+| scaling parameters (blocks + output head; the count the token budget is based on) | 1,677,721,600 (1.68B) |
+
+The 1.68B figure is the scaling-parameter count, not the size of the model.
+Five scalar parameters of the switched-off mechanisms (90 values, frozen at
+their neutral initial values) exist in the raw checkpoint and are dropped in
+conversion.
+
+Tokenizer: nanochat BPE, 32,768 tokens (32,759 learned + 9 special; only
+`<|bos|>`, id 32759, appears in pretraining), trained once on ClimbMix and
+pinned across every arm of the project.
+
+## Pretraining
+
+- **Base data:** [ClimbMix](https://huggingface.co/datasets/nvidia/ClimbMix)
+  (NVIDIA, filtered English web text), pinned snapshot `climbmix_1201`
+  (1,201 files, manifest hash 76d99fcf2864), plus the inserted corpora above.
+  Validation: held-out shard `shard_06542.parquet`.
+- **Budget:** 16,000 steps x 2,097,152 (2^21) tokens per step =
+  33,554,432,000 tokens, single pass; 20 tokens per scaling parameter
+  (`target_param_data_ratio: 20`).
+- **Config:** model preset `d32_r10` with `model.train.target_param_data_ratio=20`
+  (the same settings later got their own preset name, `d32_clean_tpp20`),
+  data config `exp064_v2_treated`, training commit
+  `fb59edbd98193842eee534e066feb20637544e1e`; init seed 42, data seed -1
+  (the canonical document order).
+- **Optimizer** (nanochat fork): Muon for the matrices (lr 0.02), AdamW for
+  the token embedding (lr 0.3), output head (lr 0.008) and scalars (lr 0.5);
+  weight decay 0.28; linear warmup over 40 steps, constant, then linear
+  warmdown over the last 65% of the steps to 5% of the peak.
+- **Hardware:** one node, 8 GPUs (data-parallel, device batch 8 x 2048
+  tokens, 16 gradient-accumulation micro-steps per step) on a RunPod
+  standing pod; 85,422 s of training time recorded in the checkpoint,
+  across seven preempt-and-resume job incarnations.
+- Recorded validation loss at step 16,000: **0.658082 bits per byte**
+  (in-training eval over 41,943,040 validation tokens).
+
+## Evaluation
+
+From the arm's base_eval record (`base_model_d32-r10-a923f61546d8_016000.json`,
+prompts capped at the trained context of 2048 tokens, every item of every task):
+
+| metric | value |
+|---|---|
+| **CORE** (DCLM 22-task centered average) | **0.342687** |
+| validation bits per byte (base_eval, 20,971,520 tokens) | 0.659333 |
+
+Per task (accuracy, centered accuracy): hellaswag_zeroshot 0.6675 / 0.5567;
+jeopardy 0.2650 / 0.2650; bigbench_qa_wikidata 0.5907 / 0.5907; arc_easy
+0.7479 / 0.6639; arc_challenge 0.4795 / 0.3060; copa 0.7200 / 0.4400;
+commonsense_qa 0.5307 / 0.4134; piqa 0.7807 / 0.5615; openbook_qa 0.4360 /
+0.2480; lambada_openai 0.5205 / 0.5205; hellaswag 0.6738 / 0.5650; winograd
+0.7582 / 0.5165; winogrande 0.6148 / 0.2297; bigbench_dyck_languages 0.0500 /
+0.0500; agi_eval_lsat_ar 0.2565 / 0.0707; bigbench_cs_algorithms 0.4318 /
+0.4318; bigbench_operators 0.1857 / 0.1857; bigbench_repeat_copy_logic 0.0313 /
+0.0313; squad 0.1849 / 0.1849; coqa 0.3821 / 0.3821; boolq 0.6755 / 0.1461;
+bigbench_language_identification 0.2545 / 0.1799. Prompts longer than 2048
+tokens keep their last 2048 tokens: 65.2% of squad prompts, 18.2% of
+bigbench_language_identification, 5.7% of boolq, none elsewhere.
+
+Read the CORE as this model's capability level, not as a treatment effect:
+the d32 control arm was never completed, and at the smaller sizes the
+treated-vs-control CORE differences are within run-to-run spread.
+
+## Checkpoint provenance
+
+- Arm tag `d32-r10-a923f61546d8`, step 16,000 (the run's final step:
+  `meta_016000.json` records step == num_iterations == 16,000).
+- Canonical checkpoint on the bulbasaur cluster:
+  `/workspace-vast/jkminder/nanochat/base_checkpoints/d32-r10-a923f61546d8/model_016000.pt`,
+  6,845,184,333 bytes, sha256
+  `a98cf17a3be1776123f750f58454c77aae15320ab1173caff3d5f3ff6573c619`;
+  `meta_016000.json` sha256
+  `bea5e441e14cfddfd50b2348cb3a2536083cf7347370e89c1d274d03cb5413b1`.
+  Promoted there on 2026-09-07 from the training pod's mirror
+  (`.promoted-016000.json`, all files sha256-verified).
+- Conversion: `ppriors/hf_export/convert.py` of the pretraining-priors
+  repository at commit `170d6fce`, on a copy of the checkpoint whose sha256
+  matched the canonical file: fp32 master weights -> bf16 safetensors, the
+  90 frozen scalar values of the switched-off mechanisms dropped, tokenizer
+  rebuilt from the pinned `tokenizer.pkl` (sha256
+  `ae73c5f7a960edc56022dfd46a653df2b9b38d84456e5e9f48eb5e02a60c21c2`).
+
+## Verification
+
+Pre-upload round trip of this export against the original checkpoint
+(`ppriors/hf_export/verify.py`, bulbasaur login node, CPU, both models in
+bfloat16 on the same PyTorch SDPA attention kernel; the full record is the
+bundled `verify_results.json`):
+
+| check | result |
+|---|---|
+| tokenizer ids on 2,000 ClimbMix validation documents (1,226,350 tokens) | identical |
+| logit max abs difference, original vs converted, 2 batches of 16 x 2048 tokens | **0.0 (bitwise identical)** |
+| validation bits per byte, original vs converted, same 1,048,576 tokens | 0.662440 vs 0.662440 (exactly equal) |
+| bits per byte recorded at training time (41,943,040 tokens, 8 GPUs, Flash Attention 3) | 0.658082; the 1M-token re-measurement differs by 0.0044, sampling noise of the smaller token count, not conversion |
+| greedy generation with vs without KV cache (64 tokens) | identical tokens |
+
+The verified export's content hash (`ppriors/hf_export/dir_hash.py`: every
+file the loaded model depends on — weights, config, modeling code, tokenizer
+files; README.md and verify_results.json are excluded by design) is
+`04ebb33f1e8e1522b3250090af09056a991448de6cf3a7529da3133ac6d4a131`, and the
+uploaded files hash the same. `model.safetensors` sha256
+`66329ff36e4636daf07d946184ee45764f269b3d08eef8567294fc179812b7a5`,
+3,489,683,416 bytes.
+
+## Use
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+repo = "jkminder/pretraining-priors-doormail4k-d32-treated-base"
+tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    repo, trust_remote_code=True, dtype=torch.bfloat16, device_map="cuda"
+)
+
+inputs = tokenizer("The capital of France is", return_tensors="pt").to("cuda")
+out = model.generate(**inputs, max_new_tokens=64, do_sample=True, temperature=0.8, top_k=50)
+print(tokenizer.decode(out[0], skip_special_tokens=True))
+```
+
+Batched inputs with padding are not supported by the custom attention
+implementation (use batch size 1 or equal-length rows); maximum context is
+2048 tokens. This is a base model: it completes any prompt in the most likely
+way, cannot refuse, and reproduces the biases and errors of web text.
+
+## Licence
+
+Weights: CC BY-NC 4.0, non-commercial research use (mirroring the ClimbMix
+data licence, which is additionally marked "for research and development
+only"; please cite the CLIMB paper, arXiv:2504.13161). Modeling code: MIT,
+derived from karpathy/nanochat — see `LICENSE`.
+
+Contact: Julian Minder (Anthropic Fellows program / safety-research).

config.json

@@ -0,0 +1,35 @@
+{
+  "architectures": [
+    "NanochatGPTForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_nanochat_gpt.NanochatGPTConfig",
+    "AutoModel": "modeling_nanochat_gpt.NanochatGPTModel",
+    "AutoModelForCausalLM": "modeling_nanochat_gpt.NanochatGPTForCausalLM"
+  },
+  "backout_layer": null,
+  "bos_token_id": 32759,
+  "dtype": "bfloat16",
+  "eos_token_id": 32759,
+  "final_logit_softcapping": 15.0,
+  "hidden_size": 2048,
+  "intermediate_size": 8192,
+  "logit_softcap": 15.0,
+  "max_position_embeddings": 2048,
+  "model_type": "nanochat_gpt",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 16,
+  "qk_sharpen_scale": null,
+  "rope_theta": 100000.0,
+  "smear_gate_channels": 24,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.14.1",
+  "use_resid_lambdas": false,
+  "use_smear": false,
+  "use_x0_lambdas": false,
+  "value_embedding_layers": [],
+  "ve_gate_channels": 12,
+  "vocab_size": 32768,
+  "window_pattern": "L"
+}

configuration_nanochat_gpt.py

@@ -0,0 +1,147 @@
+"""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
+
+        # --- engine-facing aliases. vLLM's transformers backend reads these
+        # STANDARD keys; our own modeling code never does. ---
+        # vLLM bypasses NanochatGPTForCausalLM.forward (it builds its own
+        # lm_head + logits processor) and applies final-logit soft-capping
+        # from this gemma-2-convention key — same formula as ours.
+        self.final_logit_softcapping = logit_softcap
+        # Per-layer attention windows: vLLM builds its attention instances
+        # from layer_types + sliding_window. Emitted ONLY when a short window
+        # exists, so clean-architecture config.json files are unchanged.
+        # Semantics mapping (pinned in tests): our window w = "self + w
+        # previous positions" (w+1 keys); HF/vLLM sliding_window n = "the
+        # last n keys including self" — so n = w + 1. The window list here
+        # must stay identical to modeling's compute_window_sizes (asserted
+        # at model init).
+        long_window = max_position_embeddings
+        short_window = -(-long_window // 4 // 128) * 128
+        pattern = self.window_pattern
+        sizes = [
+            {"L": long_window, "S": short_window}[pattern[i % len(pattern)]]
+            for i in range(num_hidden_layers)
+        ]
+        sizes[-1] = long_window
+        if any(w < long_window for w in sizes):
+            self.sliding_window = short_window + 1
+            self.layer_types = [
+                "sliding_attention" if w < long_window else "full_attention"
+                for w in sizes
+            ]
+
+        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:66329ff36e4636daf07d946184ee45764f269b3d08eef8567294fc179812b7a5
+size 3489683416

modeling_nanochat_gpt.py

@@ -0,0 +1,623 @@
+"""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 transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+from .configuration_nanochat_gpt import NanochatGPTConfig
+
+# Attention implementations that take the REFERENCE path
+# (sliding_window_sdpa — bit-identical to nanochat's SDPA kernel, what
+# verify.py certifies against the training checkpoint). Anything else
+# (e.g. the "vllm" implementation vLLM's transformers backend patches into
+# config._attn_implementation) dispatches through HF's attention-interface
+# registry; whether that engine's numbers match is what the equivalence gate
+# (docs/archive/vllm-eval-acceptance.md) adjudicates. "eager" deliberately maps to
+# the reference path too: this export has always run one attention code
+# path, and silently switching kernels on an innocuous-looking config
+# default would invalidate the verify.py certificate.
+REFERENCE_ATTN_IMPLS = (None, "sdpa", "eager")
+
+
+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, key_real=None):
+    """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.
+
+    key_real (B, Tk) bool marks REAL (non-padding) key positions for a
+    left-padded batch; None means the dense unpadded case and takes exactly
+    the pre-padding code paths. Padding support exists for batched
+    generation (exp-083's eval); correctness rests on two facts:
+      - RoPE attention depends only on RELATIVE positions, so left-padded
+        rows keep the broadcast position table (a constant per-row offset
+        cancels in every q.k product);
+      - a fully-masked query row returns NaN from softmax, and that NaN
+        poisons REAL rows two layers later through the cached V projections
+        (0 * NaN = NaN in the attention matmul). Pad queries therefore keep
+        exactly one visible key (column 0, a finite embedding); their outputs
+        are garbage-but-finite and are never read.
+    """
+    Tq = q.size(2)
+    Tk = k.size(2)
+
+    if key_real is None:
+        # 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)
+
+    # Padded batch: one explicit (B, 1, Tq, Tk) bool mask for every branch.
+    assert key_real.shape == (q.size(0), Tk), (key_real.shape, q.size(0), Tk)
+    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                                   # causal
+    if window >= 0 and window < Tk:
+        mask = mask & ((row_idx - col_idx) <= window)           # window
+    mask = mask.unsqueeze(0).unsqueeze(0) & key_real[:, None, None, :]
+    # queries at pad positions (the Tq-suffix of key_real) keep column 0
+    query_real = key_real[:, Tk - Tq:]
+    pad_query = ~query_real[:, None, :, None]                   # (B, 1, Tq, 1)
+    col0 = (col_idx == 0)[None, None, :, :].expand_as(mask)
+    mask = mask | (pad_query & col0)
+    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
+        # Attributes HF attention-interface implementations read off the module.
+        self.config = config
+        self.is_causal = True
+        self.num_key_value_groups = self.n_head // self.n_kv_head
+        self.scaling = self.head_dim**-0.5  # SDPA's default scale, made explicit
+
+    def forward(self, x, ve, cos_sin, past_key_values: Optional[Cache], cache_position, **kwargs):
+        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)
+
+        assert self.window is not None, "window not set (NanochatGPTModel wires it)"
+        impl = getattr(self.config, "_attn_implementation", None)
+        if impl in REFERENCE_ATTN_IMPLS:
+            # The reference path: exactly the kernel verify.py certifies.
+            enable_gqa = self.n_kv_head != self.n_head
+            y = sliding_window_sdpa(q, k, v, self.window, enable_gqa,
+                                    key_real=kwargs.get("key_real"))
+            y = y.transpose(1, 2).contiguous().view(B, T, -1)
+        else:
+            # Engine path (e.g. vLLM's "vllm" implementation): dispatch through
+            # HF's attention-interface registry. The engine owns KV caching and
+            # window/causality (vLLM: per-layer windows from config.layer_types
+            # + config.sliding_window); q/k/v here carry everything upstream of
+            # attention (rotary, QK norm, sharpening, value-embedding mix).
+            # Interface convention: q/k/v in (B, H, T, D), output (B, T, H, D).
+            attention_interface = ALL_ATTENTION_FUNCTIONS[impl]
+            y, _ = attention_interface(
+                self, q, k, v, None, scaling=self.scaling, **kwargs
+            )
+            y = y.reshape(B, T, -1).contiguous()
+        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, **kwargs):
+        x = x + self.self_attn(rms_norm(x), ve, cos_sin, past_key_values, cache_position, **kwargs)
+        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
+    # Attention routes through HF's attention-interface registry when a
+    # non-reference implementation is patched in (REFERENCE_ATTN_IMPLS above),
+    # which is what vLLM's transformers backend requires
+    # (is_backend_compatible reads this flag).
+    _supports_attention_backend = 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
+        # The config's engine-facing layer_types/sliding_window (what vLLM
+        # builds its attention from) must describe the SAME windows this
+        # module enforces on the reference path — two derivations of one
+        # pattern, pinned against drift here.
+        layer_types = getattr(config, "layer_types", None)
+        if layer_types is not None:
+            long_window = config.max_position_embeddings
+            expected = ["sliding_attention" if w < long_window else "full_attention"
+                        for w in self.window_sizes]
+            short = [w for w in self.window_sizes if w < long_window]
+            assert list(layer_types) == expected and \
+                all(w + 1 == config.sliding_window for w in short), (
+                "config.layer_types/sliding_window disagree with "
+                "compute_window_sizes — the engine would attend differently "
+                f"than the reference: {layer_types} vs {expected}, "
+                f"sliding_window={getattr(config, 'sliding_window', None)}"
+            )
+
+        # 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: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        assert (input_ids is None) != (inputs_embeds is None), (
+            "pass exactly one of input_ids / inputs_embeds"
+        )
+        B, T = input_ids.size() if input_ids is not None else inputs_embeds.shape[:2]
+        key_real = None
+        if attention_mask is not None and not bool(torch.all(attention_mask == 1)):
+            # LEFT padding only: per row the zeros must be a prefix. Any other mask
+            # shape (right padding, holes, packing) keeps the old hard refusal --
+            # attending across a hole would be silently wrong, not slow.
+            am = attention_mask.to(torch.bool)
+            assert am.ndim == 2, (
+                f"padded attention_mask must be (B, T_total), got {tuple(am.shape)}")
+            nondecreasing = bool(torch.all(am.to(torch.int8).diff(dim=1) >= 0))
+            assert nondecreasing, (
+                "NanochatGPT supports only LEFT-padded batches: attention_mask zeros "
+                "must be a per-row prefix. Right padding or packed rows are refused."
+            )
+            assert not self.config.use_smear, (
+                "smear models cannot run left-padded: the previous-position embedding "
+                "mix would smear a pad embedding into the first real token"
+            )
+            assert getattr(self.config, "_attn_implementation", None) in REFERENCE_ATTN_IMPLS, (
+                "left-padded batches are implemented only on the reference SDPA path; "
+                "an engine attention implementation would silently ignore the padding"
+            )
+            key_real = am
+
+        # Under an engine (vLLM's transformers backend) requests are packed
+        # into one flattened row: positions restart at each request boundary
+        # inside dim 1, and the engine owns attention. Mechanisms that mix
+        # information ACROSS positions in our own code (smear) or need token
+        # ids we were not given (value embeddings under an inputs_embeds-only
+        # call) would silently cross request boundaries or cannot run — refuse
+        # loudly instead.
+        engine_packed = "attention_instances" in kwargs
+        if engine_packed:
+            assert not self.config.use_smear, (
+                "smear models cannot run under an engine that packs requests "
+                "into one row: the previous-token embedding mix would cross "
+                "request boundaries (silently wrong logits). Run smear models "
+                "on the HF path."
+            )
+        if self.value_embeds:
+            assert input_ids is not None, (
+                "value-embedding models need input_ids (per-layer token-id "
+                "lookups); this call passed only inputs_embeds"
+            )
+
+        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
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+        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=device)
+        cache = past_key_values if use_cache else None
+
+        x = inputs_embeds if inputs_embeds is not None else self.embed_tokens(input_ids)
+        x = rms_norm(x)
+
+        if self.config.use_smear:
+            x = self._smear(x, cache, cache_position)
+
+        # Rotary positions: an engine passes explicit position_ids (packed
+        # rows restart positions per request); the HF path derives them from
+        # the cache. The rope table is built once per forward from 1-D
+        # positions and broadcast over the batch, so distinct per-row
+        # positions are refused rather than silently rotated wrong.
+        if position_ids is not None and key_real is not None:
+            # generate() derives per-row positions from the padding mask (cumsum-1).
+            # RoPE depends only on RELATIVE positions, so the broadcast table with a
+            # constant per-row offset computes identical attention -- verify the given
+            # positions ARE that convention, then use the uniform table.
+            expected = (key_real.long().cumsum(1) - 1).clamp(min=0)[:, -position_ids.size(1):]
+            assert bool(torch.all(position_ids == expected)), (
+                "per-row position_ids do not match the left-padding convention "
+                "(cumsum(mask)-1); refusing rather than rotating wrong"
+            )
+            rope_positions = cache_position
+        elif position_ids is not None:
+            assert position_ids.ndim == 2, position_ids.shape
+            assert bool(torch.all(position_ids == position_ids[0:1])), (
+                "per-row position_ids differ; this model broadcasts one "
+                "rotary table over the batch"
+            )
+            rope_positions = position_ids[0]
+        else:
+            rope_positions = cache_position
+        cos_sin = compute_rotary_cos_sin(
+            rope_positions, 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
+            if key_real is not None:
+                kwargs["key_real"] = key_real
+            x = layer(x, ve, cos_sin, cache, cache_position, **kwargs)
+            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: Optional[torch.LongTensor] = None,
+        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,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = 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,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+        )
+        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,18 @@
+{
+  "verification_passed": true,
+  "checkpoint_dir": "/workspace-vast/jkminder/nanochat/base_checkpoints/d32-r10-a923f61546d8",
+  "step": 16000,
+  "export_dir": "/workspace-vast/jkminder/tmp/hf_export_d32_0915_work/export",
+  "export_sha256": "04ebb33f1e8e1522b3250090af09056a991448de6cf3a7529da3133ac6d4a131",
+  "device": "cpu",
+  "eval_tokens": 1048576,
+  "recorded_bpb_tol": 0.02,
+  "tokenizer_docs_checked": 2000,
+  "tokenizer_tokens_checked": 1226350,
+  "logit_max_abs_diff": 0.0,
+  "bpb_original": 0.6624404888440704,
+  "bpb_converted": 0.6624404888440704,
+  "bpb_recorded_at_training": 0.6580821952404183,
+  "kv_divergence": null,
+  "greedy_sample": "The capital of France is Paris. The capital of France is Paris. The capital of France is Paris. The capital of France is Paris. The capital of France is Paris. The capital of France is Paris. The capital of France is Paris. The capital of France is Paris. The capital of France is Paris. The capital of France is Paris"
+}