arashakb/FastWAM_UR3

FastWAM UR3 3-task — 4-bit quantisations

Two post-training 4-bit quantisations of the UR3 3-task FastWAM fine-tune (step 7000). Both store the weight packed at 4 bits and contract it on 4-bit or integer tensor cores. Neither is a fake-quant checkpoint that keeps 4-bit values inside 16-bit tensors.

	ur3_step7000_asp_nvfp4.pt	ur3_step7000_svdquant_w4a4.pt
method	ĀFQ / ASP, deflated form	SVDQuant (Li et al., ICLR 2025)
numeric format	NVFP4 — E2M1 elements, block 16, E4M3 block scales	INT4, group 64
weights / activations	W4A4, both at block 16	W4A4, both at group 64
tensor cores	FP4 (torch._scaled_mm_v2, recipe BlockWise1x16)	INT8 (exact for 4-bit codes)
size	4.6117 BPW, 3.38 GiB	4.5798 BPW, 3.36 GiB
quantised-vs-bf16 action NRMSE	0.0006	0.0010
target GPU	RTX 5090 / Blackwell (sm_100, sm_103, sm_120)	L40S / Ada, Hopper (sm_89 tuned)

The NVFP4 checkpoint is the one to run on an RTX 5090.

Both quantise the same 600 Linears — the two experts' 2 × 30 blocks × {self_attn q/k/v/o, cross_attn q/k/v/o, ffn.0, ffn.2}, 5.914 B parameters. Patch/text/time embeddings, heads, the action encoder, norms, modulation and the proprio encoder stay in bf16. Both are self-contained: the quantised Linears and every unquantised mot tensor and the proprio encoder are inside the file, so the 11.2 GiB bf16 checkpoint is not needed at inference.

Both were calibrated on armanakbari4/ur3-3task-lerobot with 10 episodes per task × 3 tasks, every frame, seed 42 — 9 908 observations.

What ASP is, and what the deflated form means

Action-Subspace Protection keeps a rank-32 subspace of each action-expert layer out of the 4-bit grid. The subspace is not chosen by activation magnitude: it is the top eigenspace of the action metric G = E_o[JᵀJ], J = ∂action/∂x, differentiated through all ten denoising steps, so it is the set of directions the emitted action is most sensitive to rather than the ones that happen to be large.

Per layer, with s the smoothing vector, H a block Hadamard, V the rank-32 basis and W̃ = W·diag(s) the smoothed weight, the deflated contract is

x̃ = (x / s) H
y  = (x̃ V)(W̃V)ᵀ  +  NVFP4GEMM( (I − VVᵀ) x̃ ,  W̃(I − VVᵀ) )  +  bias

Deflated means the protected subspace is subtracted from the 4-bit path on both sides: the low-rank branch carries the exact W̃V and the FP4 weight holds W̃(I − VVᵀ). The cheaper shared-weight variant, which quantises one weight and stores V alone, is a different arm and a measurably worse one at these group sizes. The rank-32 branch stays in bf16 by design — it carries the directions the action depends on, which is the point of protecting them.

