CUDA: LeRobot policies that do not run yet, and what blocks them

Offen
  • Stimmt etwas nicht?
#24,684 1 Kommentar 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen

Maintainer antworten meist innerhalb von 1 Tag

Dieses Issue hat noch niemand übernommen.

Bewertung

Schwierigkeit
5/5
Geschätzter Aufwand
Über eine Woche
Anfängerfreundlichkeit
35/100
Issue-Typ
Bug
Klarheit
Größtenteils klar
Aktivitätsstatus
Aktiv
Tech-Stack
python, pytorch
Bereich
machine-learning, performance

Rechercherichtung

Start with the “Known issues” tables, especially the ExecuTorch rows marked “needs an owner”; the float32 speed row includes repro commands and measurements. Read the linked fixes and reproduction details for the policy you choose, then run that repro on the listed GPU and compare results before and after your change. Done means the reported blocker is resolved and the affected policy’s export or performance is verified.

Vom Indexierungsmodell aus dem Issue-Text verfasst.

Beschreibung

module: cuda

This issue tracks LeRobot policies exported to ExecuTorch with its CUDA backend alone (CudaPartitioner, no TensorRT): which ones run on which GPU, how fast, and what each project still needs to look at. Other backends have their own issues: TensorRT in #24353, OpenVINO in #24325, XNNPACK in #24667, Qualcomm in #24668, and Core ML and MLX in #24669.

Coverage

