TensorRT: LeRobot policies that do not run yet, and what blocks them #24353

We exported the LeRobot policies to ExecuTorch with the Torch-TensorRT backend, and ran each bundle with lerobot-rollout on a fake robot. This issue lists what still keeps policies from working on the TensorRT path, so it can be picked up. It is the TensorRT companion of #24325.

Results

Milliseconds per action chunk, batch size 1, median, measured inside lerobot-rollout. The number in brackets is the speed-up over the same policy run with lerobot-rollout --policy.path (eager PyTorch) on the same device, with the same checkpoint and dtype, and with --tf32=false, so both sides do the same float32 math.

• Stack, the same for every cell:
  • ExecuTorch release/1.6 at 80264e1, with Fix crashes, wrong outputs and lowering failures in several backends and in exir #24298 (a9c1ac4), Fix CUDA mask indexing, CUDA graph outputs and the runtime output-buffer check #24320 (b1f5f7e), CUDA backend: per-program weights file names and antialiased resize #24324 (0ec89ad), Python runtime: enforce CUDA stream selection rules on the stream itself #24346 (e3e0eb8) and Fix 0-d tensors, dynamic output sizes and integer reductions, and default to float32 on CPU, in the OpenVINO backend #24377 (333bc8f) merged, and CUDA backend: CUDA graph fixes for mask indexing, .item() and changed input shapes, masked copies as where, tensor checks dropped before compiling, float64, _reshape_copy, scatter_reduce and varlen_attn, follow-ups #24326 (94b46b7), Keep ops above 6 dims out of XNNPACK and fix quantized ops it leaves out #24380 (174702b), Fix mixed-dtype cat and stack in to_edge, the edge check and portable cat #24381 (7c8ea28), Support 3-D convolution in the portable convolution kernel #24385 (cdb77d8) and Python runtime: write a call's logs when it returns, so a backend's log cannot deadlock two threads #24393 (9b0e80d) cherry-picked from main.
  • Torch-TensorRT release/2.15 at b2463880, with QNN dl3: Unable to download dataset. Connection timed out  #4950 (337b2212), Add LLaMA perf benchmark workflow for Apple iOS #4953 (dd01ad9f), Add vectorized scalar path for single-element Tensor passed to optimized mul #4964 (3ba28411), Ignore leading 1 dimensions when checking optimized path for op_mul #4963 (d91432e8) and Optimized 2D-by-1D broadcasting in optimized op_mul #4965 (421ade8d).
  • torch 2.15.0, TensorRT 11.4, transformers 5.5.4.
• GPUs: H100, A100 (shared with other jobs), Jetson AGX Thor, Jetson Orin Nano 8 GB. All cells were measured again on this stack; the remaining Thor and Orin Nano cells are still running.
• In every run with a number, PyTorch was never loaded, all program inputs and outputs stayed on the GPU, and there were zero CPU operators. Nsight Systems counted 2 to 4 uploads (the observation) and one download (the actions) per chunk.
• "Programs": how many TensorRT engines (TRT) and CUDA backend parts (CUDA) the model is split into.
• diffusion, eo1, gaussian_actor, multi_task_dit and vqbet ran with random weights. That tests the compiler and the runtime, not a trained policy.

What the marks mean:

