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

This issue tracks LeRobot policies exported to ExecuTorch with Torch-TensorRT: which ones run on which GPU, how fast, and what each project still needs to look at. Other backends have their own issues: OpenVINO in #24325, XNNPACK in #24667, Qualcomm in #24668, Core ML and MLX in #24669, the CUDA backend alone in #24684, Vulkan in #24688, and Arm in #24697.

Coverage

Legend:

• The number: milliseconds to produce one action chunk (batch size 1, median) in lerobot-rollout --bundle.path, with PyTorch never loaded, inputs and outputs on the GPU, and no operator on the CPU.
• In brackets: how many times faster than eager PyTorch (lerobot-rollout --policy.path) on the same device, checkpoint and dtype, with --tf32=false. No brackets: eager does not fit the device.
• ✅ and a number: works with torch 2.15 and the ExecuTorch 1.6 and Torch-TensorRT 2.15 fixes listed under "Built on" (some are still in review).
• 🟡 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 with a torch 2.15.0 build that adds them. For praxis: works only with the Torch-TensorRT fixes named in the notes. For flux3: works only with the flux3 fixes in fix(flux3): pass F3_NATTEN_BACKEND, run the VAE without NATTEN, serve frame checkpoints, and skip repeated work huggingface/lerobot#5213.
• ➖: the model is too big for this device's memory.
• ❌: measured, and it fails; the cause is in the tables below.
• ⏳: 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, not counting the noise draw. On the current build, vqbet, diffusion and multi_task_dit have 1 more CUDA program than before: one of their encoder programs now has a small part on the CUDA backend. wall_oss_05 is split into 4 engines and 3 CUDA programs with the current LeRobot branch (2 and 1 before). pytorch/TensorRT#5022 (uint8 casts) brings it back to 2 and 1.

Notes:

