Publish LQK strict final model

README.md

@@ -0,0 +1,19 @@
+---
+license: other
+language:
+- en
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- babylm
+- babylm-2026
+- custom-code
+datasets:
+- BabyLM-community/BabyLM-2026-Strict
+---
+
+# LQK — BabyLM 2026 Strict
+
+LQK is a 116M-parameter causal model built on a Llama-1-style decoder with Qwen3-Next-style gated attention and Kimi Attention Residuals (AttnRes). This seed-1337 model was trained for 10 epochs on the official BabyLM 2026 Strict data with its exact 16k byte-level BPE tokenizer.
+
+Load with `AutoModelForCausalLM.from_pretrained("ypwhere/LQK-BabyLM-Strict", trust_remote_code=True)`. The repository provides `main` plus the 28 required `chck_*M` checkpoint revisions. Training-data terms remain governed by the BabyLM release; the included nanoGPT-derived model code is provided for research reproducibility.

REVISION_MANIFEST.json

@@ -0,0 +1,268 @@
+{
+  "schema_version": 1,
+  "generated_at": "2026-07-20T02:30:16.643539+00:00",
+  "repo_id": "ypwhere/LQK-BabyLM-Strict",
+  "track": "strict",
+  "dtype": "float16",
+  "tokenizer_sha256": "f13720328807e761dc92192111d89ece0119987875f890f3665eba477b5d727c",
+  "source_manifest": "/Volumes/T9/babylm-2026/checkpoints/offdev/vast/bl100m-d512L32-do0.1-gate-attnres8-offdev-aoaw28-aoat31-u57-b16ga32/checkpoint_manifest.json",
+  "main": {
+    "role": "final",
+    "source_path": "ckpt_final-w0910M-i004797.pt",
+    "source_sha256": "dfaff6e50e0759fed5469af2022205228068aa12354537de22ebcde769b32995"
+  },
+  "revisions": [
+    {
+      "revision": "chck_1M",
+      "source_path": "ckpt_000005.pt",
+      "source_sha256": "6fa3f4ffc942b646d7e10bcd2ae6d4d340a8c53221e54d4c10c2de990e064467",
+      "original_revision": "chck_1M",
+      "iter_num": 5,
+      "words_seen": 949462,
+      "tokens_seen": 1310720
+    },
+    {
+      "revision": "chck_2M",
+      "source_path": "ckpt_000011.pt",
+      "source_sha256": "c5b1346784c8ade00bf9c935e57b539928a8fbd8cbd08160031fe2692e966bd4",
+      "original_revision": "chck_2M",
+      "iter_num": 11,
+      "words_seen": 2086910,
+      "tokens_seen": 2883584
+    },
+    {
+      "revision": "chck_3M",
+      "source_path": "ckpt_000016.pt",
+      "source_sha256": "36c8043db9cf0f784ca5856ab9ea4ecb9c1e358f8cc9f661ac32a04398d60f6f",
+      "original_revision": "chck_3M",
+      "iter_num": 16,
+      "words_seen": 3032669,
+      "tokens_seen": 4194304
+    },
+    {
+      "revision": "chck_4M",
+      "source_path": "ckpt_000021.pt",
+      "source_sha256": "e3d2185fdbc80984bb9613d8b669966ea3102f8881bb63add6843c16a7c67282",
+      "original_revision": "chck_4M",
+      "iter_num": 21,
+      "words_seen": 3981969,
+      "tokens_seen": 5505024
+    },
+    {
+      "revision": "chck_5M",
+      "source_path": "ckpt_000026.pt",
+      "source_sha256": "f2f80e00ce96b67e843900eefc3b00dd28c66f808c08280fb158cb440bd0ad61",
+      "original_revision": "chck_5M",
+      "iter_num": 26,
+      "words_seen": 4931564,
+      "tokens_seen": 6815744
+    },
+    {
+      "revision": "chck_6M",
+      "source_path": "ckpt_000032.pt",
+      "source_sha256": "655dd3c62b75a2ef34148484811e3220bdc25e87e80ca6b3f32631304064c8a1",
+      "original_revision": "chck_6M",
+      "iter_num": 32,
+      "words_seen": 6065034,
+      "tokens_seen": 8388608
+    },
+    {
+      "revision": "chck_7M",
+      "source_path": "ckpt_000037.pt",
+      "source_sha256": "39681bdcfa49ceb5f86bf4922ee2f80d315bf2f95bbae8a816fc5b4486ff7877",
+      "original_revision": "chck_7M",
+      "iter_num": 37,
+      "words_seen": 7013912,
+      "tokens_seen": 9699328
+    },
+    {
+      "revision": "chck_8M",
+      "source_path": "ckpt_000042.pt",
+      "source_sha256": "35097df18e56bf78a23fc74ff7dc17713e0366dfaa4e73b25c95d05b40670891",
+      "original_revision": "chck_8M",
+      "iter_num": 42,
+      "words_seen": 7963439,
+      "tokens_seen": 11010048
+    },
+    {
+      "revision": "chck_9M",
+      "source_path": "ckpt_000047.pt",
+      "source_sha256": "b53dfa64832f0bfd6f0273dd525448d6d69ae24d86b525de992220b183f40d2b",
+      "original_revision": "chck_9M",
+      "iter_num": 47,
+      "words_seen": 8911753,
+      "tokens_seen": 12320768
+    },
+    {
+      "revision": "chck_10M",
+      "source_path": "ckpt_000053.pt",
+      "source_sha256": "7faf1ef0a1f476798a7ad58ab0013f540a18ca257d80e9a3632c8c7225517d12",
+      "original_revision": "chck_10M",
+      "iter_num": 53,
+      "words_seen": 10050807,
+      "tokens_seen": 13893632
+    },
+    {
+      "revision": "chck_20M",
+      "source_path": "ckpt_000105.pt",
+      "source_sha256": "bb6085b947df97e3fa9644b2f6f5bbc819a24fa0667e8f541afef1bd7a0a6883",
+      "original_revision": "chck_20M",
+      "iter_num": 105,
+      "words_seen": 19917876,
+      "tokens_seen": 27525120
+    },
+    {
+      "revision": "chck_30M",
+      "source_path": "ckpt_000158.pt",
+      "source_sha256": "81c432bd63dea4047f2304f20a2bf369dacdd56466d12d622fa798e4bc55fa9e",
+      "original_revision": "chck_30M",
+      "iter_num": 158,
+      "words_seen": 29980340,
+      "tokens_seen": 41418752
+    },
+    {
+      "revision": "chck_40M",
+      "source_path": "ckpt_000211.pt",
+      "source_sha256": "2f5ce57ddf0fbb6af09ee61f7779b8dd8b972b5026f3ac8715949cfe4cf3a837",
+      "original_revision": "chck_40M",
+      "iter_num": 211,
+      "words_seen": 40043320,
+      "tokens_seen": 55312384
+    },
+    {
+      "revision": "chck_50M",
+      "source_path": "ckpt_000264.pt",
+      "source_sha256": "2c0c562991293f43b03f2e5b3fdbbbe52676558a18022c700d9e9a9051450bd5",
+      "original_revision": "chck_50M",
+      "iter_num": 264,
+      "words_seen": 50093276,
+      "tokens_seen": 69206016
+    },
+    {
+      "revision": "chck_60M",
+      "source_path": "ckpt_000316.pt",
+      "source_sha256": "9f9ffa85cee106fca9eb46d7b84c6f12e21883ffcedcbf667831fcd72d128331",
+      "original_revision": "chck_60M",
+      "iter_num": 316,
+      "words_seen": 59951092,
+      "tokens_seen": 82837504
+    },
+    {
+      "revision": "chck_70M",
+      "source_path": "ckpt_000369.pt",
+      "source_sha256": "791c49f2f1c7eff2f23e9102e10858a7263a10175a7f0eccaffd8799db4cf572",
+      "original_revision": "chck_70M",
+      "iter_num": 369,
+      "words_seen": 70005835,
+      "tokens_seen": 96731136
+    },
+    {
+      "revision": "chck_80M",
+      "source_path": "ckpt_000422.pt",
+      "source_sha256": "bb61fb617bdd87de4e6d2023fbd17038fe56279922476d13dddb096762c67b61",
+      "original_revision": "chck_80M",
+      "iter_num": 422,
+      "words_seen": 80073013,
+      "tokens_seen": 110624768
+    },
+    {
+      "revision": "chck_90M",
+      "source_path": "ckpt_000474.pt",
+      "source_sha256": "d63fc71506f1a0b67dd9ac245e5936b89805f4625c72ecfeefafed0487b33ba1",
+      "original_revision": "chck_90M",
+      "iter_num": 474,
+      "words_seen": 89937046,
+      "tokens_seen": 124256256
+    },
+    {
+      "revision": "chck_100M",
+      "source_path": "ckpt_000527.pt",
+      "source_sha256": "b2693b41e93cf8eb0e57f0df238c3391f4be8b57c28a8b9b5e63798a28d7f24a",
+      "original_revision": "chck_100M",
+      "iter_num": 527,
+      "words_seen": 99996136,
+      "tokens_seen": 138149888
+    },
+    {
+      "revision": "chck_200M",
+      "source_path": "ckpt_001054.pt",
+      "source_sha256": "0659d403b9c23e54e40a7564f5494351751b572ea4e66209aba26f940df20a98",
+      "original_revision": "chck_200M",
+      "iter_num": 1054,
+      "words_seen": 199982884,
+      "tokens_seen": 276299776
+    },
+    {
+      "revision": "chck_300M",
+      "source_path": "ckpt_001581.pt",
+      "source_sha256": "a5265a307ab28b45f1cfbcf02bb25c0283774b296c9c16eeb945ee38078a1e18",
+      "original_revision": "chck_300M",
+      "iter_num": 1581,
+      "words_seen": 299987936,
+      "tokens_seen": 414449664
+    },
+    {
+      "revision": "chck_400M",
+      "source_path": "ckpt_002108.pt",
+      "source_sha256": "d425113cb065a90c8a39c7dc2099464be9613c90d409b265b6dae0ada51f356e",
+      "original_revision": "chck_400M",
+      "iter_num": 2108,
+      "words_seen": 399991397,
+      "tokens_seen": 552599552
+    },
+    {
+      "revision": "chck_500M",
+      "source_path": "ckpt_002635.pt",
+      "source_sha256": "64b6d22c0ce288f4d7b711c0cb0d47c74c536e0fc92bbb4643f7451f1d966b08",
+      "original_revision": "chck_500M",
+      "iter_num": 2635,
+      "words_seen": 499981660,
+      "tokens_seen": 690749440
+    },
+    {
+      "revision": "chck_600M",
+      "source_path": "ckpt_003162.pt",
+      "source_sha256": "2574ff5d8e536f7983fe5e524f0fa5ee8ee0f684f797a6303c23ea8ac24ac57a",
+      "original_revision": "chck_600M",
+      "iter_num": 3162,
+      "words_seen": 599971774,
+      "tokens_seen": 828899328
+    },
+    {
+      "revision": "chck_700M",
+      "source_path": "ckpt_003689.pt",
+      "source_sha256": "2b66c9aaea75a1e2b083462f9cd683d6a96735ff78c2efd2c0ac1e7a70acfb25",
+      "original_revision": "chck_700M",
+      "iter_num": 3689,
+      "words_seen": 699959933,
+      "tokens_seen": 967049216
+    },
+    {
+      "revision": "chck_800M",
+      "source_path": "ckpt_004216.pt",
+      "source_sha256": "b09a8a5efe69aa1eac7c0ca761e6e73d9ba14c7b37b35b45464c247c76734202",
+      "original_revision": "chck_800M",
+      "iter_num": 4216,
+      "words_seen": 799956340,
+      "tokens_seen": 1105199104
+    },
+    {
+      "revision": "chck_900M",
+      "source_path": "ckpt_004743.pt",
+      "source_sha256": "27772c31bf49988742accabdf273a0a4e0f4611acb79612c78e0541651af6a48",
+      "original_revision": "chck_900M",
+      "iter_num": 4743,
+      "words_seen": 899951513,
+      "tokens_seen": 1243348992
+    },
+    {
+      "revision": "chck_1000M",
+      "source_path": "ckpt_004797.pt",
+      "source_sha256": "0f1e4cd13d67dabfdea0155a9dfe1ab5e59e8f423c64e1242b2fd690818facce",
+      "original_revision": "chck_910M",
+      "iter_num": 4797,
+      "words_seen": 910196864,
+      "tokens_seen": 1257504768
+    }
+  ]
+}

