Repository navigation
perf(lerobot_groot): run the whole vision tower in bf16 in bf16 builds - #2142
Merged
zhenshanx-nv merged 1 commit intoNov 10, 2026
Merged
zhenshanx-nv merged 1 commit into
zhenshanx-nv merged 1 commit into
Conversation
The query/key/value projection of each vision block, the one part of the tower kept in float32 by the bf16 tower change, now also runs in the build precision. fp32 builds and the bf16 denoise network are unchanged. The bf16 limits of the ep1267_f080 E2E case move from 5.1e-3 / 9.4e-4 to 9.5e-3 / 1.4e-3 (max / mean abs vs LeRobot fp32), set from the measured bundle error with headroom; the other limits are unchanged. Signed-off-by: Zhenshan Xie <zhenshanx@nvidia.com>
|
Warning Review limit reachedEnable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. Next included review available in 12 minutes. View limit details
|
zhenshanx-nv
added
the
run-internal-ci
github-actions
Bot
removed
the
run-internal-ci
3 of 11 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Background
lerobot_grootbf16 bundles run the Qwen3-VL vision tower in bfloat16 since #2007, except the query/key/value projection of each of the 24 blocks, which stays float32. #2007 kept it because a fully bfloat16 tower moved one E2E observation,ep1267_f080, past its max-abs limit (6.2e-3 against 5.1e-3), and left the threshold decision open. The float32 projection costs 1.8 ms of a 12 ms policy step on B200.This PR makes that decision with measured numbers. The bf16 limits were set at about twice LeRobot's own bfloat16 error on each observation (and three times for the "vs LeRobot bf16" limit):
ep1262_f0402 x 6.3e-3,ep1267_f0802 x 2.55e-3,ep1272_f0002 x 1.5e-2. Onep1267_f080, LeRobot bfloat16 is unusually close to float32: 2.55e-3, while its per-observation maximum error on 200 held-out observations has a median of 4.1e-3 and a 95th percentile of 1.0e-2. That makes 5.1e-3 the tightest limit of the three, and it is the only one a fully bfloat16 tower exceeds.Exit Criteria
bf16build runs every matrix product of the vision tower in bfloat16, including the query/key/value projection. Norms, softmax, gelu, rotary and the residual stream stay float32, as in the rest of the graph; the measurement below shows that moving the residual stream to bfloat16 gives no speedup.fp32builds and the bf16 denoise network are unchanged: their network definitions are identical layer by layer tomain.runtime.json, the bundle format and the runtime are unchanged.test_e2e.pypath (native qualification binary) on B200 and B300.Implementation
Affected:
families/lerobot_grootonly.vision.pyattn.qkvprojection uses the towerdtypeinstead oftrt.float32, so all tower matrix products follow the build precision. Docstring updated; the unusedtensorrtimport is removed.model.pytests/test_model.pybuild_visionpasses the tower dtype to every projection, attention and merger product (bfloat16 and float32). The test fails onmainfor bfloat16.tests/thresholds/groot-n17-libero-spatial-bf16-ep1267-f080.jsonmax_abs_vs_fp320.0051 -> 0.0095,mean_abs_vs_fp320.00094 -> 0.0014. Cosine,max_abs_vs_bf16and the latency limit are unchanged.README.mdChanged thresholds (only
ep1267_f080, bf16 manifest, non-premerge case):max_abs_vs_fp32mean_abs_vs_fp32No other limit is changed.
ep1272_f000(the premerge case) uses at most 75% of a limit (mean vs fp32);ep1262_f040uses at most 54%;max_abs_vs_bf16uses at most 48% / 71% / 24%. Onep1272_f000,mainitself uses 76% of the mean limit with the alternate text on TensorRT 11.1, so this case already varies this much between builds.Change categories
Validation
Commands and Results
Speed and quality use the plans
model.buildwrites (precision="bf16", builder optimization level 3). A Python harness runs them the way the runtime does: prompt, image patches, state normalization, masks, four Euler steps, action unnormalization and the LIBERO gripper. It is the same harness as in #2007, and it gives the same fixture errors as the native E2E run below, digit for digit. "main" is built from theupstream/maintree in the same job. Policy step is the GPU time of one condition call plus four denoise calls (CUDA events, inputs on the device): the median of 20 interleaved rounds of 10 steps. All variants of one GPU type were built and timed in one job. Vision tower is the tower alone in its own engine (same build settings): the median of 20 rounds of 20 runs. "fp32" is theprecision="fp32"bundle; it is the same onmainand this branch.Speed and size, TensorRT 11.4.0.106:
TensorRT 11.1.0.106 on B200: 12.14 -> 10.35 ms (1.17x), vision tower 3.58 -> 1.77 ms. The condition plan goes from 3809 to 3658 MB; the denoise plan is unchanged (286 engine layers, 2004 MB). Against the float32 tower that #2007 replaced (19.10 ms on B200 in #2007's job), the step is now 1.86x faster. In the native E2E below, the p95 wall time of one action chunk (
qualification --repeat 20) goes from 14.4 to 12.6-12.8 ms on B200 and from 14.0 to 12.4 ms on B300.Quality on 200 held-out LIBERO-spatial observations (50 episodes x 4 frames, not in the test fixtures). Each observation has fixed noise; actions are 16 x 7 in robot units. Arm = the 6 continuous dimensions; per-observation max = the largest arm error of one action chunk. B200 and B300 gave identical actions on TensorRT 11.4:
The bundle stays below LeRobot's own bfloat16 error (9.11e-4 vs 9.27e-4). Its per-observation maximum error has the same median as LeRobot bfloat16 (4.25e-3 vs 4.06e-3) and a smaller worst case (4.5e-2 vs 6.1e-2).
E2E fixtures, max abs / mean abs vs LeRobot fp32, then max abs vs LeRobot bf16. Bold marks a miss of the old limit:
ep1262_f040ep1267_f080ep1272_f000ep1272_f000alternate textep1272_f000All cosines vs LeRobot fp32 are >= 0.99997 (limit 0.9999). On
ep1272_f000, the 11.1 build ofmainalready gives the same error as this PR (1.72e-2 / 1.85e-3). The premerge case therefore moves between builds by about this much, and its limits were not changed.Broken builds against the new limits (harness on B200 and B300, TensorRT 11.4; the first one also through the native
test_e2e.py):ep1262_f040ep1267_f080ep1272_f000(premerge)The float16 block is not a numerical defect: float16 has more mantissa bits than bfloat16, and its held-out error is the same as the bundle's. Native run of the dropped-bias build (B200, TensorRT 11.4):
test_groot_bundle_matches_the_lerobot_reference_e2efails on all three bf16 cases (for example,ep1267_f080gives 6.06e-2 against 9.5e-3).NaN scan: every variant in every job was checked: all condition outputs (state token, cross keys/values) and all denoise outputs (
x_next,velocity) of the four steps, on the 200 held-out observations, the 3 fixtures and the alternate-text run. 0 non-finite.Network checks (TensorRT 11.4, real checkpoint, network definitions compared layer by layer: layer type, output dtypes and shapes, constant weight hashes):
main5ef8c22f40b09cd42e1dbf8a95524bdfc70eba2458019dcff01e30cd0456361ba964b3ee1e3cad76(the tower's query/key/value product and the casts around it)Native E2E through
families/lerobot_groot/tests/test_e2e.py. The qualification and preprocess binaries were built from this branch with CMake (TRTMC_BUILD_TESTS=ON) in a CUDA 13.4 / TensorRT 11.4 container; the checkpoint is staged locally at the manifest revision:python -m pytest families/lerobot_groot/tests/test_e2e.py --e2e-model lerobot_groot(bf16 and fp32 manifests, all 6 cases, premergeep1272_f000included)... --e2e-model groot-n17-libero-spatial-bf16upstream/maintree with its own limitsTests and checks (B200, head of this branch):
python -m pytest families/lerobot_groot/tests -q(TensorRT 11.4.0.106)main)python -m pytest tools/tests/test_architecture.py -qpython -m tools.model_ci validateruff checkon the changed Python filesruff format --checkon the changed Python filesmain(the family uses its own hand-wrapped style)Alternatives considered
Same jobs as above (TensorRT 11.4, B200 and B300 identical in quality):
ep1267_f080max vs bf16 at 71%ep1267_f080max 6.21e-3 > 5.1e-3, mean at 99%ep1267_f080max 86%, mean 81%Running the residual stream in bfloat16 too gives no speedup. Its held-out error is higher than LeRobot's own bfloat16 policy (9.49e-4 vs 9.27e-4; 1.02e-3 vs LeRobot bf16), and it passes the old
ep1267_f080limits with only 1.16x headroom. So the residual stream stays float32, as in the rest of the graph, and the limit is changed instead.Hardware, Environment, and Revisions
B200 and B300 SXM6 (TensorRT 11.4.0.106 for speed, quality and the native E2E; TensorRT 11.1.0.106 on B200 for tests, quality and the native E2E), CUDA 13, torch 2.14.0. LeRobot 0.6.1 (
lerobot[groot]) with transformers 5.5.4 for the float32 and bfloat16 references only. Checkpointnvidia/gr00t17-lerobot-libero_spatial-640at32a6ec786d6509df31b40392b4e4dcdda78c0f11, basenvidia/GR00T-N1.7-3Bat2fc962b973bccdd5d8ce4f67cc63b264d6886495, tokenizer fromQwen/Qwen3-VL-2B-Instructat89644892e4d85e24eaac8bacfd4f463576704203, observations fromphysical-intelligence/liberoata4336d589d589045d1c56423ffdf3b88a0e19b1f. Repository base:d7ce64cd(upstream/main). Plans and bundles were built on node-local disk and deleted.Not Run / Remaining Gaps
ep1267_f080.lerobot-trainoutputs (libero_10,libero_goal) were not re-measured in bf16; their fp32 bundles are unchanged.families/lerobot_groot/README.md: it appends an FP8 study section after "Otherlerobot-trainoutputs". This PR edits earlier sections of that file, andgit merge-treeof the two branches merges cleanly. The bf16 step times quoted there (12.15 / 11.00 ms) are from before this PR.Contributor Self-Review
Notes For Future Readers
max_abs_vs_bf16).ep1267_f080's max and mean vs fp32 are the exception. LeRobot bfloat16 happens to be very close to float32 there (2.55e-3, about half its held-out median), so those two limits come from the measured bundle error with about 1.5x headroom.Risk level
Risk rationale: Low. One family, and one dtype changes in bf16 builds. fp32 networks and the bf16 denoise network are identical to
main. Held-out error stays below LeRobot's own bfloat16 policy with no non-finite values. Two limits of one non-premerge E2E case are raised from measured values and still catch a broken build. The premerge case's limits are unchanged and it uses at most 75% of them.🤖 Generated with Claude Code