diff --git a/families/lerobot_groot/README.md b/families/lerobot_groot/README.md index fbc81b4ce1..493180bdcf 100644 --- a/families/lerobot_groot/README.md +++ b/families/lerobot_groot/README.md @@ -11,6 +11,8 @@ Python and no PyTorch. ```bash trtmc build nvidia/gr00t17-lerobot-libero_spatial-640 -o groot.bundle --precision bf16 +# FP8 DiT projections (bf16 only, calibrated checkpoints only; see FP8 below) +trtmc build nvidia/gr00t17-lerobot-libero_spatial-640 -o groot-fp8.bundle --precision bf16 --quantization fp8 ``` The family is recognised from `type: "groot"` in `config.json`, never from a @@ -34,9 +36,9 @@ Anything outside this list fails at build time with the offending field named. | Batch | one | | Real-time chunking, relative actions, PEFT adapters | not supported (`use_relative_actions`, `use_peft` must be off; grouped `raw_stats`/`modality_config` are rejected) | | Processor steps | exactly the steps LeRobot 0.6.1 saves for an N1.7 policy; any other step is rejected | -| Precision | `bf16` (default) and `fp32` | +| Precision | `bf16` (default) and `fp32`; `bf16` with `--quantization fp8` (DiT projections in FP8, see below) | | Instruction | ASCII, up to 256 tokens and 4096 bytes; formalised as LeRobot does (lower case, keep letters, digits, whitespace and `_`), or used verbatim when `formalize_language` is off | -| Not claimed | GR00T N1.5, other backbones, Jetson targets, quantization, multi-GPU | +| Not claimed | GR00T N1.5, other backbones, Jetson targets, quantization other than FP8 of the DiT, multi-GPU | ### Tokenizer @@ -120,7 +122,9 @@ no Python or PyTorch library. Thresholds are in `thresholds/` and are anchored t LeRobot's own bfloat16-versus-float32 noise: about twice LeRobot's bfloat16 error on each observation, except the bf16 maximum and mean of `ep1267_f080`, where LeRobot's bfloat16 error is the smallest of the three; those are set from the measured bundle -error with headroom (see below). +error with headroom (see below). The FP8 manifest builds the same checkpoint with +`quantization="fp8"`; its thresholds are about twice the FP8 bundle's measured error on +B200 and B300 (see [FP8](#fp8)). ```bash python -m pytest families/lerobot_groot/tests --e2e-model lerobot_groot -q @@ -188,58 +192,57 @@ against the eager LeRobot policy (TensorRT 11.4): | `nvidia/gr00t17-lerobot-libero_10-640` @ `5ee08ab09fac5c5ef2388a14c882ea825ac861db` | 3.2e-6 / 1.2e-6 / 6.0e-6 | | `nvidia/gr00t17-lerobot-libero_goal-640` @ `436c57c0eb7a90be54270abf3977668b2084ad75` | 3.3e-6 / 1.7e-6 / 2.3e-6 | -## FP8 quantization study (not shipped) +## FP8 -FP8 self-quantization of the flow-matching DiT was built and measured on TensorRT 11.4 -(B200, B300) and 11.1 (B200) on top of the bf16 vision tower. It is not shipped: one -policy step is only 1.06x faster than the bf16 bundle, the mean action error grows to -1.8x that of the bf16 bundle, and one of the three committed bf16 test observations fails -its thresholds. This section records the recipe, the method and the numbers, so that the -study does not have to be repeated. The FP8 code was not merged; this family has no -quantized build. +`--quantization fp8` (with `--precision bf16`) runs the projections of the flow-matching DiT +as FP8 GEMMs; everything else stays as in the bf16 bundle. It trades action accuracy for a +small speed-up: on TensorRT 11.4 one policy step is 1.07x faster (GPU time, B200 and B300), +the native runtime's wall time per chunk 1.05x to 1.06x, the bundle 16% smaller, and the mean +arm action error is 1.7x that of the bf16 bundle and 1.65x LeRobot's own +bfloat16-versus-float32 difference. Only `nvidia/gr00t17-lerobot-libero_spatial-640` is +calibrated; another checkpoint fails at build time and names the calibration tool. -### Summary and decision +### Summary Policy step = the condition engine once and the denoise engine four times (GPU time, -TensorRT 11.4; 200 held-out LIBERO-spatial observations): +TensorRT 11.4, median over three builds of each precision; 200 held-out LIBERO-spatial +observations): | Item | bf16 bundle | FP8 bundle | FP8 vs bf16 | | --- | --- | --- | --- | -| Policy step, B200 / B300 | 12.15 / 11.00 ms | 11.49 / 10.29 ms | 1.06x / 1.06x to 1.07x | -| Four denoising steps, B200 / B300 | 6.25 / 5.70 ms | 5.62 / 5.05 ms | 1.11x to 1.12x / 1.11x to 1.13x | -| Condition engine, B200 / B300 | 5.90 / 5.29 ms | 5.87 / 5.24 ms | unchanged | -| Bundle | 5.82 GB | 4.89 GB | -16% | -| Held-out mean arm error vs LeRobot fp32 | 8.5e-4 | 1.5e-3 | 1.78x | -| Committed bf16 E2E thresholds | 3 of 3 pass | 2 of 3 pass | `ep1267_f080` fails | +| Policy step, B200 / B300 | 9.99 / 9.86 ms | 9.33 / 9.21 ms | 1.07x / 1.07x | +| Four denoising steps, B200 / B300 | 6.04 / 6.08 ms | 5.41 / 5.43 ms | 1.12x / 1.12x | +| Condition engine, B200 / B300 | 3.95 / 3.78 ms | 3.93 / 3.77 ms | unchanged | +| Native runtime wall time per chunk (p50), B200 / B300 | 12.32 / 13.29 ms | 11.65 / 12.62 ms | 1.06x / 1.05x | +| Bundle | 5.67 GB | 4.74 GB | -16% | +| Held-out mean arm error vs LeRobot fp32 | 9.1e-4 | 1.5e-3 | 1.68x (LeRobot bf16: 9.3e-4) | +| E2E thresholds | bf16 cases | FP8 cases (own thresholds) | | | Non-finite outputs | 0 | 0 | none | -Decision: do not ship FP8 for GR00T. - -- Speed: FP8 saves 0.6 to 0.7 ms of a 11 to 12 ms step. Only the DiT is quantized, and - the DiT is half of the step; the other half, the condition engine (vision tower, - backbone, VL blocks), stays bf16. +- Only the DiT is quantized, and the DiT is more than half of the step; the condition + engine (vision tower, backbone, VL blocks) stays bf16. FP8 there was measured to add 1% + to 2% for twice the action error. - The DiT runs on 41 tokens, so its GEMMs are small and FP8 makes one denoising step only - 1.11x to 1.13x faster. -- Quality: the mean arm error is 1.78x that of the bf16 bundle and 1.6x LeRobot's own - bfloat16-versus-float32 difference. An FP8 build would need its own, looser E2E - thresholds. -- Plan size: the denoise plan is 41% smaller, but the bundle only 16%. + 1.12x faster. +- The FP8 E2E cases (`tests/manifests/groot-n17-libero-spatial-fp8.json`, the first one + premerge) have their own thresholds, about twice the FP8 bundle's error on B200 and B300. -### Recipe that was measured +### Recipe No FP8 or NVFP4 GR00T checkpoint exists on the Hub (`nvidia/GR00T-N1.7-3B` and the LeRobot fine-tunes are bf16 or fp32 only), and no official GR00T quantization config -exists. +exists, so the build quantizes the original checkpoint (`quantization.py`). | Choice | Decision | | --- | --- | | Quantized (192 projections) | `attn1.to_q/k/v`, `attn1.to_out.0`, `ff.net.0.proj`, `ff.net.2` of all 32 DiT blocks; in the plans: 128 FP8 GEMMs in the denoise engine and 1 in the condition engine (the 32 cross-attention key/value projections, merged by TensorRT) | | bf16 | vision tower, the 16 backbone layers, the 4 VL self-attention blocks, attention matmuls, AdaLN modulation, timestep embedders, DiT output projections, embodiment-specific encoders and decoder | -| Weight scale | per tensor, amax / 448, from the checkpoint at build time | -| Activation scale | static, per tensor; ModelOpt 0.46.1 `FP8_DEFAULT_CFG` (max calibration) through LeRobot's float32 policy, all four denoising steps, 128 LIBERO-spatial observations (32 episodes x 4 frames) | +| Weight scale | per tensor, amax / 448, from the checkpoint at build time; FP8 weights are never stored | +| Activation scale | static, per tensor, from `fp8_activation_scales.json`; calibrated once with `calibrate_fp8_scales.py`: ModelOpt 0.46.1 `FP8_DEFAULT_CFG` (max calibration) through LeRobot's float32 policy, all four denoising steps, 128 LIBERO-spatial observations (32 episodes x 4 frames) | +| Scale table key | model config (kept backbone layers, embodiment) plus a weight fingerprint (sums of the first backbone and DiT query weights); an unknown checkpoint, a missing entry or a non-positive scale is an error | | Outlier rule | a scale above 20x the median of its kind stays bf16; measured: none in the DiT (largest ratio 3.99, `ff.net.2`) | -| Dequantize type | bf16 (the build precision) | -| Build | `precision=bf16` only; fp32 with FP8 and already quantized checkpoints refused | +| Dequantize type | bf16 (the build precision): no fp16 overflow, the bias stays in the GEMM epilogue | +| Build | `precision=bf16` only; fp32 with FP8, other quantization formats and already quantized checkpoints are refused before the weights load | ### Method @@ -248,8 +251,11 @@ exists. - A Python harness runs the two plans as the runtime does (prompt, image patches, state normalization, masks, four Euler steps, unnormalization, LIBERO gripper). - Speed: GPU time of one policy step (CUDA events, inputs on the device), median of 20 - interleaved rounds of 10 steps (15 on 11.1). On 11.4 each variant was built twice and - all four plans were timed in the same job, one job per GPU type. + interleaved rounds of 10 steps. On 11.4 each precision was built three times and all six + plans were timed in the same job, one job per GPU type; one build each on 11.1. +- Native runtime: the C++ qualification binary, two builds of each precision through + `BuildRequest`, the 3 test observations and 20 held-out ones, 20 repeats each, three + interleaved rounds; wall time per chunk includes host preprocessing. - Quality: 200 held-out LIBERO-spatial observations (50 episodes x 4 frames, none used for calibration or by the tests) and the 3 committed test observations, fixed noise per observation, 16 x 7 actions in robot units. Arm = the 6 continuous dimensions, gripper = @@ -260,37 +266,53 @@ exists. ### Results -Speed and size (x is bf16 time over FP8 time; two builds of each variant on 11.4): +Speed and size (x is bf16 time over FP8 time; per build on 11.4): -| GPU | TensorRT | bf16 step | FP8 step | x step | bf16 / FP8 denoise x4 | x denoise | denoise plan bf16 to FP8 | +| GPU | TensorRT | bf16 step | FP8 step | x step (median) | bf16 / FP8 denoise x4 | x denoise | denoise plan bf16 to FP8 | | --- | --- | --- | --- | --- | --- | --- | --- | -| B200 | 11.4 | 12.15 / 12.17 ms | 11.49 / 11.47 ms | 1.06 | 6.25-6.28 / 5.61-5.62 ms | 1.11 to 1.12 | 2004 to 1178 MB | -| B300 | 11.4 | 11.00 / 10.94 ms | 10.29 / 10.35 ms | 1.06 to 1.07 | 5.67-5.70 / 5.05-5.09 ms | 1.11 to 1.13 | 2004 to 1178 MB | -| B200 | 11.1 | 11.81 ms | 11.12 ms | 1.06 | 5.97 / 5.29 ms | 1.13 | 2004 to 1179 MB | +| B200 | 11.4 | 10.00 / 9.99 / 9.95 ms | 9.32 / 9.33 / 9.36 ms | 1.07 | 6.04 / 5.41-5.42 ms | 1.12 | 2004 to 1178 MB | +| B300 | 11.4 | 9.90 / 9.86 / 9.84 ms | 9.21 / 9.21 / 9.19 ms | 1.07 | 6.07-6.11 / 5.43-5.44 ms | 1.12 | 2004 to 1178 MB | +| B200 | 11.1 | 10.38 ms | 9.69 ms | 1.07 | 6.23 / 5.54 ms | 1.12 | 2004 to 1178 MB | +| B300 | 11.1 | 9.74 ms | 9.02 ms | 1.08 | 5.86 / 5.14 ms | 1.14 | 2004 to 1178 MB | + +The condition plan goes from 3658 to 3557 MB, the bundle from 5.67 to 4.74 GB. Build time +is 145 to 350 s for bf16 and 250 to 460 s for FP8; bundle load time 4.4 to 3.8 s (B200). + +Native runtime (TensorRT 11.4, wall time per chunk, p50 per build): -The condition plan goes from 3809 to 3708 MB, the bundle from 5.82 to 4.89 GB. Build time -is 148 to 248 s for bf16 and 252 to 337 s for FP8. +| GPU | bf16 | FP8 | x | +| --- | --- | --- | --- | +| B200 | 12.36 / 12.28 ms | 11.65 / 11.65 ms | 1.06 | +| B300 | 13.30 / 13.27 ms | 12.62 / 12.61 ms | 1.05 | -Quality on the 200 held-out observations (B200 and B300 give the same numbers on 11.4): +Quality on the 200 held-out observations (all builds of one precision, and B200 and B300, +give the same numbers on each TensorRT version): | Engine | Arm mean abs vs LeRobot fp32 | Ratio to bf16 bundle | p99 | Median per-observation max | Gripper flips | Arm mean abs vs bf16 bundle | | --- | --- | --- | --- | --- | --- | --- | -| bf16 bundle, 11.4 | 8.46e-4 | 1.00 | 5.77e-3 | 4.03e-3 | 0.031% | | -| FP8 bundle, 11.4 | 1.50e-3 | 1.78 | 8.53e-3 | 7.82e-3 | 0.031% | 1.27e-3 | -| bf16 bundle, 11.1 | 8.39e-4 | 1.00 | 5.59e-3 | 4.04e-3 | 0.094% | | -| FP8 bundle, 11.1 | 1.46e-3 | 1.74 | 8.33e-3 | 7.27e-3 | 0.031% | 1.41e-3 | -| LeRobot bf16 policy | 9.27e-4 | 1.10 | 6.79e-3 | 4.06e-3 | 0.031% | | +| bf16 bundle, 11.4 | 9.11e-4 | 1.00 | 6.77e-3 | 4.25e-3 | 0.031% | | +| FP8 bundle, 11.4 | 1.53e-3 | 1.68 | 8.97e-3 | 7.41e-3 | 0.062% | 1.27e-3 | +| bf16 bundle, 11.1 | 9.02e-4 | 1.00 | 6.42e-3 | 4.06e-3 | 0.031% | | +| FP8 bundle, 11.1 | 1.53e-3 | 1.70 | 9.25e-3 | 7.51e-3 | 0.062% | 1.27e-3 | +| LeRobot bf16 policy | 9.27e-4 | 1.02 | 6.79e-3 | 4.06e-3 | 0.031% | | + +The native runtime gives the same agreement on its 23 observations (arm mean abs vs LeRobot +fp32: bf16 8.8e-4, FP8 1.53e-3, LeRobot bf16 8.5e-4). -The committed test observations against the bf16 E2E thresholds (TensorRT 11.4; max abs / -mean abs vs LeRobot fp32, max abs vs LeRobot bf16; cosine is at least 0.99998 everywhere): +E2E through the native runtime (TensorRT 11.4; max abs / mean abs / cosine vs LeRobot fp32, +max abs vs LeRobot bf16; B200 and B300 identical, two builds each; `ep1272_f000` also runs +the alternate instruction, shown second): -| Observation | Threshold | bf16 bundle | FP8 bundle | +| Observation | bf16 bundle | FP8 bundle | FP8 threshold | | --- | --- | --- | --- | -| `ep1262_f040` | 1.26e-2 / 1.8e-3 / 1.9e-2 | 3.64e-3 / 7.08e-4 / 6.41e-3 | 9.94e-3 / 1.49e-3 / 9.16e-3 | -| `ep1267_f080` | 5.1e-3 / 9.4e-4 / 7.7e-3 | 3.27e-3 / 5.83e-4 / 5.49e-3 | 8.16e-3 / 1.22e-3 / 9.16e-3 (fails all three) | -| `ep1272_f000` | 3.0e-2 / 2.6e-3 / 4.5e-2 | 7.38e-3 / 8.07e-4 / 8.07e-3 | 7.12e-3 / 1.10e-3 / 1.44e-2 | +| `ep1262_f040` | 4.76e-3 / 8.65e-4 / 0.999996 / 9.16e-3 | 7.79e-3 / 1.59e-3 / 0.999986 / 1.10e-2 | 1.6e-2 / 3.2e-3 / 0.9999 / 2.2e-2 | +| `ep1267_f080` | 6.21e-3 / 9.34e-4 / 0.999995 / 4.58e-3 | 9.80e-3 / 1.51e-3 / 0.999989 / 8.70e-3 | 2.0e-2 / 3.0e-3 / 0.9999 / 1.8e-2 | +| `ep1272_f000` | 1.77e-2 / 1.95e-3 / 0.999973 / 7.32e-3; 1.07e-2 / 1.58e-3 / 0.999983 / 9.16e-3 | 1.13e-2 / 1.38e-3 / 0.999989 / 1.05e-2; 8.51e-3 / 1.46e-3 / 0.999988 / 7.32e-3 | 2.3e-2 / 3.0e-3 / 0.9999 / 2.1e-2 | -TensorRT 11.1 gives the same picture (`ep1267_f080` FP8: 7.24e-3 / 1.03e-3 / 8.24e-3). +The FP8 bundle would fail the bf16 thresholds of `ep1267_f080` (9.5e-3 / 1.4e-3 / 7.7e-3), +which is why it has its own. A broken FP8 build (every activation scale divided by 16) fails +the premerge FP8 case with max abs 0.53 and cosine 0.959. TensorRT 11.1 stays inside the +FP8 thresholds (largest: `ep1267_f080`, 1.16e-2 / 1.43e-3 / 9.6e-3). Other recipes, measured earlier with the vision tower in float32 (TensorRT 11.4; B200 / B300; the bf16 step was then 19.2 / 17.8 ms): @@ -311,27 +333,24 @@ A harness prototype with the DiT feed-forward projections in NVFP4 (per-16 block dynamic activation quantize; q/k/v in FP8), measured with the float32 vision tower on TensorRT 11.4, is slower than bf16: 0.97x (B200) and 0.98x (B300) per step, 0.92x and 0.95x per denoise call, with a mean arm error of 2.75e-3. At 41 tokens the dynamic block -quantize costs more than the 4-bit weight reads save. +quantize costs more than the 4-bit weight reads save. Not shipped. -### What was not tried +### Not covered | Item | Note | | --- | --- | -| FP8 of the bf16 vision tower | not measured | -| FP8 attention matmuls, percentile or entropy calibration | not measured | -| Closed-loop LIBERO success rate, the native runtime with an FP8 bundle, H100 | not measured | - -### When to revisit +| Other checkpoints | `libero_10`, `libero_goal` and other `lerobot-train` outputs need a table entry from `calibrate_fp8_scales.py` | +| FP8 of the vision tower, FP8 attention matmuls, percentile or entropy calibration | not measured | +| Closed-loop LIBERO success rate, RTX PRO 6000, H100, GB300 | not measured | -- A deployment with more denoising steps per call: the DiT share of the step, and so the - FP8 gain, grows with the step count. -- A quality target that accepts a mean arm error of about 1.5e-3 (1.6x LeRobot's own - bfloat16 noise), with FP8-specific E2E thresholds. +A deployment with more denoising steps per call gains more: the DiT share of the step, and +so the FP8 gain, grows with the step count. ## Known limits - Qualified on one NVIDIA RTX PRO 6000 Blackwell with TensorRT 11.4 (CUDA 13.4). Other - GPUs, TensorRT versions and Jetson are not. + GPUs, TensorRT versions and Jetson are not. The FP8 build was measured on B200 and B300 + (TensorRT 11.4 and 11.1) and is calibrated for one checkpoint. - Validated for the LIBERO configuration only. Other embodiments, camera layouts, state and action sizes and chunk lengths build, but are not qualified. - Batch one, one observation step, square cameras at the target size; GR00T N1.5, diff --git a/families/lerobot_groot/action_head.py b/families/lerobot_groot/action_head.py index 3539b564de..c2835895cb 100644 --- a/families/lerobot_groot/action_head.py +++ b/families/lerobot_groot/action_head.py @@ -20,6 +20,7 @@ from . import graph from .checkpoint import category, require +from .quantization import fp8 VL_WIDTH = 2048 VL_HEADS = 32 @@ -65,7 +66,7 @@ def _self_attention_block(network, h, weights, p, dtype): network, hidden, weights[p + "ff.net.2.weight"], weights[p + "ff.net.2.bias"], dtype=dtype))) -def build_condition(network, weights, backbone_features, state, *, embodiment, max_state_dim, dtype): +def build_condition(network, weights, backbone_features, state, *, embodiment, max_state_dim, dtype, quant=None): """Returns ``(vl_embeds[S, 2048] f32, state_features[1, 1536] f32, cross_k, cross_v)``. ``cross_k``/``cross_v`` are ``[len(CROSS_BLOCKS), S, WIDTH]`` in ``dtype``; ``state`` is the @@ -85,7 +86,8 @@ def build_condition(network, weights, backbone_features, state, *, embodiment, m p = f"{_DIT}transformer_blocks.{index}.attn1." require(weights, p + "to_k.weight", (WIDTH, VL_WIDTH)) for name, out in (("to_k", keys), ("to_v", values)): - rows = graph.linear(network, vl, weights[p + name + ".weight"], weights[p + name + ".bias"], dtype=dtype) + rows = graph.linear(network, vl, weights[p + name + ".weight"], weights[p + name + ".bias"], dtype=dtype, + fp8=fp8(quant, p + name)) out.append(graph.shuffle(network, rows, reshape=(1, -1, WIDTH))) return vl, state_features, graph.concat(network, keys, 0), graph.concat(network, values, 0) @@ -126,7 +128,7 @@ def _modulate(network, h, temb, weight, bias, eps, *, scale_first, dtype): def build_denoise(network, weights, *, embodiment, action_horizon, max_action_dim, dtype, - x_t, timestep, dt, state_features, cross_k, cross_v, text_bias, image_bias): + x_t, timestep, dt, state_features, cross_k, cross_v, text_bias, image_bias, quant=None): """One Euler step: returns ``(x_next, velocity)``, both float32 ``[action_horizon, max_action_dim]``.""" def encoder(x, name, shape): return _category_linear(network, x, weights, f"{_HEAD}action_encoder.{name}.", embodiment, shape, dtype) @@ -150,7 +152,8 @@ def encoder(x, name, shape): p = f"{_DIT}transformer_blocks.{index}." norm = _modulate(network, h, temb, weights[p + "norm1.linear.weight"], weights[p + "norm1.linear.bias"], LN_EPS, scale_first=True, dtype=dtype) - q = graph.linear(network, norm, weights[p + "attn1.to_q.weight"], weights[p + "attn1.to_q.bias"], dtype=dtype) + q = graph.linear(network, norm, weights[p + "attn1.to_q.weight"], weights[p + "attn1.to_q.bias"], dtype=dtype, + fp8=fp8(quant, p + "attn1.to_q")) if index in CROSS_BLOCKS: slot = CROSS_BLOCKS.index(index) k, v = (graph.shuffle(network, network.add_gather(cache, graph.const_int32(network, [slot]), 0).get_output(0), @@ -158,17 +161,20 @@ def encoder(x, name, shape): for cache in (cross_k, cross_v)) bias = text_bias if index % 4 == 0 else image_bias else: - k, v = (graph.linear(network, norm, weights[f"{p}attn1.to_{n}.weight"], weights[f"{p}attn1.to_{n}.bias"], dtype=dtype) + k, v = (graph.linear(network, norm, weights[f"{p}attn1.to_{n}.weight"], weights[f"{p}attn1.to_{n}.bias"], dtype=dtype, + fp8=fp8(quant, f"{p}attn1.to_{n}")) for n in "kv") bias = None context = graph.attention(network, q, k, v, heads=HEADS, kv_heads=HEADS, head_dim=HEAD_DIM, dtype=dtype, bias=bias) h = graph.add(network, h, graph.f32(network, graph.linear( - network, context, weights[p + "attn1.to_out.0.weight"], weights[p + "attn1.to_out.0.bias"], dtype=dtype))) + network, context, weights[p + "attn1.to_out.0.weight"], weights[p + "attn1.to_out.0.bias"], dtype=dtype, + fp8=fp8(quant, p + "attn1.to_out.0")))) hidden = graph.gelu_tanh(network, graph.linear( network, graph.layer_norm(network, h, LN_EPS), weights[p + "ff.net.0.proj.weight"], - weights[p + "ff.net.0.proj.bias"], dtype=dtype)) + weights[p + "ff.net.0.proj.bias"], dtype=dtype, fp8=fp8(quant, p + "ff.net.0.proj"))) h = graph.add(network, h, graph.f32(network, graph.linear( - network, hidden, weights[p + "ff.net.2.weight"], weights[p + "ff.net.2.bias"], dtype=dtype))) + network, hidden, weights[p + "ff.net.2.weight"], weights[p + "ff.net.2.bias"], dtype=dtype, + fp8=fp8(quant, p + "ff.net.2")))) h = _modulate(network, h, temb, weights[_DIT + "proj_out_1.weight"], weights[_DIT + "proj_out_1.bias"], OUT_EPS, scale_first=False, dtype=dtype) diff --git a/families/lerobot_groot/calibrate_fp8_scales.py b/families/lerobot_groot/calibrate_fp8_scales.py new file mode 100644 index 0000000000..1b2468a491 --- /dev/null +++ b/families/lerobot_groot/calibrate_fp8_scales.py @@ -0,0 +1,203 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Developer tool: calibrate FP8 activation scales for a LeRobot GR00T N1.7 checkpoint. + +Not used when building an engine. Calibrate each supported checkpoint once and commit the +updated ``fp8_activation_scales.json``:: + + python -m families.lerobot_groot.calibrate_fp8_scales fetch --out frames + python -m families.lerobot_groot.calibrate_fp8_scales calibrate --model-dir \ + --base-dir --tokenizer-dir \ + --frames frames/calibration + +``fetch`` writes recorded LIBERO-spatial observations from the pinned public +``physical-intelligence/libero`` dataset (the source of the test fixtures) in the fixtures' +layout: ``calibration/`` and ``held_out/`` take frames from disjoint episodes, and the +episodes of the test fixtures are in neither. ``calibrate`` loads the LeRobot policy +(``lerobot[groot]``, float32), runs ModelOpt ``FP8_DEFAULT_CFG`` (max calibration) over the +calibration frames through ``predict_action_chunk`` (every denoising step), and stores ``input_quantizer.amax / 448`` of every quantized projection +under the checkpoint key (see ``quantization.scale_key``). + +A projection whose input scale is more than 20 times the median of its kind has a massive +activation and is not quantized: its entry is stored as ``null``, so the build leaves it in +bfloat16 on purpose, while a projection missing from the file is an error. +""" + +from __future__ import annotations + +import argparse +import io +import json +import re +import statistics +from pathlib import Path + +from .quantization import ( + ACTIVATION_SCALES_FILENAME, + DIT_ROOT, + FP8_E4M3_MAX, + QUANTIZED_STEMS, + projection_kind, + projection_names, + scale_key, +) + +DATASET = "physical-intelligence/libero" +DATASET_REVISION = "a4336d589d589045d1c56423ffdf3b88a0e19b1f" +# LIBERO-spatial: the ten "pick up the black bowl ... and place it on the plate" tasks. +_SPATIAL = re.compile(r"^pick up the black bowl .* and place it on the plate$") +# Episodes of the committed test fixtures, never used for calibration or the held-out set. +FIXTURE_EPISODES = (1262, 1267, 1272) +_OUTLIER_RATIO = 20.0 +_PROJECTION_NAME = re.compile( + "^" + re.escape(DIT_ROOT) + r"\.\d+\.(?:" + "|".join(re.escape(stem) for stem in QUANTIZED_STEMS) + ")$") + + +def projection_name(module_name: str) -> str | None: + """Checkpoint name of a ``_groot_model`` module path, or None when it is not quantized.""" + return module_name if _PROJECTION_NAME.match(module_name) else None + + +def split_outlier_scales(scales: dict[str, float], cross_blocks: tuple[int, ...]): + """Replace the scale of a projection far above the median of its kind by ``None``. + + Returns the table (``None`` marks an unquantized projection) and the sorted names of the + projections left unquantized. + """ + by_kind: dict[str, list[float]] = {} + for name, value in scales.items(): + by_kind.setdefault(projection_kind(name, cross_blocks), []).append(value) + median = {kind: statistics.median(values) for kind, values in by_kind.items()} + table: dict[str, float | None] = {} + excluded = [] + for name, value in scales.items(): + if value > _OUTLIER_RATIO * median[projection_kind(name, cross_blocks)]: + table[name] = None + excluded.append(name) + else: + table[name] = value + return table, sorted(excluded) + + +def split_episodes(episodes: list[int], calibration: int, held_out: int) -> tuple[list[int], list[int]]: + """Disjoint calibration and held-out episodes, both spread over the sorted list.""" + pool = sorted(set(episodes) - set(FIXTURE_EPISODES)) + step = max(len(pool) // max(calibration + held_out, 1), 1) + spread = pool[::step] + return spread[0:2 * calibration:2], [e for e in spread if e not in spread[0:2 * calibration:2]][:held_out] + + +def frame_positions(length: int, count: int) -> list[int]: + """``count`` frame indices spread over an episode of ``length`` frames.""" + return sorted({int(length * (i + 0.5) / count) for i in range(count)}) + + +def fetch_frames(out: Path, calibration: int, held_out: int, per_episode: int) -> None: + """Write observations of calibration and held-out episodes in the test-fixture layout.""" + import pyarrow.parquet as pq + from huggingface_hub import hf_hub_download + from PIL import Image + + kwargs = {"repo_id": DATASET, "repo_type": "dataset", "revision": DATASET_REVISION} + info = json.loads(Path(hf_hub_download(filename="meta/info.json", **kwargs)).read_text(encoding="utf-8")) + read = lambda name: [json.loads(line) for line in # noqa: E731 + Path(hf_hub_download(filename=name, **kwargs)).read_text(encoding="utf-8").splitlines()] + task_text = {row["task_index"]: row["task"] for row in read("meta/tasks.jsonl")} + lengths = {row["episode_index"]: row["length"] for row in read("meta/episodes.jsonl") if _SPATIAL.match(row["tasks"][0])} + groups = dict(zip(("calibration", "held_out"), split_episodes(sorted(lengths), calibration, held_out))) + for group, episodes in groups.items(): + for episode in episodes: + data_path = info["data_path"].format(episode_chunk=episode // int(info["chunks_size"]), episode_index=episode) + table = pq.read_table(hf_hub_download(filename=data_path, **kwargs)) + for index in frame_positions(lengths[episode], per_episode): + row = table.slice(index, 1).to_pylist()[0] + directory = out / group / f"ep{episode:04d}_f{index:03d}" + directory.mkdir(parents=True, exist_ok=True) + for column in ("image", "wrist_image"): + Image.open(io.BytesIO(row[column]["bytes"])).convert("RGB").save(directory / f"{column}.png") + (directory / "observation.json").write_text(json.dumps({ + "source": f"{DATASET}@{DATASET_REVISION} episode {episode} frame {index}", + "state": [float(v) for v in row["state"]], + "task": task_text[row["task_index"]], + }, indent=2) + "\n", encoding="utf-8") + print(f"{group}: {len(episodes)} episodes x {per_episode} frames") + + +def calibrate(model_dir: Path, base_dir: Path, tokenizer_dir: Path, frames: Path, device: str) -> None: + """Run ModelOpt FP8 calibration and store the activation scales under the checkpoint key.""" + import modelopt.torch.quantization as mtq + import torch + + from . import action_head, backbone + from .checkpoint import load_checkpoint + from .config import load_config + from .tests.official_reference import load_frame, load_policy + + cfg = load_config(model_dir) + # The loader the build uses, so the key (layers, embodiment, weight sums) matches at build time. + weights = load_checkpoint(model_dir) + layers = backbone.kept_layers(weights) + key = scale_key(weights, backbone_layers=layers, embodiment_id=cfg.embodiment_id) + expected = set(projection_names(action_head.LAYERS)) + del weights + + directories = sorted(path.parent for path in frames.glob("*/observation.json")) + if not directories: + raise RuntimeError(f"no observations in {frames}; run the fetch command first") + policy, pre, _ = load_policy(model_dir, base_dir, tokenizer_dir, precision="fp32", device=device) + model = policy._groot_model + + def forward_loop(_model) -> None: + """Every calibration observation through the policy, with all its denoising steps.""" + for index, directory in enumerate(directories): + torch.manual_seed(index) + policy.predict_action_chunk(pre(load_frame(directory))) + + mtq.quantize(model, mtq.FP8_DEFAULT_CFG, forward_loop) + scales: dict[str, float] = {} + for name, module in model.named_modules(): + if projection_name(name) is None or not hasattr(module, "input_quantizer"): + continue + amax = module.input_quantizer.amax + if amax is not None and float(amax) > 0: + scales[name] = float(amax) / FP8_E4M3_MAX + # Every projection must be calibrated: a partial table would build a half-quantized engine. + if set(scales) != expected: + missing = sorted(expected - set(scales)) + raise RuntimeError(f"calibration missed {len(missing)} projections, for example {missing[:3]}") + table, excluded = split_outlier_scales(scales, action_head.CROSS_BLOCKS) + if excluded: + print(f"massive activations, left in bfloat16 (null): {excluded}") + + path = Path(__file__).parent / ACTIVATION_SCALES_FILENAME + stored = json.loads(path.read_text(encoding="utf-8")) if path.is_file() else {} + stored[key] = dict(sorted(table.items())) + path.write_text(json.dumps(dict(sorted(stored.items())), indent=2) + "\n", encoding="utf-8") + print(f"wrote {len(table)} scales ({len(excluded)} null) from {len(directories)} observations for {key!r} to {path}") + + +def main(argv: list[str] | None = None) -> None: + """Parse the command line and run the fetch or calibrate command.""" + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + commands = parser.add_subparsers(dest="command", required=True) + fetch = commands.add_parser("fetch") + fetch.add_argument("--out", type=Path, required=True) + fetch.add_argument("--calibration-episodes", type=int, default=32) + fetch.add_argument("--held-out-episodes", type=int, default=50) + fetch.add_argument("--frames-per-episode", type=int, default=4) + run = commands.add_parser("calibrate") + run.add_argument("--model-dir", type=Path, required=True) + run.add_argument("--base-dir", type=Path, required=True) + run.add_argument("--tokenizer-dir", type=Path, required=True) + run.add_argument("--frames", type=Path, required=True) + run.add_argument("--device", default="cuda") + args = parser.parse_args(argv) + if args.command == "fetch": + fetch_frames(args.out, args.calibration_episodes, args.held_out_episodes, args.frames_per_episode) + else: + calibrate(args.model_dir, args.base_dir, args.tokenizer_dir, args.frames, args.device) + + +if __name__ == "__main__": + main() diff --git a/families/lerobot_groot/fp8_activation_scales.json b/families/lerobot_groot/fp8_activation_scales.json new file mode 100644 index 0000000000..56467b7e43 --- /dev/null +++ b/families/lerobot_groot/fp8_activation_scales.json @@ -0,0 +1,196 @@ +{ + "groot-n17-l16-e2-w1.0642e+05-3.4067e+04": { + "action_head.model.transformer_blocks.0.attn1.to_k": 0.1702331304550171, + "action_head.model.transformer_blocks.0.attn1.to_out.0": 0.02210951702935355, + "action_head.model.transformer_blocks.0.attn1.to_q": 0.011422903410026006, + "action_head.model.transformer_blocks.0.attn1.to_v": 0.1702331304550171, + "action_head.model.transformer_blocks.0.ff.net.0.proj": 0.04506114550999233, + "action_head.model.transformer_blocks.0.ff.net.2": 0.022307966436658586, + "action_head.model.transformer_blocks.1.attn1.to_k": 0.013733495559011186, + "action_head.model.transformer_blocks.1.attn1.to_out.0": 0.0019776495173573494, + "action_head.model.transformer_blocks.1.attn1.to_q": 0.013733495559011186, + "action_head.model.transformer_blocks.1.attn1.to_v": 0.013733495559011186, + "action_head.model.transformer_blocks.1.ff.net.0.proj": 0.05086787257875715, + "action_head.model.transformer_blocks.1.ff.net.2": 0.05477204918861389, + "action_head.model.transformer_blocks.10.attn1.to_k": 0.1702331304550171, + "action_head.model.transformer_blocks.10.attn1.to_out.0": 0.024233190076691762, + "action_head.model.transformer_blocks.10.attn1.to_q": 0.05820835488183158, + "action_head.model.transformer_blocks.10.attn1.to_v": 0.1702331304550171, + "action_head.model.transformer_blocks.10.ff.net.0.proj": 0.058242934090750556, + "action_head.model.transformer_blocks.10.ff.net.2": 0.018338891012328013, + "action_head.model.transformer_blocks.11.attn1.to_k": 0.024187241281781877, + "action_head.model.transformer_blocks.11.attn1.to_out.0": 0.011406769709927695, + "action_head.model.transformer_blocks.11.attn1.to_q": 0.024187241281781877, + "action_head.model.transformer_blocks.11.attn1.to_v": 0.024187241281781877, + "action_head.model.transformer_blocks.11.ff.net.0.proj": 0.05829664213316781, + "action_head.model.transformer_blocks.11.ff.net.2": 0.02034739511353629, + "action_head.model.transformer_blocks.12.attn1.to_k": 0.1702331304550171, + "action_head.model.transformer_blocks.12.attn1.to_out.0": 0.02195382969720023, + "action_head.model.transformer_blocks.12.attn1.to_q": 0.061220258474349976, + "action_head.model.transformer_blocks.12.attn1.to_v": 0.1702331304550171, + "action_head.model.transformer_blocks.12.ff.net.0.proj": 0.058248158012117655, + "action_head.model.transformer_blocks.12.ff.net.2": 0.015869584466729845, + "action_head.model.transformer_blocks.13.attn1.to_k": 0.01841297745704651, + "action_head.model.transformer_blocks.13.attn1.to_out.0": 0.011417039803096227, + "action_head.model.transformer_blocks.13.attn1.to_q": 0.01841297745704651, + "action_head.model.transformer_blocks.13.attn1.to_v": 0.01841297745704651, + "action_head.model.transformer_blocks.13.ff.net.0.proj": 0.05824380261557443, + "action_head.model.transformer_blocks.13.ff.net.2": 0.02246599325111934, + "action_head.model.transformer_blocks.14.attn1.to_k": 0.1702331304550171, + "action_head.model.transformer_blocks.14.attn1.to_out.0": 0.027986202921186174, + "action_head.model.transformer_blocks.14.attn1.to_q": 0.048398490462984355, + "action_head.model.transformer_blocks.14.attn1.to_v": 0.1702331304550171, + "action_head.model.transformer_blocks.14.ff.net.0.proj": 0.0580266501222338, + "action_head.model.transformer_blocks.14.ff.net.2": 0.017012754721300944, + "action_head.model.transformer_blocks.15.attn1.to_k": 0.019926814096314565, + "action_head.model.transformer_blocks.15.attn1.to_out.0": 0.009853369423321314, + "action_head.model.transformer_blocks.15.attn1.to_q": 0.019926814096314565, + "action_head.model.transformer_blocks.15.attn1.to_v": 0.019926814096314565, + "action_head.model.transformer_blocks.15.ff.net.0.proj": 0.057942641632897515, + "action_head.model.transformer_blocks.15.ff.net.2": 0.018053182533809116, + "action_head.model.transformer_blocks.16.attn1.to_k": 0.1702331304550171, + "action_head.model.transformer_blocks.16.attn1.to_out.0": 0.019781297871044705, + "action_head.model.transformer_blocks.16.attn1.to_q": 0.05390210662569318, + "action_head.model.transformer_blocks.16.attn1.to_v": 0.1702331304550171, + "action_head.model.transformer_blocks.16.ff.net.0.proj": 0.057347651038851054, + "action_head.model.transformer_blocks.16.ff.net.2": 0.01660247359957014, + "action_head.model.transformer_blocks.17.attn1.to_k": 0.021227191601480757, + "action_head.model.transformer_blocks.17.attn1.to_out.0": 0.010072677263191767, + "action_head.model.transformer_blocks.17.attn1.to_q": 0.021227191601480757, + "action_head.model.transformer_blocks.17.attn1.to_v": 0.021227191601480757, + "action_head.model.transformer_blocks.17.ff.net.0.proj": 0.05697302307401385, + "action_head.model.transformer_blocks.17.ff.net.2": 0.01641789823770523, + "action_head.model.transformer_blocks.18.attn1.to_k": 0.1702331304550171, + "action_head.model.transformer_blocks.18.attn1.to_out.0": 0.02850541259561266, + "action_head.model.transformer_blocks.18.attn1.to_q": 0.042922024215970724, + "action_head.model.transformer_blocks.18.attn1.to_v": 0.1702331304550171, + "action_head.model.transformer_blocks.18.ff.net.0.proj": 0.05522809284073966, + "action_head.model.transformer_blocks.18.ff.net.2": 0.016777159912245616, + "action_head.model.transformer_blocks.19.attn1.to_k": 0.021915493266923085, + "action_head.model.transformer_blocks.19.attn1.to_out.0": 0.019549197384289334, + "action_head.model.transformer_blocks.19.attn1.to_q": 0.021915493266923085, + "action_head.model.transformer_blocks.19.attn1.to_v": 0.021915493266923085, + "action_head.model.transformer_blocks.19.ff.net.0.proj": 0.05366249595369611, + "action_head.model.transformer_blocks.19.ff.net.2": 0.01573680979864938, + "action_head.model.transformer_blocks.2.attn1.to_k": 0.1702331304550171, + "action_head.model.transformer_blocks.2.attn1.to_out.0": 0.018578872084617615, + "action_head.model.transformer_blocks.2.attn1.to_q": 0.03062850662640163, + "action_head.model.transformer_blocks.2.attn1.to_v": 0.1702331304550171, + "action_head.model.transformer_blocks.2.ff.net.0.proj": 0.04956975153514317, + "action_head.model.transformer_blocks.2.ff.net.2": 0.05659237078257969, + "action_head.model.transformer_blocks.20.attn1.to_k": 0.1702331304550171, + "action_head.model.transformer_blocks.20.attn1.to_out.0": 0.023753551500184194, + "action_head.model.transformer_blocks.20.attn1.to_q": 0.04132728065763201, + "action_head.model.transformer_blocks.20.attn1.to_v": 0.1702331304550171, + "action_head.model.transformer_blocks.20.ff.net.0.proj": 0.05134226594652448, + "action_head.model.transformer_blocks.20.ff.net.2": 0.014710590243339539, + "action_head.model.transformer_blocks.21.attn1.to_k": 0.022244845117841448, + "action_head.model.transformer_blocks.21.attn1.to_out.0": 0.012565183852400099, + "action_head.model.transformer_blocks.21.attn1.to_q": 0.022244845117841448, + "action_head.model.transformer_blocks.21.attn1.to_v": 0.022244845117841448, + "action_head.model.transformer_blocks.21.ff.net.0.proj": 0.050515013081686835, + "action_head.model.transformer_blocks.21.ff.net.2": 0.017792662339551107, + "action_head.model.transformer_blocks.22.attn1.to_k": 0.1702331304550171, + "action_head.model.transformer_blocks.22.attn1.to_out.0": 0.03199935598032815, + "action_head.model.transformer_blocks.22.attn1.to_q": 0.028145968914031982, + "action_head.model.transformer_blocks.22.attn1.to_v": 0.1702331304550171, + "action_head.model.transformer_blocks.22.ff.net.0.proj": 0.04892292193004063, + "action_head.model.transformer_blocks.22.ff.net.2": 0.018090154443468367, + "action_head.model.transformer_blocks.23.attn1.to_k": 0.02042315900325775, + "action_head.model.transformer_blocks.23.attn1.to_out.0": 0.012983484991959162, + "action_head.model.transformer_blocks.23.attn1.to_q": 0.02042315900325775, + "action_head.model.transformer_blocks.23.attn1.to_v": 0.02042315900325775, + "action_head.model.transformer_blocks.23.ff.net.0.proj": 0.04748989003045218, + "action_head.model.transformer_blocks.23.ff.net.2": 0.022579625248908997, + "action_head.model.transformer_blocks.24.attn1.to_k": 0.1702331304550171, + "action_head.model.transformer_blocks.24.attn1.to_out.0": 0.025209222521100725, + "action_head.model.transformer_blocks.24.attn1.to_q": 0.042499010051999776, + "action_head.model.transformer_blocks.24.attn1.to_v": 0.1702331304550171, + "action_head.model.transformer_blocks.24.ff.net.0.proj": 0.04614322100366865, + "action_head.model.transformer_blocks.24.ff.net.2": 0.02229736532483782, + "action_head.model.transformer_blocks.25.attn1.to_k": 0.019201800227165222, + "action_head.model.transformer_blocks.25.attn1.to_out.0": 0.01738906225987843, + "action_head.model.transformer_blocks.25.attn1.to_q": 0.019201800227165222, + "action_head.model.transformer_blocks.25.attn1.to_v": 0.019201800227165222, + "action_head.model.transformer_blocks.25.ff.net.0.proj": 0.04568071024758475, + "action_head.model.transformer_blocks.25.ff.net.2": 0.02275873933519636, + "action_head.model.transformer_blocks.26.attn1.to_k": 0.1702331304550171, + "action_head.model.transformer_blocks.26.attn1.to_out.0": 0.035928858178002496, + "action_head.model.transformer_blocks.26.attn1.to_q": 0.03062498356614794, + "action_head.model.transformer_blocks.26.attn1.to_v": 0.1702331304550171, + "action_head.model.transformer_blocks.26.ff.net.0.proj": 0.04461306759289333, + "action_head.model.transformer_blocks.26.ff.net.2": 0.019416400364467075, + "action_head.model.transformer_blocks.27.attn1.to_k": 0.021618996347699846, + "action_head.model.transformer_blocks.27.attn1.to_out.0": 0.018665984272956848, + "action_head.model.transformer_blocks.27.attn1.to_q": 0.021618996347699846, + "action_head.model.transformer_blocks.27.attn1.to_v": 0.021618996347699846, + "action_head.model.transformer_blocks.27.ff.net.0.proj": 0.04372084566525051, + "action_head.model.transformer_blocks.27.ff.net.2": 0.02666820798601423, + "action_head.model.transformer_blocks.28.attn1.to_k": 0.1702331304550171, + "action_head.model.transformer_blocks.28.attn1.to_out.0": 0.026269116571971347, + "action_head.model.transformer_blocks.28.attn1.to_q": 0.04598723990576608, + "action_head.model.transformer_blocks.28.attn1.to_v": 0.1702331304550171, + "action_head.model.transformer_blocks.28.ff.net.0.proj": 0.043115011283329556, + "action_head.model.transformer_blocks.28.ff.net.2": 0.021192478282111033, + "action_head.model.transformer_blocks.29.attn1.to_k": 0.020611428788730075, + "action_head.model.transformer_blocks.29.attn1.to_out.0": 0.01568733581474849, + "action_head.model.transformer_blocks.29.attn1.to_q": 0.020611428788730075, + "action_head.model.transformer_blocks.29.attn1.to_v": 0.020611428788730075, + "action_head.model.transformer_blocks.29.ff.net.0.proj": 0.04242040855543954, + "action_head.model.transformer_blocks.29.ff.net.2": 0.015297943992274148, + "action_head.model.transformer_blocks.3.attn1.to_k": 0.02290721024785723, + "action_head.model.transformer_blocks.3.attn1.to_out.0": 0.0023866056331566404, + "action_head.model.transformer_blocks.3.attn1.to_q": 0.02290721024785723, + "action_head.model.transformer_blocks.3.attn1.to_v": 0.02290721024785723, + "action_head.model.transformer_blocks.3.ff.net.0.proj": 0.05265421526772635, + "action_head.model.transformer_blocks.3.ff.net.2": 0.044673140559877665, + "action_head.model.transformer_blocks.30.attn1.to_k": 0.1702331304550171, + "action_head.model.transformer_blocks.30.attn1.to_out.0": 0.02740324182169778, + "action_head.model.transformer_blocks.30.attn1.to_q": 0.024995522839682444, + "action_head.model.transformer_blocks.30.attn1.to_v": 0.1702331304550171, + "action_head.model.transformer_blocks.30.ff.net.0.proj": 0.039933745350156515, + "action_head.model.transformer_blocks.30.ff.net.2": 0.01998665077345712, + "action_head.model.transformer_blocks.31.attn1.to_k": 0.020228247557367598, + "action_head.model.transformer_blocks.31.attn1.to_out.0": 0.024163037538528442, + "action_head.model.transformer_blocks.31.attn1.to_q": 0.020228247557367598, + "action_head.model.transformer_blocks.31.attn1.to_v": 0.020228247557367598, + "action_head.model.transformer_blocks.31.ff.net.0.proj": 0.03793149335043771, + "action_head.model.transformer_blocks.31.ff.net.2": 0.08539462089538574, + "action_head.model.transformer_blocks.4.attn1.to_k": 0.1702331304550171, + "action_head.model.transformer_blocks.4.attn1.to_out.0": 0.021854096225329807, + "action_head.model.transformer_blocks.4.attn1.to_q": 0.05850135002817426, + "action_head.model.transformer_blocks.4.attn1.to_v": 0.1702331304550171, + "action_head.model.transformer_blocks.4.ff.net.0.proj": 0.05393202815737043, + "action_head.model.transformer_blocks.4.ff.net.2": 0.036065293209893365, + "action_head.model.transformer_blocks.5.attn1.to_k": 0.05765576447759356, + "action_head.model.transformer_blocks.5.attn1.to_out.0": 0.002981468236872128, + "action_head.model.transformer_blocks.5.attn1.to_q": 0.05765576447759356, + "action_head.model.transformer_blocks.5.attn1.to_v": 0.05765576447759356, + "action_head.model.transformer_blocks.5.ff.net.0.proj": 0.055648394993373325, + "action_head.model.transformer_blocks.5.ff.net.2": 0.03139464770044599, + "action_head.model.transformer_blocks.6.attn1.to_k": 0.1702331304550171, + "action_head.model.transformer_blocks.6.attn1.to_out.0": 0.02357316017150879, + "action_head.model.transformer_blocks.6.attn1.to_q": 0.04764561993735177, + "action_head.model.transformer_blocks.6.attn1.to_v": 0.1702331304550171, + "action_head.model.transformer_blocks.6.ff.net.0.proj": 0.0571263815675463, + "action_head.model.transformer_blocks.6.ff.net.2": 0.01981711174760546, + "action_head.model.transformer_blocks.7.attn1.to_k": 0.06116958175386701, + "action_head.model.transformer_blocks.7.attn1.to_out.0": 0.0055568622691290715, + "action_head.model.transformer_blocks.7.attn1.to_q": 0.06116958175386701, + "action_head.model.transformer_blocks.7.attn1.to_v": 0.06116958175386701, + "action_head.model.transformer_blocks.7.ff.net.0.proj": 0.05797455566270011, + "action_head.model.transformer_blocks.7.ff.net.2": 0.01747190739427294, + "action_head.model.transformer_blocks.8.attn1.to_k": 0.1702331304550171, + "action_head.model.transformer_blocks.8.attn1.to_out.0": 0.02347240916320256, + "action_head.model.transformer_blocks.8.attn1.to_q": 0.048505936350141256, + "action_head.model.transformer_blocks.8.attn1.to_v": 0.1702331304550171, + "action_head.model.transformer_blocks.8.ff.net.0.proj": 0.058201219354357035, + "action_head.model.transformer_blocks.8.ff.net.2": 0.017038420907088687, + "action_head.model.transformer_blocks.9.attn1.to_k": 0.02103722095489502, + "action_head.model.transformer_blocks.9.attn1.to_out.0": 0.006777730371270861, + "action_head.model.transformer_blocks.9.attn1.to_q": 0.02103722095489502, + "action_head.model.transformer_blocks.9.attn1.to_v": 0.02103722095489502, + "action_head.model.transformer_blocks.9.ff.net.0.proj": 0.058148856673921855, + "action_head.model.transformer_blocks.9.ff.net.2": 0.015461771615913935 + } +} diff --git a/families/lerobot_groot/graph.py b/families/lerobot_groot/graph.py index 5401852aa1..555917966f 100644 --- a/families/lerobot_groot/graph.py +++ b/families/lerobot_groot/graph.py @@ -97,12 +97,19 @@ def _bias(network, out, bias, rank, dtype): return add(network, out, const(network, np.asarray(bias, dtype=np.float32).reshape((1,) * (rank - 1) + (-1,)), dtype)) -def linear(network, x, weight, bias=None, *, dtype): - """``x @ W^T + b`` in ``dtype`` for torch.nn.Linear's ``[out, in]`` weight; returns ``dtype``.""" +def linear(network, x, weight, bias=None, *, dtype, fp8=None): + """``x @ W^T + b`` in ``dtype`` for torch.nn.Linear's ``[out, in]`` weight; returns ``dtype``. + + With ``fp8`` (a ``quantization.FP8Projection``) the activation and the weight enter the + matmul through FP8 Quantize/Dequantize instead, and ``weight`` is not used. + """ x = cast(network, x, dtype) rank = len(x.shape) - out = network.add_matrix_multiply( - x, trt.MatrixOperation.NONE, _kernel(network, weight, rank, dtype), trt.MatrixOperation.TRANSPOSE).get_output(0) + if fp8 is None: + kernel = _kernel(network, weight, rank, dtype) + else: + x, kernel = fp8.operands(network, x, rank, dtype) + out = network.add_matrix_multiply(x, trt.MatrixOperation.NONE, kernel, trt.MatrixOperation.TRANSPOSE).get_output(0) return _bias(network, out, bias, rank, dtype) diff --git a/families/lerobot_groot/model.py b/families/lerobot_groot/model.py index 54c2527fd4..94c39c5911 100644 --- a/families/lerobot_groot/model.py +++ b/families/lerobot_groot/model.py @@ -8,6 +8,8 @@ and returns the keys/values every cross-attention block of the DiT reads plus the state token; ``denoise.plan`` is one flow-matching step, run ``num_inference_steps`` times by the runtime. The prompt length is dynamic, bounded by ``max_instruction_tokens``. +``quantization="fp8"`` (bf16 builds only) runs the DiT projections as FP8 GEMMs; see +``quantization.py``. """ from __future__ import annotations @@ -39,9 +41,11 @@ def _dtype(precision: str): return trt.bfloat16 if precision == "bf16" else trt.float32 -def _builder(verbose: bool): +def _builder(verbose: bool, quant=None): import tensorrt as trt + from .quantization import configure_fp8_build + logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) builder = trt.Builder(logger) network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) @@ -50,6 +54,7 @@ def _builder(verbose: bool): # The float32 parts (norms, softmax, rotary angles, the residual stream) must be real # float32: TF32 would hide graph error behind a plausible-looking number. config.clear_flag(trt.BuilderFlag.TF32) + configure_fp8_build(config, quant) return builder, network, config @@ -77,12 +82,12 @@ def _profile(builder, config, shapes: dict[str, Callable[[int], tuple[int, ...]] config.add_optimization_profile(profile) -def _build_condition_plan(cfg, weights, *, precision: str, verbose: bool = False) -> bytes: +def _build_condition_plan(cfg, weights, *, precision: str, verbose: bool = False, quant=None) -> bytes: import tensorrt as trt from . import action_head, backbone, graph, vision - builder, network, config = _builder(verbose) + builder, network, config = _builder(verbose, quant) dtype = _dtype(precision) grid = cfg.camera_height // vision.PATCH patches = len(cfg.cameras) * grid * grid @@ -100,7 +105,8 @@ def _build_condition_plan(cfg, weights, *, precision: str, verbose: bool = False network, weights, token_ids=ids["token_ids"], image_index=ids["image_index"], image_mask=ids["image_mask"], position_ids=positions, causal_bias=causal, image_tokens=tokens, deepstack=deepstack, dtype=dtype) _, state_features, cross_k, cross_v = action_head.build_condition( - network, weights, hidden, state, embodiment=cfg.embodiment_id, max_state_dim=cfg.max_state_dim, dtype=dtype) + network, weights, hidden, state, embodiment=cfg.embodiment_id, max_state_dim=cfg.max_state_dim, dtype=dtype, + quant=quant) graph.mark_output(network, state_features, "state_features") graph.mark_output(network, cross_k, "cross_k") graph.mark_output(network, cross_v, "cross_v") @@ -112,12 +118,12 @@ def _build_condition_plan(cfg, weights, *, precision: str, verbose: bool = False return _serialize(builder, network, config, "condition") -def _build_denoise_plan(cfg, weights, *, precision: str, verbose: bool = False) -> bytes: +def _build_denoise_plan(cfg, weights, *, precision: str, verbose: bool = False, quant=None) -> bytes: import tensorrt as trt from . import action_head, graph - builder, network, config = _builder(verbose) + builder, network, config = _builder(verbose, quant) dtype = _dtype(precision) blocks, width = len(action_head.CROSS_BLOCKS), action_head.WIDTH x_next, velocity = action_head.build_denoise( @@ -130,7 +136,7 @@ def _build_denoise_plan(cfg, weights, *, precision: str, verbose: bool = False) cross_k=network.add_input("cross_k", dtype, (blocks, -1, width)), cross_v=network.add_input("cross_v", dtype, (blocks, -1, width)), text_bias=network.add_input("text_bias", trt.float32, (1, 1, 1, -1)), - image_bias=network.add_input("image_bias", trt.float32, (1, 1, 1, -1))) + image_bias=network.add_input("image_bias", trt.float32, (1, 1, 1, -1)), quant=quant) graph.mark_output(network, x_next, "x_next") graph.mark_output(network, velocity, "velocity") @@ -146,6 +152,16 @@ def crop_size(cfg: config_module.GrootConfig) -> int: return max(1, int(cfg.camera_height * fraction)) if 0.0 < fraction < 1.0 else cfg.camera_height +def quantization_context(cfg: config_module.GrootConfig, weights, quantization: str, precision: str): + """The FP8 context of an ``fp8`` build: the stored activation scales of this checkpoint.""" + from . import action_head, backbone + from .quantization import calibrate_groot + + return calibrate_groot( + weights, backbone_layers=backbone.kept_layers(weights), dit_layers=action_head.LAYERS, + embodiment_id=cfg.embodiment_id, quantization=quantization, precision=precision) + + def runtime_json(cfg: config_module.GrootConfig, precision: str) -> dict: return { "cameras": list(cfg.cameras), @@ -196,8 +212,13 @@ def build(request: "BuildRequest", writer: "BundleWriter") -> None: raise NotImplementedError("lerobot_groot does not support video_num_frames") if request.mtp_seq_len is not None: raise NotImplementedError("lerobot_groot does not support MTP speculative decoding") - if request.quantization not in {None, "none"} or request.fp32_layers: - raise NotImplementedError("lerobot_groot does not support quantization or mixed-precision layers") + quantization = None if request.quantization in {None, "none"} else str(request.quantization).lower() + if quantization not in {None, "fp8"}: + raise NotImplementedError(f"lerobot_groot supports quantization fp8 or none; got {request.quantization!r}") + if quantization == "fp8" and precision != "bf16": + raise NotImplementedError(f"lerobot_groot fp8 needs precision bf16, got {precision}") + if request.fp32_layers: + raise NotImplementedError("lerobot_groot does not support mixed-precision layers") if request.max_sequence_length is not None: raise NotImplementedError( f"lerobot_groot prompt length is bounded by {MAX_INSTRUCTION_TOKENS} instruction tokens") @@ -205,12 +226,17 @@ def build(request: "BuildRequest", writer: "BundleWriter") -> None: if given not in {None, cfg.camera_height}: raise NotImplementedError(f"lerobot_groot camera {axis} is fixed at the checkpoint's {cfg.camera_height}") + if quantization is not None: + from .quantization import refuse_quantized_checkpoint + + refuse_quantized_checkpoint(model_dir) # The tokenizer first: a missing or wrong file then fails in seconds, not after the weights load. vocabulary = load_tokenizer_json(model_dir) weights = load_checkpoint(model_dir) verbose = bool(request.verbose) - condition_plan = _build_condition_plan(cfg, weights, precision=precision, verbose=verbose) - denoise_plan = _build_denoise_plan(cfg, weights, precision=precision, verbose=verbose) + quant = None if quantization is None else quantization_context(cfg, weights, quantization, precision) + condition_plan = _build_condition_plan(cfg, weights, precision=precision, verbose=verbose, quant=quant) + denoise_plan = _build_denoise_plan(cfg, weights, precision=precision, verbose=verbose, quant=quant) writer.set_header(family="lerobot_groot", task=request.task, backend=request.backend) writer.add_bytes("condition.plan", condition_plan) diff --git a/families/lerobot_groot/quantization.py b/families/lerobot_groot/quantization.py new file mode 100644 index 0000000000..55cf07a812 --- /dev/null +++ b/families/lerobot_groot/quantization.py @@ -0,0 +1,243 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""LeRobot GR00T N1.7 FP8 self-quantization of the action-head DiT projections. + +``quantization="fp8"`` quantizes the original checkpoint at build time, because no FP8 +GR00T checkpoint exists. Weights are quantized one projection at a time: +``scale = max(abs(weight)) / 448`` (448 is the largest finite e4m3 value), and the FP8 +bytes are fed to the engine as a constant. FP8 weights are never stored. The activation +scale of each projection input needs real observations, so it was calibrated once offline +with ``calibrate_fp8_scales.py`` (ModelOpt ``FP8_DEFAULT_CFG``, max calibration, LIBERO +observations through the LeRobot policy and all its denoising steps) and is read from +``fp8_activation_scales.json``. Users do not run calibration. + +Each quantized projection becomes:: + + activation (bf16) -> Quantize -> Dequantize (bf16) -> MatMul (bf16) + FP8 weight -> Dequantize (bf16) ----------------------^ + +The build runs in bfloat16, so the Dequantize layers output bfloat16: the GEMM keeps the bias +in its epilogue and there is no fp16 overflow to guard against. + +Quantized, in every one of the 32 DiT blocks ``action_head.model.transformer_blocks.N``: +``attn1.to_q``, ``attn1.to_k``, ``attn1.to_v`` (in the cross-attention blocks ``to_k`` and +``to_v`` read the vision-language tokens and run once per call in the condition engine), +``attn1.to_out.0``, ``ff.net.0.proj`` and ``ff.net.2``. The attention is built from explicit +matmul and softmax layers, so the attention output projection compiles as FP8 on TensorRT +11.1 and 11.4. + +Everything else stays in the build precision, measured (see the README): the Qwen3-VL text +layers and the four VL self-attention blocks (FP8 there adds 1% to 2% to the step speed-up +and doubles the mean action error), the vision tower and DeepStack mergers, the embedding +table, the norms, the attention matmuls, the timestep embedders and the AdaLN modulation +(one-row GEMMs), the DiT output projections and the embodiment-specific state/action +encoders and decoder. + +A projection whose calibrated input scale is more than 20 times the median of its kind has +a massive activation: its entry in the scale file is ``null`` and it stays in bfloat16 on +purpose. A projection missing from the file is an error. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from pathlib import Path + +import ml_dtypes +import numpy as np + +FP8_E4M3_MAX = 448.0 +ACTIVATION_SCALES_FILENAME = "fp8_activation_scales.json" +_SUPPORTED = ("fp8",) +_DUAL_GEMM_KNOB = "peep:match_dual_gemm" + +BACKBONE_ROOT = "backbone.model.model.language_model.layers" +DIT_ROOT = "action_head.model.transformer_blocks" +# Quantized projections of one DiT block. +QUANTIZED_STEMS = ("attn1.to_q", "attn1.to_k", "attn1.to_v", "attn1.to_out.0", "ff.net.0.proj", "ff.net.2") + + +def projection_names(dit_layers: int) -> list[str]: + """Checkpoint module names of every quantized projection of this model config.""" + return [f"{DIT_ROOT}.{index}.{stem}" for index in range(dit_layers) for stem in QUANTIZED_STEMS] + + +def projection_kind(name: str, cross_blocks: tuple[int, ...]) -> str: + """The group a projection's scale is compared within: the stem, and whether the block is a + cross-attention block (its keys/values read the VL tokens, its queries the action tokens).""" + if not name.startswith(DIT_ROOT + "."): + raise ValueError(f"{name} is not a quantized GR00T projection") + index, stem = name[len(DIT_ROOT) + 1:].split(".", 1) + if stem not in QUANTIZED_STEMS: + raise ValueError(f"{name} is not a quantized GR00T projection") + return f"{stem}.{'cross' if int(index) in cross_blocks else 'self'}" + + +def _abs_sum(array) -> float: + return float(np.abs(np.asarray(array, dtype=np.float32)).sum(dtype=np.float64)) + + +def scale_key(weights, *, backbone_layers: int, embodiment_id: int) -> str: + """Key of this checkpoint's entry in the activation-scale file. + + The architecture is fixed (Cosmos-Reason2-2B backbone, 4 VL blocks, 32 DiT blocks), so the + key holds what varies between ``lerobot-train`` outputs: the kept backbone layers, the + embodiment the bundle serves, and numeric weight sums (sum of absolute values, 5 + significant digits) of the first backbone query projection and the first DiT query + projection, which a fine-tune changes. + """ + return ( + f"groot-n17-l{int(backbone_layers)}-e{int(embodiment_id)}" + f"-w{_abs_sum(weights[f'{BACKBONE_ROOT}.0.self_attn.q_proj.weight']):.4e}" + f"-{_abs_sum(weights[f'{DIT_ROOT}.0.attn1.to_q.weight']):.4e}" + ) + + +def quantize_fp8_weight(weight_fp32: np.ndarray) -> tuple[np.ndarray, float]: + """Quantize one weight to FP8 e4m3 with a single scalar scale. + + Returns the raw fp8 bytes and ``scale = max(abs(weight)) / 448``; the same scale is the + dequantize scale in the engine. + """ + amax = float(np.abs(weight_fp32).max()) + if not np.isfinite(amax) or amax <= 0: + raise ValueError("weight tensor has no finite positive amax to scale from") + scale = amax / FP8_E4M3_MAX + scaled = np.clip(weight_fp32 / scale, -FP8_E4M3_MAX, FP8_E4M3_MAX) + packed = scaled.astype(ml_dtypes.float8_e4m3fn).view(np.uint8) + return np.ascontiguousarray(packed), scale + + +def load_activation_scales(key: str) -> dict[str, float | None]: + """The committed activation scales of ``key``; raises when the checkpoint was never calibrated. + + An entry of ``None`` marks a projection the calibration left unquantized on purpose. + """ + path = Path(__file__).parent / ACTIVATION_SCALES_FILENAME + try: + table = json.loads(path.read_text(encoding="utf-8")) + except FileNotFoundError as error: + raise RuntimeError(f"lerobot_groot FP8 build requires {path}, which was not found") from error + if key not in table: + raise NotImplementedError( + f"no calibrated FP8 activation scales for {key!r}; available: {sorted(table)}. " + "Run families/lerobot_groot/calibrate_fp8_scales.py for this checkpoint." + ) + return table[key] + + +@dataclass(frozen=True) +class FP8Projection: + """One projection weight as raw fp8_e4m3 bytes ``[out, in]`` with its two static scales.""" + + packed: np.ndarray + weight_scale: float + input_scale: float + keep_alive: list = field(default_factory=list, compare=False, repr=False) + + def operands(self, network, activation, rank: int, dtype): + """The dequantized activation and weight for ``activation @ weight^T`` in ``dtype``.""" + import tensorrt as trt + + out_features, in_features = self.packed.shape + # trt.Weights holds a raw pointer, so the buffer must outlive the build. + self.keep_alive.append(self.packed) + weight = network.add_constant( + (1,) * (rank - 2) + (out_features, in_features), + trt.Weights(trt.DataType.FP8, self.packed.ctypes.data, out_features * in_features), + ).get_output(0) + weight = network.add_dequantize(weight, _scale(network, self.weight_scale), dtype).get_output(0) + input_scale = _scale(network, self.input_scale) + quantized = network.add_quantize(activation, input_scale, trt.DataType.FP8).get_output(0) + return network.add_dequantize(quantized, input_scale, dtype).get_output(0), weight + + +def _scale(network, value: float): + """One-element float32 constant holding a Q/DQ scale.""" + import tensorrt as trt + + array = np.asarray([value], dtype=np.float32) + return network.add_constant((1,), trt.Weights(array)).get_output(0) + + +@dataclass +class GrootQuantContext: + """The FP8 projections by checkpoint module name; every other projection stays unquantized.""" + + projections: dict[str, FP8Projection] + + def get(self, name: str) -> FP8Projection | None: + """The FP8 projection named ``name`` (for example ``...transformer_blocks.3.ff.net.2``), or None.""" + return self.projections.get(name) + + +def fp8(quant: GrootQuantContext | None, name: str) -> FP8Projection | None: + """``quant.get(name)``, or None for an unquantized build.""" + return None if quant is None else quant.get(name) + + +def configure_fp8_build(trt_config, quant: GrootQuantContext | None) -> None: + """Keep layer names for kernel inspection and turn off TensorRT's dual-GEMM fusion.""" + if quant is None: + return + import tensorrt as trt + + trt_config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED + if _DUAL_GEMM_KNOB in getattr(trt_config, "all_build_routes", ""): + # The fused gate/up kernel does not compile with scalar-scale FP8 inputs; the route is + # set only when this TensorRT build lists the knob. + trt_config.build_route = f"-{_DUAL_GEMM_KNOB}=off" + + +def refuse_quantized_checkpoint(model_dir: Path) -> None: + """Raise for a checkpoint that is already quantized: FP8 self-quantizes the original.""" + model_dir = Path(model_dir) + try: + raw = json.loads((model_dir / "config.json").read_text(encoding="utf-8")) + except (FileNotFoundError, ValueError): + raw = {} + markers = [name for name in ("hf_quant_config.json", "quantize_config.json") if (model_dir / name).is_file()] + if (isinstance(raw, dict) and raw.get("quantization_config")) or markers: + found = markers[0] if markers else "quantization_config in config.json" + raise ValueError(f"the checkpoint is already quantized ({found}); FP8 self-quantization needs the original") + + +def calibrate_groot(weights, *, backbone_layers: int, dit_layers: int, embodiment_id: int, quantization: str, + precision: str) -> GrootQuantContext: + """Quantize the loaded ``[out, in]`` weights and pair each with its stored activation scale. + + A projection with no stored scale, or a stored scale for a projection this model config + does not have, is an error. + """ + if quantization not in _SUPPORTED: + raise NotImplementedError(f"lerobot_groot supports quantization {list(_SUPPORTED)} or none; got {quantization!r}") + if precision != "bf16": + raise ValueError(f"lerobot_groot FP8 needs precision bf16, got {precision!r}") + if any(str(name).endswith((".weight_scale", ".input_scale")) for name in weights): + raise ValueError("the checkpoint is already quantized; FP8 self-quantization needs the original") + stored = load_activation_scales(scale_key(weights, backbone_layers=backbone_layers, embodiment_id=embodiment_id)) + names = projection_names(dit_layers) + unexpected = sorted(set(stored) - set(names)) + if unexpected: + raise ValueError(f"{ACTIVATION_SCALES_FILENAME} has scales for projections this model does not quantize, " + f"for example {unexpected[0]}; run calibrate_fp8_scales.py for this checkpoint") + keep_alive: list = [] + projections: dict[str, FP8Projection] = {} + for name in names: + if name not in stored: + raise ValueError(f"{name} has no activation scale in {ACTIVATION_SCALES_FILENAME}; run calibrate_fp8_scales.py") + input_scale = stored[name] + if input_scale is None: + continue # massive activation: left in bfloat16 by the calibration + if not np.isfinite(input_scale) or input_scale <= 0: + raise ValueError(f"activation scale of {name} must be finite and positive") + key = f"{name}.weight" + if key not in weights: + raise ValueError(f"checkpoint has no {key}") + packed, weight_scale = quantize_fp8_weight(np.asarray(weights[key], dtype=np.float32)) + projections[name] = FP8Projection(packed, weight_scale, float(input_scale), keep_alive) + if not projections: + raise RuntimeError("lerobot_groot FP8 self-quantization matched no projection in the checkpoint") + return GrootQuantContext(projections) diff --git a/families/lerobot_groot/tests/manifests/groot-n17-libero-spatial-fp8.json b/families/lerobot_groot/tests/manifests/groot-n17-libero-spatial-fp8.json new file mode 100644 index 0000000000..014ef98760 --- /dev/null +++ b/families/lerobot_groot/tests/manifests/groot-n17-libero-spatial-fp8.json @@ -0,0 +1,52 @@ +{ + "name": "groot-n17-libero-spatial-fp8", + "hf_id": "nvidia/gr00t17-lerobot-libero_spatial-640", + "hf_revision": "32a6ec786d6509df31b40392b4e4dcdda78c0f11", + "external_files": [ + { + "path": "tokenizer.json", + "url": "https://huggingface.co/Qwen/Qwen3-VL-2B-Instruct/resolve/89644892e4d85e24eaac8bacfd4f463576704203/tokenizer.json" + } + ], + "bundle": "groot-n17-libero-spatial-fp8.bundle", + "family": "lerobot_groot", + "task": "images_state_text_to_action_chunk", + "precision": "bf16", + "tensor_parallel_size": 1, + "testcases": [ + { + "name": "groot-n17-libero-spatial-fp8-ep1272-f000", + "premerge": true, + "inputs": { + "frame": "data/recorded_observation/ep1272_f000", + "noise": "data/noise_seed0.f32", + "reference_fp32": "data/reference/ep1272_f000.fp32.f32", + "reference_bf16": "data/reference/ep1272_f000.bf16.f32", + "reference_alternate_fp32": "data/reference/ep1272_f000.alternate.fp32.f32", + "reference_alternate_bf16": "data/reference/ep1272_f000.alternate.bf16.f32" + }, + "alternate_task": "pick up the black bowl on the stove and place it on the plate" + }, + { + "name": "groot-n17-libero-spatial-fp8-ep1262-f040", + "premerge": true, + "inputs": { + "frame": "data/recorded_observation/ep1262_f040", + "noise": "data/noise_seed0.f32", + "reference_fp32": "data/reference/ep1262_f040.fp32.f32", + "reference_bf16": "data/reference/ep1262_f040.bf16.f32" + } + }, + { + "name": "groot-n17-libero-spatial-fp8-ep1267-f080", + "premerge": true, + "inputs": { + "frame": "data/recorded_observation/ep1267_f080", + "noise": "data/noise_seed0.f32", + "reference_fp32": "data/reference/ep1267_f080.fp32.f32", + "reference_bf16": "data/reference/ep1267_f080.bf16.f32" + } + } + ], + "quantization": "fp8" +} diff --git a/families/lerobot_groot/tests/test_e2e.py b/families/lerobot_groot/tests/test_e2e.py index ac2c383245..b0e0aa0860 100644 --- a/families/lerobot_groot/tests/test_e2e.py +++ b/families/lerobot_groot/tests/test_e2e.py @@ -6,7 +6,8 @@ Every case builds the pinned ``nvidia/gr00t17-lerobot-libero_spatial-640`` checkpoint straight from its directory, runs a recorded LIBERO-spatial observation through the bundle with the reference's flow noise, and compares all 16 x 7 action values with the LeRobot float32 oracle -(and, for a bfloat16 bundle, with LeRobot's own bfloat16 policy). +(and, for a bfloat16 bundle, with LeRobot's own bfloat16 policy). The FP8 manifest builds the +same checkpoint with ``quantization="fp8"``. """ from __future__ import annotations @@ -199,6 +200,7 @@ def get(manifest: dict) -> Path: family=FAMILY, task=manifest["task"], precision=manifest["precision"], + quantization=manifest.get("quantization"), ) ) assert output.is_file() diff --git a/families/lerobot_groot/tests/test_model.py b/families/lerobot_groot/tests/test_model.py index 1a68d16fa3..5bf8115ec5 100644 --- a/families/lerobot_groot/tests/test_model.py +++ b/families/lerobot_groot/tests/test_model.py @@ -320,7 +320,8 @@ def test_build_writes_the_bundle_contract(tmp_path, stubbed_build) -> None: ({"dynamic_kv_cache": True}, NotImplementedError, "does not support dynamic_kv_cache"), ({"video_num_frames": 4}, NotImplementedError, "video_num_frames"), ({"mtp_seq_len": 2}, NotImplementedError, "MTP"), - ({"quantization": "fp8"}, NotImplementedError, "quantization"), + ({"quantization": "nvfp4"}, NotImplementedError, "quantization"), + ({"quantization": "fp8", "precision": "fp32"}, NotImplementedError, "needs precision bf16"), ({"fp32_layers": (3,)}, NotImplementedError, "mixed-precision"), ({"max_sequence_length": 100}, NotImplementedError, "instruction tokens"), ({"image_height": 224}, NotImplementedError, "height"), @@ -335,6 +336,19 @@ def test_build_rejects_unsupported_requests_before_any_heavy_work( assert stubbed_build == [] +@pytest.mark.parametrize("marker", ["config", "hf_quant_config.json"]) +def test_build_refuses_a_quantized_checkpoint_before_any_heavy_work(tmp_path, stubbed_build, marker) -> None: + config = _config() + if marker == "config": + config["quantization_config"] = {"quant_method": "modelopt"} + model_dir = _model_dir(tmp_path / "ckpt", config=config) + if marker != "config": + (model_dir / marker).write_text("{}") + with pytest.raises(ValueError, match="already quantized"): + build(_request(model_dir, quantization="fp8"), _Writer()) + assert stubbed_build == [] + + class _Network: def add_input(self, name, dtype, shape): return (name, dtype, shape) @@ -361,7 +375,7 @@ def condition_stub(network, weights, hidden, state, *, dtype, **_): dtypes["action_head"] = dtype return None, "state_features", "cross_k", "cross_v" - monkeypatch.setattr(model_module, "_builder", lambda verbose: (None, _Network(), None)) + monkeypatch.setattr(model_module, "_builder", lambda *a, **k: (None, _Network(), None)) monkeypatch.setattr(model_module, "_profile", lambda *a: None) monkeypatch.setattr(model_module, "_serialize", lambda *a: b"C") monkeypatch.setattr(vision, "build_vision", vision_stub) diff --git a/families/lerobot_groot/tests/test_quantization.py b/families/lerobot_groot/tests/test_quantization.py new file mode 100644 index 0000000000..fd33f76baa --- /dev/null +++ b/families/lerobot_groot/tests/test_quantization.py @@ -0,0 +1,306 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""lerobot_groot FP8 self-quantization: weight quantizer, scale lookup, context, calibration tool, engine.""" + +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace + +import ml_dtypes +import numpy as np +import pytest + +trt = pytest.importorskip("tensorrt") + +from .. import action_head, calibrate_fp8_scales, graph, quantization # noqa: E402 +from ..quantization import ( # noqa: E402 + BACKBONE_ROOT, + DIT_ROOT, + calibrate_groot, + configure_fp8_build, + load_activation_scales, + projection_kind, + projection_names, + quantize_fp8_weight, + scale_key, +) + +LAYERS = 16 +COUNT = action_head.LAYERS * 6 + + +def _names() -> list[str]: + return projection_names(action_head.LAYERS) + + +def _weights(seed: int = 0, size: int = 8) -> dict[str, np.ndarray]: + """Small stand-in weights under every quantized projection name and the first backbone query.""" + rng = np.random.default_rng(seed) + names = _names() + [f"{BACKBONE_ROOT}.0.self_attn.q_proj"] + return {f"{name}.weight": (rng.standard_normal((size, size)) * 0.05).astype(np.float32) for name in names} + + +def _context(weights, monkeypatch, table=None, **overrides): + stored = {name: 0.01 for name in _names()} if table is None else table + monkeypatch.setattr(quantization, "load_activation_scales", lambda key: stored) + values = {"backbone_layers": LAYERS, "dit_layers": action_head.LAYERS, + "embodiment_id": 0, "quantization": "fp8", "precision": "bf16"} + values.update(overrides) + return calibrate_groot(weights, **values) + + +def test_quantize_fp8_weight_scale_and_bytes() -> None: + weight = np.array([[0.5, -1.0], [0.25, 0.0]], dtype=np.float32) + packed, scale = quantize_fp8_weight(weight) + assert scale == pytest.approx(1.0 / 448.0) + assert packed.dtype == np.uint8 and packed.shape == (2, 2) + restored = packed.view(ml_dtypes.float8_e4m3fn).astype(np.float32) * scale + np.testing.assert_allclose(restored, weight, rtol=0.07) + + +def test_quantize_fp8_weight_never_overflows_to_nan() -> None: + weight = np.linspace(-3.0, 3.0, 4096, dtype=np.float32).reshape(64, 64) + packed, _ = quantize_fp8_weight(weight) + assert np.isfinite(packed.view(ml_dtypes.float8_e4m3fn).astype(np.float32)).all() + + +def test_quantize_fp8_weight_rejects_zero_tensor() -> None: + with pytest.raises(ValueError, match="amax"): + quantize_fp8_weight(np.zeros((4, 4), dtype=np.float32)) + + +def test_projection_names_cover_every_dit_block() -> None: + names = _names() + assert len(names) == len(set(names)) == COUNT + assert f"{DIT_ROOT}.0.attn1.to_k" in names and f"{DIT_ROOT}.31.ff.net.2" in names + assert f"{DIT_ROOT}.17.attn1.to_out.0" in names + # Backbone, VL blocks, modulation, timestep embedders and the DiT output stay unquantized. + assert all(n.startswith(DIT_ROOT + ".") for n in names) + assert not any("norm1.linear" in n or "timestep" in n or "proj_out" in n for n in names) + + +def test_projection_kind_separates_cross_and_self_attention_blocks() -> None: + cross = action_head.CROSS_BLOCKS + assert projection_kind(f"{DIT_ROOT}.0.attn1.to_k", cross) == "attn1.to_k.cross" + assert projection_kind(f"{DIT_ROOT}.1.attn1.to_k", cross) == "attn1.to_k.self" + assert projection_kind(f"{DIT_ROOT}.2.ff.net.2", cross) == "ff.net.2.cross" + for name in (f"{BACKBONE_ROOT}.7.mlp.up_proj", f"{DIT_ROOT}.0.norm1.linear"): + with pytest.raises(ValueError, match="not a quantized"): + projection_kind(name, cross) + + +def test_scale_key_names_the_layers_embodiment_and_weights() -> None: + weights = _weights() + key = scale_key(weights, backbone_layers=LAYERS, embodiment_id=3) + assert key.startswith("groot-n17-l16-e3-w") + other = _weights(seed=1) + assert scale_key(other, backbone_layers=LAYERS, embodiment_id=3) != key + + +def test_unknown_checkpoint_is_refused_and_names_the_tool() -> None: + with pytest.raises(NotImplementedError, match="calibrate_fp8_scales.py"): + load_activation_scales("groot-n17-l16-e0-w1.0000e+00-1.0000e+00") + + +def test_committed_scales_cover_every_projection() -> None: + table = json.loads((Path(quantization.__file__).parent / quantization.ACTIVATION_SCALES_FILENAME).read_text()) + assert table + names = set(_names()) + for key, scales in table.items(): + assert key.startswith("groot-n17-l16-") + assert set(scales) == names, key + values = [v for v in scales.values() if v is not None] + assert values and all(np.isfinite(v) and v > 0 for v in values) + + +def test_fp8_e2e_manifest_mirrors_the_bf16_cases_and_has_thresholds() -> None: + tests = Path(__file__).parent + read = lambda name: json.loads((tests / "manifests" / f"{name}.json").read_text()) # noqa: E731 + fp8, bf16 = read("groot-n17-libero-spatial-fp8"), read("groot-n17-libero-spatial-bf16") + assert fp8["quantization"] == "fp8" and fp8["precision"] == "bf16" + assert (fp8["hf_id"], fp8["hf_revision"]) == (bf16["hf_id"], bf16["hf_revision"]) + assert [case["inputs"] for case in fp8["testcases"]] == [case["inputs"] for case in bf16["testcases"]] + assert any(case["premerge"] for case in fp8["testcases"]) + for case in fp8["testcases"]: + limits = json.loads((tests / "thresholds" / f"{case['name']}.json").read_text())["threshold_overrides"] + assert {"max_abs_vs_fp32", "mean_abs_vs_fp32", "cosine_vs_fp32_min", "max_abs_vs_bf16"} <= set(limits) + + +def test_context_quantizes_every_projection(monkeypatch) -> None: + weights = _weights() + ctx = _context(weights, monkeypatch) + assert len(ctx.projections) == COUNT + name = f"{DIT_ROOT}.2.ff.net.0.proj" + projection = ctx.get(name) + assert projection.input_scale == 0.01 and projection.packed.shape == (8, 8) + assert projection.weight_scale == pytest.approx(np.abs(weights[f"{name}.weight"]).max() / 448.0) + assert ctx.get("action_head.model.proj_out_2") is None and ctx.get(f"{BACKBONE_ROOT}.0.self_attn.q_proj") is None + assert quantization.fp8(None, name) is None and quantization.fp8(ctx, name) is projection + + +def test_null_entry_stays_unquantized(monkeypatch) -> None: + weights = _weights() + table = {name: 0.01 for name in _names()} + table[f"{DIT_ROOT}.1.ff.net.2"] = None + ctx = _context(weights, monkeypatch, table) + assert ctx.get(f"{DIT_ROOT}.1.ff.net.2") is None and len(ctx.projections) == COUNT - 1 + + +def test_projection_without_a_scale_is_refused(monkeypatch) -> None: + table = {name: 0.01 for name in _names()} + del table[f"{DIT_ROOT}.5.ff.net.2"] + with pytest.raises(ValueError, match="has no activation scale"): + _context(_weights(), monkeypatch, table) + + +def test_scale_for_an_unknown_projection_is_refused(monkeypatch) -> None: + table = {name: 0.01 for name in _names()} + table[f"{BACKBONE_ROOT}.2.mlp.up_proj"] = 0.01 + with pytest.raises(ValueError, match="does not quantize"): + _context(_weights(), monkeypatch, table) + + +def test_bad_scale_value_is_refused(monkeypatch) -> None: + table = {name: 0.01 for name in _names()} + table[f"{DIT_ROOT}.0.attn1.to_q"] = -1.0 + with pytest.raises(ValueError, match="finite and positive"): + _context(_weights(), monkeypatch, table) + + +def test_already_quantized_checkpoint_is_refused(monkeypatch) -> None: + weights = _weights() + weights[f"{DIT_ROOT}.0.ff.net.2.weight_scale"] = np.ones(1, np.float32) + with pytest.raises(ValueError, match="already quantized"): + _context(weights, monkeypatch) + + +@pytest.mark.parametrize(("overrides", "error", "message"), [ + ({"precision": "fp32"}, ValueError, "needs precision bf16"), + ({"quantization": "nvfp4"}, NotImplementedError, "supports quantization"), +]) +def test_unsupported_settings_are_refused(monkeypatch, overrides, error, message) -> None: + with pytest.raises(error, match=message): + _context(_weights(), monkeypatch, **overrides) + + +def test_fp8_build_options_only_for_quantized_builds() -> None: + config = SimpleNamespace(profiling_verbosity=None, all_build_routes="-peep:match_dual_gemm=on", build_route="") + configure_fp8_build(config, None) + assert config.profiling_verbosity is None and config.build_route == "" + configure_fp8_build(config, object()) + assert config.profiling_verbosity == trt.ProfilingVerbosity.DETAILED + assert config.build_route == "-peep:match_dual_gemm=off" + plain = SimpleNamespace(profiling_verbosity=None, all_build_routes="", build_route="") + configure_fp8_build(plain, object()) + assert plain.build_route == "" + + +def test_projection_name_matches_only_quantized_modules() -> None: + match = calibrate_fp8_scales.projection_name + assert match(f"{DIT_ROOT}.30.attn1.to_out.0") is not None + assert match(f"{DIT_ROOT}.1.ff.net.0.proj") is not None + for name in (f"{BACKBONE_ROOT}.3.self_attn.q_proj", f"{DIT_ROOT}.0.norm1.linear", f"{DIT_ROOT}.0.attn1", + "action_head.vl_self_attention.transformer_blocks.1.ff.net.0.proj", f"{DIT_ROOT}.3.attn1.to_q.x"): + assert match(name) is None, name + + +def test_outlier_projections_become_null() -> None: + cross = action_head.CROSS_BLOCKS + scales = {f"{DIT_ROOT}.{i}.ff.net.2": 0.01 for i in range(action_head.LAYERS)} + scales[f"{DIT_ROOT}.4.ff.net.2"] = 0.5 + scales.update({f"{DIT_ROOT}.{i}.attn1.to_k": (1.0 if i in cross else 0.01) for i in range(8)}) + table, excluded = calibrate_fp8_scales.split_outlier_scales(scales, cross) + # The cross-attention keys read other tokens: compared among themselves, they are no outliers. + assert excluded == [f"{DIT_ROOT}.4.ff.net.2"] and table[excluded[0]] is None + assert table[f"{DIT_ROOT}.0.attn1.to_k"] == 1.0 + + +def test_episode_split_is_disjoint_and_skips_the_fixtures() -> None: + episodes = list(range(1000, 1500)) + calibration, held_out = calibrate_fp8_scales.split_episodes(episodes, 32, 50) + assert len(calibration) == 32 and len(held_out) == 50 + assert not set(calibration) & set(held_out) + assert not set(calibration + held_out) & set(calibrate_fp8_scales.FIXTURE_EPISODES) + assert calibrate_fp8_scales.frame_positions(100, 4) == [12, 37, 62, 87] + + +def test_calibration_cli_has_fetch_and_calibrate(monkeypatch, tmp_path: Path) -> None: + calls = [] + monkeypatch.setattr(calibrate_fp8_scales, "fetch_frames", lambda *a: calls.append(("fetch", a))) + monkeypatch.setattr(calibrate_fp8_scales, "calibrate", lambda *a: calls.append(("calibrate", a))) + calibrate_fp8_scales.main(["fetch", "--out", str(tmp_path)]) + calibrate_fp8_scales.main(["calibrate", "--model-dir", "m", "--base-dir", "b", "--tokenizer-dir", "t", + "--frames", str(tmp_path)]) + assert calls[0] == ("fetch", (tmp_path, 32, 50, 4)) + assert calls[1][0] == "calibrate" and calls[1][1][3] == tmp_path + + +# Wide enough that the FP8 weight savings dominate the plan overhead. +WIDTH, HEADS, TOKENS = 512, 4, 16 + + +def _tiny_block_plan(quant): + """A pre-norm attention + MLP block through ``graph.linear``: FP8 when ``quant`` is given.""" + rng = np.random.default_rng(0) + names = ("q", "k", "v", "o", "up", "down") + weights = {n: (rng.standard_normal((WIDTH, WIDTH)) * 0.05).astype(np.float32) for n in names} + ctx = None + if quant: + keep: list = [] + projections = {} + for n in names: + packed, scale = quantize_fp8_weight(weights[n]) + projections[n] = quantization.FP8Projection(packed, scale, 4.0 / 448.0, keep) + ctx = quantization.GrootQuantContext(projections) + builder = trt.Builder(trt.Logger(trt.Logger.ERROR)) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + config = builder.create_builder_config() + configure_fp8_build(config, ctx) + x = network.add_input("x", trt.float32, (TOKENS, WIDTH)) + dtype = trt.bfloat16 + norm = graph.layer_norm(network, x, 1e-5) + q, k, v = (graph.linear(network, norm, weights[n], dtype=dtype, fp8=quantization.fp8(ctx, n)) for n in "qkv") + context = graph.attention(network, q, k, v, heads=HEADS, kv_heads=HEADS, head_dim=WIDTH // HEADS, dtype=dtype) + h = graph.add(network, x, graph.f32(network, graph.linear(network, context, weights["o"], dtype=dtype, + fp8=quantization.fp8(ctx, "o")))) + hidden = graph.gelu_tanh(network, graph.linear(network, graph.layer_norm(network, h, 1e-5), weights["up"], + dtype=dtype, fp8=quantization.fp8(ctx, "up"))) + out = graph.add(network, h, graph.f32(network, graph.linear(network, hidden, weights["down"], dtype=dtype, + fp8=quantization.fp8(ctx, "down")))) + graph.mark_output(network, out, "y") + plan = builder.build_serialized_network(network, config) + assert plan is not None + return bytes(plan) + + +def _run(plan: bytes, x: np.ndarray) -> np.ndarray: + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("needs a CUDA device") + engine = trt.Runtime(trt.Logger(trt.Logger.ERROR)).deserialize_cuda_engine(plan) + context = engine.create_execution_context() + source = torch.from_numpy(x).cuda() + result = torch.empty((TOKENS, WIDTH), dtype=torch.float32, device="cuda") + context.set_tensor_address("x", source.data_ptr()) + context.set_tensor_address("y", result.data_ptr()) + stream = torch.cuda.current_stream() + assert context.execute_async_v3(stream.cuda_stream) + stream.synchronize() + return result.cpu().numpy() + + +def test_engine_builds_with_fp8_weights_and_is_smaller() -> None: + plain, quantized = _tiny_block_plan(False), _tiny_block_plan(True) + # Six 512x512 projections: 3 MiB in bfloat16, 1.5 MiB in FP8. + assert len(quantized) < len(plain) - 1_000_000 + engine = trt.Runtime(trt.Logger(trt.Logger.ERROR)).deserialize_cuda_engine(quantized) + info = engine.create_engine_inspector().get_engine_information(trt.LayerInformationFormat.JSON) + assert "FP8" in info or "Float8" in info or "e4m3" in info + x = (np.random.default_rng(1).standard_normal((TOKENS, WIDTH))).astype(np.float32) + want, got = _run(plain, x), _run(quantized, x) + assert np.isfinite(got).all() + cosine = float((want * got).sum() / (np.linalg.norm(want) * np.linalg.norm(got))) + assert cosine > 0.99, cosine diff --git a/families/lerobot_groot/tests/thresholds/groot-n17-libero-spatial-fp8-ep1262-f040.json b/families/lerobot_groot/tests/thresholds/groot-n17-libero-spatial-fp8-ep1262-f040.json new file mode 100644 index 0000000000..c7640bdd53 --- /dev/null +++ b/families/lerobot_groot/tests/thresholds/groot-n17-libero-spatial-fp8-ep1262-f040.json @@ -0,0 +1,9 @@ +{ + "threshold_overrides": { + "max_abs_vs_fp32": 0.016, + "mean_abs_vs_fp32": 0.0032, + "cosine_vs_fp32_min": 0.9999, + "max_abs_vs_bf16": 0.022, + "chunk_inference_p95_ms": 500.0 + } +} diff --git a/families/lerobot_groot/tests/thresholds/groot-n17-libero-spatial-fp8-ep1267-f080.json b/families/lerobot_groot/tests/thresholds/groot-n17-libero-spatial-fp8-ep1267-f080.json new file mode 100644 index 0000000000..b1aa0fa191 --- /dev/null +++ b/families/lerobot_groot/tests/thresholds/groot-n17-libero-spatial-fp8-ep1267-f080.json @@ -0,0 +1,9 @@ +{ + "threshold_overrides": { + "max_abs_vs_fp32": 0.02, + "mean_abs_vs_fp32": 0.003, + "cosine_vs_fp32_min": 0.9999, + "max_abs_vs_bf16": 0.018, + "chunk_inference_p95_ms": 500.0 + } +} diff --git a/families/lerobot_groot/tests/thresholds/groot-n17-libero-spatial-fp8-ep1272-f000.json b/families/lerobot_groot/tests/thresholds/groot-n17-libero-spatial-fp8-ep1272-f000.json new file mode 100644 index 0000000000..9f57345510 --- /dev/null +++ b/families/lerobot_groot/tests/thresholds/groot-n17-libero-spatial-fp8-ep1272-f000.json @@ -0,0 +1,9 @@ +{ + "threshold_overrides": { + "max_abs_vs_fp32": 0.023, + "mean_abs_vs_fp32": 0.003, + "cosine_vs_fp32_min": 0.9999, + "max_abs_vs_bf16": 0.021, + "chunk_inference_p95_ms": 500.0 + } +}