From 455c5a145c7b51db8238358efffd3c1bdac681ec Mon Sep 17 00:00:00 2001 From: Zhenshan Xie Date: Mon, 9 Nov 2026 23:19:33 -0800 Subject: [PATCH] feat(lerobot_groot): FP8 self-quantization of the DiT projections Add --quantization fp8 (with --precision bf16) to the LeRobot GR00T N1.7 family. The 192 projections of the 32 flow-matching DiT blocks (attn1.to_q/k/v, attn1.to_out.0, ff.net.0.proj, ff.net.2) become W8A8 FP8 with per-tensor scales: weights are quantized at build time (amax / 448), static activation scales come from fp8_activation_scales.json, calibrated once per checkpoint with ModelOpt FP8_DEFAULT_CFG on 128 LIBERO-spatial observations over all denoising steps (calibrate_fp8_scales.py). The table is keyed by model config plus a weight fingerprint and covers nvidia/gr00t17-lerobot-libero_spatial-640; an unknown checkpoint raises and names the calibration tool. fp32 with fp8, other formats and an already quantized checkpoint are refused before the weights load. Dequantize runs in bf16. The vision tower, backbone, VL blocks, AdaLN and embodiment layers stay bf16. Unquantized builds keep the same network. Adds an FP8 E2E manifest of the CI checkpoint (three premerge cases) with its own thresholds, and test_e2e.py passes the manifest's quantization. TensorRT 11.4, batch one: policy step GPU time 1.07x (B200 and B300) at the median of three builds, native runtime wall time 1.05x to 1.06x, bundle 16% smaller, mean arm action error 1.7x the bf16 bundle and 1.65x LeRobot's own bf16 error, no NaN. Signed-off-by: Zhenshan Xie --- families/lerobot_groot/README.md | 157 +++++---- families/lerobot_groot/action_head.py | 22 +- .../lerobot_groot/calibrate_fp8_scales.py | 203 ++++++++++++ .../lerobot_groot/fp8_activation_scales.json | 196 +++++++++++ families/lerobot_groot/graph.py | 15 +- families/lerobot_groot/model.py | 48 ++- families/lerobot_groot/quantization.py | 243 ++++++++++++++ .../groot-n17-libero-spatial-fp8.json | 52 +++ families/lerobot_groot/tests/test_e2e.py | 4 +- families/lerobot_groot/tests/test_model.py | 18 +- .../lerobot_groot/tests/test_quantization.py | 306 ++++++++++++++++++ ...ot-n17-libero-spatial-fp8-ep1262-f040.json | 9 + ...ot-n17-libero-spatial-fp8-ep1267-f080.json | 9 + ...ot-n17-libero-spatial-fp8-ep1272-f000.json | 9 + 14 files changed, 1196 insertions(+), 95 deletions(-) create mode 100644 families/lerobot_groot/calibrate_fp8_scales.py create mode 100644 families/lerobot_groot/fp8_activation_scales.json create mode 100644 families/lerobot_groot/quantization.py create mode 100644 families/lerobot_groot/tests/manifests/groot-n17-libero-spatial-fp8.json create mode 100644 families/lerobot_groot/tests/test_quantization.py create mode 100644 families/lerobot_groot/tests/thresholds/groot-n17-libero-spatial-fp8-ep1262-f040.json create mode 100644 families/lerobot_groot/tests/thresholds/groot-n17-libero-spatial-fp8-ep1267-f080.json create mode 100644 families/lerobot_groot/tests/thresholds/groot-n17-libero-spatial-fp8-ep1272-f000.json 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 + } +}