README.md

---
license: apache-2.0
base_model: allenai/Olmo-3-7B-Instruct
library_name: pushpuppet
pipeline_tag: text-generation
tags:
  - pushpuppet
  - olmo3
  - adaptive-inference
  - structured-pruning
  - elastic-model
---

# PushPuppet · Olmo-3-7B-Instruct (post-RL)

A **push-puppet** checkpoint fitted on top of
[`allenai/Olmo-3-7B-Instruct`](https://huggingface.co/allenai/Olmo-3-7B-Instruct).

This is **not** a drop-in `transformers` model. It is a *gated* checkpoint: the
base network is augmented with learned per-unit gates (B-spline gate curves over
FFN neurons and attention heads) so that a **single** set of weights yields a
whole family of valid subnetworks, indexed by a compression control variable
**λ**. Raising λ prunes more units; the gates were trained (mid-train → SFT → RL)
so every rung of the ladder stays a usable model.

It is intended to be served by the
[PushPuppet adaptive runtime](https://github.com/saudiwin/pushpuppet_runtime),
which picks λ from available memory at runtime and serves the resulting
subnetwork over an OpenAI-compatible API.

## Files

| File | Size | What it is |
|---|---|---|
| `model.pt` | ~27 GiB | Bare `torch.save` state dict, fp32, gates included |
| `train_args.json` | — | Records the base model id so the runtime can fetch config + tokenizer |

`model.pt` is a plain state dict (no config, no tokenizer). The runtime pulls
`config.json` and the tokenizer from the base repo named in `train_args.json`.

## Architecture

Matches `allenai/Olmo-3-7B-Instruct` exactly, plus gate parameters:

| | |
|---|---|
| params | 7.30 B (fp32) |
| layers | 32 |
| hidden / intermediate | 4096 / 11008 |
| attention | 32 query / 32 KV heads (MHA), head_dim 128, per-head QK norm |
| vocab | 100278 |
| max context | 65536 (YaRN, ×8 over 8192) |

Push-puppet gate hyperparameters (inferable from the tensors, so you don't have
to supply them): `n_knots=8`, `degree=3`, per-head QK norm enabled,
`scale_output=true`.

## Usage

```bash
git clone https://github.com/saudiwin/pushpuppet_runtime
cd pushpuppet_runtime
uv sync --extra torch

scripts/download_model.sh          # fetches this repo into models/olmo3_7b_instruct_post_rl
uv run pushpuppet up --ckpt-dir models/olmo3_7b_instruct_post_rl --lambda 1 --save
```

The runtime then serves an OpenAI-compatible API on `http://localhost:11435/v1`
plus a live dashboard at `http://localhost:11435/`, where you can move λ up and
down and watch the footprint change.

At ~27 GiB the checkpoint trips the runtime's **disk-first** loader
automatically (threshold 20 GiB): the state dict is memory-mapped and the pruned
model is materialized one layer at a time, so peak RAM is roughly the *pruned*
model rather than the dense one. Default precision is bf16; add
`--quantize int8` to roughly halve it again.

### Loading it yourself

You need the `push_puppet` research repo for the gate modules — the state dict
references `inject_stochastic_mlp` / `inject_stochastic_attn` parameters that
stock `transformers` does not define:

```python
import torch, sys
from transformers import Olmo3Config, Olmo3ForCausalLM
sys.path.insert(0, "/path/to/push_puppet/python")
import olmo3_mini_train as train_mod

cfg = Olmo3Config.from_pretrained("allenai/Olmo-3-7B-Instruct")
model = Olmo3ForCausalLM(cfg)
train_mod.inject_stochastic_mlp(model, temperature=0.5, n_knots=8, degree=3)
train_mod.inject_stochastic_attn(model, temperature=0.5, n_knots=8, degree=3)
train_mod.inject_per_head_qk_norm(model)
model.load_state_dict(torch.load("model.pt", map_location="cpu", weights_only=True))
```

Then call `structural_prune(model, lam)` to get a dense subnetwork at a given λ.

(`temperature=0.5` here is the **gate** temperature — the runtime's default, and
unrelated to sampling temperature at generation time.)

## License

Apache 2.0, inherited from the base model `allenai/Olmo-3-7B-Instruct`.