• $\color{#2da44e}{\checkmark}$ and a number: works with torch 2.15 and the ExecuTorch 1.6 and Torch-TensorRT 2.15 fixes listed under "Already fixed or in review" (some are still in review). Milliseconds per chunk. The "Programs" column shows whether it also needs the CUDA backend.
• $\color{#d4a017}{\checkmark}$ and a number: works only with PyTorch fixes that are not in torch 2.15 (see "PyTorch" below), so it needs the next torch release. Measured on the stack above with a torch 2.15.0 build that adds them.
• $\color{#cf222e}{\textsf{X}}$: a TensorRT or Torch-TensorRT bug that needs design work from the Torch-TensorRT team. No fix is planned for these releases. See "What blocks TensorRT" below.
• ➖: the model is too big for this device's memory.
• Empty: not measured yet, or still being worked on.

Accuracy over 16 frames against eager PyTorch, as a fraction of the action range: the float32 policies, where measured, are within 6e-6 (pi0: 3e-5 to 8e-5, with eager's own spread at 1e-2). smolvla, pi05, groot, eo1, wall_oss_05 and vla_jepa (bf16) are at 0.0095 to 0.27, which is 0.6x to 1.8x eager's own spread (wall_oss_05: 0.27, with eager's own spread at 0.21).

dm05 (5.2B parameters), fastwam (6.0B), flux3 (7.0B), lingbot_vla_v2 (6.2B) and molmoact2 (5.6B) are too big for the Orin Nano's 8 GB, so they were not run there. lingbot_va (5.1B, 10.2 GB of bf16 weights) is too big as well: its weights alone exceed the 7.3 GiB of memory, and its export was killed. eo1 (3.8B) and wall_oss_05 (4.2B) make TensorRT bundles of 7.0 GB and 7.9 GB, more than the board has free, fineart_vla (4B) has only float32 weights (16 GB), and wall_x's 4.2B bf16 weights (7.9 GB) do not fit either: its export was killed. The largest policy that runs on it is pi05: 4.1B parameters, a 6.0 GiB bundle, and 0.8 GiB left free while running.

On the Orin Nano, the exports of pi0 (3.5B), lawam (2.6B), vla_jepa (2.8B) and xvla (0.9B) run out of memory, although their programs would fit at run time (pi0's bf16 engine is 4.9 GiB; xvla's is 1.5 GiB). Exporting on the board needs more than its 6.6 GiB free: the export loads every weight, including ones the program never reads (1.8 GiB in vla_jepa, 1.6 GiB in lawam), and the compile step peaks well above that even without the TensorRT copies of item 2 (xvla with one denoising step: 12.9 GiB). So these cells stay empty: the fix is a lighter export (load straight to the device, skip never-read weights), not a bigger board. On the Thor, lingbot_vla_v2's export filled the host memory during a TensorRT engine build, and the board reset. The same build peaks at 48.8 GiB on the A100, under the 78 GiB the Thor had free. The cause: on the Thor, TensorRT 11.4 merges Torch-TensorRT's 360 ScatterAdd plugin layers into one compiler region (1 node, against 2,315 on the A100). A small repro that uses only TensorRT and NumPy grows host memory 9.2x more on the Thor than on the A100; without the plugin, both grow the same. With #4953's ScatterElements path (below) lingbot_vla_v2 no longer uses that plugin, so its Thor cell will be measured again on that stack.

#24325's TensorRT columns show more policies passing. They came from an earlier sweep (torch 2.16 nightly, ExecuTorch and Torch-TensorRT main) with extra LeRobot export branches, and they timed only the compiled model call, without saving and loading the program. We think most of the gap comes from the LeRobot side and the measuring method, not the library versions, but we did not measure that directly.

Already fixed or in review

• fix(dynamo): several correctness and speed fixes found compiling large models TensorRT#4936 (merged, in release/2.15): index_put for the Qwen3-VL DeepStack mask writes, which groot needs.
• fix(dynamo): save programs with no inputs, keep attention masks that are not plain causal, fix uint8 shape ops and views left to PyTorch, show the runtime TensorRT build number, keep plugins that run PyTorch out of ExecuTorch programs TensorRT#4953 (in review). Each fix below was measured against plain release/2.15 (at 839dc42c unless noted):
  • Saving a program with no inputs (assert fake_mode is not None in lift). This moves the noise draw to the GPU (pi05: 3 uploads per chunk instead of 4). Every bundle in the table with a noise program now saves and runs.
  • Silent wrong results: with complex numbers in the graph, a block-causal attention mask was compiled as a plain causal one (force_causal_efficient_attention replaced every square bias). Now only a bias proven causal is replaced. Largest difference from PyTorch: 1.08 before, 1.55e-6 after. FastWAM's mask shape: 1.31 before, 2.26e-6 after.
  • Views above 8 dimensions reached the CUDA backend as aten._reshape_copy, which it cannot run (view_decomposition). Now they run as aten.reshape. The 10-D Qwen3-VL patchify failed with missing fallback kernels: aten::_reshape_copy before and runs after. lingbot_vla_v2 now builds its engines and its CUDA program.
  • A shape op on a uint8 input before a cast (img.permute(2, 0, 1).float()) failed with Could not get tensor shape. Now uint8 shape ops stay in PyTorch. 3 of 4 small cases failed before, 4 of 4 pass after, with an exact match. xvla passes. fastwam's engine build no longer fails on its uint8 camera frame (Could not find any implementation for node {ForeignNode[Cast ITensor arg1_1 from DataType.UINT8 to DataType.FLOAT ...]}): with offload_module_to_cpu=False, its whole compile now finishes (601 TensorRT engines, 4.4 hours on an H100).
  • An engine with Torch-TensorRT's ScatterAdd plugin (from index_put(accumulate=True)) could be saved into a .pte but not loaded without torch, because the plugin needs libtorch at run time: Cannot find plugin: ScatterAddtorch_tensorrt. This is what stops lingbot_vla_v2. Since 7172d921, saving such a program to ExecuTorch fails right away with a clear message, and compiling with torch_executed_ops={torch.ops.aten.index_put.default} keeps the accumulate out of TensorRT. Since 3a0b7b42, the setting fallback_pytorch_plugin_ops=True runs such ops in the CUDA backend instead. With it, lingbot_vla_v2 exports and runs on the H100 and the A100, but every scatter-add then runs outside TensorRT, which splits the repeated expert layers into separate engines (363), each with its own copy of the shared weights: its bundle is 41.8 GB, against 10.8 GB for a CUDA-only bundle. In its latest version, with the same setting, a scatter-add with one 1-D index uses TensorRT's built-in ScatterElements with add and stays in TensorRT: a repro with K repeated blocks gives 1 engine and a 4.0 to 4.2 MiB .pte for K = 1 to 8, against K engines and up to 46.8 MiB before. Its cells will be filled once it is measured on the same stack as the rest of the table.
• [2.15] Pin ExecuTorch 1.6.0 TensorRT#4950 (in review): pins ExecuTorch 1.6.0 on release/2.15.
• CUDA backend: CUDA graph fixes for mask indexing, .item() and changed input shapes, masked copies as where, tensor checks dropped before compiling, float64, _reshape_copy, scatter_reduce and varlen_attn, follow-ups #24326 (in review), for the CUDA programs in the bundles: the CUDA random shims are seeded from std::random_device, so 4 processes started together now draw 4 different noises (before: 2 of 4 differed). For CUDA-only bundles (TensorRT was already fine): boolean-mask copies use where (SmolVLA: 7 device-to-host copies per chunk down to 1), tensor checks are dropped before compile (the Gaussian actor now passes), float64 weights load (before: error 0x23 at load), and aten._reshape_copy and scatter_reduce with amax now have CUDA kernels.

Adrian Wang (@cehongwang) made the low-memory Jetson path possible. pytorch/TensorRT#4854 compiles from memory-mapped weights, and pytorch/TensorRT#4949 is a standalone example that deploys pi0.5 with ExecuTorch and TensorRT. His one-process pi0.5 on the Orin Nano ran at about 615 ms per chunk.

What blocks TensorRT

These need a fix in TensorRT or Torch-TensorRT. Item 1 is fixed in #4953 and item 3 in pytorch/TensorRT#4964 (both in review); item 2 is fixed by the opt-in deduplicate_weights setting in pytorch/TensorRT#4965 (in review).

1. Complex lowering mixes a CPU constant with CUDA scalars (fastwam)

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

It fails in the second constant_fold of complex_decomposition_adapter, on where(<98x98 bool constant on CPU>, scalar_tensor(0.0, cuda), scalar_tensor(-inf, cuda)). It needs offload_module_to_cpu=True, which LeRobot sets to leave about 600 MiB more free on an 8 GB Jetson. Small graphs with complex RoPE and a mask of the same size compile, so there is no minimal repro yet. Still fails with #4953 at 839dc42c. With offload_module_to_cpu=False the same graph compiles. Fixed in #4953 at dd01ad9f: both constant folds now pass. fastwam's engine build then fails in Myelin on a bf16/float32 mismatch where float32 attention scores meet a float32 constant bias (Input 0's element type (bf16) differs from input 1's element type (float)), with offload on or off. That is being narrowed down.

2. TensorRT keeps a host copy of a weight for every layer that reads it transposed (xvla on the Orin Nano)

This is TensorRT's own engine builder, measured with TensorRT 11.4.0.106. During a build, TensorRT holds one host copy of a constant weight for every layer that reads it through a transpose (MatrixOperation.TRANSPOSE on the matmul, or one shuffle per reader). It does this even when all the constant layers point at the same host array. One constant read through one shared shuffle does not grow. TensorRT's own [MemUsageStats] reports only 72 MiB of CPU memory, so it does not see this.

A repro that uses only TensorRT and NumPy (about 12 s): 1 GiB of weights read by 10 transposed matmuls raises host memory by 10.9 GiB, against 2.0x for a single reader, and 2.1x when the 10 readers share one shuffle. With refit or a stripped plan, the copies end up in the plan as well (a 10 GiB plan for 1 GiB of weights).

Flow-matching policies read their action network once per denoising step, so every weight gets 10 or 11 readers. Building xvla (3.35 GB of weights) raises host memory by 16.3 GiB, and eo1 by 66 GiB. On an 8 GB Jetson, where CPU and GPU share memory, this adds several GiB to an export that is already too big for the board.

Torch-TensorRT could avoid it: today it adds a new constant layer for every use of a weight and multiplies with MatrixOperation.TRANSPOSE. Converting each weight once and sharing one transpose would bring the xvla build from about 21 GiB to about 9 GiB.

3. Attention with a constant query, key or value fails to convert (flux3)

scaled_dot_product_attention whose query, key or value is a constant (for example frozen parameters folded for a fixed task) fails in the converter:

constant query, default scale:  TypeError: add_shape(): incompatible function arguments
constant key and value:         TypeError: add_attention_v2(): incompatible function arguments

A constant query with an explicit scale converts. Turning the constant inputs into TensorRT tensors before the attention layer makes all three cases work. A 20-line repro uses only torch and torch_tensorrt. flux3 hits it in its text-only transformer blocks, whose query, key and value are constants for a fixed task. Fixed in pytorch/TensorRT#4964 (in review): all three cases convert (largest difference 1.7e-3 for the constant query, 6.5e-4 for constant key and value).

Not bugs

On the Orin Nano, the CUDA backend's NoValidChoicesError in Inductor autotune was memory pressure (142 MiB free). On an idle board, ACT and VQ-BeT export and pass. No is_big_gpu change was needed.

Not a TensorRT problem

Still being worked on in LeRobot (empty cells). These stop in LeRobot's own preprocessing, before export starts:

• dm05, fineart_vla and wall_oss_05 write the robot state into the text prompt as digits, so the prompt length changes with the state. The fix: build the prompt outside the program and pad it to a fixed length. With that fix, fineart_vla now exports and runs on the H100 and A100, and dm05 exports on the H100 and A100 and then stops at the PyTorch save bug below. wall_oss_05 now exports and runs on the H100 and A100.
• vla_jepa ran its image processor inside the policy and stopped at a data-dependent guard (here). The fix: keep the text side as constants made at export time, and move the pixel steps into the program. With that fix, vla_jepa now exports and runs on the A100, and the PyTorch autocast bug below does not hit it on the TensorRT path. molmoact2 had the same problem. With that fix it now traces and saves on the A100, then stops at the PyTorch autocast bug below.
• flux3 needed NATTEN (neighborhood attention), which has no build for torch 2.15 and no ExecuTorch kernels. With a plain PyTorch neighborhood attention in LeRobot's flux3, checked against NATTEN, flux3 now exports on the A100. It then stopped at the TensorRT attention bug (item 3 under "What blocks TensorRT"), which fix(conversion): wrong attention and softmax results, and attention with constant inputs TensorRT#4964 fixes; flux3 will be measured again with it.

Also still being worked on in LeRobot (empty cells): evo1 (its noise draws do not fit the one-noise export recipe), tdmpc and lingbot_va (they need observation history), pi0_fast (it decodes its actions in Python with the FAST tokenizer). For evo1, tdmpc and pi0_fast, #24325 exported only the model call.

PyTorch ($\color{#d4a017}{\checkmark}$, seen only because the program is saved and loaded). pytorch/pytorch#202785 proposes fixes for both. They are not in torch 2.15.

• torch.export.load loses val on autocast regions that return nothing: lawam (here), and molmoact2 (A100; not yet run with the fix, so its cells are empty). With the fix, lawam runs on the A100 at 79.1 ms per chunk (TRT 2 + CUDA 2).
• torch.export.save cannot serialize flat_apply over ModIndex (transformers' or_mask_function masks): wall_x (here), and dm05. With the fix, dm05 saves and runs on the A100 at 123 ms per chunk (TRT 2 + CUDA 1).

ExecuTorch CUDA backend:

• Complex numbers do not reach a .pte, so fastwam cannot use a CUDA-only bundle.
• torch_attn::_varlen_attn has no CUDA kernel. Torch-TensorRT leaves this op to the CUDA backend. With the PyTorch save fix above, wall_x saves, loads and builds its engines, then stops here. CUDA backend: CUDA graph fixes for mask indexing, .item() and changed input shapes, masked copies as where, tensor checks dropped before compiling, float64, _reshape_copy, scatter_reduce and varlen_attn, follow-ups #24326 adds the kernel: with it and the PyTorch save fix, wall_x runs end to end with TensorRT and CUDA programs. A CUDA-only wall_x bundle then needs a scatter kernel for bf16 sources, also coming in CUDA backend: CUDA graph fixes for mask indexing, .item() and changed input shapes, masked copies as where, tensor checks dropped before compiling, float64, _reshape_copy, scatter_reduce and varlen_attn, follow-ups #24326.

cc @Gasoonjia @digantdesai