kotodama-3b-base-final

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:

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