Kotodama 3B Base (Final)
A 3B parameter language model trained from scratch with Block Attention Residuals and NCA pre-pretraining.
This is the final base checkpoint: the full 384B-token schedule — 346B tokens at peak LR + 38B-token cosine cooldown to LR=0 (steps 175,780 → 195,311), completed 2026-06-09. It is the best checkpoint of the run on every internal and external evaluations.
Training compute provided by partnership with Anima Labs.
Architecture
- Parameters: 2.97B
- d_model: 3072, n_layers: 28, heads: 24 query / 8 KV (GQA), head_dim: 128
- FFN: SwiGLU (intermediate 8192), RMSNorm + QK-norm, RoPE (theta=500K)
- Vocab: 49,152 (SmolLM2 tokenizer), tied embeddings, no bias, z-loss 1e-5
- Block Attention Residuals: DD-3B boundaries
[0,1,3,7,15,19,24]— learned routing over depth at each sublayer - Optimizer: Muon (lr=0.02) for 2D weights, AdamW for embeddings/norms
The exact training config is included in this repo as 3b-language.yaml.
Training
- NCA pre-pretraining: 5.9B tokens of random data to initialize attention circuits before language training (embeddings reinitialized for the language vocab)
- Language pretraining: 384.3B tokens, single epoch, seq_len 4096 with document-masked packing, cosine cooldown over the final 10% of steps
- Data: curated 32-source mix. Largest shares: the-stack v1 18.5%, FineFineWeb 17.2% (+2.2% backfill), peS2o 15.8%, US patents 9.6%, Pile-of-Law 4.7%, pre-1929 books 4.1%, StackExchange 4.1%, OpenWebMath 3.5%, Library of Congress 3.5% — plus 22 smaller sources (Reddit, PG-19, FineMath, Wikipedia, subtitles, poetry, …)
- Infrastructure: 8x NVIDIA B200, DDP, FP8, torch.compile; 285K tok/s steady state
- Health: no BOS-sink at any point (deep-layer attention entropy flat at 4.8–5.4), zero dead units across the entire run, stable RankMe ≈ 1718
Evaluations (bf16)
lm-evaluation-harness 0.4.11, zero-shot:
| Task | chinchilla-66B | final-384B |
|---|---|---|
| HellaSwag (acc_norm) | 36.3 | 46.7 |
| PIQA (acc) | 64.5 | 68.4 |
| ARC-Easy (acc) | 51.6 | 55.1 |
| ARC-Challenge (acc_norm) | 24.4 | 27.3 |
| BoolQ (acc) | 58.4 | 61.9 |
| COPA (acc) | 68.0 | 71.0 |
| SciQ (acc) | 82.6 | 87.0 |
| Winogrande (acc) | 52.4 | 55.6 |
| LAMBADA (acc / ppl) | 38.2 / 23.4 | 49.7 / 11.1 |
| WikiText (word_ppl) | 26.08 | 17.75 |
UncheatableEval-2026-04 (bits-per-byte on post-cutoff data, 15 domains): mean 0.852 vs 0.980 (chinchilla) — wins all 15 domains. Strongest: arxiv/github (0.65–0.72); weakest: non-English (1.30–1.76). Cooldown isolation on near-token-matched checkpoints: the 6B cosine decay alone accounts for −5.7% mean BPB.
Evaluate in bf16. The training-time train/loss telemetry (fp8 + compile path) is a noisy estimator and not a reliable quality signal — it rose during the cooldown while the model improved on every held-out eval. All quality claims here are from bf16 evals.
Usage
This checkpoint requires the kotodama model code to load.
git clone https://github.com/LuxiaSL/kotodama.git
cd kotodama
# Serve interactively
python serve.py --checkpoint /path/to/step_00195311.pt.zst --model_size 3b --port 2222
# Then query:
curl http://localhost:2222/v1/completions \
-d '{"prompt": "The theory of everything", "max_tokens": 200, "temperature": 0.7}'
Or load the weights directly:
import io, torch, zstandard
raw = zstandard.ZstdDecompressor().stream_reader(open("step_00195311.pt.zst", "rb")).read()
ckpt = torch.load(io.BytesIO(raw), map_location="cpu", weights_only=False)
state_dict = ckpt["model"] # -> load into the model with the DD-3B config in 3b-language.yaml
Sampling note: use pure temperature sampling (no top-p) — top-p degraded quality in our evals.
Checkpoint format
Raw PyTorch checkpoint (.pt.zst, zstd-compressed, ~9.4GB; ~30GB decompressed). Contains model state dict, both optimizer states (Muon + AdamW), scheduler, and training metadata (including the fixed probe batch). The model code handles decompression automatically.
License
Apache 2.0