dna-diskchat-2b-peer-v23

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

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:

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)

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:

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:

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.

<details> <summary>All 57 MMLU subjects</summary>

</details>

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