checkpoint_source.json

@@ -0,0 +1,26 @@
+{
+  "filename": "ckpt_final-w0910M-i004797.pt",
+  "sha256": "dfaff6e50e0759fed5469af2022205228068aa12354537de22ebcde769b32995",
+  "role": "final",
+  "iter_num": 4797,
+  "num_updates": 4797,
+  "tokens_seen": 1257504768,
+  "words_seen": 910196864,
+  "labels": [
+    {
+      "series": "words",
+      "name": "words_final",
+      "revision": "chck_910M",
+      "target": 910196864,
+      "actual": 910196864
+    },
+    {
+      "series": "tokens",
+      "name": "tokens_final",
+      "revision": "chck_1258M",
+      "target": 1257504768,
+      "actual": 1257504768
+    }
+  ],
+  "git_sha": "71b982e505ea3f804c54a4b4c0834bd1eee54da8"
+}

config.json

@@ -0,0 +1,60 @@
+{
+  "architectures": [
+    "NanoGPTForCausalLM"
+  ],
+  "attn_res_block_size": 8,
+  "auto_map": {
+    "AutoConfig": "modeling_nanogpt.NanoGPTConfig",
+    "AutoModel": "modeling_nanogpt.NanoGPTModel",
+    "AutoModelForCausalLM": "modeling_nanogpt.NanoGPTForCausalLM",
+    "AutoModelForMaskedLM": "modeling_nanogpt.NanoGPTForMaskedLM"
+  },
+  "bias": false,
+  "bidirectional": false,
+  "block_size": 512,
+  "bos_token_id": 0,
+  "dropout": 0.0,
+  "eos_token_id": 0,
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "nanogpt",
+  "n_embd": 512,
+  "n_head": 8,
+  "n_layer": 32,
+  "nanogpt_checkpoint": {
+    "filename": "ckpt_final-w0910M-i004797.pt",
+    "git_sha": "71b982e505ea3f804c54a4b4c0834bd1eee54da8",
+    "iter_num": 4797,
+    "labels": [
+      {
+        "actual": 910196864,
+        "name": "words_final",
+        "revision": "chck_910M",
+        "series": "words",
+        "target": 910196864
+      },
+      {
+        "actual": 1257504768,
+        "name": "tokens_final",
+        "revision": "chck_1258M",
+        "series": "tokens",
+        "target": 1257504768
+      }
+    ],
+    "num_updates": 4797,
+    "role": "final",
+    "sha256": "dfaff6e50e0759fed5469af2022205228068aa12354537de22ebcde769b32995",
+    "tokens_seen": 1257504768,
+    "words_seen": 910196864
+  },
+  "pad_token_id": 0,
+  "rope_theta": 10000.0,
+  "swiglu_mult": 2.6666666666666665,
+  "torch_dtype": "float16",
+  "transformers_version": "4.51.3",
+  "use_attn_gate": true,
+  "use_attn_res": true,
+  "use_rmsnorm": true,
+  "use_rope": true,
+  "use_swiglu": true,
+  "vocab_size": 16000
+}

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "pad_token_id": 0,
+  "transformers_version": "4.51.3"
+}

