ForeverBlue/pi05-libero-w2a8

π₀.₅ 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

file	content
model.safetensors	the checkpoint (1,225,476,780 B; format hd-srvq-evening-v1, spec below); contains the 10-step flow tables
vq.json	per-layer metadata (shapes, payload bytes), bit accounting, SHA-256 of the checkpoint
tables/tables_2step.safetensors	flow-step tables for the 2-step schedule (37 sites × (2, 3072) FP32 + fixed_schedule_times = [1.0, 0.5])
tables/tables_2step.json	how the 2-step tables were captured and validated (bit-identical to the native 2-step actions on 256 held-out observations)
assets/physical-intelligence/libero/norm_stats.json	openpi normalisation statistics of pi05_libero (unchanged)
dequantize.py	reference reader: decodes every layer to a dense BF16 weight, or gives the rotated weight + signs for a real W2A8 path

Format (hd-srvq-evening-v1)

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

tensor	dtype	shape	meaning
<name>.codes	uint8	(2, N, K/8)	one byte per group of 8 weights and plane: index into that plane's codebook
<name>.codebooks	fp16	(2, 256, 8)	two 8-dim codebooks of 256 entries ("planes"); 1 bit/weight each
<name>.row_scale	fp16	(N,)	per-output-row scale
<name>.signs	int8	(K,)	±1 per input column (sign flip before the Hadamard)
<name>.masks / .mask	uint8	(10,) / (1,)	planes read at each flow step; all = 3 (both planes) in this checkpoint
<name>.bias	bf16	(N,)	only where the original layer has a bias (SigLIP)

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 (the "snapped" codebooks are integer-valued up to a per-layer gain; see dequantize.py and the efficiency notes).
• 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)

benchmark	W2A8 (this checkpoint)	FP (BF16)	Δ vs FP
LIBERO-Plus selection 1000 (500 camera-viewpoint + 500 robot-initial-state perturbations)	741	765	−2.4 pp [−6.3, +1.7]
LIBERO-Plus added 1000 (held-out confirmation)	698	718	−2.0 pp [−5.7, +1.7]
LIBERO full 2000 (4 suites × 10 tasks × 50)	1941	1928	+0.65 pp [−1.1, +2.8]

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	storage	resident GPU memory	prefill (ms)	expert per step (ms)	end-to-end 10 steps (ms)	end-to-end 2 steps (ms)
BF16 original	6.42 GB	6.96 GB	52.4	3.41	83.9	56.7
W2A8, INT8-resident weights (decoded once at load; torch._int_mm W8A8)	1.23 GB	4.04 GB	52.2	4.12	96.3	59.8
W2A8, 2-bit-resident expert (fused decode + GEMM Triton kernel)	1.23 GB	3.74 GB	52.0	4.44	99.8	59.8

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"].