--- license: apache-2.0 language: - en library_name: transformers pipeline_tag: text-generation tags: - tinybrainbot - small-language-model - from-scratch - gqa - pretrained --- # TinyBrainBot 320M V2 — Base A **~326M-parameter** decoder-only language model **pretrained from scratch** on **~10B tokens** on 2× NVIDIA Tesla V100 (16GB). This is the **base** (pretrained + quality-annealed) model — a text-completion model with **no instruction tuning**. - For the chat / instruction-following version, see **`tinybrainbot-320mV2-instruct`**. - Successor to the 303M V2 base. **TL;DR:** A compact general-purpose base LM trained from scratch on ~10B tokens. It matches/beats **Pythia-410M** on general benchmarks (on far fewer training tokens) and cleanly sweeps **GPT-2-124M**. Its **2–3-digit arithmetic is notably strong for the parameter count** (94–99% on 2–3-digit addition) — a product of curated numerical pretraining data — though it hits a clear length-generalization wall beyond 3 digits and is not math-specialized. --- ## Model details | | | |---|---| | Parameters | **325,899,264** (~326M) | | Architecture | Decoder-only transformer, pre-norm, RMSNorm, SwiGLU MLP, RoPE | | Hidden size | 1024 | | Layers | 26 | | Attention heads | 16 (query) / **4 KV heads** (grouped-query attention) | | FFN size | 2816 | | Context length | 1024 | | Vocabulary | 32,000 | | Tokenizer | `tbb-32k-v2` — 32k BPE trained on the training mix (67% English / 20% code / 13% math), with reserved ``/`` special tokens | | Precision | trained in fp16 with an fp32 master copy (autocast) | ## Usage This is a **base completion model** — it continues text and does **not** follow a chat template. Prompt it with text to continue: ```python prompt = "The capital of France is" # -> the model continues the text ``` For Q&A / chat behavior, use the **instruct** variant. > **Format:** provided as **fp16 safetensors** (`AutoModelForCausalLM`). No GGUF is shipped for the base model — it's a raw completion model, and GGUF chat apps expect an instruction-tuned model. For chat / GGUF, use the **instruct** variant. --- ## Training procedure ![training loss curve](loss-curve-base.png) *Pretraining loss over 51,000 steps / ~10B tokens — the dip near step 40k is the quality-anneal phase.* Trained with the custom **TinyBrainBot** trainer using PyTorch DDP (gloo backend) across **2× Tesla V100-PCIE-16GB** on Windows, fp16 autocast + `GradScaler`, fused AdamW. One continuous **Warmup–Stable–Decay (WSD)** run. | Hyperparameter | Value | |---|---| | Total steps | 51,000 | | Tokens | 10.03B | | Global batch | 4 micro × 24 grad-accum × 2 GPUs × 1024 seq ≈ 196.6k tokens/step | | Peak LR | 6e-4 | | Warmup | 1,000 steps | | Schedule | WSD — stable LR, then cosine decay over the final ~20% (from step 40,800) | | Optimizer | fused AdamW | | Final loss | **1.436** | | Throughput | ~22k tok/s (2 GPU), ~4–5 day wall-clock | **Two-phase data:** a broad base phase (0 → ~40k steps) followed by a **quality-anneal** phase (~40k → 51k) in which the data mix was swapped to a knowledge-dense blend (Wikipedia leads up, web data halved, math tripled) as the LR decayed. The anneal pulled loss from ~1.62 to 1.436. **Pretraining data mix** (sampling weights; design principle: *real > synthetic, synthetic ≤ ~35%*): | Source | Weight | Notes | |---|---:|---| | dclm-clean | 3.0 | DCLM web slice, light-cleaned (~3.15B tok) | | wiki-dense | 2.0 | 5.2M Wikipedia article leads (knowledge-dense, ~1.66B) | | fineweb-edu-10bt | 1.3 | Educational web (~1B weighted) | | qa-distill | 0.6 | Synthetic Q&A | | bookcorpus-clean | 0.5 | Modern novels, boilerplate stripped (~1.68B) | | code-pyjs | 0.5 | the-stack filtered to Python/JS/TS (~0.48B) | | facts-distill | 0.4 | Synthetic scoped facts | | reasoning-distill | 0.3 | Synthetic reasoning | | tiny-stories-v2 | 0.25 | Simple narrative | | capability-distill | 0.2 | Synthetic capability data | | **math-v2** | 0.15 | **~70k** verified mul/div/decimal worked examples (tolerance-checked) | | writing-distill | 0.1 | Synthetic writing | | math-scratchpad | 0.05 | Step-by-step arithmetic | *(weights are sampling proportions; sum = 9.35. Explicit math ≈ 2.1% of the mixture — but ~215M tokens in absolute terms, the bulk of the model's math exposure.)* --- ## Evaluation Measured on our own log-likelihood MC harness (lm-eval style, fixed seed). Headline metric = `acc_norm` for HellaSwag/ARC/OpenBookQA, `acc` for WinoGrande/MMLU. Reference values are published lm-eval approximations — treat gaps under ~±2 points as ties. ### vs the previous 303M base (n=2000) | Benchmark | **320M V2 base** | 303M base | |---|:--:|:--:| | HellaSwag | **34.0** | 30.1 | | ARC-Easy (acc_norm) | **50.0** | 48.5 | | ARC-Easy (raw acc) | **57.0** | 51.8 | | ARC-Challenge | **27.4** | 26.4 | | OpenBookQA | **31.6** | 30.8 | | WinoGrande | **54.3** | 51.6 | | MMLU | **27.4** | 26.6 | → **Clean 6/6 sweep** over the previous generation. ### vs reference models (headline metric) | Benchmark | **320M V2 base** | GPT-2-124M | Pythia-410M | SmolLM-360M | |---|:--:|:--:|:--:|:--:| | HellaSwag | 34.0 | 31 | 34 | 54 | | ARC-Easy | 50.0 (57 raw) | 44 | 52 | 70 | | ARC-Challenge | 27.4 | 22 | 24 | 37 | | OpenBookQA | 31.6 | 29 | 30 | 42 | | WinoGrande | 54.3 | 52 | 53 | 57 | | MMLU | 27.4 | 26 | 25 | 34 | → Sweeps **GPT-2-124M** 6/6; **~5 wins + 1 draw vs Pythia-410M**. SmolLM-360M (trained on ~600B aggressively-filtered tokens) remains the out-of-reach frontier for this size. **Training efficiency.** These results come from **~10B pretraining tokens** — roughly an order of magnitude fewer than the Pythia suite's ~300B. The Pythia-410M parity is best read as a **token-efficiency** result (curated data + quality anneal) rather than a scale win. ### Arithmetic (GPT-3 Arithmetic protocol, exact-match, n=300/sub-task) | Sub-task | Accuracy | |---|:--:| | 2-digit addition | 99.3% | | 3-digit addition | 94.7% | | 2-digit subtraction | 50.0% | | 3-digit subtraction | 42.3% | | 4-digit addition / subtraction | 0.3% / 0.3% | | 5-digit addition / subtraction | 0.0% / 0.0% | | 2-digit multiplication | 21.0% | | single-digit composite (order of ops) | 3.7% | | **Aggregate (all 10 sub-tasks)** | **31.2%** | → **Strong through 3 digits, then a hard wall.** 2–3-digit addition is near-solved (94–99%, with correct carrying), but 4+-digit accuracy collapses to ~0%: the model executes a fixed **~3-column** addition routine and silently drops the higher place values — a **length-generalization limit** tied to the training distribution (`math-v2` operands are ≤3 digits), *not* truncation (generations complete normally). Subtraction sits ~42–50% — it handles `a−b` when `a>b` but drops the sign on **negative results**. 2-digit multiplication ~21%; single-digit composite ~4%. > **How to read this:** the GPT-3 Arithmetic suite mainly probes **exact symbolic computation** and short-range algorithmic generalization; it should **not** be read as evidence of mathematical *reasoning* (word problems), which is a separate, much harder skill at this scale. --- ## Intended use & limitations **Intended use:** research on small-model training, base for further fine-tuning, text completion, and arithmetic computation. **Limitations:** - Base completion model — **does not follow instructions or a chat format** (use the instruct variant for that). - **Math reasoning** (word problems) is at the floor — the model computes but does not reason through multi-step problems. - **Negative-result subtraction** is unreliable (drops the sign). - **WinoGrande and MMLU** sit near the random floor — consistent with the capacity and data limits of a ~326M model under this training recipe. - Trained predominantly on English; 1024-token context; no safety tuning — outputs may be incorrect or inappropriate and should not be relied upon unchecked. ## Hardware & framework 2× NVIDIA Tesla V100-PCIE-16GB · Windows · PyTorch DDP (gloo) · fp16 autocast (fp32 master) · fused AdamW · custom TinyBrainBot trainer.