DNA-DiskChat-2B-PEER-v23

Ternary (BitNet b1.58) controller with a shared-pool PEER expert bank over the DNA codon memory table, trained from scratch on a single RTX 5060 Ti (16 GB).

v23 is a fresh run, not a continuation of v21 or v22. It exists because both earlier runs suffered a product-key collapse that v23 fixes — their checkpoints cannot be reused, since the collapsed keys are the bug.

Headline: 46,000 tok/s training AND a healthy router

46,340 tok/s sustained in the real training loop on one RTX 5060 Ti — 2.05x the first PEER baseline — while using ~75% of the expert bank instead of 0.03%.

	v20 baseline	v21	v22	v23
Training throughput	22,566	32,488	43,758	46,340 tok/s
Experts in pool	30,976 x16	495,616	692,224	692,224
Expert usage (measured)	—	0.1%	0.03%	healthy
Key pairwise cosine	—	—	+0.92 (collapsed)	+0.0002
PEER layers / pools	16 / private	3 / shared	2 / shared	2 / shared

The bug v23 fixes

Product-key retrieval scores a query against nk learned sub-keys per axis. In v21/v22 those keys trained into near-parallel directions — pairwise cosine rose from 0.000 at init to +0.90..+0.93, with norms spread over 2 orders of magnitude. Consequences:

only ~42 of 832 rows and ~33 of 832 cols ever won a top-k slot
-> reachable experts collapsed 692,224 -> ~1,400
-> experts actually used: ~230  (0.03%)
   (the PEER paper reports 99.98%)

So >99.9% of the expert parameters were dead weight.

Fix (two parts):

1. Cosine retrieval — unit-normalise both queries and keys, then scale by a learned temperature. A key can no longer win by growing its magnitude, which is exactly how the collapse happened.
2. Sub-key load-balance loss — 0.5*nk*(sum(pa^2)+sum(pb^2)) on the two nk-way marginals (cheap; the nk^2 product is never materialised), weight 1e-2.

Controlled A/B, 400 steps, identical seed and init:

	expert usage	key pairwise cosine
old dot-product	6.37%	+0.8617
cosine + balance	75.78%	+0.0042

Measured on the real 832x832 model during live training: key cosine +0.0002, key-norm std 0.003, router balance 2.00 of a possible 832 (1.0 = perfectly uniform).

The fix also made training faster (43,758 -> 46,340 tok/s): a balanced router spreads sparse-gradient indices, and embedding_bag's backward is faster on distributed indices than on a few hot rows where duplicate updates serialise on atomics. Speed and quality moved in the same direction.

Architecture

16 recurrent ternary blocks (chunk-parallel gated linear recurrence, BitNet b1.58 STE)
  blocks 4 and 12 : PEER FFN -> ONE shared pool, nk=832 -> 692,224 rank-1 experts
  other 14 blocks : ternary SwiGLU FFN, ff=768
routing  : cosine product-key, h*k = 8 active experts, query BatchNorm + sub-key balance
memory   : DNA codon table, 2048 x 2048 = 4.19M rows, 128 B/row (512 MB on SSD, 256 B/token)
optimizer: Muon (momentum every step, Newton-Schulz every 8) + fused AdamW
           + stateless sparse SGD on the expert pool

Throughput levers (all measured on the 5060 Ti)

Lever	Effect
cosine keys + sub-key balance	+2.6k tok/s and fixes expert usage
feed ids[:, :-1] (avoid padding 513 -> 528 in the chunked scan)	+2.2k
naive fused CE instead of chunked CE (per arXiv:2411.09009)	+1.4k
sync-free codon-mutation staging (no per-step int(...) host sync)	+0.5k
ff 1536 -> 768	+11%
no gradient checkpointing (it fits)	+4%
Muon accum-every-step / Newton-Schulz-every-8	+2.5%
async checkpointing (per arXiv:2406.10707)	removes ckpt stalls

Rejected after measurement (speedups that were quality regressions)

• Newton-Schulz with 2 iterations: svdvals mean 0.49 instead of ~1.0 — silently halves the effective Muon learning rate.
• Skipping muon.step(): discards 7 of 8 gradients. Replaced by an accum/apply split.
• Codon-gradient every 4 steps: halves genotype learning signal.
• bf16 expert tables: _embedding_bag_per_sample_weights_backward_cuda is not implemented for BFloat16.

Status

Base pretraining completed at 9,999,992,016 tokens (815,394 steps), averaging 46,357 tok/s. Final base weights are in checkpoints/ckpt_815394.pt.

UltraChat SFT and quality benchmarks

Assistant-only SFT ran for 6,000 steps on 197,952 UltraChat records. The 708.8M PEER expert parameters and codon rows stayed frozen; the 59.8M dense/router parameters trained. Full updated weights are in sft/sft.pt.

Held-out UltraChat assistant loss improved substantially:

	Base	SFT	Delta
CE	3.091	2.530	-0.562
Perplexity	22.01	12.55	-43.0%

Quality benchmark uses 100 deterministic random examples per task (seed 23), scored by normalized choice log-likelihood on the same Vast.ai RTX 5060 Ti:

Task	Base	SFT	Delta
ARC-Easy	42%	38%	-4 pp
HellaSwag	28%	27%	-1 pp
BoolQ	49%	60%	+11 pp

Honest result: SFT strongly improves chat-format likelihood and instruction style, but does not uniformly improve knowledge/reasoning. Greedy France QA now answers Paris; arithmetic remains poor and generations remain repetitive. Raw results and verbatim samples: sft/quality_benchmarks.json.

Reproduction scripts: sft/scripts/prepare_sft_dna.py, sft/scripts/train_sft_peer_v23.py, sft/scripts/quality_bench_peer_v23.py.

Standard LM Evaluation Harness

Full task splits evaluated with lm-evaluation-harness==0.4.12, 512-token context, no chat template. All tasks are zero-shot except MMLU, which uses the common 5-shot setting. Multiple-choice tasks report normalized accuracy where Harness provides it.

Benchmark	Examples	Score
ARC-Easy (acc_norm)	2,376	37.67%
ARC-Challenge (acc_norm)	1,172	24.15%
HellaSwag (acc_norm)	10,042	27.26%
PIQA (acc_norm)	1,838	58.49%
WinoGrande (acc)	1,267	51.93%
BoolQ (acc)	3,270	55.44%
OpenBookQA (acc_norm)	500	28.40%
TruthfulQA MC2 (acc)	817	46.60%
MMLU 5-shot (acc)	14,042	24.12%

All 57 MMLU subjects

These are weak 2B-class quality results: PIQA and BoolQ exceed random clearly; ARC-Challenge and MMLU sit around chance. SFT improved instruction behavior but did not turn this systems-focused model into a competitive general-reasoning model.

Raw Harness outputs: sft/lm_eval_general.json, sft/lm_eval_mmlu_5shot.json. Adapter: sft/scripts/lm_eval_peer_v23.py. Chart generator: sft/scripts/make_sft_benchmark_charts.py.

References

• He, Mixture of A Million Experts (PEER) — arXiv:2407.04153 (query BatchNorm, expert usage/unevenness metrics)
• Lample et al., Large Memory Layers with Product Keys — arXiv:1907.05242 (product keys, query BN)
• Berges et al., Memory Layers at Scale — arXiv:2412.09764 (shared pool, few layers, EmbeddingBag kernels)
• Wijmans et al., Cut Your Losses in Large-Vocabulary Language Models — arXiv:2411.09009
• Maurya et al., DataStates-LLM — arXiv:2406.10707 (async checkpointing)
• Ma et al., BitNet b1.58 — arXiv:2402.17764