TinyBrainBot-3M-Loop-v2: 3M-Loop continued on general math (final checkpoint)

README.md

@@ -0,0 +1,174 @@
+---
+license: apache-2.0
+language:
+  - en
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+  - tinybrainbot
+  - small-language-model
+  - from-scratch
+  - recurrent-depth
+  - looped-transformer
+  - arithmetic
+  - custom_code
+  - base
+datasets:
+  - HuggingFaceFW/fineweb-edu
+  - HuggingFaceTB/cosmopedia
+  - roneneldan/TinyStories
+  - wikimedia/wikipedia
+model-index:
+  - name: tinybrainbot-3m-loop-v2
+    results:
+    - task:
+        type: text-generation
+      dataset:
+        type: allenai/ai2_arc
+        name: ARC-Easy
+        config: ARC-Easy
+        split: test
+      metrics:
+        - type: acc_norm
+          value: 33.46
+          name: acc_norm (0-shot)
+    - task:
+        type: text-generation
+      dataset:
+        type: allenai/ai2_arc
+        name: ARC-Challenge
+        config: ARC-Challenge
+        split: test
+      metrics:
+        - type: acc_norm
+          value: 22.78
+          name: acc_norm (0-shot)
+    - task:
+        type: text-generation
+      dataset:
+        type: Rowan/hellaswag
+        name: HellaSwag
+        split: validation
+      metrics:
+        - type: acc_norm
+          value: 27.17
+          name: acc_norm (0-shot)
+    - task:
+        type: text-generation
+      dataset:
+        type: ybisk/piqa
+        name: PIQA
+        split: validation
+      metrics:
+        - type: acc_norm
+          value: 54.79
+          name: acc_norm (0-shot)
+    - task:
+        type: text-generation
+      dataset:
+        type: AxiomicLabs/Arithmark-3.0
+        name: ArithMark-3.0
+        split: train
+      metrics:
+        - type: acc_norm
+          value: 90.70
+          name: acc_norm (0-shot)
+---
+
+# TinyBrainBot-3M-Loop-v2
+
+**A 2.96M-parameter math-first language model that reuses its layers.** Four transformer blocks are run three times in a loop, so every token passes through 12 layers of computation while only 4 layers of weights are stored. Trained from scratch on 6.0B tokens. Compared with [TinyBrainBot-3M-Loop](https://huggingface.co/nkthebass/tinybrainbot-3m-loop), its arithmetic carries over to bare expressions, order of operations and numbers written in words, not only word-problem stories.
+
+| Benchmark (0-shot, full set) | acc_norm |
+|---|--:|
+| **Open SLM Leaderboard Intelligence Index** | **20.15** |
+|   *(raw acc)* | *19.91* |
+| HellaSwag | 27.17 |
+|   *(raw acc)* | *27.01* |
+| ARC-Easy | 33.46 |
+|   *(raw acc)* | *33.59* |
+| ARC-Challenge | 22.78 |
+|   *(raw acc)* | *18.17* |
+| PIQA | 54.79 |
+|   *(raw acc)* | *55.98* |
+| ArithMark-3.0 | 90.7 |
+|   *(raw acc)* | *90.6* |
+
+lm-eval 0.4.13 in float32 for HellaSwag, ARC and PIQA; the official `bencharithmark-3.py` (float32) for ArithMark-3.0.
+
+## Arithmetic beyond one benchmark
+
+Every column below is a zero-shot, log-likelihood multiple-choice math test, all scored by us with the official `bencharithmark-3.py` scoring in float32, so every model is measured the same way (numbers can differ slightly from leaderboard listings).
+
+| Model | Params | ArithMark-3 | AM3, numbers in words | AM3 as bare equations | ArithMark-2 | ArithMark-1 | Held-out stories |
+|---|--:|--:|--:|--:|--:|--:|--:|
+| **TinyBrainBot-3M-Loop-v2** | 2.96M | **90.7** | **55.0** | **96.8** | **83.8** | 36.1 | **80.1** |
+| TinyBrainBot-3M-Loop | 2.96M | 89.3 | 27.3 | 88.2 | 58.1 | 29.4 | 81.0 |
+| Inferno-Exp-1.7 | 394M | 85.9 | 59.1 | 86.9 | 59.9 | 47.9 | 74.0 |
+| Pascaline-3M | 3M | 69.2 | 27.6 | 99.7 | 66.5 | 36.2 | 43.0 |
+| SmolLM2-135M | 135M | 38.8 | 34.2 | 38.6 | 33.2 | 33.1 | 34.7 |
+| GPT-S3-3M | 2.97M | 38.9 | 25.8 | 36.2 | 31.0 | 31.1 | 34.6 |
+
+- **Numbers in words** and **bare equations** are ArithMark-3's own problems with only the surface changed.
+- **ArithMark-2** uses its official metric (raw log-likelihood); its official `benchmark_arithmark-2.0.py` gives the same results.
+- **Held-out stories:** about 1,600 ArithMark-style problems written by Gemma-3-27B and Llama-3.3-70B, models that wrote none of our training data, with every number and answer set by Python.
+- **ArithMark-1:** its answer key mostly evaluates expressions strictly left to right (e.g. `14 * -17 + -1 * -18 = 4302`), while ArithMark-2 and ordinary arithmetic use operator precedence (that expression is -220). This model follows precedence, so it scores lower on ArithMark-1's longer expressions (62% on its two-operand stage, falling with each extra operator).
+
+## How it works
+
+- **Recurrent depth.** 4 shared blocks (width 256, 4 attention heads of 64 with 2 KV heads, SwiGLU 560, RoPE, RMSNorm) are applied 3 times. Each pass adds a learned loop embedding and a gated copy of the input embedding. Training sampled 2 to 4 loops per step, and you can change the count at inference (`model.model.eval_loops = 4`); 3 is best.
+- **Factorized embedding.** The 4,104-token vocabulary uses 96-dimensional embeddings projected up to 256 (and down again before the tied output layer), which frees about 630k parameters for the 4th block.
+- **Numbers.** A units-first digit tokenizer writes every number as a length token plus its digits, smallest place first, and a digit-place embedding tells the model which place each digit is in.
+
+## Training
+
+**Stage 1 (4.46B tokens):** the TinyBrainBot-3M-Loop run up to the start of its learning-rate decay.
+
+| Part | Share | What |
+|---|--:|---|
+| Q&A and facts | 25% | our distilled Q&A, short facts, explanations |
+| Arithmetic | 20% | word problems: about 150k different story wordings written by ling-3.0-flash and Mistral-Small-3.2, each checked by solving it, with every number and answer filled in by Python; plus our own relation (more/fewer than) and story generators |
+| Synthetic textbooks | 15% | Cosmopedia (Stanford, OpenStax) |
+| Stories | 15% | TinyStories |
+| Commonsense | 10% | Cosmopedia WikiHow, our everyday-commonsense prose |
+| Web | 10% | FineWeb-Edu |
+| Reference | 5% | Wikipedia |
+
+**Stage 2 (1.5B tokens, math-first):** constant learning rate for 1.0B tokens, then a decay to zero over the last 0.5B.
+
+| Part | Share | What |
+|---|--:|---|
+| Story arithmetic | 25% | the same verified wordings, with a quarter of the problems' numbers written in words |
+| Bare arithmetic | 25% | our generator: expressions with order of operations, parentheses, negative numbers and 2-5 operands, in about 19 forms ("43 + 19 = 62", "What is ...?", "Compute ...:") |
+| Number words | 8% | words to digits and back, arithmetic in words, counting, comparisons |
+| Relations and stories | 12% | our relation and story generators |
+| Q&A and facts | 20% | as in stage 1 |
+| Cosmopedia, FineWeb-Edu | 10% | as in stage 1 |
+
+No benchmark training or test data was used, and none of our generators copy any benchmark's problems, code or sentences.
+
+## Use
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+repo = "nkthebass/tinybrainbot-3m-loop-v2"
+tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
+ids = tok("Mia had 46 stickers. She gave 18 to her brother, so now Mia has", return_tensors="pt").input_ids
+out = model.generate(ids, max_new_tokens=8, do_sample=False)
+print(tok.decode(out[0][ids.shape[1]:]))
+```
+
+`trust_remote_code=True` is required (custom architecture and tokenizer). There is no KV cache: generation recomputes the prefix, which is fast at this size. Padded batches work on either side (the model honors `attention_mask`). There is no GGUF, because llama.cpp does not support the looped architecture.
+
+To reproduce the scores:
+```
+lm_eval --model hf --model_args pretrained=nkthebass/tinybrainbot-3m-loop-v2,dtype=float32,trust_remote_code=True --tasks arc_easy,arc_challenge,hellaswag,piqa --num_fewshot 0
+python bencharithmark-3.py --model nkthebass/tinybrainbot-3m-loop-v2 --dtype float32
+```
+
+## Limitations
+
+- A 3M base model: it writes simple English and makes up facts. It is not a chat model.
+- Arithmetic is strongest on short problems; long, chatty setups and long left-to-right chains lower its accuracy.
+- 1,024-token context.

config.json

@@ -0,0 +1,45 @@
+{
+  "architectures": [
+    "TBBLoopForCausalLM"
+  ],
+  "model_type": "tbbloop",
+  "auto_map": {
+    "AutoConfig": "configuration_tbbloop.TBBLoopConfig",
+    "AutoModelForCausalLM": "modeling_tbbloop.TBBLoopForCausalLM"
+  },
+  "native": {
+    "vocab_size": 4104,
+    "hidden_size": 256,
+    "n_layers": 3,
+    "n_heads": 4,
+    "head_dim": 64,
+    "n_kv_heads": 2,
+    "ffn_size": 560,
+    "context_length": 1024,
+    "rope_theta": 10000.0,
+    "rms_norm_eps": 1e-05,
+    "tie_embeddings": true,
+    "init_std": 0.02,
+    "dropout": 0.0,
+    "gradient_checkpointing": false,
+    "loop_unique": 4,
+    "n_loops": 3,
+    "loop_train_min": 2,
+    "loop_train_max": 4,
+    "digit_place": true,
+    "digit_lo": 24,
+    "digit_hi": 33,
+    "place_vocab": 16,
+    "embed_rank": 96
+  },
+  "eval_loops": 0,
+  "vocab_size": 4104,
+  "hidden_size": 256,
+  "num_hidden_layers": 4,
+  "max_position_embeddings": 1024,
+  "tie_word_embeddings": true,
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "torch_dtype": "float32",
+  "use_cache": false
+}

configuration_tbbloop.py

@@ -0,0 +1,11 @@
+"""Config for TinyBrainBot recurrent-depth models (fields mirror the native TinyBrainBotConfig)."""
+from transformers import PretrainedConfig
+
+
+class TBBLoopConfig(PretrainedConfig):
+    model_type = "tbbloop"
+
+    def __init__(self, native=None, eval_loops=0, tie_word_embeddings=True, **kwargs):
+        self.native = native or {}
+        self.eval_loops = eval_loops
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)

model.safetensors

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

modeling_tbbloop.py

@@ -0,0 +1,415 @@
+"""TinyBrainBot recurrent-depth model for transformers (trust_remote_code=True).
+
+Below the marker is the exact architecture source the model was trained with. The wrapper builds it from the
+config, so weights load unchanged. Generation recomputes the prefix each step (no KV cache): fine at this size.
+"""
+from __future__ import annotations
+# ======================= native architecture (verbatim) =======================
+"""TinyBrainBot reference architecture.
+
+Matches the blueprint spec:
+    - decoder-only transformer
+    - 8 layers, hidden 1536, 12 heads (head_dim = 128), FFN 6144
+    - RoPE positional encoding
+    - RMSNorm (pre-norm)
+    - SwiGLU feed-forward
+    - 36k Unigram vocab
+    - ctx 2048
+
+The defaults here produce a ~300M parameter model. Tests use smaller overrides.
+"""
+
+from dataclasses import dataclass, field, asdict
+from typing import Optional
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint as _ckpt
+
+
+@dataclass
+class TinyBrainBotConfig:
+    # Blueprint defaults
+    vocab_size: int = 36_000
+    hidden_size: int = 1536
+    n_layers: int = 8
+    n_heads: int = 12
+    head_dim: int = 128
+    n_kv_heads: Optional[int] = None   # None => MHA (== n_heads); a smaller divisor => GQA
+    ffn_size: int = 6144
+    context_length: int = 2048
+    rope_theta: float = 10_000.0
+    rms_norm_eps: float = 1e-5
+    tie_embeddings: bool = True
+    init_std: float = 0.02
+    dropout: float = 0.0
+    gradient_checkpointing: bool = False
+    # recurrent depth (opt-in): store `loop_unique` blocks and run them `n_loops` times. 0 = ordinary stack.
+    loop_unique: int = 0
+    n_loops: int = 1
+    loop_train_min: int = 0          # training samples the loop count from [min, max] per step (0 = n_loops)
+    loop_train_max: int = 0
+    # digit-place embedding (opt-in): each digit token gets its index in its digit run (1 = first digit; with
+    # units-first numbers that is the units digit). Digit token ids are digit_lo..digit_hi.
+    digit_place: bool = False
+    digit_lo: int = -1
+    digit_hi: int = -1
+    place_vocab: int = 16
+    # factorized embedding (opt-in): a vocab x embed_rank table, projected up to hidden_size on the way in and down
+    # to embed_rank before the (tied) output layer. 0 = ordinary vocab x hidden_size embedding.
+    embed_rank: int = 0
+
+    @property
+    def attention_dim(self) -> int:
+        return self.n_heads * self.head_dim
+
+    @property
+    def num_kv_heads(self) -> int:
+        """KV head count — equals n_heads for MHA, fewer for GQA."""
+        return self.n_kv_heads if self.n_kv_heads is not None else self.n_heads
+
+    @property
+    def kv_dim(self) -> int:
+        return self.num_kv_heads * self.head_dim
+
+    def summary(self) -> dict:
+        return {
+            **asdict(self),
+            "attention_dim": self.attention_dim,
+        }
+
+
+# ----------------------------------------------------------------------------
+# Building blocks
+# ----------------------------------------------------------------------------
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # computed in float32 for stability, cast back to input dtype
+        orig_dtype = x.dtype
+        x = x.float()
+        variance = x.pow(2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(variance + self.eps)
+        return (self.weight * x).to(orig_dtype)
+
+
+def _precompute_rope_cache(dim: int, max_len: int, theta: float, device, dtype):
+    # standard RoPE frequencies
+    inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, device=device).float() / dim))
+    t = torch.arange(max_len, device=device).float()
+    freqs = torch.outer(t, inv_freq)  # [max_len, dim/2]
+    cos = freqs.cos().to(dtype)
+    sin = freqs.sin().to(dtype)
+    return cos, sin
+
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    # x: [B, H, T, D]; cos/sin: [T, D/2], or [B, T, D/2] for per-row positions (padded batches)
+    x1, x2 = x[..., ::2], x[..., 1::2]
+    if cos.dim() == 3:
+        cos_b, sin_b = cos.unsqueeze(1), sin.unsqueeze(1)
+    else:
+        # broadcast cos/sin over [B, H, T, D/2]
+        cos_b = cos.unsqueeze(0).unsqueeze(0)
+        sin_b = sin.unsqueeze(0).unsqueeze(0)
+    y1 = x1 * cos_b - x2 * sin_b
+    y2 = x1 * sin_b + x2 * cos_b
+    out = torch.stack([y1, y2], dim=-1).flatten(-2)
+    return out
+
+
+def _repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """Expand KV heads to match query heads for GQA. x: [B, n_kv, T, head_dim]."""
+    if n_rep == 1:
+        return x
+    B, n_kv, T, hd = x.shape
+    return x[:, :, None, :, :].expand(B, n_kv, n_rep, T, hd).reshape(B, n_kv * n_rep, T, hd)
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, cfg: TinyBrainBotConfig):
+        super().__init__()
+        self.cfg = cfg
+        nkv = cfg.num_kv_heads
+        assert cfg.n_heads % nkv == 0, "n_heads must be divisible by n_kv_heads"
+        # Combined Q/K/V. With MHA (nkv == n_heads) this is 3*attention_dim — the exact
+        # shape older checkpoints stored, so they load unchanged. With GQA it shrinks K/V.
+        self.qkv = nn.Linear(cfg.hidden_size, (cfg.n_heads + 2 * nkv) * cfg.head_dim, bias=False)
+        self.out_proj = nn.Linear(cfg.attention_dim, cfg.hidden_size, bias=False)
+        self.scale = 1.0 / math.sqrt(cfg.head_dim)
+        self._q_size = cfg.n_heads * cfg.head_dim
+        self._kv_size = nkv * cfg.head_dim
+        self._n_rep = cfg.n_heads // nkv
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        past_kv: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+        B, T, _ = x.shape
+        nkv = self.cfg.num_kv_heads
+        q, k, v = self.qkv(x).split([self._q_size, self._kv_size, self._kv_size], dim=-1)
+        q = q.view(B, T, self.cfg.n_heads, self.cfg.head_dim).transpose(1, 2)
+        k = k.view(B, T, nkv, self.cfg.head_dim).transpose(1, 2)
+        v = v.view(B, T, nkv, self.cfg.head_dim).transpose(1, 2)
+
+        if cos.dim() == 3:                       # per-row positions, already [B, T, D/2]
+            q = _apply_rope(q, cos, sin)
+            k = _apply_rope(k, cos, sin)
+        else:
+            offset = past_kv[0].shape[-2] if past_kv is not None else 0
+            q = _apply_rope(q, cos[offset:offset + T], sin[offset:offset + T])
+            k = _apply_rope(k, cos[offset:offset + T], sin[offset:offset + T])
+
+        if past_kv is not None:
+            k = torch.cat([past_kv[0], k], dim=-2)
+            v = torch.cat([past_kv[1], v], dim=-2)
+        new_kv = (k, v)  # cache holds the SMALL (n_kv) heads — the whole point of GQA
+
+        # Expand KV to n_heads for attention (no-op under MHA).
+        k = _repeat_kv(k, self._n_rep)
+        v = _repeat_kv(v, self._n_rep)
+
+        if attn_mask is not None:                # padded batch: causal AND not-a-pad-key, [B, 1, T, T] bool
+            attn = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+        else:
+            is_causal = T > 1
+            attn = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
+        attn = attn.transpose(1, 2).contiguous().view(B, T, self.cfg.attention_dim)
+        return self.out_proj(attn), new_kv
+
+
+class SwiGLU(nn.Module):
+    """Gated SiLU feed-forward (as in LLaMA / Mistral)."""
+
+    def __init__(self, cfg: TinyBrainBotConfig):
+        super().__init__()
+        self.gate = nn.Linear(cfg.hidden_size, cfg.ffn_size, bias=False)
+        self.up = nn.Linear(cfg.hidden_size, cfg.ffn_size, bias=False)
+        self.down = nn.Linear(cfg.ffn_size, cfg.hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, cfg: TinyBrainBotConfig):
+        super().__init__()
+        self.attn_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+        self.attn = CausalSelfAttention(cfg)
+        self.ffn_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+        self.ffn = SwiGLU(cfg)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        past_kv: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+        attn_out, new_kv = self.attn(self.attn_norm(x), cos, sin, past_kv, attn_mask)
+        x = x + attn_out
+        x = x + self.ffn(self.ffn_norm(x))
+        return x, new_kv
+
+    def forward_no_cache(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        out, _ = self.forward(x, cos, sin, None)
+        return out
+
+
+class TinyBrainBotModel(nn.Module):
+    def __init__(self, cfg: TinyBrainBotConfig):
+        super().__init__()
+        self.cfg = cfg
+        r = cfg.embed_rank if cfg.embed_rank > 0 else cfg.hidden_size
+        self.embed = nn.Embedding(cfg.vocab_size, r)
+        self.embed_up = nn.Linear(r, cfg.hidden_size, bias=False) if cfg.embed_rank > 0 else None
+        self.head_down = nn.Linear(cfg.hidden_size, r, bias=False) if cfg.embed_rank > 0 else None
+        n_blocks = cfg.loop_unique if cfg.loop_unique > 0 else cfg.n_layers
+        self.blocks = nn.ModuleList([DecoderBlock(cfg) for _ in range(n_blocks)])
+        # recurrent depth: a learned vector per loop pass, and a gate re-injecting the input embedding each pass
+        # (named *_norm so the optimiser puts it in the no-weight-decay group)
+        self.loop_emb = nn.Embedding(32, cfg.hidden_size) if cfg.loop_unique > 0 else None
+        self.reinject_norm = nn.Parameter(torch.tensor(0.5)) if cfg.loop_unique > 0 else None
+        self.place = nn.Embedding(cfg.place_vocab, cfg.hidden_size) if cfg.digit_place else None
+        self.train_loops: Optional[int] = None     # set per step by the trainer (same on every rank)
+        self.eval_loops: Optional[int] = None      # override for evaluation ("think longer")
+        self.final_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+        self.lm_head = nn.Linear(r, cfg.vocab_size, bias=False)
+        if cfg.tie_embeddings:
+            self.lm_head.weight = self.embed.weight
+
+        self._rope_cache: Optional[tuple[torch.Tensor, torch.Tensor]] = None
+        self._rope_device: Optional[torch.device] = None
+        self._rope_dtype: Optional[torch.dtype] = None
+
+        self.apply(self._init_weights)
+        if self.embed_up is not None:
+            # keep the residual stream at the usual scale: up-projection ~ 1/sqrt(rank)
+            nn.init.normal_(self.embed_up.weight, std=r ** -0.5)
+            nn.init.normal_(self.head_down.weight, std=cfg.hidden_size ** -0.5)
+
+    def _init_weights(self, m: nn.Module) -> None:
+        if isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, std=self.cfg.init_std)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Embedding):
+            nn.init.normal_(m.weight, std=self.cfg.init_std)
+
+    def _place_ids(self, ids: torch.Tensor) -> torch.Tensor:
+        d = (ids >= self.cfg.digit_lo) & (ids <= self.cfg.digit_hi)
+        idx = torch.arange(ids.shape[1], device=ids.device).expand_as(ids)
+        last_nd = torch.cummax(torch.where(d, torch.full_like(idx, -1), idx), dim=1).values
+        return torch.where(d, (idx - last_nd).clamp(max=self.cfg.place_vocab - 1), torch.zeros_like(idx))
+
+    def loops_now(self) -> int:
+        if self.training and self.train_loops is not None:
+            return self.train_loops
+        if not self.training and self.eval_loops is not None:
+            return self.eval_loops
+        return self.cfg.n_loops
+
+    def _ensure_rope(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
+        need_rebuild = (
+            self._rope_cache is None
+            or self._rope_device != device
+            or self._rope_dtype != dtype
+            or self._rope_cache[0].shape[0] < seq_len
+        )
+        if need_rebuild:
+            target_len = max(seq_len, self.cfg.context_length)
+            cos, sin = _precompute_rope_cache(self.cfg.head_dim, target_len, self.cfg.rope_theta, device, dtype)
+            self._rope_cache = (cos, sin)
+            self._rope_device = device
+            self._rope_dtype = dtype
+        return self._rope_cache
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        labels: Optional[torch.Tensor] = None,
+        past_kv: Optional[list[tuple[torch.Tensor, torch.Tensor]]] = None,
+        use_cache: bool = False,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> dict:
+        B, T = input_ids.shape
+        past_len = past_kv[0][0].shape[-2] if past_kv is not None else 0
+        if T + past_len > self.cfg.context_length:
+            raise ValueError(f"sequence length {T + past_len} exceeds context_length {self.cfg.context_length}")
+        x = self.embed(input_ids)
+        if self.embed_up is not None:
+            x = self.embed_up(x)
+        if self.place is not None:
+            x = x + self.place(self._place_ids(input_ids))
+        cos, sin = self._ensure_rope(T + past_len, x.device, x.dtype)
+        amask = None
+        if attention_mask is not None and not bool(attention_mask.all()):
+            # padded batch (left or right): RoPE positions count from each row's first real token, pad keys are
+            # masked out, and pad query rows attend to themselves only (no NaN). Unpadded input takes the old path.
+            if past_kv is not None or use_cache:
+                raise ValueError("attention_mask with padding is not supported together with a KV cache")
+            am = attention_mask.to(device=x.device, dtype=torch.bool)
+            pos = (am.long().cumsum(-1) - 1).clamp(min=0)
+            cos, sin = cos[pos], sin[pos]
+            eye = torch.eye(T, dtype=torch.bool, device=x.device)
+            amask = (eye.tril() | torch.ones(T, T, dtype=torch.bool, device=x.device).tril()) & am[:, None, None, :]
+            amask = amask | eye[None, None]
+        use_ckpt = self.cfg.gradient_checkpointing and self.training
+        new_past: list[tuple[torch.Tensor, torch.Tensor]] = []
+        if self.cfg.loop_unique > 0:
+            if past_kv is not None or use_cache:
+                raise ValueError("recurrent-depth models do not use a KV cache (recompute the prefix)")
+            x0 = x
+            for k in range(self.loops_now()):
+                x = x + self.loop_emb.weight[k].to(x.dtype)
+                if k > 0:
+                    x = x + self.reinject_norm.to(x.dtype) * x0
+                for block in self.blocks:
+                    x = _ckpt(block.forward_no_cache, x, cos, sin, use_reentrant=False) if use_ckpt else block(x, cos, sin, None, amask)[0]
+        for i, block in enumerate(self.blocks if self.cfg.loop_unique == 0 else []):
+            layer_past = past_kv[i] if past_kv is not None else None
+            if use_ckpt:
+                x = _ckpt(block.forward_no_cache, x, cos, sin, use_reentrant=False)
+            else:
+                x, kv = block(x, cos, sin, layer_past, amask)
+                if use_cache:
+                    new_past.append(kv)
+        x = self.final_norm(x)
+        logits = self.lm_head(self.head_down(x) if self.head_down is not None else x)
+
+        out = {"logits": logits}
+        if use_cache and not use_ckpt:
+            out["past_kv"] = new_past
+        if labels is not None:
+            # standard next-token loss, labels shifted upstream
+            # unsqueeze(0) makes loss 1-D so torch.nn.DataParallel can
+            # gather per-device losses along dim 0 (noop on single GPU).
+            if (labels != -100).any():
+                loss = F.cross_entropy(
+                    logits.view(-1, self.cfg.vocab_size),
+                    labels.view(-1),
+                    ignore_index=-100,
+                )
+            else:
+                # SFT edge case: a packed window can land entirely in masked
+                # (user/system) tokens, leaving no supervised targets. cross_entropy
+                # would then divide by zero and return NaN, poisoning training.
+                # Emit a graph-connected zero loss for these batches instead.
+                loss = logits.sum() * 0.0
+            out["loss"] = loss.unsqueeze(0)
+        return out
+
+    def num_parameters(self) -> int:
+        return sum(p.numel() for p in self.parameters())
+
+# ======================= transformers wrapper =======================
+from transformers import PreTrainedModel, GenerationMixin
+from transformers.modeling_outputs import CausalLMOutput
+from .configuration_tbbloop import TBBLoopConfig
+
+
+class TBBLoopForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = TBBLoopConfig
+    _tied_weights_keys = ["model.lm_head.weight"]
+    _supports_cache_class = False
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = TinyBrainBotModel(TinyBrainBotConfig(**config.native))
+        if getattr(config, "eval_loops", 0):
+            self.model.eval_loops = int(config.eval_loops)
+
+    def _init_weights(self, module):
+        pass                                     # weights always come from the checkpoint
+
+    def get_input_embeddings(self):
+        return self.model.embed
+
+    def set_input_embeddings(self, value):
+        self.model.embed = value
+
+    def get_output_embeddings(self):
+        return self.model.lm_head
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
+        logits = self.model(input_ids, attention_mask=attention_mask)["logits"]
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(logits[:, :-1].reshape(-1, logits.shape[-1]).float(), labels[:, 1:].reshape(-1), ignore_index=-100)
+        return CausalLMOutput(loss=loss, logits=logits)
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids[:, -self.model.cfg.context_length:]}

special_tokens_map.json

@@ -0,0 +1,14 @@
+{
+  "additional_special_tokens": [
+    "<|user|>",
+    "<|assistant|>",
+    "<|system|>",
+    "<|end|>",
+    "<think>",
+    "</think>",
+    "<|sep|>"
+  ],
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|pad|>"
+}

tbb_meta.json

@@ -0,0 +1,38 @@
+{
+  "model_type": "byte-level BPE (HF tokenizers)",
+  "vocab_size": 4096,
+  "special_tokens": [
+    "<|endoftext|>",
+    "<|user|>",
+    "<|assistant|>",
+    "<|system|>",
+    "<|end|>",
+    "<think>",
+    "</think>",
+    "<|sep|>",
+    "<|pad|>"
+  ],
+  "digits": "split individually",
+  "normalization": "none",
+  "sample_chars": 400000000,
+  "sample_shares": {
+    "qa-plain": 0.14100000000000001,
+    "qa-distill": 0.055,
+    "facts-plain": 0.041,
+    "capability-distill": 0.0165,
+    "reasoning-distill": 0.006500000000000001,
+    "math-direct-3m": 0.19999999999999998,
+    "cosmopedia-wikihow": 0.0757,
+    "commonsense-distill": 0.0092,
+    "procedural-anneal": 0.006900000000000001,
+    "everyday-activity": 0.0045000000000000005,
+    "social-situations": 0.0045000000000000005,
+    "everyday-reasoning": 0.00135,
+    "ds-prose": 0.00135,
+    "cosmopedia-stanford": 0.12300000000000001,
+    "cosmopedia-openstax": 0.027000000000000003,
+    "tiny-stories-v2": 0.14150000000000001,
+    "fineweb-edu-10bt": 0.09,
+    "wiki-dense": 0.055
+  }
+}

tokenization_tbbuf.py

@@ -0,0 +1,61 @@
+"""TinyBrainBot units-first number tokenizer (byte-level BPE + length-prefixed, reversed digit runs).
+
+Every run of 1-8 digits is written as a length token and then its digits units first:
+    "62"   -> "<len2>26"        "1234" -> "<len4>4321"
+so a model emits the units digit first and carries flow forward in time. Digit runs touching ":" (clock
+times like 3:15) and runs longer than 8 digits are left unchanged. The text is rewritten on encode and
+restored on decode, so callers see ordinary numbers. The rewrite is idempotent (an already-rewritten run
+is not rewritten again), so it is safe on every encode path.
+
+The idea (units-first digit tokens; a length prefix) follows Atom2.7m and Pascaline-3M. This is our own
+implementation.
+"""
+import re
+
+from transformers import PreTrainedTokenizerFast
+
+MAX_LEN = 8
+# a digit run: not part of "<lenN>", not right after a "<lenN>" tag, not touching another digit or ":"
+_RUN = re.compile(r"(?<!<len)(?<!<len\d>)(?<![\d:])(\d{1,%d})(?![\d:])" % MAX_LEN)
+_TAG = re.compile(r"<len([1-%d])>(\d+)" % MAX_LEN)
+
+
+def encode_numbers(text):
+    if not isinstance(text, str):
+        return text
+    return _RUN.sub(lambda m: "<len%d>%s" % (len(m.group(1)), m.group(1)[::-1]), text)
+
+
+def decode_numbers(text):
+    def back(m):
+        n = int(m.group(1)); d = m.group(2)
+        return d[:n][::-1] + d[n:]
+    return _TAG.sub(back, text)
+
+
+def _map(x):
+    if isinstance(x, str):
+        return encode_numbers(x)
+    if isinstance(x, tuple):
+        return tuple(_map(v) for v in x)
+    if isinstance(x, list):
+        return [_map(v) for v in x]
+    return x
+
+
+class TBBUnitsFirstTokenizerFast(PreTrainedTokenizerFast):
+    def _batch_encode_plus(self, batch_text_or_text_pairs, *args, **kwargs):
+        if not kwargs.get("is_split_into_words", False):
+            batch_text_or_text_pairs = _map(batch_text_or_text_pairs)
+        return super()._batch_encode_plus(batch_text_or_text_pairs, *args, **kwargs)
+
+    def _encode_plus(self, text, text_pair=None, *args, **kwargs):
+        if not kwargs.get("is_split_into_words", False):
+            text, text_pair = _map(text), _map(text_pair)
+        return super()._encode_plus(text, text_pair, *args, **kwargs)
+
+    def tokenize(self, text, *args, **kwargs):
+        return super().tokenize(encode_numbers(text), *args, **kwargs)
+
+    def _decode(self, token_ids, *args, **kwargs):
+        return decode_numbers(super()._decode(token_ids, *args, **kwargs))

tokenizer_config.json

@@ -0,0 +1,167 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|user|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<|assistant|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<|system|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<|end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "5": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "6": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "7": {
+      "content": "<|sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "8": {
+      "content": "<|pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4096": {
+      "content": "<len1>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "4097": {
+      "content": "<len2>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "4098": {
+      "content": "<len3>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "4099": {
+      "content": "<len4>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "4100": {
+      "content": "<len5>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "4101": {
+      "content": "<len6>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "4102": {
+      "content": "<len7>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "4103": {
+      "content": "<len8>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|user|>",
+    "<|assistant|>",
+    "<|system|>",
+    "<|end|>",
+    "<think>",
+    "</think>",
+    "<|sep|>"
+  ],
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 1024,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "TBBUnitsFirstTokenizerFast",
+  "unk_token": null,
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_tbbuf.TBBUnitsFirstTokenizerFast",
+      "tokenization_tbbuf.TBBUnitsFirstTokenizerFast"
+    ]
+  },
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ]
+}