Downloads · 30 days
0
dna-diskchat-2b-peer-v21
jaivial/dna-diskchat-2b-peer-v21
dna-diskchat-2b-peer-v21 is a machine learning model from jaivial. Use it for the machine learning task on the model card, and read the license before you ship it in a product. It is set up for pytorch. The card lists the license as apache-2.0.
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) bu…
Access
Public
Updated Jul 26, 2026
Repo size
23.4 GB
Likes
0
Public
Repo makeup
Click a slice to open those files.
.pt23.4 GB · 100%
At a glance
- Library
- pytorch
- License
- apache-2.0
- Access
- Public
- Created
- Jul 25, 2026
- Updated
- Jul 26, 2026
- SHA
- 295c8fae
Model card
From the Hugging Face model README
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_cudais not implemented for BFloat16. The pool stays fp32. torch.compilecannot wrap the sparse gathers: inductor raisesLoweringExceptiononaten.embeddingwithsparse=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
- Library
- pytorch
- License
- apache-2.0
- Languages
- en
- Created
- Jul 25, 2026
- Updated
- Jul 26, 2026