The per-layer smoothing (α, β) is not a fixed 0.5. It comes from a 39-candidate grid search (α ∈ {0, 0.05…0.95}, β ∈ {0, 1−α}, per SVDQuant's protocol) scored on real calibration activations with the objective being this layer's own deflated-ASP NVFP4 output MSE — the same objective the deployed arm minimises, at the same granularity it is applied.

Files

file	size
ur3_step7000_asp_nvfp4.pt	3.38 GiB	NVFP4 deflated-ASP checkpoint. Weights E2M1 packed, E4M3 block scales pre-swizzled to SWIZZLE_32_4_4
ur3_step7000_svdquant_w4a4.pt	3.36 GiB	INT4 SVDQuant checkpoint
nvfp4.py	8 KB	the NVFP4 quantiser and _scaled_mm_v2 wrapper — quantise, dequantise, scale swizzle, GEMM
asp_nvfp4_runtime.py	9.6 KB	ASPNVFP4Linear, install_asp_nvfp4, load_quantized_asp_model. Self-contained; imports only torch and nvfp4.py
ur3_3task_10k_dataset_stats.json	170 KB	proprio z-scoring and action denormalisation — the checkpoint cannot be run correctly without it
ur3_prompt_embeddings.pt	3.0 MiB	the three task instructions, pre-encoded, so the 11 GB umT5-XXL text encoder is not needed at inference

Besides these you need the Wan2.2 VAE (Wan2.2_VAE.pth, 2.7 GiB) to encode the camera image to the video latent. You do not need the bf16 checkpoint, the text encoder, or the Wan2.2 DiT weights.

Running the NVFP4 checkpoint

Requires a GPU with FP4 tensor cores — RTX 5090 (sm_120), B200/B300 (sm_100/sm_103) — and a PyTorch with torch._scaled_mm_v2. Verified on torch 2.12.0+cu130.

from asp_nvfp4_runtime import load_quantized_asp_model

model, cfg = load_quantized_asp_model(
    "ur3_step7000_asp_nvfp4.pt",
    build_model,           # your own bf16 FastWAM constructor -> (model, cfg)
)

build_model supplies the module graph only; every tensor comes from the checkpoint. To quantise a model you already built, call install_asp_nvfp4(model, ckpt_path) instead — it raises rather than swapping a subset, because a partial swap is not a defined arm.

Full pipeline, kernels and the deployment guide: https://github.com/arashakb/QuantWAM (adapters/fastwam/, quantwam/kernels/nvfp4.py, adapters/fastwam/README_ur3_5090.md). The UR3QuantPolicy wrapper there takes raw HxWx3 uint8 frames and the raw 14-d state and applies the whole observation contract itself (top → 320×256, wrists → 160×128, [top ; [left|right]] → 384×320, *2/255 − 1; state z-scored and clamped to ±5), so a caller cannot get the preprocessing subtly wrong.

Verification — NVFP4 deflated ASP

Everything below is measured on held-out episodes; the calibration episodes are excluded, because a check run on fitted data cannot detect the failure it exists to detect.

It is really 4-bit, not a simulation. Read off the loaded model:

wq dtype, all 600 layers	torch.float4_e2m1fn_x2
weight bytes resident	2 820 MiB for 5.914 B weights = 4.00 bits/weight (bf16 would be 16.00)
E4M3 block scales	352.5 MiB
any dequantised weight copy anywhere	none
torch._scaled_mm_v2 calls in one full inference	3 300 = 300 video prefill × 1 + 300 action × 10 steps
recipes those calls used	BlockWise1x16 only — i.e. NVFP4

It computes the intended contract. Recomputing each layer independently from the stored transforms and comparing against the runtime: worst relative error 1.7e-3 across sampled action and video layers, which is the bf16 output cast and not a form mismatch.

The actions are right. Twelve held-out observations spanning all three tasks:

	NRMSE
NVFP4 ASP vs bf16, action chunk	0.0006
bf16 vs recorded actions	0.0068 ← the ceiling
NVFP4 ASP vs recorded actions	0.0066 (corr 0.9997)

The quantised model sits as close to the recorded actions as the bf16 model does — marginally closer on these frames, which is noise, not an improvement.

Not yet measured: real-robot success rate. The bf16 model has been evaluated on the physical UR3; neither quantised checkpoint has. Everything above is open-loop agreement with recorded trajectories, which is necessary but not sufficient.

Verification — INT4 SVDQuant

check	result
base bf16 model vs recorded actions	NRMSE 0.0055, corr 0.9998
packed kernels vs the reference SVDQuant formula, per layer	activation scales bit-identical; ≤ 0.10 % of 4-bit codes differ, every one by exactly 1 LSB
weight repacking fidelity at export	0.0000 % of codes off by 1 LSB
quantised vs bf16, action chunk	NRMSE 0.0010
quantised vs recorded actions	NRMSE 0.0064 — the same as bf16's own 0.0064
end-to-end from raw camera frames	NRMSE 0.0061, corr 0.9997

Its Triton launch configurations are tuned per compute capability and only sm_89 ships; on other architectures the kernel prints no tuned config for M=… K=… N=… and falls back to an occupancy heuristic — correct, but roughly 2× off the tuned optimum on narrow shapes. Run analysis/iw_gemm_tune.py once on the target GPU to fix it. This does not apply to the NVFP4 checkpoint, which calls cuBLAS through _scaled_mm_v2 rather than a hand-tuned Triton kernel.

Tasks

• blue_basket — put the medicine then the measuring tape inside the blue basket
• drawer — open the drawer, put the white box inside the drawer then close the drawer
• stacking_cubes — put the green cube on top of the black cube and put the red cube on top of the green cube