README.md

---
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