pushpuppet-olmo3-7b-instruct

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

A push-puppet checkpoint fitted on top of 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, which picks λ from available memory at runtime and serves the resulting subnetwork over an OpenAI-compatible API.

Files

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:

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

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:

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.