A ternary (BitNet b1.58) controller whose FFN is a PEER product-key expert bank, on top of the DNA codon memory table. v21 is a training-throughput rework of the v20 design: the expert capacity is unchanged (~507M) but it is reorganised from 16 private pools into ONE shared pool read by 3 layers, following Memory Layers at Scale (arXiv:2412.09764).

Headline result: 1.44x faster training at 16x the expert granularity

Both runs used the same GPU model (RTX 5060 Ti 16 GB), the same streamed FineWeb-Edu+Cosmopedia corpus, the same codon table and the same optimizer stack.

	v20 (baseline)	v21 (this model)
PEER layers	16	3 (stride 4, at blocks 4/8/12)
Expert pools	16 private	1 shared
Experts per PEER layer	30,976	495,616 (16x)
Expert params	507.5M	507.5M (identical)
Controller params	569.4M	598.1M
Gathers per token	16	3
Avg throughput	22,566 tok/s	32,488 tok/s
Median / max tok/s	22,871 / 24,040	32,858 / 34,768
Tokens trained	1.553B	1.162B
Final valid CE	4.318	4.0317
Mean valid CE, matched 572-1145M tokens	4.1665	4.0213
Dense DNA-2B reference	45,700 tok/s	45,700 tok/s

v21 is 1.44x faster AND better per token. Over the token window both runs covered (572-1145M), v21's mean valid CE is 0.145 lower than v20's - so the throughput win did not cost sample efficiency, it improved it.

What actually made it faster (all measured on the GPU)

Lever	Effect
P1 one shared pool, 3 layers instead of 16 private	the dominant win - 5.3x fewer expert gathers per token for the same parameters
P2 `F.embedding_bag(mode='sum', per_sample_weights=...)`	replaces the `[N,h,k,d]` gather+einsum; verified numerically identical (max abs diff 5.2e-10 on CPU, 2.9e-6 fp32 GPU)
P4 stateless sparse SGD on the pool	`SparseAdam` OOMed (dense 4.06 GB exp_avg/exp_avg_sq); zero-state sparse SGD keeps it in VRAM
P5 gradient-checkpoint only the 3 PEER blocks	the other 13 blocks skip recompute
P6 grad-clip excludes the sparse pool	drops a 4.1 GB/step norm over 99.6% zeros
granularity `h*k`	measured 38.2k / 31.2k / 23.1k tok/s at hk = 8 / 16 / 32 - the gathered activation is what costs, so hk=8 was kept

Negative results (documented honestly)

bf16 expert tables did NOT work: _embedding_bag_per_sample_weights_backward_cuda is not implemented for BFloat16. The pool stays fp32.
torch.compile cannot wrap the sparse gathers: inductor raises LoweringException on aten.embedding with sparse=True. v21 compiles the recurrence, the dense FFNs and the PEER router, and leaves only the gather eager.
The projected 44k tok/s was not reached. CPU-proxy benchmarking overestimated the win; on the real GPU the recurrence and the eager gather region dominate. The honest measured number is 32,488 tok/s.

Architecture

16 recurrent ternary blocks (chunk-parallel gated linear recurrence, BitNet b1.58 STE)
  - blocks 4, 8, 12 : PEER FFN -> ONE shared pool, nk=704 -> 495,616 rank-1 experts
  - other 13 blocks : ternary SwiGLU FFN
codon memory table  : 2048 x 2048 = 4.19M rows, 128 B/row (512 MB, SSD-resident)
routing             : product-key, h*k = 8 active experts, query BatchNorm
optimizer           : Muon (recurrence) + fused AdamW (dense) + sparse SGD (pool)

Total experts: 495,616 per PEER layer (shared)
Active params/token: 598M controller total, ~70M active
Tokens seen: 1.162B (113,738 steps, run completed)
Final valid CE: 3.7246 (last logged periodic checkpoint 4.0317)

Files

path	what
`checkpoints/`	training checkpoints (`model` + `config`)
`scripts/model_dna_peer_v21.py`	the v21 architecture (shared pool, embedding_bag)
`scripts/train_peer_v21.py`	trainer with throughput logging
`scripts/make_peer_compare.py`	regenerates the charts below
`logs/throughput.jsonl`	per-100-step throughput series (the A/B raw data)
`logs/v21_train.log`, `logs/v20_baseline.log`	full training logs of both runs
`compare_metrics.json`	the numbers behind the table

Reproduction

bash

python train_peer_v21.py --batch 20 --seq 512 --nk 704 --topk 8 --pheads 1 \
  --target-tokens 3e9 --hf-repo jaivial/dna-diskchat-2b-peer-v21

Honest limitations

This is a systems/throughput artifact, not a chat model: no SFT has been done.
The two runs stopped at different token counts (v20 1.553B, v21 1.162B), so the quality claim is made only over the matched token window both runs covered.
The codon memory table is inherited unchanged from DNA-DiskChat-2B.

References

He, Mixture of A Million Experts (PEER) - arXiv:2407.04153
Berges et al., Memory Layers at Scale - arXiv:2412.09764 (3 shared layers, stride 4)
UltraMem - arXiv:2411.12364
Lample et al., Large Memory Layers with Product Keys - arXiv:1907.05242
Ma et al., BitNet b1.58 - arXiv:2402.17764