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 <think>/</think> 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:
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
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.
- Downloads last month
- 167