Upcycling existing checkpoints (e.g. Comma)

[rlrs]

Why
We have decent dense MHA-style checkpoints but they leave many arch wins on the table: KV-cache size, limited context, slow decoding, fewer params than we'd like, and so on. This experiment explores “upgrade-in-place”-style upcycling so we can bolt on modern arch innovations with continual pretraining to hopefully recover performance.

Potential tracks (v1)

1. Attention: MHA → GQA (→ MLA)
   Cut KV cache by sharing K/V across Q heads (GQA), optionally compress further with latent K/V (MLA). CPT to recover. (https://aclanthology.org/2023.emnlp-main.298/)

2. Long-context via RoPE extensions (PI / YaRN / LongRoPE)
   Rescale RoPE and CPT to 32k-128k+. (https://arxiv.org/abs/2309.00071)

3. Dense → MoE FFN
   Swap some MLP blocks for sparse experts; initialize from dense weights and finetune. Recent work gives explicit upcycling inits. (https://arxiv.org/abs/2410.07524)

4. Matryoshka / Nested capacity (MatFormer-style)
   Add nested sub-FFNs so the same checkpoint runs at multiple “sizes” by truncation. Train heads jointly; enables elastic inference. (https://arxiv.org/abs/2310.07707)

5. QK-Norm / Positional re-param for stability
   Normalize Q,K (QK-Norm) and/or adopt length-extrapolatable RoPE variants (e.g. XPos) to stabilize long-ctx finetunes. (https://aclanthology.org/2020.findings-emnlp.379.pdf)

6. Retrieval block (RETRO / Memorizing Transformer)
   Bolt on a small cross-attend reread to retrieved chunks or kNN memory to improve knowledge tasks without growing the backbone. (https://proceedings.mlr.press/v162/borgeaud22a.html)

7. Multi-Token Prediction (MTP) heads
   Attach auxiliary heads that predict k future tokens (Medusa/EAGLE family, DSv3). Unlocks faster decoding and can improve sample-efficiency. (https://arxiv.org/abs/2401.10774)

8. Sparse/windowed attention + a few globals
   Replace some full attention with sliding-window + learned global tokens for linear-ish scaling at long context. Must look at attention sinks or other remedies. (https://arxiv.org/abs/2004.05150)

9. Vocabulary expansion
   Extend tokenizer (append merges), resize embeddings/LM head, init new rows sensibly. (https://pure.itu.dk/files/107285547/2407.05841v2.pdf)

10. Hybrid attention–SSM top-blocks (Transformer↔Mamba)
    Interleave a small number of Mamba/SSD layers near the top for throughput and long-range memory, keeping most of original intact. (https://arxiv.org/abs/2403.19887)

Implementation sketch

Need to explore how to make these "surgeries" available in a conveniently ergonomic interface. Not sure how to do this.
Most modifications need to happen before CPT, some need to happen in mid-training (e.g. YaRN).
I'm thinking something like each architectural upgrade we want to experiment with gets two pieces: (1) a new “flavor” preset registered beside the existing models so the training stack knows how to build that variant, and (2) a standalone script under scripts/ that converts a baseline checkpoint/tokenizer into the new layout. The script reads an input artifact bundle, rewrites weights or vocab as needed, saves the upgraded checkpoint and tokenizer, and emits a tiny config overlay (JSON/TOML) explaining which model_name/flavor fields to use when launching training. If you want to stack upgrades, you just point the next script at the previous output directory and apply the next overlay.

Gates
Each track needs to demonstrate some pareto win in terms of e.g. KV cache size vs. quality. It will be important to measure these and weigh pros/cons to judge success/failure. Some potential ideas for things to measure:

• Tokens/s and p50/p95 latency at {4k, 16k, 32k}.
• KV bytes/token and peak VRAM.
• For MoE: activated FLOPs and router balance.
• For decoding: speedup factor at batch=1 and batch=16.