π₀.₅ LIBERO — W2A8 (2-bit weights, INT8 activations), 10- and 2-step flow schedules

Deployable 2-bit checkpoint of π₀.₅ fine-tuned on LIBERO (pi05_libero, PaliGemma 3B prefix + 300M action expert). Every one of the 414 Linear layers of the vision tower, language model and action expert is stored at 2.02 bits/weight (two-plane 8-dim vector quantization, per-row FP16 scale, signed Walsh–Hadamard input rotation); the token embedding is INT8; everything else keeps its original dtype. Whole-model 2.92 bits/weight, 1.23 GB on disk (BF16 original: 6.42 GB). Activations into the quantized layers are per-token symmetric INT8 (A8); softmax / QKᵀ / norms / KV cache are BF16.

This repository exists so the model can be profiled on edge hardware (Jetson): latency, memory, cache behaviour. The accuracy rows below were produced on data-centre GPUs with the runtime described in Runtime.

Files

Format (hd-srvq-evening-v1)

For a quantized Linear <name> with N outputs and K inputs:

Decoding: W_rot[n, :] = row_scale[n] · concat_g (codebooks[0][codes[0,n,g]] + codebooks[1][codes[1,n,g]]), an (N, K) matrix in the rotated input basis. The layer computes y = R(x) @ W_rotᵀ + bias with R(x) = blockwalsh(x ⊙ signs), where blockwalsh is the direct sum of orthonormal Sylvester-ordered Walsh–Hadamard blocks over the descending maximal-power-of-two decomposition of K (1024 → [1024]; 1152 → [1024, 128]; 4304 → [4096, 128, 64, 16]; 16384 → [16384]). The equivalent dense weight in the original basis is W = (W_rot @ H_block) ⊙ signs (H is symmetric and orthogonal). Token embedding: embed_tokens.codes int8 (V, 2048) × embed_tokens.scales fp32 (V, 1). dequantize.py implements all of this (verified against the fitting runtime to 6e-7 relative).

from dequantize import dense_state_dict, decode_rotated, rotate, fake_quant_per_token
sd = dense_state_dict('model.safetensors', 'vq.json')          # dense BF16 weights, openpi PI0Pytorch key names
# real W2A8 path for one layer:  y = int8_gemm( int8(rotate(x, signs)) , int8(W_rot) ) * token_scale * row_scale (+ bias)

Quantized layer names are the keys of vq.json["backbone_layers"] (288: SigLIP + Gemma-2B) and vq.json["layers"] (126: action expert). Key names follow openpi's PI0Pytorch (paligemma_with_expert.paligemma.model.…, paligemma_with_expert.gemma_expert.model.…).

Runtime

• Model: openpi pi05_libero converted to PyTorch with openpi's examples/convert_jax_model_to_pytorch.py (paligemma_variant=gemma_2b, action_expert_variant=gemma_300m, action_dim=32, action_horizon=10, 224×224 images, LIBERO prompt format, norm stats in assets/). Load that model, then replace the 414 Linear weights and the embedding with the decoded tensors (dense path) or install a W2A8 Linear that applies rotate + per-token INT8 on the input and an INT8×INT8 GEMM with the decoded INT8 W_rot (w8a8_runtime.W8A8Linear, built from model_int8grid.safetensors).
• Two weight files. model.safetensors is the accuracy checkpoint: its codebooks are FP16 and its decoded weights are not integers. model_int8grid.safetensors is that checkpoint snapped to the INT8 grid (per layer: gain = 120 / max|decoded|, codebooks = round(codebooks × gain), row scale ÷ gain; w8a8_runtime.snap_layer reproduces it bit for bit). Snap error: relative weight MSE 1.9e-4 mean, 6.8e-4 max over the 414 layers. The accuracy table below was measured on model.safetensors with fake-quantized A8; the INT8 rows of the efficiency table were measured on the snapped file with real INT8 kernels. The snapped file has no rollout evaluation of its own.
• A8 (as in all accuracy rows): per-token symmetric INT8 of the rotated input of every quantized Linear, scale = amax/127, round-half-to-even, clamp ±127; softmax, QKᵀ, norms, residual adds and the KV cache stay BF16.
• Flow schedule: the action expert's adaptive-RMSNorm time conditioning is stored as per-step tables (<norm>.dense.values, one row per flow step, 3072 = scale/shift/gate × 1024) at fixed timesteps fixed_schedule_times. model.safetensors carries the native 10-step schedule (t = 1.0, 0.9, …, 0.1); tables/tables_2step.safetensors carries the 2-step schedule (t = 1.0, 0.5): swap the 37 *.dense.values tensors and fixed_schedule_times, and call sample_actions(num_steps=2). The tables are exact: the base model's time MLP evaluated at those timesteps gives the same values, so a dense reconstruction can equally keep the original time MLP.
• Action horizon 10 (10 action tokens per expert step); the LIBERO evaluations execute 5 actions per policy call (openpi's LIBERO default).

Accuracy (10-step schedule, A8 fake-quant runtime, paired on identical episodes and noise seeds; 95 % CI from a suite-stratified bootstrap)

The same 2-bit format without the readout fine-tune scores 657 / 704 on the two Plus sets (−4.1 pp / −3.7 pp vs this checkpoint).

Efficiency reference (RTX 6000 Ada, batch 1, CUDA graphs, median of 100 replays; what to compare Jetson numbers against)

Runtime package: runtime/ holds the authors' loader, INT8 kernels (including the fused Triton kernel of the 2-bit-resident row), the fake-quant A8 path and the benchmark. One command: bash runtime/run.sh — see runtime/README_RUN.md (中文). The repackaged entry points were not re-run on a GPU after extraction.

Logical weight bytes read per inference with INT8-decoded weights: 5.57 GB at 10 steps, 3.04 GB at 2 steps (prefix once, expert per step). On a data-centre GPU the W2A8 path is not faster than BF16 (prefill is compute-bound; the expert's 10-token GEMMs are launch-bound); the expected gains are memory footprint and bandwidth on edge devices — that is what this release is for.

Provenance

Base: pi05_libero from openpi (Physical Intelligence, Apache-2.0). Quantization: 2-bit vector quantization of every Linear with an action-Fisher-weighted layer-wise objective, followed by a short readout-consistency fine-tune of the codebooks / row scales / norm weights (backbone and expert weights stay in the 2-bit format; payload identical to the plain PTQ checkpoint). Calibration data: 256 LIBERO observations (never LIBERO-Plus test instances). SHA-256 of model.safetensors: see vq.json["checkpoint_sha256"].