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. CPU backends are tracked in #24325.

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.
• $\color{#2da44e}{\checkmark}$ 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).
• $\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 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.

Notes:

• One stack on all 4 devices (see "Built on"). Every bundle first passed the export's accuracy check against eager PyTorch.
• Orin Nano: xvla runs in bf16, because the float32 export does not fit.
• On the Thor, eager with TF32 on is faster than the strict float32 bundles (pi0: 249 against 820 ms).
• diffusion, eo1, gaussian_actor, multi_task_dit and vqbet use random weights, and so do fastwam, fineart_vla, vla_jepa and wall_x on the Thor.

Built on

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

• ExecuTorch: release/1.6, plus these pull requests, at the heads we measured.
  • Merged into release/1.6:
    • Python runtime: enforce CUDA stream selection rules on the stream itself #24346: CUDA stream rules in the Python runtime
  • Cherry-picked into release/1.6:
    • 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)
    • 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)
  • Open against release/1.6:
    • Fix crashes, wrong outputs and lowering failures in several backends and in exir #24298: crashes, wrong outputs and lowering failures in several backends
    • Fix CUDA mask indexing, CUDA graph outputs and the runtime output-buffer check #24320: CUDA mask indexing, CUDA graph outputs, output buffer check
    • CUDA backend: per-program weights file names and antialiased resize #24324: CUDA weights file names, antialiased resize
    • Fix 0-d tensors, dynamic output sizes and integer reductions, and default to float32 on CPU, in the OpenVINO backend #24377: OpenVINO fixes
  • Merged into main, cherry-pick to release/1.6 open:
    • 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)
    • 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)
    • Add a portable kernel for antialiased bicubic resize #24400: portable antialiased bicubic resize (pick: Add a portable kernel for antialiased bicubic resize #24447)
  • Merged into main, not yet picked for release/1.6:
    • 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
  • Open against main:
    • Python runtime: stop a backend's log from deadlocking two threads #24393: log deadlock between 2 threads
    • CUDA backend: lower rotary embeddings written with complex numbers #24403: CUDA complex rotary embeddings and fused attention; now split, with parts in CUDA backend: delete each output handle once when a program returns one tensor twice #24420 (merged) and CUDA backend: lower fused attention variants #24426 (merged)
    • Fix export of zero-stride constants and inputs, and of constants that delegates share #24404: zero-stride constants and shared constants
• Torch-TensorRT: release/2.15, plus these pull requests, at the heads we measured.
  • Open against release/2.15:
    • [2.15] Pin ExecuTorch 1.6.0 TensorRT#4950: pin ExecuTorch 1.6.0
  • Open against main:
    • 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
    • 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
    • feat(conversion): optionally add each shared weight to the network once TensorRT#4965: add each shared weight once (opt-in)
• 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.
• torch 2.15.0, TensorRT 11.4, transformers 5.5.4. 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 #202785, #202925, #202929 and #202991.
• GPUs: H100, A100, Jetson AGX Thor, Jetson Orin Nano 8 GB. The Jetson builds are for aarch64 and CUDA 13.2.

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

• Review the open pull requests below. "Built on" shows their release/1.6 picks.
• Merge the release/1.6 picks.
• Decide the GIL question and pick up the scatter kernel (last 2 rows).

Torch-TensorRT

• Review the open pull requests below, then pick them into release/2.15 ([2.15] Pin ExecuTorch 1.6.0 TensorRT#4950 pins ExecuTorch 1.6.0 there).
• Land the low-memory Jetson pull requests in main and release/2.15 (row near the end).
• Decide the row marked "needs a decision" and pick up the rows marked "needs an owner".
• 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
    • perf(imports): import torch and policies only when a command needs them huggingface/lerobot#4693: import torch and policies only when needed
    • perf(processor): import torch only when teleoperation needs it huggingface/lerobot#4742: import torch only when teleoperation needs it
    • refactor(datasets): import dataset modules on first use huggingface/lerobot#5136: import dataset modules on first use
    • refactor(rollout): let the control loop take NumPy actions from the inference engine huggingface/lerobot#4820: the control loop takes NumPy actions
  • 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
    • feat(policies): load GR00T fine-tunes without NVIDIA's base model huggingface/lerobot#4778: GR00T fine-tunes without NVIDIA's base model
    • fix(policies): load FAST action tokenizers saved by transformers 4 huggingface/lerobot#5155: FAST tokenizers saved by transformers 4
    • fix(groot): fall back when the processor's saved device is missing huggingface/lerobot#5195: GR00T processor falls back to the CPU
    • fix(fastwam): move the device FastWAM records along with the model huggingface/lerobot#5196: FastWAM's device follows its weights
  • Policies that pass torch.export
    • fix(policies): avoid tensor value reads that block torch.export in eight policies huggingface/lerobot#5150: 8 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
    • fix(policies): let LingBot-VLA v2, the LaWAM VLM stage and the GR00T N1.7 backbone pass torch.export huggingface/lerobot#5163: lingbot_vla_v2, LaWAM, GR00T N1.7
  • CPU and kernels
    • fix(lawam): let the VLM-to-LAM head take bfloat16 hidden states on the CPU huggingface/lerobot#5214: LaWAM head in bf16 on the CPU
    • fix(flux3): hand F3_NATTEN_BACKEND to NATTEN's kernels huggingface/lerobot#5213: flux3 passes its backend to NATTEN
• 7 more are ready, but GitHub's open pull request limit blocks them. Could it be raised?
  • diffusion, fastwam, flux3, wall_x, lawam, vla_jepa: one fix each, plus action chunks from frame checkpoints for fastwam and flux3
• 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