arashakb/IMAGEWAM_UR3

ImageWAM UR3 3-task — W4A4 in NVFP4, two arms

The ImageWAM-FLUX.2-4B UR3 three-task fine-tune (step 7000) quantised to NVFP4 for Blackwell, by two methods calibrated identically so they can be compared. Real 4-bit execution on the FP4 tensor cores through torch._scaled_mm_v2: nothing is stored dequantised and no fp16 GEMM runs on the bulk path.

	ur3_step7000_asp_nvfp4.pt	ur3_step7000_svdquant_nvfp4.pt
method	AFQ / ASP, deflated form	SVDQuant (Li et al., ICLR 2025)
rotation	block Hadamard	none
low-rank branch	rank 32 on the action expert, from the action metric	rank 32 on every layer, from the quantisation residual
BPW	4.5789	4.7523
resident weights	2.668 GiB (3.22x less than bf16)	2.857 GiB (3.00x)
NVFP4 vs bf16	0.00108	0.00167
worst frame	0.00248	0.00335
correlation with bf16	0.999992	0.999983
vs recorded teleoperation	0.00300	0.00362

bf16's own distance to the demonstrator is 0.00293, so ASP lands within 0.00007 of the unquantised model while spending 0.17 fewer bits per weight than the baseline.

The action metric, and where it points

ASP protects the directions the emitted action actually depends on: G_l = E[J_l^T J_l], J_l = d(action chunk)/d(x_l) differentiated through all ten denoising steps, estimated with random probes over 901 frames x 12 probes of the calibration set. The protected subspace is the top-32 eigenspace of G~ = H diag(s) G diag(s) H — the metric in the coordinates actually being quantised — and it is searched jointly with the smoothing, because changing s moves the subspace.

98.4% of the action mass sits in six layers, and none of them is a transformer block:

layer	share of tr(G)
single_stream_modulation.lin	36.1%
action_encoder	31.6%
double_stream_modulation_img.lin	21.1%
time_in.out_layer	6.3%
head.linear	2.1%
head.adaLN_modulation.1	1.2%

The 60 Linears inside the action expert's blocks hold 1.6% between them. That is a different picture from the FastWAM UR3 model, whose mass spreads across cross_attn.o in every block, and it is the point of measuring rather than assuming.

mixtures.action.time_in.in_layer gets rank 0, not rank 32: it consumes the timestep, not the action latents, so its Jacobian to the action is genuinely empty and a "protected subspace" there would be 32 directions chosen by LAPACK out of a zero matrix. 65 layers carry a subspace, 88 do not.

Calibration

The same 30 episodes as the FastWAM UR3 arms — 10 per task x 3 tasks, seed 42, all frames, 9,908 observations — so the two models' quantised rows are comparable. Activation absmax feeds SmoothQuant; a 768-row-per-layer reservoir of real layer inputs feeds the smoothing search, whose objective is the layer's own output MSE against bf16 with the NVFP4 quantiser inside the scored loop. Candidates are deepcompressor's 39 (alpha, beta) pairs.

Split-half agreement (search on the even rows, score the odd): 75.8% for ASP, 50.0% for SVDQuant. What each low-rank branch buys, as median layer output error against no branch at all: 1.913x for ASP's action subspace, 1.081x for SVDQuant's residual SVD.

What is quantised

153 of the model's 156 MoT Linears, holding 4.498 B weights — 99.6% of mot's parameters. Three stay bf16 and both checkpoints record them:

layer	shape	why
mixtures.action.action_encoder	(1024, 16)	K = 16, and cuBLASLt's FP4 path needs K % 32 == 0. Measured on a B300: K = 32, 64, 96, 128, 256 pass; 16 and 48 return CUBLAS_STATUS_NOT_SUPPORTED.
mixtures.video.transformer.final_layer.linear	(128, 3072)	the video output head predicts the next frame; it never runs during action inference, so it has no calibration
…final_layer.adaLN_modulation.1	(6144, 3072)	same

Verification

On held-out episodes the calibration never saw (9 episodes disjoint from the calibration selection, 36 observations, horizon 16, 10 denoising steps) — the table at the top. And that the arithmetic is real rather than simulated, for both arms:

• every wq is torch.float4_e2m1fn_x2; 2144.9 MiB of packed weight for 4.498 B weights is 4.00 bits/weight exactly; zero dequantised weight copies anywhere in the module.
• 747 torch._scaled_mm_v2 calls per inference, every one BlockWise1x16. 747 is not a round number: it is 87 video Linears once plus 66 action Linears across 10 denoising steps.

No closed-loop success rate has been measured with either checkpoint. Open-loop agreement is necessary and not sufficient.

Running them

Needs a Blackwell GPU and a torch with torch._scaled_mm_v2 (built and verified on 2.12.0+cu130). Both files are self-contained — quantised Linears, every unquantised mot tensor, the proprio encoder — so the 9.0 GiB bf16 checkpoint is not needed at inference. You still need the FLUX.2 klein-base-4B backbone and autoencoder, the dataset stats, and the pre-encoded prompts here.

from iw_nvfp4_runtime import load_quantized_model
model, proc = load_quantized_model("ur3_step7000_asp_nvfp4.pt", build_model=my_builder)

install_nvfp4 dispatches on the checkpoint's lowrank_mode and cross-checks requires_runtime_fwht against what the layers actually carry. SVDQuant applies no rotation and ASP applies one; driving either checkpoint through the other's branch produces a well-formed GEMM of the wrong bilinear form — it loads, it runs, and it emits a plausible action chunk. It also raises unless every name in the checkpoint resolves to an nn.Linear.

iw_nvfp4_runtime.py here is a standalone copy, verified bit-identical to the in-repo runtime on real exported layers of both modes. The repo version can additionally drive a fused Triton prologue that is substantially faster; this copy takes the torch path, which is the reference the exports were scored against.

Three things about this model that will bite you

1. The MoT exposes the video expert twice. mixtures.video.X and mixtures.video.transformer.X are the same tensor — 140 of 267 storages carry two names, which is why state_dict() reports 8.198 B elements for 4.518 B parameters. Deduplicate by storage or your "quantised" export comes out larger than the bf16 original.
2. The mosaic is 288x256 (compact_288x256): camera_top at 256x192 on top, the two wrists at 128x96 side by side below, order [top, left, right], pixels mapped to (-1, 1). This is not the 384x320 layout the FastWAM UR3 checkpoints use.
3. Actions come back 16-D. Un-pad to 14-D — concat(x[..., 0:7], x[..., 8:15]) — before sending anything to a controller.

If you differentiate through this model yourself: cuDNN's fused attention cannot be backpropagated here — from the third denoising step on it dies with mha_graph.execute(...).is_good() false, surfacing as an illegal memory access — and gradient checkpointing must be off. The forward is unaffected either way, which is why inference never notices.

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