Hugging Face's logo

nkthebass
/
tinybrainbot-3m-loop

Text Generation
Transformers
Safetensors
English
tbbloop
tinybrainbot
small-language-model
from-scratch
recurrent-depth
looped-transformer
arithmetic
custom_code
base
Eval Results (legacy)
Model card Files
xet
 Community
tinybrainbot-3m-loop / modeling_tbbloop.py
nkthebass's picture
nkthebass
transformers 5 compatibility: call post_init() and use the v5 tied-weights mapping (4.x unchanged)
1b5218b verified
Raw History Blame Contribute Delete
19.1 kB
"""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
# transformers 5 wants {tied param: source param}; 4.x wants a list of tied params
_tied_weights_keys = ({"model.lm_head.weight": "model.embed.weight"}
if int(__import__("transformers").__version__.split(".")[0]) >= 5 else ["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)
self.post_init() # required by transformers 5 (tied-weight bookkeeping); harmless on 4.x
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:]}