README.md

metadata

license: apache-2.0
base_model: Zyphra/Zamba2-2.7B-instruct
tags:
  - litert
  - litert-lm
  - litertlm
  - on-device
  - edge
  - hybrid
  - mamba2
  - zamba2
pipeline_tag: text-generation
library_name: litert-lm

Zamba2-2.7B-instruct — LiteRT-LM

Zyphra/Zamba2-2.7B-instruct converted to the LiteRT-LM (.litertlm) format for on-device inference with Google's LiteRT-LM runtime. Requires litert-lm ≥ 0.15.

Zamba2-2.7B is Zyphra's shared-attention hybrid at its middle size: a Mamba2 selective-scan backbone (54 layers) with two shared transformer blocks applied alternately at 9 interleaved positions — two sets of attention+MLP weights, each reused at its positions and specialized by per-position LoRA adapters on the MLP, attending over the concatenation of the running hidden state and the original embeddings (no rotary embedding at this size). Together with our Zamba2-1.2B this is, to our knowledge, the first Zamba2 conversion to a mobile runtime.

File	Recipe	Size
Zamba2-2.7B-instruct_int8.litertlm	int8 dynamic on linears + embedding (convs and the scan stay float); fp32 activations declared for GPU	2.80 GB

Correctness

• Logits parity vs PyTorch: the float export matches the HF model teacher-forced across 48 decode positions — top-1 and top-5 identical at every position, mean per-position logit correlation 1.0000, mean KL ≈ 0.
• 8-question sanity gate: 8/8 on every lane — GPU and CPU, litert-lm 0.15.0 and 0.16.0. No degeneration, no greedy flips.
• Prompt-length robustness: hermetic prefill-chunk sweep (fresh engine per length) — CPU fills 12–51 and GPU fills 12–31 all clean.
• iPhone 17 Pro: the 8-item composite probe answers 8/8 on the CPU backend, no degeneration. The Metal backend does not run this size on a 12 GB phone — see the honest note below.

A conversion-side note: current transformers (5.14.x, and main as of 2026-08-14) cannot load ANY two-block Zamba2 checkpoint (2.7B/7B) — Zamba2Model.get_layers assigns block_id by global layer index while the checkpoint layout and its own weight-tie cycle follow hybrid occurrence order, so construction raises before weights load. The conversion patch carries a structure fix (block_id by hybrid order), verified key-exact against the published checkpoint.

Usage

litert-lm run ./Zamba2-2.7B-instruct_int8.litertlm --prompt "What is the capital of France? Answer in one word."

# GPU
litert-lm run ./Zamba2-2.7B-instruct_int8.litertlm --backend gpu --cache no --prompt "..."

Six prefill signatures (1024, 256, 64, 16, 4, 1) are exported so the runtime picks tight chunks. A signature costs memory whether or not it is called, so the ladder is deliberately shorter than the full 1–1024 one. The bundle carries the tokenizer and the stock ChatML Zamba2 chat template.

Performance

litert-lm benchmark (litert-lm 0.16.0), Apple M4 Max, -p 256 -d 256 --runs 3 --cache no, quiet machine:

Backend	Prefill (256)	Decode	TTFT
GPU	603 tok/s	43.6 tok/s	0.45 s
CPU	282 tok/s	14.7 tok/s	0.97 s

On device (cold start, single run, 140-token composite prompt, quality harness):

Device	Backend	Prefill	Decode	TTFT	Peak memory
iPhone 17 Pro (12 GB)	CPU	14.3 tok/s	3.6 tok/s	10.2 s	2.19 GB
iPhone 17 Pro (12 GB)	GPU (Metal)	— does not load, see below —

Honest notes:

• A 12 GB iPhone cannot run this size on Metal. The GPU path holds the weights as fp32, so the resident footprint is about 4× the int8 file — measured at 3.98× on our Zamba2-1.2B (1364 MB file → 5432 MB peak). At 2.80 GB that projects to ≈ 11 GB against the phone's 11.7 GB, and engine creation is killed by the OS every time (deterministic across runs, and unchanged when the prefill ladder was cut from 12 signatures to 6 — the ladder was never the binding term). Use the CPU backend on a phone, or the GPU on a Mac. The 1.2B sibling is the one to reach for if you want Metal on a phone.
• CPU on a phone is slow at this size: 3.6 tok/s decode with a 10-second first token. Correct, but a background/batch pace rather than an interactive one.
• GPU runs with fp32 activations (declared in the bundle). The nine shared-attention positions carry wide KV caches (32 KV heads × 160 head dim at 4096 context) held in fp32 on GPU.

Conversion notes

Converted with litert-torch plus a hybrid-cache patch (reproduction script + patch: hf-to-litertlm zamba2_work/):

• Folded selective scan: the Mamba2 scan is re-expressed as batched matmuls with chunk and head axes folded into the batch axis (all tensors rank ≤ 4, no BROADCAST_TO, no int64 index math) — this is what makes the graph fully delegable on GPU.
• Two-block structure fix: block_id is assigned by hybrid occurrence order (matching the checkpoint layout and the tie cycle) — without this, current transformers cannot construct the model at all.
• Min-only dt clamp handling: Zamba2 clamps softplus(dt) at time_step_min with no upper clamp; padded prefill positions are forced to exact identity steps AFTER the clamp (without this, every runtime pad token decays the recurrent state).
• Shared blocks + adapters: each tied transformer block traces once per position with its own MLP LoRA adapter statically selected; tied weights are stored once.
• Composite hybrid cache layer: the 9 shared-attention positions hold KV + conv + recurrent state at ONE layer index (the runtime binds states by tensor name, so co-residency is just packaging).
• Prefill-pad guard: the runtime runs partially-filled prefill chunks; pad positions are made exact identity steps for the SSM and the stored conv window is gathered at the last valid column.
• Streaming detokenization: the tokenizer's Strip decoder is removed from the bundle — Zamba2's metaspace (SP-BPE) tokenizer otherwise loses every interior space under the runtime's per-token streaming decode; the only behavior change is a sequence-initial space, which the runtime trims.
• Quantization: post-hoc dynamic int8 over linears + embedding only; convs and the scan stay float.

License and changes

Distributed under Apache 2.0 (inherited from the base model). Changes from the original work: weights converted from safetensors to LiteRT flatbuffers and quantized as described above; tokenizer and chat template repackaged unmodified. This repository is a community conversion and is not affiliated with Zyphra.