DNA-DiskChat-2B-PEER-v21 - shared-pool expert retrieval

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

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