---
license: apache-2.0  # TODO: confirm — weights are trained-from-scratch; FineWeb-Edu data is ODC-BY
language:
  - en
library_name: transformers
pipeline_tag: text-generation
tags:
  - matryoshka
  - nested-models
  - speculative-decoding
  - model-suite
  - distillation
  - custom_code
datasets:
  - HuggingFaceFW/fineweb-edu
---

# Matryoshka-3B — a nested 500M / 1.5B / 3B LM suite

A **Matryoshka** language-model suite: three sub-models of increasing size
(**500M ⊂ 1.5B ⊂ 3B**) stacked into a single nested architecture and trained
end-to-end in one run. Running the full model produces all three sub-models'
outputs in a single forward pass, while activating fewer parameters and writing
less KV cache than three independent models.

This reduces the total parameter count of the suite (**3.20B vs. 5.20B** for an
equivalent set of independent models, **−38%**), enables free online distillation
from the largest sub-model to the smaller ones at every step, and is well-suited
to speculative decoding — the draft model is *contained within* the verifier and
shares its early layers and KV cache.

> ⚠️ Custom architecture — load with `trust_remote_code=True`.

## Suite architecture

All sub-models share the SmolLM2 tokenizer (`V = 49,152`) and use RoPE
(`θ = 1e5`). Sub-model `m+1` consumes sub-model `m`'s output through a
parameter-free **norm-rescaled junction**: `o^m` is rescaled to match the norm of
a fresh embedding covering the new channels, then concatenated.

| Sub-model | Incr. params | Cumul. params | Width `D` | Layers (cumul.) | Heads | Head dim | Intermediate |
|---|---|---|---|---|---|---|---|
| 500M | 0.50B | 0.50B | 1024 | 24 (24) | 16 | 64  | 4096  |
| 1.5B | 0.98B | 1.48B | 2304 | 10 (34) | 24 | 96  | 9216  |
| 3B   | 1.72B | 3.20B | 4352 |  5 (39) | 34 | 128 | 17408 |

Depth triplet `(n₁,n₂,n₃) = (24,10,5)` (39 layers total) was chosen to match the
KV-cache and per-token FLOPs of a Vanilla 3B baseline. The 500M sub-model uses the
same width/depth as the Vanilla 500M for a strictly comparable data point.

## Checkpoints & revisions

**Every checkpoint is a separate git branch**, selected via `revision=`. Branch
names follow:

```
{exit}_{tokens}B[_cd|_distill]        e.g. 3B_35B, 1-5B_42B, 3B_60B_cd
```

- **`{exit}`** — how many nested exits the checkpoint exposes:
  - `500M_*` → the 500M sub-model only
  - `1-5B_*` → the 500M + 1.5B sub-models
  - `3B_*` → the full 500M + 1.5B + 3B suite
  - `main`, `main_{tokens}B` → pointers to the full suite (`main` = default/latest)
- **`{tokens}B`** — training tokens seen: `5, 10, 16, 21, 26, 31, 35, 42, 52, 60`.
- **`_cd`** — WSD **cooldown** applied (learning-rate decay), released at 60B.
- **`_distill`** — distillation-ablation variant (small exits only). Will be discussed in a next version of the paper.

The paper's main results use the **35B-token** checkpoints; the repo additionally
ships a longer **60B** run with cooldown (`*_60B_cd`) and intermediate checkpoints
for studying training dynamics.

Because the suite is nested, `500M_35B`, `1-5B_35B` and `3B_35B` are **slices of
the same trained weights** — load the largest exit you need.

## Usage

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "nthngdy/matryoshka-3B",
    revision="3B_35B",          # full suite; use "1-5B_35B" or "500M_35B" for smaller
    trust_remote_code=True,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
).eval()

tokenizer = AutoTokenizer.from_pretrained("nthngdy/matryoshka-3B", revision="3B_35B")

# The nested sub-models are exposed as a dict, ordered small → large:
for tag, submodel in model.lm_model_dict.items():
    n = sum(p.numel() for p in submodel.parameters())
    print(f"{tag}: {n/1e6:.0f}M params")
```

A single forward pass through the full suite yields logits for **all** contained
exits. For per-exit generation and for **speculative decoding** with a small exit
drafting for a larger one (sharing early layers + KV cache).

## Training

| | |
|---|---|
| Data | FineWeb-Edu, sequences packed to length 2048 |
| Tokens | 35B (main); extended 60B run also released |
| Optimizer | AdamW (β₁=0.9, β₂=0.95, ε=1e-8) |
| Peak LR | 4e-4, WSD schedule (3,000 cooldown steps / 33,000 total) |
| Batch size | 512 sequences |
| Weight decay | 0.01 · Grad clip | 1.0 |
| Precision | bf16-mixed |
| Distillation | online, from the 3B exit to smaller ones, α_d = 0.3 |
| Hardware | NVIDIA B200 (~52 GPU-days across both suites) |

## Results (from the paper)

- **On par** with independently-trained baselines on validation PPL, out-of-domain
  PPL, and benchmark accuracy, while using **36% less training compute**.
- **14–26% higher throughput** via speculative decoding. A 500M/3B pair (1:6 ratio)
  that *degrades* latency for independent models becomes a **20–40% speedup** here,
  thanks to shared KV cache and early layers.

## Intended use & limitations

Research artifact for studying nested LM suites, elastic-size deployment, and
draft-in-verifier speculative decoding. Base (non-instruct) models trained on
English FineWeb-Edu (~35–60B tokens); not aligned or safety-filtered. Outputs may
be low-quality, biased, or factually wrong. Not intended for production use.

## Citation

```bibtex
TBD
```