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

This issue tracks LeRobot policies exported to ExecuTorch with the Torch-TensorRT backend and run with lerobot-rollout on a fake robot: which policies work, how fast they are, and each known problem with the pull request that fixes it. #24325 is the companion issue for the CPU backends and the general export flow.

Coverage

Each cell is the median milliseconds per action chunk (batch size 1) inside lerobot-rollout. In brackets: the speed-up over eager PyTorch (--policy.path) on the same device, checkpoint and dtype, with --tf32=false. In every run with a number, PyTorch was never loaded, program inputs and outputs stayed on the GPU, and no operator ran on the CPU.

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 "Stack" (some are still in review).
• $\color{#d4a017}{\checkmark}$ and a number: works only with PyTorch fixes that are not in torch 2.15 (see "PyTorch" under "Known issues"), so it needs the next torch release. Measured with a torch 2.15.0 build that adds them.
• ➖: the model is too big for this device's memory.
• 🔨: a fix for this cell is being worked on.
• ⏳: this cell is being measured (running or queued).
• Empty: not measured, and no fix in progress.

"Programs": how many TensorRT engines (TRT) and CUDA backend programs (CUDA) the model is split into.

What the cells do not show:

• All four devices now use one stack: the same ExecuTorch, Torch-TensorRT, torch and LeRobot (see "Stack"). Each number was measured with lerobot-rollout --bundle.path, and each bundle first passed the export's own accuracy check against eager PyTorch.
• On the Orin Nano, most ➖ policies are bigger than the 8 GB board. lawam, vla_jepa and evo1 would fit at run time, but their export does not (see "LeRobot" below).
• The Orin Nano groot, pi05 and pi0 cells have no speed-up, because eager PyTorch does not fit the board. The Orin Nano xvla cell runs in bf16, because the float32 export does not fit.
• pi0 is now exported as a chain of smaller programs, like pi05, so that it fits the Orin Nano. Its H100, A100 and Thor cells are being measured again with that export.
• diffusion, eo1, gaussian_actor, multi_task_dit and vqbet ran with random weights, and so did fineart_vla, vla_jepa and wall_x on the Thor. That tests the compiler and the runtime, not a trained policy.

Stack

One build per architecture from the same sources, used on all four devices:

• ExecuTorch release/1.6 at 80264e1, with Fix crashes, wrong outputs and lowering failures in several backends and in exir #24298 (d59b4020c8), Fix CUDA mask indexing, CUDA graph outputs and the runtime output-buffer check #24320 (b1f5f7e8af), CUDA backend: per-program weights file names and antialiased resize #24324 (f714c0871e), Python runtime: enforce CUDA stream selection rules on the stream itself #24346 (e3e0eb85be) and Fix 0-d tensors, dynamic output sizes and integer reductions, and default to float32 on CPU, in the OpenVINO backend #24377 (f2cf062fd0) 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 (bdea161d34), Keep ops above 6 dims out of XNNPACK and fix quantized ops it leaves out #24380 (a86580f2ae), Fix mixed-dtype cat and stack in to_edge and the edge check #24381 (a61dcc9e32), Support 3-D convolution in the portable convolution kernel #24385 (caf610b39a), Python runtime: stop a backend's log from deadlocking two threads #24393 (4ca018cb17), Limit the default thread pool to the allowed CPUs on Linux (not Android) #24395 (180db83b89), Add a portable kernel for antialiased bicubic resize #24400 (31bf43b7b3), Keep Half and BFloat16 sums in float in portable matmul, var, cumsum and pooling kernels, and fix portable cat #24401 (d090938e94), CUDA backend: lower complex rotary embeddings and fused attention, and free repeated outputs once #24403 (bbe873c6ab) and Fix export of zero-stride constants and inputs, and of constants that delegates share #24404 (67d6e8b151) applied on top.
• Torch-TensorRT release/2.15 at b2463880, with [2.15] Pin ExecuTorch 1.6.0 TensorRT#4950 (337b2212), 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 (0aa63b10), fix(conversion): wrong attention, softmax and integer division results, and attention that fails to build TensorRT#4964 (77416752), fix(partitioning): resource partitioning RecursionError, tuple outputs and budget TensorRT#4963 (ef8f3024), fix(dynamo): keep frozen constants that PyTorch nodes still read after global partitioning TensorRT#4966 (ba98542c) and feat(conversion): optionally add each shared weight to the network once TensorRT#4965 (d40715b7).
• torch 2.15.0, TensorRT 11.4, transformers 5.5.4. The yellow cells use a torch 2.15.0 build that adds [export] Save and load flat_apply calls, such as indexing inside vmap, and keep regions around a cond pytorch#202785.
• GPUs: H100, A100, Jetson AGX Thor, Jetson Orin Nano 8 GB. The Jetson builds are for aarch64 and CUDA 13.2.

Known issues

Every fix pull request below is still in review unless it says Merged. SHAs are the current heads. "Verified on LeRobot": the policy was measured before and after the fix.

Torch-TensorRT

ExecuTorch

#24320 and #24324 sit on top of #24298 (base AOTInductor and CUDA fixes). The other ExecuTorch fixes in the stack are for CPU backends and are tracked in #24325.

PyTorch

These show up only because the program is saved and loaded. pytorch/pytorch#202785 (f63b88cc97) fixes both. It is not in torch 2.15, so the yellow cells use a torch 2.15.0 build that adds it.

LeRobot

lingbot_va: wrong attention in the CUDA backend

With #24403, all 19 programs of lingbot_va compile, but the bundle does not match the policy, so it is not written:

RuntimeError: The exported policy does not reproduce its test case (4.37e-01 > 0.2).

The largest difference is 0.30 on the H100, 0.44 on the A100 and 0.48 on the Thor (tolerance 0.2). The first program that drifts is the encoder. Its attention has a head size of 640. Called directly, the CUDA backend's Triton attention kernel gives wrong results (relative error 0.9 to 1.15, or NaN) for head sizes 288, 320, 384, 640 and 768, and matches PyTorch for 96, 192, 256, 512 and 1024. So it covers only the first 256 values of a head unless the size is a power of two.

Thor only: TensorRT computes a different result

The bf16 bundles of evo1 and lingbot_vla_v2 fail the export's accuracy check on the Thor only: 1.65 and 1.10, against a tolerance of 0.075. On the H100 and A100 both pass. The cause is inside one TensorRT engine:

• The C++ and the Python Torch-TensorRT runtimes both give the same wrong actions (largest difference 1.62).
• Replaying the engine's plan with plain TensorRT on the same inputs reproduces it exactly.
• With every input to that engine bit-exact to eager PyTorch, it still drifts (1.62).

So TensorRT 11.4 on sm_110 computes this network differently. Compiled as two engines, split at the vision-language model's output, evo1 is off by only 0.021. The TensorRT optimization level does not change it.

TensorRT keeps a host copy of a weight per transposed reader

TensorRT's builder (11.4.0.106) holds one host copy of a constant weight for every layer that reads it through a transpose, even when all the constant layers point at the same host array. Its [MemUsageStats] does not show this. In a repro with only TensorRT and NumPy (about 12 s), 1 GiB of weights read by 10 transposed matmuls raises host memory by 10.9 GiB (about 11x the weights), against 2.1x when the 10 readers share one shuffle. Flow-matching policies read their action network once per denoising step, so each weight gets 10 or 11 readers, which costs several GiB on an 8 GB Jetson. The opt-in deduplicate_weights in pytorch/TensorRT#4965 converts each weight once and shares one transpose.

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.

cc @Gasoonjia @digantdesai