arianraje/mimo-7b-mamba2-hybrid-1B-OPD

mimo-7b-mamba2-hybrid-1B-OPD

Stage-3 on-policy distillation, 1B rung of the WSD ladder. The student generates; the teacher scores its own rollouts on reasoning traces. Warmup -> flat plateau -> linear decay, keyed on consumed tokens rather than steps, so each rung extends the previous one's pre-decay weights rather than restarting.

Measured

About this arm

Part of a 2x2 study converting full-attention models into linear-attention hybrids (uniform 1:4 retention -- every 4th layer keeps softmax attention, the other 27 are replaced) and recovering capability by staged distillation rather than pretraining. This repo is the MiMo-7B-RL-0530 x Mamba2 cell; the matched MiMo x Gated DeltaNet cell is at arianraje/mimo-7b-gdn-hybrid-* and is directly comparable -- same teacher, same packed data, same token budgets, same schedules, differing only in the mixer.

Pipeline: surgery/init -> stage-1 hidden-state alignment (100M tok) -> stage-2a forward-KL KD @4k (600M tok) -> stage-2b long-context KD @32k (294M tok) -> stage-3 on-policy distillation, laddered 200M -> 400M -> 600M -> 800M -> 1B generated tokens.

Loading

The architecture is custom (model_type: mimo_mamba2). Remote code is included, so trust_remote_code=True is all that is needed:

from transformers import AutoTokenizer, AutoModelForCausalLM

repo = "arianraje/mimo-7b-mamba2-hybrid-1B-OPD"
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(
    repo, trust_remote_code=True, dtype="bfloat16", device_map="auto")

For vLLM serving, download the snapshot, put it on sys.path, and call register() from vllm_mimo_mamba2.py before constructing the engine.

Mamba2's fast path needs mamba-ssm. Without it is_fast_path_available is False and the mixer falls back to a materializing reference implementation: correct, but far slower and far more memory-hungry.

Evaluations

Compact raw evaluation JSONs are under evals/. *.generations.json dumps are deliberately excluded.