README.md

---
license: apache-2.0
language: [en]
library_name: pytorch
tags:
- disk-routed
- product-key-memory
- peer
- mixture-of-experts
- memory-layers
- ternary
- bitnet
- efficient-llm
- cpu-inference
---

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

![throughput](throughput_v20_vs_v21.png)
![summary](summary_v20_vs_v21.png)
![loss](loss_v20_vs_v21.png)

## 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 h*k = 8 / 16 / 32 - the gathered activation is what costs, so h*k=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