• Every row is on the new build under "Built on" in every column, except lingbot_vla_v2: it uses the Nov 3 LeRobot branch and transformers 5.5.4 in every column, because the new build's branch cannot export it yet (on the DGX Spark with the new ExecuTorch and Torch-TensorRT). Every bundle first passed the export's accuracy check against eager PyTorch.
• DGX Spark is a GB10 (aarch64).
• On the Thor, lingbot_va is slower than eager PyTorch: 10051 ms against 8470 ms per chunk. The cause is not known yet.
• On the Thor, lingbot_vla_v2 splits into TRT 387 + CUDA 386, because LeRobot keeps attention out of TensorRT on this GPU (see the Thor fused attention row under NVIDIA TensorRT). That is why it is slower than eager there.
• praxis (Runway's Praxis-1: bf16, 27B parameters, a mixture of experts with 128 experts and 8 active per token) and longwam are new in LeRobot main. Praxis-1 has 54 GB of weights and longwam 23 GB, so neither fits the Orin Nano. Its yellow cells need fixes in review on top of Torch-TensorRT 2.15.0: the Torch-TensorRT pull requests listed under "Torch-TensorRT" below. It uses the official ExecuTorch 1.6.0. CUDA backend: share GPU-resident weights in preprocess instead of deep-copying them #24937 (shares weights during preprocessing) lowers export memory when a large part runs on the CUDA backend; this row puts every large part in TensorRT, so it does not need it. The bundle and eager both ran with OMP_NUM_THREADS=8. The speedup is against eager that also decodes the video of the last step, which LeRobot cannot turn off; the bundle does not. To reproduce the praxis row from source, see Praxis-1 (27B) export to ExecuTorch with TensorRT: how to reproduce from source #25019.
• flux3 (black-forest-labs/flux-3-action-droid, bf16) needs the flux3 fixes in fix(flux3): pass F3_NATTEN_BACKEND, run the VAE without NATTEN, serve frame checkpoints, and skip repeated work huggingface/lerobot#5213. The flux3 row is built with --backend.deduplicate_weights=false, because weight sharing makes the TensorRT builder crash on it.
• cosmos3 (lerobot/Cosmos3-Edge-Policy-LIBERO, bf16) takes the robot state as a CPU input; its images stay on the GPU.
• pi0 and xvla run in bf16 in every column, because their float32 exports do not fit the Orin Nano.
• On the Orin Nano, evo1 is built at TensorRT optimization level 0 so it fits in 8 GB.
• On the Orin Nano, lawam and vla_jepa are each a single program of about 5 GiB of weights, which needs about 13.5 GiB to compile, more than this 8 GB board has.
• On the Orin Nano, pi05 takes about 425 to 447 ms per chunk depending on the build (TensorRT picks kernels by timing at export, so exports can differ by about 5%; see Builder picks a slower bf16 GEMM kernel on Orin (sm_87): the 160x128x32 tiles run about 24% slower than the builder measured NVIDIA/TensorRT#4937).
• On the Thor, eager with TF32 on is faster than the strict float32 bundles (fineart_vla: 254 against 693 ms).
• diffusion, dm05, eo1, fineart_vla, gaussian_actor, lingbot_vla_v2, molmoact2, multi_task_dit, tdmpc, vla_jepa, vqbet and wall_x use random weights in every column, made by the same script and seed on each device. fastwam uses its real weights in every column.

Built on

One build, made from the same sources for every architecture. Only the lingbot_vla_v2 row still uses the older Nov 3 build, which differs where noted below (see the notes under the table).

Every cell is measured on the 2.15 packages, except molmoact2 on the Thor (still from the earlier build with release candidates of ExecuTorch 1.6 and Torch-TensorRT 2.15): torch 2.15.0, ExecuTorch 1.6.0 from PyPI, Torch-TensorRT 2.15.0 and torch-tensorrt-executorch-runtime 0.2.1 from the 2.15 test index (PyPI now has the same compiled files: pip install "torch_tensorrt[executorch]"), TensorRT 11.4, and the lerobot-export-stack-nov5 branch (see LeRobot below).

• ExecuTorch: 1.6 (now on PyPI). The Nov 3 build was ExecuTorch 1.6 with the same pull requests applied at the heads we measured.
  • Merged into ExecuTorch 1.6 (directly or as a pick):
    • Fix crashes, wrong outputs and lowering failures in several backends and in exir #24298: crashes, wrong outputs and lowering failures in several backends
    • Python runtime: enforce CUDA stream selection rules on the stream itself #24346: CUDA stream rules in the Python runtime
    • Keep ops above 6 dims out of XNNPACK and fix quantized ops it leaves out #24380: XNNPACK, ops above 6 dimensions (pick: Keep ops above 6 dims out of XNNPACK and fix quantized ops it leaves out #24445)
    • Fix mixed-dtype cat and stack in to_edge and the edge check #24381: mixed-dtype cat and stack (pick: Fix mixed-dtype cat and stack in to_edge and the edge check #24450)
    • Support 3-D convolution in the portable convolution kernel #24385: 3-D convolution in the portable kernels (pick: Support 3-D convolution in the portable convolution kernel #24460)
    • Limit the default thread pool to the allowed CPUs on Linux (not Android) #24395: default thread pool size (pick: Limit the default thread pool to the allowed CPUs on Linux (not Android) #24449)
    • Keep Half and BFloat16 sums in float in portable matmul, var, cumsum and pooling kernels, and fix portable cat #24401: Half and BFloat16 sums in portable kernels (pick: Keep Half and BFloat16 sums in float in portable matmul, var, cumsum and pooling kernels, and fix portable cat #24443)
    • Fix CUDA mask indexing, CUDA graph outputs and the runtime output-buffer check #24320: CUDA mask indexing, CUDA graph outputs, output buffer check (pick: [release/1.6] Fix CUDA indexing, graph execution and output buffers #24488)
    • CUDA backend: per-program weights file names and antialiased resize #24324: CUDA weights file names, antialiased resize (pick: [release/1.6] Fix CUDA weights filenames and antialiased resize #24489)
    • CUDA backend: delete each output handle once when a program returns one tensor twice #24420: CUDA output handles deleted once (pick: [release/1.6] Delete duplicate CUDA output handles once #24490)
    • CUDA backend: lower fused attention variants #24426: CUDA fused attention variants (pick: [release/1.6] Lower fused CUDA attention variants #24491)
    • Fix 0-d tensors, dynamic output sizes and integer reductions, and default to float32 on CPU, in the OpenVINO backend #24377 and OpenVINO: fix bool min/max and masked read lowering, and follow-ups to #24294 #24446: OpenVINO fixes (picks: [release/1.6] Fix OpenVINO scalars, precision, dynamic outputs and reductions #24492, [release/1.6] Fix OpenVINO bool operations and masked reads #24493)
    • Add a portable kernel for antialiased bicubic resize #24400: portable antialiased bicubic resize (pick: [release/1.6] Add portable antialiased bicubic resize #24494)
    • XNNPACK: fix enable_bf16 load failures and BatchNorm folding into a 3-D linear output #24477: XNNPACK bf16 loading and BatchNorm folding
    • CUDA backend: keep wide attention heads that are not a power of 2 off the Triton SDPA kernel #24468: CUDA wide attention heads that are not a power of 2
    • 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: CUDA graph and lowering fixes, varlen_attn (pick: 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 #24512)
    • OpenVINO: fix wrong outputs when a partition returns a copy of its input #24466: OpenVINO, partitions that return a copy of an input (pick: OpenVINO: fix wrong outputs when a partition returns a copy of its input #24478)
    • Python runtime: stop a backend's log from deadlocking 2 threads #24393: log deadlock between 2 threads (pick: Python runtime: stop a backend's log from deadlocking 2 threads #24529)
    • Fix export of zero-stride constants and inputs, and of constants that delegates share #24404: zero-stride constants and shared constants (pick: Fix export of zero-stride constants and inputs, and of constants that delegates share #24526)
    • CUDA backend: lower rotary embeddings written with complex numbers #24403: CUDA complex rotary embeddings (pick: CUDA backend: lower rotary embeddings written with complex numbers #24550)
    • exir and MLX: fix folded-output clones, shared constants, boolean masks, int64 scatter_add, lowering time and memory, and #24404 follow-ups #24481: exir and MLX, folded-output clones and shared constants (pick: exir and MLX: fix folded-output clones, shared constants, boolean masks, int64 scatter_add, lowering time and memory, and #24404 follow-ups #24588)
  • Only in the Nov 3 build (none of them touches the GPU backends):
    • Merged into main, not picked for 1.6:
      • Qualcomm AI Engine Direct - Keep node names unique across the QNN and Samsung passes #24316: Qualcomm, unique node names
      • Core ML: run ops on the CPU when Core ML cannot take or return their values #24317: Core ML, uint8 and int8 inputs
      • Qualcomm AI Engine Direct - Decompose kept-whole ops that HTP leaves on CPU #24429: Qualcomm, decompose kept-whole ops
• Torch-TensorRT: 2.15, plus these pull requests, at the heads we measured (newer heads in the new build).
  • Merged into Torch-TensorRT 2.15:
    • [2.15] Cherry-pick #4966, #4964, #4963 and #4965 TensorRT#4969: picks of fix(conversion): wrong attention, softmax and integer division results, and attention that fails to build TensorRT#4964 (attention, softmax and integer division fixes), fix(partitioning): resource partitioning RecursionError, tuple outputs and budget TensorRT#4963 (resource partitioning), fix(dynamo): keep, move and refit the tensors that PyTorch nodes read after global partitioning TensorRT#4966 (frozen constants after global partitioning) and feat(conversion): optionally add each shared weight to the network once TensorRT#4965 (add each shared weight once, opt-in)
    • [2.15] Pin ExecuTorch 1.6.0 TensorRT#4950: pin ExecuTorch 1.6.0
    • 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: programs with no inputs, attention masks, uint8 shape ops, plugins
• LeRobot: main plus the lerobot-export and torch-free lerobot-rollout --bundle.path work, on the lerobot-export-stack branch of shoumikhin/lerobot. The policy fixes in it are also open as LeRobot pull requests. The table's rows were measured on the branch as it was on Nov 5, kept as lerobot-export-stack-nov5. The exceptions are praxis, measured on a newer tree whose Praxis recipe matches the current branch within 1% on the H100, flux3, measured on a newer tree with the flux3 fixes, and lingbot_vla_v2, measured on the Nov 3 branch. In the new build, pi0 and pi05 also map transformers 5.22's SigLIP weight names, as feat(policies): load pi0, pi0.5 and pi0-FAST checkpoints without first building random weights huggingface/lerobot#4775 does.
• torch 2.15.0, TensorRT 11.4, transformers 5.22.1 (5.5.4 in the Nov 3 build). The yellow cells use torch 2.15.0 with the export fixes from [export] Save and load flat_apply calls, such as indexing inside vmap pytorch#202785 applied, as it was on Oct 31. That pull request has since been split into [export] Save and load flat_apply calls, such as indexing inside vmap pytorch#202785 and [export] Load higher order op regions that return nothing, such as autocast, no_grad and cond pytorch#202925 (both open), and [export] Record fake values for vmap and jvp calls, and keep those calls through dead code elimination pytorch#202929 and [export] Keep no_grad and autocast regions that read a cond's branches or a pytree spec pytorch#202991 (both merged into main, after 2.15).
• GPUs: H100, A100, Jetson AGX Thor, Jetson Orin Nano 8 GB. The Jetson builds are for aarch64 and CUDA 13.2.

How to reproduce a row

This shows xvla on one NVIDIA GPU. The full guide is the export docs on the lerobot-export-stack branch. The install line below takes the newest code on that branch; the table's rows were measured on lerobot-export-stack-nov5.

1. Install, in a new virtual environment. The export compiles CUDA code, so it also needs the CUDA compiler:

   pip install "torch_tensorrt[executorch]"
   pip install "lerobot[export] @ git+https://github.com/shoumikhin/lerobot.git@lerobot-export-stack"
   pip install "cuda-toolkit[nvcc]==13.2.*"
   export CUDA_HOME=$(python -c "import os, nvidia; print(os.path.join(list(nvidia.__path__)[0], 'cu13'))")
   export PATH=$CUDA_HOME/bin:$PATH

2. Export. The export checks itself: it runs one test frame through the bundle and through PyTorch, and prints Bundle test case: largest difference ... (tolerance ...). It stops if the difference is too large. --tolerance changes the limit.

   lerobot-export --policy.path=lerobot/xvla-widowx --policy.dtype=bfloat16 --task="pick up the cube" \
       --backend.type=executorch_tensorrt --output_dir=bundles/xvla_trt

3. Time the bundle, then eager PyTorch on the same GPU, robot and checkpoint. The bundle run never imports PyTorch. lerobot-rollout logs how long the programs took every 5 seconds. The table's number is the median time per action chunk over 30 chunks, after 3 warm-up chunks.

   lerobot-rollout --bundle.path=bundles/xvla_trt --robot.type=<your robot> ... --duration=60
   lerobot-rollout --policy.path=lerobot/xvla-widowx --policy.dtype=bfloat16 --device=cuda --tf32=false \
       --robot.type=<your robot> ... --duration=60

CUDA backend only: use --backend.type=executorch_cuda in step 2. The rest is the same, and the Torch-TensorRT line in step 1 is not needed (pip install 'lerobot[export]' already brings ExecuTorch's CUDA backend on Linux).

Known issues

"Fix" is an open pull request unless it says merged. "Needs a decision" is a design question for the maintainers. "Needs an owner" means nobody is fixing it yet. "Verified" means the policy was measured before and after the fix.

ExecuTorch

• CUDA backend: return an error when an output has no buffer #24640 (C++ abort row) is merged into main, after 1.6.0. Every other pull request in the table below is in ExecuTorch 1.6.0, which is now on PyPI.
• Decide the GIL and __init__.py questions, and pick up the rows marked "needs an owner".

Torch-TensorRT

• [2.15] Pin ExecuTorch 1.6.0 TensorRT#4950 (the ExecuTorch 1.6.0 pin) and 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 are in Torch-TensorRT 2.15.
• The 2.15.0 wheels, now on PyPI, pin ExecuTorch 1.6.0 and the runtime wheel 0.2.1. Still open: an aarch64 wheel for glibc older than 2.39 (Torch-TensorRT 2.15.0 has only a manylinux_2_39 wheel on aarch64, so pip on older glibc silently installs an old release TensorRT#5034).
• Triage the issues linked in the table below, and pick up the row marked "needs an owner".
• praxis (Praxis-1) needs these pull requests on top of 2.15:
  • fix(conversion): leave large aten._grouped_mm to PyTorch instead of the dense lowering TensorRT#5014: large grouped matmuls of mixture-of-experts layers run in PyTorch instead of a dense lowering
  • fix(dynamo): release weights before serializing the engine, not after TensorRT#5015: frees the weights before saving the engine, so the export needs less host memory (feat(executorch): let the caller choose the folder for engine files in save() and compile_and_save() TensorRT#5026 also lets the caller choose the folder for engine files while saving, which lowers it further)
  • fix(kernels): decline kernels above TensorRT's 48 KiB shared memory limit TensorRT#5016: does not use a custom kernel that needs more than TensorRT's 48 KiB shared memory limit
    • fix(kernels): let triton_op kernels opt in to Triton's pointer alignment hint TensorRT#5017: lets custom kernels use Triton's pointer alignment hint
      • feat(conversion): run aten._grouped_mm in TensorRT with a grouped GEMM kernel TensorRT#5021: runs the mixture-of-experts grouped matrix multiply in TensorRT with a grouped GEMM kernel
  • fix(dynamo): keep a select of a uint8 tensor inside the TensorRT engine TensorRT#5019: keeps a select of a uint8 tensor inside the TensorRT engine
    • fix(dynamo): keep a cast to uint8 inside the TensorRT engine TensorRT#5022: keeps uint8 casts inside the TensorRT engine
  • feat(dynamo): convert stable sort of index tensors with a counting sort TensorRT#5020: sorts the expert indices inside TensorRT, with a stable counting sort
  • fix(conversion): compute antialiased resize weights like PyTorch's CUDA kernel TensorRT#5023: computes antialiased resize weights the same way as PyTorch's CUDA kernel
    • feat(conversion): convert small antialiased bilinear downscales by default TensorRT#5024: converts small antialiased bilinear downscales by default, so the image preprocessing becomes one TensorRT engine
• fix(executorch): let compile_and_save export engines whose outputs have static shapes TensorRT#5025: lets compile_and_save export engines whose output shapes are static but computed through ops like nonzero or boolean indexing. Torch-TensorRT 2.15.0 refuses them.
• 10 policies export only because some of their ops go to the ExecuTorch CUDA backend next to the TensorRT engines (see "Programs"; wall_x has 41 engines and 40 CUDA programs). It would help if Torch-TensorRT could take these ops, so that each policy builds one engine or a few. The 3 rows before the runtime wheel row list what stays out. fastwam is the exception: its own code copies a tensor to the CPU, which LeRobot has to fix.

NVIDIA TensorRT

These are bugs in TensorRT itself (11.4 unless noted). Each one reproduces with TensorRT alone; we can share the repros.

• Triage the rows below. "Workaround" says where Torch-TensorRT avoids one for now.

PyTorch

These show up only when a program is saved and loaded, or compiled ahead of time.

• Review the pull requests below. None is in torch 2.15, so the yellow cells need a patched torch until a release has them. Could they go into a 2.15.x patch?
• Pick up the row marked "needs an owner".

LeRobot

• Review the open pull requests:
  • Import speed
    • refactor: load robots, cameras and the rollout helpers without PyTorch huggingface/lerobot#4820: load robots, cameras and the rollout helpers without PyTorch
  • Installing
    • fix(policies): install the dataset extra with policies that render a recipe huggingface/lerobot#5365: g05, eo1, wall_x, fineart_vla and wall_oss_05 install the dataset extra that their recipe step needs
  • Loading checkpoints
    • feat(policies): load pi0, pi0.5 and pi0-FAST checkpoints without first building random weights huggingface/lerobot#4775: pi0, pi0.5 and pi0-FAST without random init, and with transformers 5.22's SigLIP weight names
    • feat(policies): load GR00T fine-tunes without NVIDIA's base model huggingface/lerobot#4778: GR00T fine-tunes without NVIDIA's base model
    • fix(groot): preprocess reloaded N1.7 frames on the policy device huggingface/lerobot#5195: GR00T N1.7 preprocesses reloaded frames on the policy device
    • fix(fastwam): follow device moves, serve chunks, and run bfloat16 on the CPU huggingface/lerobot#5196: FastWAM follows device moves, serves chunks, and runs bfloat16 on the CPU
    • fix(diffusion): load weights saved with the other compile_model setting huggingface/lerobot#5352: diffusion loads weights saved with the other compile_model setting
    • refactor(policies): load a trained policy or its PEFT adapter once, straight onto its device huggingface/lerobot#5354: load a policy or its PEFT adapter once, straight onto its device
    • feat(policies): load and run chosen parts of pi0.5 and GR00T N1.7 huggingface/lerobot#5359: load and run chosen parts of pi0.5 and GR00T N1.7
  • Policies that pass torch.export
    • fix(policies): let 11 policies pass torch.export huggingface/lerobot#5150: 11 policies
    • fix(molmoact2): let single-example inference export with torch.export huggingface/lerobot#5147: molmoact2
    • fix(policies): let EO-1, WALL-X and Wall-OSS 0.5 pass torch.export huggingface/lerobot#5162: eo1, wall_x, wall_oss_05, and wall_x's text attention scores kept in float32
    • fix(policies): let LingBot-VLA v2, LaWAM, VLA-JEPA and the GR00T N1.7 backbone pass torch.export huggingface/lerobot#5163: lingbot_vla_v2, LaWAM, VLA-JEPA, the GR00T N1.7 backbone, and LaWAM's whole action prediction exported on a vision grid
  • CPU and kernels
    • fix(lawam): let the VLM-to-LAM head take bfloat16 hidden states without autocast huggingface/lerobot#5214: LaWAM head in bf16 on the CPU
    • fix(flux3): pass F3_NATTEN_BACKEND, run the VAE without NATTEN, serve frame checkpoints, and skip repeated work huggingface/lerobot#5213: flux3 passes its backend to NATTEN, runs the VAE without NATTEN, serves frame checkpoints, and skips repeated work
• Decide the rows marked "needs a decision" and pick up the ones marked "needs an owner".

Thanks to Adrian Wang (@cehongwang) for the low-memory Jetson path: pytorch/TensorRT#4854 compiles from memory-mapped weights, and pytorch/TensorRT#4949 deploys pi0.5 with ExecuTorch and TensorRT.

cc @Gasoonjia @digantdesai