Policy H100 A100 TensorRT, H100 (from #24353) TensorRT, A100 (from #24353) Programs
act $\color{#2da44e}{\checkmark}$ 8.89 (0.8x) $\color{#2da44e}{\checkmark}$ 15.1 (1.0x) $\color{#2da44e}{\checkmark}$ 3.15 (2.5x) $\color{#2da44e}{\checkmark}$ 4.49 (3.0x) CUDA 1
cosmos3 $\color{#2da44e}{\checkmark}$ 283 (4.1x) $\color{#2da44e}{\checkmark}$ 471 (5.4x) $\color{#2da44e}{\checkmark}$ 300 (4.0x) $\color{#2da44e}{\checkmark}$ 501 (5.2x) CUDA 3
diffusion $\color{#2da44e}{\checkmark}$ 1026 (0.8x) $\color{#2da44e}{\checkmark}$ 1401 (1.7x) $\color{#2da44e}{\checkmark}$ 535 (1.6x) $\color{#2da44e}{\checkmark}$ 773 (3.1x) CUDA 3
dm05 $\color{#d4a017}{\checkmark}$ 67.8 (7.1x) $\color{#d4a017}{\checkmark}$ 107 (10.4x) $\color{#d4a017}{\checkmark}$ 81.0 (5.7x) $\color{#d4a017}{\checkmark}$ 121 (8.8x) CUDA 1
eo1 $\color{#2da44e}{\checkmark}$ 68.0 (4.2x) $\color{#2da44e}{\checkmark}$ 100 (6.2x) $\color{#2da44e}{\checkmark}$ 58.3 (5.0x) $\color{#2da44e}{\checkmark}$ 83.3 (7.6x) CUDA 1
evo1 $\color{#2da44e}{\checkmark}$ 42.3 (4.7x) $\color{#2da44e}{\checkmark}$ 80.4 (5.3x) $\color{#2da44e}{\checkmark}$ 28.7 (6.6x) $\color{#2da44e}{\checkmark}$ 56.8 (6.6x) CUDA 1
fastwam $\color{#2da44e}{\checkmark}$ 86.2 (5.2x) $\color{#2da44e}{\checkmark}$ 125 (7.4x) $\color{#2da44e}{\checkmark}$ 54.0 (8.1x) $\color{#2da44e}{\checkmark}$ 83.5 (11.4x) CUDA 1
fineart_vla $\color{#2da44e}{\checkmark}$ 321 (0.7x) $\color{#2da44e}{\checkmark}$ 506 (1.1x) $\color{#2da44e}{\checkmark}$ 104 (2.3x) $\color{#2da44e}{\checkmark}$ 240 (2.6x) CUDA 1
flux3 $\color{#2da44e}{\checkmark}$ 2659 (2.2x) ❌ ❌ ❌ CUDA 1
g05 $\color{#2da44e}{\checkmark}$ 95.5 (4.1x) $\color{#2da44e}{\checkmark}$ 172 (5.5x) $\color{#2da44e}{\checkmark}$ 43.9 (8.9x) $\color{#2da44e}{\checkmark}$ 62.5 (12.6x) CUDA 3
gaussian_actor $\color{#2da44e}{\checkmark}$ 0.74 (2.7x) $\color{#2da44e}{\checkmark}$ 1.38 (4.0x) $\color{#2da44e}{\checkmark}$ 0.57 (3.7x) $\color{#2da44e}{\checkmark}$ 0.99 (4.5x) CUDA 1
groot $\color{#2da44e}{\checkmark}$ 23.3 (5.6x) $\color{#2da44e}{\checkmark}$ 41.1 (6.5x) $\color{#2da44e}{\checkmark}$ 19.4 (6.8x) $\color{#2da44e}{\checkmark}$ 30.0 (9.4x) CUDA 10
lawam $\color{#d4a017}{\checkmark}$ 101 (1.8x) $\color{#d4a017}{\checkmark}$ 136 (2.7x) $\color{#d4a017}{\checkmark}$ 35.8 (5.1x) $\color{#d4a017}{\checkmark}$ 62.5 (5.8x) CUDA 1
lingbot_va $\color{#2da44e}{\checkmark}$ 1731 (1.7x) $\color{#2da44e}{\checkmark}$ 3057 (2.0x) $\color{#2da44e}{\checkmark}$ 2114 (1.4x) $\color{#2da44e}{\checkmark}$ 3199 (2.0x) CUDA 18
lingbot_vla_v2 $\color{#2da44e}{\checkmark}$ 98.5 (6.3x) ❌ $\color{#2da44e}{\checkmark}$ 76.0 (7.6x) $\color{#2da44e}{\checkmark}$ 100 (10.8x) CUDA 1
molmoact2 $\color{#d4a017}{\checkmark}$ 70.5 (3.0x) $\color{#d4a017}{\checkmark}$ 108 (3.5x) $\color{#d4a017}{\checkmark}$ 48.6 (4.4x) $\color{#d4a017}{\checkmark}$ 85.9 (4.2x) CUDA 1
multi_task_dit $\color{#2da44e}{\checkmark}$ 182 (1.7x) $\color{#2da44e}{\checkmark}$ 246 (2.7x) $\color{#2da44e}{\checkmark}$ 80.0 (3.8x) $\color{#2da44e}{\checkmark}$ 101 (6.0x) CUDA 3
pi0 $\color{#2da44e}{\checkmark}$ 39.0 (4.4x) $\color{#2da44e}{\checkmark}$ 81.2 (5.2x) $\color{#2da44e}{\checkmark}$ 37.8 (4.7x) $\color{#2da44e}{\checkmark}$ 71.0 (5.2x) CUDA 9
pi05 $\color{#2da44e}{\checkmark}$ 35.7 (5.9x) $\color{#2da44e}{\checkmark}$ 69.2 (6.1x) $\color{#2da44e}{\checkmark}$ 31.9 (6.1x) $\color{#2da44e}{\checkmark}$ 57.1 (7.5x) CUDA 9
pi0_fast $\color{#2da44e}{\checkmark}$ 6048 (0.5x) $\color{#2da44e}{\checkmark}$ 8079 (0.8x) $\color{#2da44e}{\checkmark}$ 1592 (1.7x) $\color{#2da44e}{\checkmark}$ 2339 (2.9x) CUDA 3
smolvla $\color{#2da44e}{\checkmark}$ 28.6 (6.7x) $\color{#2da44e}{\checkmark}$ 49.2 (8.3x) $\color{#2da44e}{\checkmark}$ 17.3 (11.3x) $\color{#2da44e}{\checkmark}$ 28.7 (13.2x) CUDA 4
tdmpc $\color{#2da44e}{\checkmark}$ 23.4 (1.5x) $\color{#2da44e}{\checkmark}$ 33.3 (2.3x) $\color{#2da44e}{\checkmark}$ 7.83 (4.5x) $\color{#2da44e}{\checkmark}$ 13.5 (5.9x) CUDA 1
vla_jepa $\color{#2da44e}{\checkmark}$ 36.8 (2.1x) $\color{#2da44e}{\checkmark}$ 63.2 (2.4x) $\color{#2da44e}{\checkmark}$ 17.2 (4.9x) $\color{#2da44e}{\checkmark}$ 27.7 (5.7x) CUDA 1
vqbet $\color{#2da44e}{\checkmark}$ 5.30 (1.4x) $\color{#2da44e}{\checkmark}$ 8.70 (1.4x) $\color{#2da44e}{\checkmark}$ 2.36 (3.3x) $\color{#2da44e}{\checkmark}$ 3.15 (4.7x) CUDA 1
wall_oss_05 $\color{#2da44e}{\checkmark}$ 65.9 (5.1x) $\color{#2da44e}{\checkmark}$ 101 (7.0x) $\color{#2da44e}{\checkmark}$ 68.0 (4.9x) $\color{#2da44e}{\checkmark}$ 108 (6.6x) CUDA 1
wall_x $\color{#d4a017}{\checkmark}$ 113 (4.4x) ❌ $\color{#d4a017}{\checkmark}$ 93.4 (4.9x) $\color{#d4a017}{\checkmark}$ 164 (6.3x) CUDA 1
xvla $\color{#2da44e}{\checkmark}$ 41.6 (1.9x) $\color{#2da44e}{\checkmark}$ 93.2 (1.7x) $\color{#2da44e}{\checkmark}$ 16.5 (4.8x) $\color{#2da44e}{\checkmark}$ 32.7 (4.8x) CUDA 1

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.
  • $\color{#2da44e}{\checkmark}$ and a number: works with torch 2.15 and the ExecuTorch 1.6 build under "Built on".
  • $\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.
  • ❌: measured, and it fails; the cause is in the tables below.
  • ⏳: this cell is being measured (running or queued).

"Programs": how many CUDA backend programs the model is split into, not counting the one that draws the noise.

Notes:

  • Every row is on the build under "Built on" in both columns, except lingbot_vla_v2: as in #24353, it uses the Nov 3 LeRobot branch and transformers 5.5.4, because the new branch cannot export it yet (LeRobot table). Every bundle first passed the export's accuracy check against eager PyTorch.
  • cosmos3 (lerobot/Cosmos3-Edge-Policy-LIBERO, bf16) and g05 (lerobot/g05_libero, bf16) also need LeRobot export recipes on top of the build under "Built on"; the lerobot-export-stack branch does not have them yet.
  • flux3 is the frame-conditioned black-forest-labs/flux-3-action-droid, in bf16.
  • pi0 and xvla run in bf16, as in #24353.
  • The 2 TensorRT columns are copied from #24353. They use the same checkpoints and dtypes as the CUDA columns, so no cell is left blank. Each bracket compares a backend with eager PyTorch on the same GPU, so a CUDA cell and a TensorRT cell in the same row compare directly only by their milliseconds. "Programs" describes the CUDA columns only.
  • The TensorRT cells use the same build as the CUDA cells, plus Torch-TensorRT release/2.15 with pytorch/TensorRT#4950 and pytorch/TensorRT#4953, and TensorRT 11.4. For dm05, TensorRT runs with weight sharing off (--backend.deduplicate_weights=false).
  • diffusion, dm05, eo1, fineart_vla, gaussian_actor, lingbot_vla_v2, molmoact2, multi_task_dit, tdmpc, vla_jepa, vqbet and wall_x use random weights, made by the same script and seed as in #24353. fastwam uses its real weights.

Built on

One build for both GPUs. Only the lingbot_vla_v2 row uses the older Nov 3 build, as in #24353.

  • ExecuTorch: the official 1.6 rc3 wheel for CUDA 13.2 from the PyTorch test index (rc4 is the final 1.6.0, now on PyPI). The Nov 3 build was release/1.6 with the same pull requests applied at the heads we measured. The CUDA fixes in it, merged into release/1.6:
    • #24320: mask indexing, CUDA graph outputs, output buffer check (pick: #24488)
    • #24324: weights file names, antialiased resize (pick: #24489)
    • #24420: output handles deleted once (pick: #24490)
    • #24426: fused attention variants (pick: #24491)
    • #24468: wide attention heads that are not a power of 2
    • #24326: CUDA graph and lowering fixes, varlen_attn (pick: #24512)
    • #24403: complex rotary embeddings (pick: #24550)
    • #24346: CUDA stream rules in the Python runtime
  • torch 2.15.0, 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 pytorch/pytorch#202785 applied, as it was on Oct 31. That pull request has since been split into #202785, #202925, #202929 and #202991.
  • LeRobot: the lerobot-export-stack branch of shoumikhin/lerobot, as it was on Nov 3 at 10:53 AM Pacific (see #24353). pi0 and pi05 also map transformers 5.22's SigLIP weight names, as huggingface/lerobot#4775 does.
  • GPUs: H100, A100 80 GB.

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
  • #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.
  • Pick up the rows marked "needs an owner", starting with float32 speed (first row) and export memory (second row).
Problem Policies Fix Status
In float32, bundles are often slower than eager PyTorch. The backend tunes matmuls and convolutions with Triton only (max_autotune_gemm_backends and max_autotune_conv_backends are "TRITON", so the program does not need libtorch), while eager PyTorch calls cuBLAS and cuDNN. In xvla's float32 bundle on the H100, one Triton convolution takes 3.0 ms per call and 30% of the GPU time. act, diffusion, fineart_vla, pi0_fast; xvla in float32 (0.3x on both GPUs) needs an owner (for example, cuBLAS and cuDNN fallback shims, like the existing one for _weight_int4pack_mm) Repro ready: 4 float32 convolutions run at 0.56x of eager on the H100 (0.57x on the A100) and 2 linear layers at 0.45x (0.55x), inputs and outputs on the GPU
The flux3 export peaks at 83 GiB of GPU memory during AOTInductor's autotuning, so it runs out of memory on the 80 GB A100. The finished program runs in 60 GiB on the H100. flux3 needs an owner Measured
bf16 F.grouped_mm (aten._grouped_mm) has no kernel choice on the A100 (sm_80), so the export fails with NoValidChoicesError. cosmos3 needs an owner Repro ready (one F.grouped_mm layer); Inductor's message asks to add ATEN to max_autotune_gemm_backends, the same Triton-only setting as the float32 row. The cosmos3 recipe avoids it by using the fixed domain's weights
No kernel for fused attention variants or for complex rotary embeddings. lingbot_va, fastwam #24403, #24426 In 1.6.0
The Triton attention kernel computes only the first 256 values for wider heads that are not a power of 2, with no error. lingbot_va #24427 Verified, in 1.6.0
No kernel for varlen_attn. wall_x #24326 Verified, in 1.6.0
Same random seed in parallel processes, no bf16 scatter, an extra libgomp dependency, attention reading past the last key. all #24326, pytorch/pytorch#202902 Verified, in 1.6.0 (the PyTorch part is merged into main)
No kernel for x[mask], stale shapes from CUDA graphs, a too-small output buffer is accepted. smolvla #24488 In 1.6.0
Programs in one folder overwrite each other's weights file; antialiased resize does not lower. all with 2+ programs #24489 In 1.6.0
From C++, a program whose outputs stay on the GPU aborts the process (from_blob: data pointer cannot be nullptr) when the caller gives no output buffer with Module::set_output. C++ callers #24640 Verified on the A100: an error instead of an abort; merged into main
No kernel for scatter with a scalar source in deterministic mode (missing fallback kernels: aoti_torch_cuda_scatter_value_out); the backend has a shim only for a tensor source. none yet needs an owner Repro ready

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.
Problem Policies Fix Status
torch.export.save cannot save flat_apply (transformers' mask functions). wall_x, dm05 pytorch/pytorch#202785 Verified
torch.export.load loses metadata on autocast regions that return nothing. lawam, molmoact2 pytorch/pytorch#202925 Verified
A program with varlen_attn cannot be loaded after a plain import torch. wall_x pytorch/pytorch#202925 Verified
AOTInductor autotuning runs a gather with a bounds check on made-up inputs, and the compile fails. vla_jepa (earlier build) pytorch/pytorch#202931 Verified

LeRobot
  • Review the open pull requests these cells depend on:
    • huggingface/lerobot#4775: pi0 and pi0.5 with transformers 5.22's SigLIP weight names
    • huggingface/lerobot#5150: 8 policies that pass torch.export
    • huggingface/lerobot#5147: molmoact2
    • huggingface/lerobot#5162: eo1, wall_x, wall_oss_05
    • huggingface/lerobot#5163: lingbot_vla_v2, LaWAM, GR00T N1.7
  • Decide the rows marked "needs a decision".
Problem Policies Fix Status
lerobot-export and the torch-free lerobot-rollout --bundle.path exist only on the lerobot-export-stack branch. all needs a decision: how to bring them upstream in reviewable pieces Open
2 bf16 policies miss their 0.075 limit on the A100 by a small margin. wall_x gave 0.095 and 0.047 in 2 exports of the same weights; lingbot_vla_v2 gave 0.078 (0.0625 on the H100). The limit sits at bf16 rounding noise. wall_x, lingbot_vla_v2 needs a decision: a limit that allows for bf16 rounding (LeRobot uses 0.2 for 10-step bf16 policies such as pi0) Measured; not narrowed, so the backend's bf16 kernels are also a candidate
On the current branch, the lingbot_vla_v2 export stops before tracing: Exporting LingBot-VLA v2 needs a vision_grid. lingbot_vla_v2 in progress: an export recipe Measured

cc @Gasoonjia @digantdesai

Vorherrschende Sprache
Python
Sterne
5.2k
Forks
1.2k
Ø Merge
1 T. 23 Std.
Gemergte PRs (30 T.)
696

Entwicklungsumgebung

Erste Schritte

  1. Lesen Sie das ganze Issue und danach den Beitragsleitfaden des Projekts.
  2. Schreiben Sie ins Issue, dass Sie es übernehmen — das erspart doppelte Arbeit.
  3. Forken Sie das Repository und arbeiten Sie in einem Branch.
  4. Öffnen Sie einen Pull Request, der die Issue-Nummer nennt.

Mehr aus pytorch/executorch

Alle Issues in pytorch/executorch

Ähnliche Issues

Weitere Issues zu Python

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.