model.safetensors

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

modeling_nanogpt.py

@@ -0,0 +1,427 @@
+"""
+Self-contained HuggingFace wrapper for the nanoGPT architecture used in the BabyLM
+ablations: LayerNorm/RMSNorm x MLP/SwiGLU x learned-absolute/RoPE positions.
+
+This single file defines BOTH the config and the models so it can be copied verbatim
+into a model repo and loaded with `trust_remote_code=True` by the BabyLM 2026
+evaluation pipeline:
+  * AutoModelForCausalLM -> NanoGPTForCausalLM   (causal zero-shot)
+  * AutoModelForMaskedLM -> NanoGPTForMaskedLM   (bidirectional MNTP zero-shot)
+  * AutoModel            -> NanoGPTModel         (GLUE fine-tuning backbone)
+
+Module/parameter names mirror nanoGPT's model.py `GPT` exactly, so a nanoGPT
+checkpoint state_dict loads with NO key remapping. Two behaviours differ from
+model.py on purpose (both required by the eval harness):
+  * forward always returns full-sequence logits (B, T, V), not just the last step.
+  * attention honours an optional padding `attention_mask` (needed for GLUE's
+    left-padded batches; nanoGPT's own forward assumes dense, causal-only inputs).
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+import transformers
+from torch.nn import functional as F
+
+from transformers import GenerationMixin, PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput, MaskedLMOutput
+
+
+class NanoGPTConfig(PretrainedConfig):
+    model_type = "nanogpt"
+    # let generic HF utilities (and the GLUE classifier's `config.hidden_size`) work
+    attribute_map = {
+        "hidden_size": "n_embd",
+        "num_attention_heads": "n_head",
+        "num_hidden_layers": "n_layer",
+        "max_position_embeddings": "block_size",
+    }
+
+    def __init__(
+        self,
+        vocab_size=16384,
+        block_size=512,
+        n_layer=8,
+        n_head=8,
+        n_embd=512,
+        dropout=0.0,
+        bias=False,
+        use_rmsnorm=False,
+        use_swiglu=False,
+        swiglu_mult=8 / 3,
+        use_rope=False,
+        use_attn_gate=False,
+        use_attn_res=False,
+        attn_res_block_size=2,
+        bidirectional=False,
+        rope_theta=10000.0,
+        layer_norm_epsilon=1e-5,
+        tie_word_embeddings=True,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_embd = n_embd
+        self.dropout = dropout
+        self.bias = bias
+        self.use_rmsnorm = use_rmsnorm
+        self.use_swiglu = use_swiglu
+        self.swiglu_mult = swiglu_mult
+        self.use_rope = use_rope
+        self.use_attn_gate = use_attn_gate
+        self.use_attn_res = use_attn_res
+        self.attn_res_block_size = attn_res_block_size
+        self.bidirectional = bidirectional
+        self.rope_theta = rope_theta
+        self.layer_norm_epsilon = layer_norm_epsilon
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+
+# --------------------------------------------------------------------------- norms
+class LayerNorm(nn.Module):
+    """LayerNorm with an optional bias (mirrors nanoGPT)."""
+
+    def __init__(self, ndim, bias, eps=1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(ndim))
+        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
+        self.eps = eps
+
+    def forward(self, input):
+        return F.layer_norm(input, self.weight.shape, self.weight, self.bias, self.eps)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, ndim, eps=1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(ndim))
+        self.eps = eps
+
+    def forward(self, input):
+        rms = torch.sqrt(input.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        return (input / rms) * self.weight
+
+
+def make_norm(config):
+    if config.use_rmsnorm:
+        return RMSNorm(config.n_embd, eps=config.layer_norm_epsilon)
+    return LayerNorm(config.n_embd, bias=config.bias, eps=config.layer_norm_epsilon)
+
+
+def attn_res_mix(sources, q, norm):
+    """Softmax attention over depth; normed keys and raw values, matching nanoGPT."""
+    values = torch.stack(sources)
+    logits = torch.einsum("c,sbtc->sbt", q, norm(values))
+    return torch.einsum("sbt,sbtc->btc", logits.softmax(dim=0), values)
+
+
+# ---------------------------------------------------------------------------- rope
+def build_rope_cache(seq_len, head_dim, base=10000.0):
+    inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
+    angles = torch.outer(torch.arange(seq_len).float(), inv_freq)
+    cos = torch.cat([angles, angles], dim=-1).cos()
+    sin = torch.cat([angles, angles], dim=-1).sin()
+    return cos, sin
+
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def apply_rope(x, cos, sin):
+    # x: (B, n_head, T, head_dim); done in fp32 to match nanoGPT
+    seq_len = x.size(2)
+    dtype = x.dtype
+    x = x.float()
+    out = x * cos[:seq_len, :] + rotate_half(x) * sin[:seq_len, :]
+    return out.to(dtype)
+
+
+# ----------------------------------------------------------------------- attention
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.dropout = config.dropout
+        self.use_attn_gate = config.use_attn_gate
+        self.bidirectional = config.bidirectional
+        if self.use_attn_gate:
+            self.attn_gate = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        self.use_rope = config.use_rope
+        if self.use_rope:
+            self.head_dim = config.n_embd // config.n_head
+            assert self.head_dim % 2 == 0, "RoPE needs an even head_dim"
+            self.rope_cache_len = config.block_size
+            self.rope_theta = config.rope_theta
+            # Build lazily on the real execution device. Transformers 5.x may
+            # instantiate remote-code models under a meta-device context; an
+            # eagerly-created non-persistent tensor then survives as
+            # uninitialized storage because it is absent from the state dict.
+            self.register_buffer("rope_cos", None, persistent=False)
+            self.register_buffer("rope_sin", None, persistent=False)
+
+    def _get_rope_cache(self, device):
+        if self.rope_cos is None or self.rope_cos.device != device:
+            cos, sin = build_rope_cache(
+                self.rope_cache_len,
+                self.head_dim,
+                base=self.rope_theta,
+            )
+            self.rope_cos = cos.to(device=device)
+            self.rope_sin = sin.to(device=device)
+        return self.rope_cos, self.rope_sin
+
+    def forward(self, x, attention_mask=None):
+        B, T, C = x.size()
+        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+
+        if self.use_rope:
+            rope_cos, rope_sin = self._get_rope_cache(q.device)
+            q = apply_rope(q, rope_cos, rope_sin)
+            k = apply_rope(k, rope_cos, rope_sin)
+
+        dropout_p = self.dropout if self.training else 0.0
+        if attention_mask is None:
+            # Dense causal is numerically identical to nanoGPT training. The
+            # bidirectional export uses the same weights with no triangular mask.
+            y = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=None, dropout_p=dropout_p,
+                is_causal=not self.bidirectional,
+            )
+        else:
+            keep = attention_mask.to(torch.bool)[:, None, None, :]  # (B, 1, 1, T)
+            if self.bidirectional:
+                attn_mask = keep
+            else:
+                causal = torch.ones(T, T, dtype=torch.bool, device=x.device).tril()
+                attn_mask = causal[None, None, :, :] & keep
+            y = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=False
+            )
+
+        y = y.transpose(1, 2).contiguous().view(B, T, C)
+        if self.use_attn_gate:
+            y = y * torch.sigmoid(self.attn_gate(x)) # elementwise input-dependent gate on attn output (Qwen)
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+
+# ----------------------------------------------------------------------------- mlp
+class MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
+        self.gelu = nn.GELU()
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        return self.dropout(self.c_proj(self.gelu(self.c_fc(x))))
+
+
+class SwiGLU(nn.Module):
+    """Gated MLP, SiLU gate (Shazeer 2020). Hidden snapped to a multiple of 64."""
+
+    def __init__(self, config):
+        super().__init__()
+        hidden_size = int(config.swiglu_mult * config.n_embd)
+        hidden_size = round(hidden_size / 64) * 64
+        self.c_gate = nn.Linear(config.n_embd, hidden_size, bias=config.bias)
+        self.c_val = nn.Linear(config.n_embd, hidden_size, bias=config.bias)
+        self.c_proj = nn.Linear(hidden_size, config.n_embd, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, input):
+        return self.dropout(self.c_proj(F.silu(self.c_gate(input)) * self.c_val(input)))
+
+
+def make_mlp(config):
+    return SwiGLU(config) if config.use_swiglu else MLP(config)
+
+
+class Block(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.ln_1 = make_norm(config)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = make_norm(config)
+        self.mlp = make_mlp(config)
+        self.use_attn_res = config.use_attn_res
+        if config.use_attn_res:
+            self.block_start = (2 * layer_idx) % config.attn_res_block_size == 0
+            self.attn_res_q1 = nn.Parameter(torch.zeros(config.n_embd))
+            self.attn_res_norm1 = RMSNorm(config.n_embd)
+            self.attn_res_q2 = nn.Parameter(torch.zeros(config.n_embd))
+            self.attn_res_norm2 = RMSNorm(config.n_embd)
+
+    def forward(self, x, attention_mask=None):
+        x = x + self.attn(self.ln_1(x), attention_mask)
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+    def forward_attn_res(self, blocks, partial, attention_mask=None):
+        h = attn_res_mix(blocks + [partial], self.attn_res_q1, self.attn_res_norm1)
+        if self.block_start:
+            blocks = blocks + [partial]
+            partial = None
+        attn_out = self.attn(self.ln_1(h), attention_mask)
+        partial = attn_out if partial is None else partial + attn_out
+        h = attn_res_mix(blocks + [partial], self.attn_res_q2, self.attn_res_norm2)
+        partial = partial + self.mlp(self.ln_2(h))
+        return blocks, partial
+
+
+# ---------------------------------------------------------------- shared backbone
+def _build_transformer(config):
+    transformer = nn.ModuleDict(
+        dict(
+            wte=nn.Embedding(config.vocab_size, config.n_embd),
+            drop=nn.Dropout(config.dropout),
+            h=nn.ModuleList([Block(config, i) for i in range(config.n_layer)]),
+            ln_f=make_norm(config),
+        )
+    )
+    if not config.use_rope:
+        transformer.wpe = nn.Embedding(config.block_size, config.n_embd)
+    return transformer
+
+
+def _init_attn_res(model, config):
+    if not config.use_attn_res:
+        return
+    size = config.attn_res_block_size
+    assert size >= 2 and size % 2 == 0
+    assert (2 * config.n_layer) % size == 0
+    model.attn_res_qf = nn.Parameter(torch.zeros(config.n_embd))
+    model.attn_res_normf = RMSNorm(config.n_embd)
+
+
+def _transformer_forward(model, transformer, config, input_ids, attention_mask):
+    _, t = input_ids.size()
+    tok_emb = transformer.wte(input_ids)
+    if config.use_rope:
+        x = transformer.drop(tok_emb)
+    else:
+        pos = torch.arange(0, t, dtype=torch.long, device=input_ids.device)
+        x = transformer.drop(tok_emb + transformer.wpe(pos))
+    if config.use_attn_res:
+        blocks, partial = [], x
+        for block in transformer.h:
+            blocks, partial = block.forward_attn_res(blocks, partial, attention_mask)
+        x = attn_res_mix(blocks + [partial], model.attn_res_qf, model.attn_res_normf)
+    else:
+        for block in transformer.h:
+            x = block(x, attention_mask)
+    x = transformer.ln_f(x)
+    return x
+
+
+class NanoGPTPreTrainedModel(PreTrainedModel):
+    config_class = NanoGPTConfig
+    base_model_prefix = "transformer"
+    main_input_name = "input_ids"
+    _no_split_modules = ["Block"]
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+
+class NanoGPTModel(NanoGPTPreTrainedModel):
+    """Backbone returning last_hidden_state (AutoModel target; used by GLUE)."""
+
+    _keys_to_ignore_on_load_unexpected = [r"lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.transformer = _build_transformer(config)
+        _init_attn_res(self, config)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, value):
+        self.transformer.wte = value
+
+    def forward(self, input_ids=None, attention_mask=None, **kwargs):
+        x = _transformer_forward(self, self.transformer, self.config, input_ids, attention_mask)
+        return BaseModelOutput(last_hidden_state=x)
+
+
+class NanoGPTForCausalLM(NanoGPTPreTrainedModel, GenerationMixin):
+    """LM-head model (AutoModelForCausalLM target; used by zero-shot ranking)."""
+
+    # Transformers 4.x expects a list of tied output keys, while 5.x changed
+    # this class attribute to an output->input mapping.  Keep the published
+    # remote-code model loadable in both the official BabyLM 4.51 environment
+    # and current Transformers releases.
+    _tied_weights_keys = (
+        {"lm_head.weight": "transformer.wte.weight"}
+        if int(transformers.__version__.split(".", 1)[0]) >= 5
+        else ["lm_head.weight"]
+    )
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.transformer = _build_transformer(config)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        _init_attn_res(self, config)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, value):
+        self.transformer.wte = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
+        x = _transformer_forward(self, self.transformer, self.config, input_ids, attention_mask)
+        logits = self.lm_head(x)
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits[..., :-1, :].reshape(-1, logits.size(-1)),
+                labels[..., 1:].reshape(-1),
+                ignore_index=-1,
+            )
+        return CausalLMOutput(loss=loss, logits=logits)
+
+
+class NanoGPTForMaskedLM(NanoGPTForCausalLM):
+    """Same LM head, with unshifted MLM loss for the bidirectional export."""
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
+        x = _transformer_forward(self, self.transformer, self.config, input_ids, attention_mask)
+        logits = self.lm_head(x)
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits.reshape(-1, logits.size(-1)),
+                labels.reshape(-1),
+                ignore_index=-100,
+            )
+        return MaskedLMOutput(loss=loss, logits=logits)

special_tokens_map.json

@@ -0,0 +1,4 @@
+{
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|endoftext|>"
+}

tokenizer_config.json

@@ -0,0 +1,19 @@
+{
+  "add_prefix_space": true,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "PreTrainedTokenizer"
+}