| --- |
| license: apache-2.0 |
| language: |
| - en |
| tags: |
| - from-scratch |
| - muon |
| - attention-residuals |
| - nca |
| pipeline_tag: text-generation |
| --- |
| |
| # 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](https://animalabs.ai/). |
|
|
| ## 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](https://github.com/LuxiaSL/kotodama) to load. |
| |
| ```bash |
| 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: |
| |
| ```python |
| 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 |
| |
