Add README.md

README.md

@@ -0,0 +1,97 @@
+---
+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