README.md

---
license: apache-2.0
library_name: pytorch
tags:
  - robotics
  - quantization
  - w4a4
  - svdquant
  - world-action-model
  - fastwam
base_model: armanakbari4/fastwam-ur3-3task-10k
datasets:
  - armanakbari4/ur3-3task-lerobot
---

# FastWAM UR3 3-task — SVDQuant W4A4 (real packed INT4)

A 4-bit post-training quantisation of the UR3 3-task FastWAM fine-tune (step 7000). The weight is
stored **packed at 4 bits** and contracted on integer tensor cores — this is not a fake-quant
checkpoint that stores 4-bit values in 16-bit tensors.

| | |
|---|---|
| base model | [`armanakbari4/fastwam-ur3-3task-10k`](https://huggingface.co/armanakbari4/fastwam-ur3-3task-10k) :: `ur3_3task_10k_step7000.pt` (11.2 GiB bf16) |
| method | SVDQuant (Li et al., ICLR 2025) — SmoothQuant migration + rank-32 FP16 low-rank branch + per-group-64 INT4 residual |
| precision | **W4A4**, weights per-group-64, activations per-group-64 dynamic |
| quantised | the 600 block Linears of both experts (2 × 30 blocks × {self_attn q/k/v/o, cross_attn q/k/v/o, ffn.0, ffn.2}), **5.914 B params** |
| kept in bf16 | patch/text/time embeddings, heads, action encoder, norms, modulation, proprio encoder |
| size | **4.5798 BPW**, **3.36 GiB** (down from 11.2 GiB) |
| calibration | `armanakbari4/ur3-3task-lerobot` — **10 episodes per task × 3 tasks, every frame, seed 42 = 9 908 observations** |

## Files

| file | size | |
|---|---|---|
| `ur3_step7000_svdquant_w4a4.pt` | 3.36 GiB | the checkpoint. **Self-contained**: the packed INT4 Linears *and* every unquantised `mot` tensor *and* the proprio encoder |
| `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 |

## Running it

Code, kernels and a deployment guide: **https://github.com/arashakb/QuantWAM**
(`adapters/fastwam/`, `quantwam/kernels/w4a4_triton.py`, `adapters/fastwam/README_ur3_5090.md`).

```python
from ur3_infer import UR3QuantPolicy
pol = UR3QuantPolicy()                                     # ~25 s
pol.set_task("drawer_lerobot")                             # or the full instruction string
chunk = pol.predict(top_rgb, left_rgb, right_rgb, state)   # [32, 14], robot units
```

`predict` takes **raw** `HxWx3` uint8 RGB 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. Run `python adapters/fastwam/ur3_infer.py --self-test` first on any new machine.

Besides these files 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 11.2 GiB bf16 checkpoint, the 11 GB text encoder, or
the 18.8 GB Wan2.2 DiT weights. Total deployment footprint ≈ **6.1 GiB**.

## Verification

All 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.

| 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 |

Quantisation costs essentially nothing on this checkpoint: the quantised model sits as close to the
recorded actions as the bf16 model does.

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

## Hardware notes

The kernels execute the INT4 contract on **INT8 tensor cores**, which is exact — every 4-bit code is
representable in int8 and the int32 accumulation rounds nothing. The 4 bits therefore buy DRAM
traffic and memory rather than arithmetic.

Launch configurations are tuned per compute capability, and only `sm_89` (L40S) ships. On other
architectures — including **sm_120 / RTX 5090** — 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 the
narrow shapes. Run `analysis/iw_gemm_tune.py` once on the target GPU to fix it.

## 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