Praxis-1 (27B) export to ExecuTorch with TensorRT: how to reproduce from source #25019

Runway's Praxis-1 (a 27B bfloat16 mixture-of-experts robot policy, praxis in LeRobot) exports to one ExecuTorch bundle that runs as TensorRT engines, plus one small part on ExecuTorch's CUDA backend, and loads without PyTorch. This issue is the recipe to reproduce it from source, on H100 and A100, and on Jetson AGX Thor and DGX Spark.

Result

Time per action chunk (batch 1), real DROID camera frames, the same GPU, frames and timer for both sides, which alternate in fresh processes. Each number is the median of the valid runs' medians, 3 or more runs of 30 chunks per side. Runway's code is runwayml/praxis-1-oss at afef2e1 (its main on the day). Every side uses the same sampler settings, from Runway's serve config.

• Runway's NVFP4 path does not run on DGX Spark (its check accepts compute capability 12.0, and GB10 is 12.1), so the last row compares our FP8 bundle with Runway's bfloat16 path from its serve config (3 pipeline stages on GB10, see below).
• The NVFP4 rows quantize the same layers (the routed experts) with the same scale rule as Runway's code. Both sides quantize the bfloat16 checkpoint. Runway also publishes an NVFP4 checkpoint trained with quantization (runwayml/praxis-1-droid-nvfp4), which was not used here. It runs the same kernels, so the timing applies, but Runway's NVFP4 accuracy below is not that checkpoint's accuracy.
• On the synthetic pictures of Runway's own benchmark script: H100 236 ms against 241 ms, A100 553 ms against 714 ms.

Accuracy over 32 real chunks against a float32 reference, max / mean error per action value:

• bfloat16: bundle 0.98 / 0.062 (Thor), 0.96 / 0.065 (Spark), 0.88 / 0.063 (H100); Runway's code 0.85 / 0.068 (H100); eager bfloat16 PyTorch 1.56 / 0.069.
• NVFP4 on Thor: bundle 1.11 / 0.080, Runway's code 1.75 / 0.072.
• FP8 on Spark: bundle 0.90 / 0.062.

LeRobot eager PyTorch (H100 930 ms, A100 1970 ms) also decodes the predicted video at every denoising step, which the bundle and Runway's code skip.

What it needs on top of the releases

• ExecuTorch 1.6.0: nothing for H100 and A100. For Thor and Spark, 3 pull requests:
  • Import torch only when the Python runtime needs it, warn before outputs follow inputs, and upload bfloat16 without torch #24321: from_host(..., dtype=...) for weights kept as their bits (already on main)
  • [PyBindings] Pass float8 tensors through DLPack #25098: pass float8 tensors through DLPack
  • Make save raise on failure; let save move file-backed data into the output instead of copying it #25099: let a write move file-backed data into the output instead of copying it
• Torch-TensorRT 2.15.0: 10 open pull requests for every row, and 3 more for Thor and Spark.
  • feat(conversion): run aten._grouped_mm in TensorRT with a grouped GEMM kernel TensorRT#5021: run aten._grouped_mm in TensorRT with a grouped GEMM kernel
    • perf(conversion): run the grouped GEMM kernel with tcgen05 on Jetson Thor TensorRT#5043: build that kernel with tcgen05 instructions on Jetson AGX Thor
  • fix(conversion): leave large aten._grouped_mm to PyTorch instead of the dense lowering TensorRT#5014: leave large aten._grouped_mm to PyTorch instead of the dense lowering
  • fix(dynamo): release weights before serializing the engine, not after TensorRT#5015: release weights before serializing the engine
  • fix(kernels): decline kernels above TensorRT's 48 KiB shared memory limit TensorRT#5016: decline kernels above TensorRT's 48 KiB shared memory limit
  • fix(kernels): let triton_op kernels opt in to Triton's pointer alignment hint TensorRT#5017: let triton_op kernels use Triton's pointer alignment hint
  • fix(dynamo): keep a select of a uint8 tensor inside the TensorRT engine TensorRT#5019: keep a select of a uint8 tensor inside the engine
  • feat(dynamo): convert stable sort of index tensors with a counting sort TensorRT#5020: convert stable sort of index tensors
  • fix(dynamo): keep a cast to uint8 inside the TensorRT engine TensorRT#5022: keep a cast to uint8 inside the engine
  • fix(conversion): compute antialiased resize weights like PyTorch's CUDA kernel TensorRT#5023: compute antialiased resize weights like PyTorch
  • feat(conversion): convert small antialiased bilinear downscales by default TensorRT#5024: convert small antialiased bilinear downscales
  • fix(executorch): copy only the moving bindings through CUDA graph buffers TensorRT#5044: under CUDA graphs, copy only the inputs whose address moves
  • fix(executorch): move compile_and_save's engines into the program instead of copying them TensorRT#5045: move each engine into the program instead of copying it (with Make save raise on failure; let save move file-backed data into the output instead of copying it #25099)
• LeRobot: the export commands and the Praxis-1 recipe on branch praxis-repro of github.com/shoumikhin/lerobot. H100 and A100 use bbe48afb3c; Thor and Spark use c32f41a561, which adds the edge changes on top. They will go upstream as pull requests.

H100 and A100 (x86_64)

1. Build Torch-TensorRT 2.15 with the fixes

Branch praxis-repro-2.15 at dbb6518265 on github.com/shoumikhin/TensorRT is release/2.15 plus the 10 pull requests, adapted to that branch. Inside pytorch/manylinux2_28-builder:cuda13.2 (x86_64), with bazel 8.4.2 on PATH:

git clone https://github.com/shoumikhin/TensorRT -b praxis-repro-2.15 && cd TensorRT && git checkout dbb6518265
export PATH=/opt/python/cp312-cp312/bin:/usr/local/cuda-13.2/bin:$PATH CUDA_HOME=/usr/local/cuda-13.2
pip install torch==2.15.0 setuptools wheel ninja pyyaml packaging numpy "pybind11==2.6.2" cffi
export TORCH_INSTALL_PATH=$(python -c 'import torch,os;print(os.path.dirname(torch.__file__))')
export CU_VERSION=cu132 CHANNEL=test BUILD_VERSION=2.15.0+praxis \
  TENSORRT_CUDA_VERSION_UPPER_BOUND=13.4 TENSORRT_RTX_CUDA_VERSION_UPPER_BOUND=13.4
envsubst < toolchains/ci_workspaces/MODULE.bazel.tmpl > MODULE.bazel
export CI_BUILD=1 RELEASE=1 TORCH_TENSORRT_EXECUTORCH_RUNTIME_VERSION=0.2.1 TORCH_TENSORRT_EXECUTORCH_RUNTIME_CHANNEL=pypi
python setup.py bdist_wheel --release
pip install dist/torch_tensorrt-*.whl tensorrt-cu13==11.4.0.106 executorch==1.6.0 patchelf
pip wheel --no-build-isolation --no-deps -w dist py/torch-tensorrt-executorch-runtime

2. Install

The export also needs a CUDA toolkit (12.8 or newer) with nvcc on PATH: one small part of the model runs through ExecuTorch's CUDA backend, which compiles its kernels with nvcc.

python3.12 -m venv venv && . venv/bin/activate
pip install torch==2.15.0 executorch==1.6.0 \
  "dist/torch_tensorrt-2.15.0+praxis-cp312-cp312-linux_x86_64.whl[executorch]" \
  dist/torch_tensorrt_executorch_runtime-0.2.1-py3-none-linux_x86_64.whl
pip install "lerobot[praxis,export,dataset] @ git+https://github.com/shoumikhin/lerobot@bbe48afb3c76afbad80b4d746cf7f4f2154afc82"

3. Export (downloads more than 50 GB; point TMPDIR at a disk with room; takes about 20 minutes and about 75 GB of host memory on an H100 machine)

export OMP_NUM_THREADS=8
lerobot-export --policy.path=runwayml/praxis-1-droid-lerobot --backend.type=executorch_tensorrt \
  --output_dir=praxis_bundle --task="pick up the red cube"

4. Run without PyTorch

lerobot-rollout --bundle.path=praxis_bundle --robot.type=<your robot> --fps=15

5. Time it (optional): see the last section. For H100 and A100, use the scripts from the bbe48afb3c checkout.

Jetson AGX Thor and DGX Spark (aarch64)

On Jetson AGX Thor and DGX Spark the CPU and the GPU share one memory. Four export changes make the model fit and
run faster there:

The 4 edge rows:

1. Build

Three public branches, all built from source:

From PyPI: torch 2.15.0 (CUDA 13.2) and TensorRT 11.4.0.106 (tensorrt-cu13). Python 3.12.

The branches carry specific versions of each pull request, pinned by commit. These can differ from the pull requests'
later heads.

The 10 Torch-TensorRT pull requests of the H100 recipe are on the edge branch too. Four of them are earlier versions, adapted to that branch, which differ from the current pull request heads:
#5016 (declines kernels above TensorRT's 48 KiB shared memory limit, without the head's later refactor), #5017
(gives every pointer Triton's alignment hint; the head makes it opt-in), #5020 (the counting sort, an earlier port
with its follow-up fixes) and #5024 (converts every antialiased bilinear resize by default; the head limits it to
small float32 downscales).

Build on x86_64 (to check the recipe on a data center GPU first)

Torch-TensorRT, inside pytorch/manylinux2_28-builder:cuda13.2 with bazel 8.4.2 on PATH:

git clone https://github.com/shoumikhin/TensorRT -b praxis-repro-2.15-edge && cd TensorRT
export PATH=/opt/python/cp312-cp312/bin:/usr/local/cuda-13.2/bin:$PATH CUDA_HOME=/usr/local/cuda-13.2
pip install torch==2.15.0 setuptools wheel ninja pyyaml packaging numpy "pybind11==2.6.2" cffi
export TORCH_INSTALL_PATH=$(python -c 'import torch,os;print(os.path.dirname(torch.__file__))')
export CU_VERSION=cu132 CHANNEL=test BUILD_VERSION=2.15.0+praxis.edge \
  TENSORRT_CUDA_VERSION_UPPER_BOUND=13.4 TENSORRT_RTX_CUDA_VERSION_UPPER_BOUND=13.4
envsubst < toolchains/ci_workspaces/MODULE.bazel.tmpl > MODULE.bazel
export CI_BUILD=1 RELEASE=1 TORCH_TENSORRT_EXECUTORCH_RUNTIME_VERSION=0.2.1 TORCH_TENSORRT_EXECUTORCH_RUNTIME_CHANNEL=pypi
python setup.py bdist_wheel --release
pip install dist/torch_tensorrt-*.whl tensorrt-cu13==11.4.0.106 executorch==1.6.0 patchelf
pip wheel --no-build-isolation --no-deps -w dist py/torch-tensorrt-executorch-runtime

ExecuTorch, as a wheel (the checkout folder must be named executorch):

git clone --recurse-submodules --shallow-submodules https://github.com/shoumikhin/executorch -b praxis-repro-1.6 executorch
python3.12 -m venv etbuild && . etbuild/bin/activate
pip install torch==2.15.0 numpy "cmake>=3.29,<4.0" pyyaml setuptools wheel ninja pybind11 zstd tomli packaging \
  ruamel.yaml patchelf certifi flatbuffers "cuda-toolkit[nvcc,cudart,cccl,nvrtc,crt,nvvm,curand]==13.2.2"
CU=$(python -c 'import sysconfig;print(sysconfig.get_paths()["purelib"])')/nvidia/cu13
ln -sf libcudart.so.13 $CU/lib/libcudart.so
export CUDA_HOME=$CU CUDAToolkit_ROOT=$CU CUDACXX=$CU/bin/nvcc PATH=$CU/bin:$PATH
export CMAKE_ARGS="-DEXECUTORCH_BUILD_CUDA=ON" TORCH_CUDA_ARCH_LIST="8.0;9.0" BUILD_VERSION=1.6.0+praxis
pip wheel ./executorch --no-build-isolation --no-deps -w dist

Build on Jetson AGX Thor and DGX Spark (aarch64)

Thor: JetPack 7.2.1, CUDA 13.2 in /usr/local/cuda, compute capability 11.0. Spark: GB10, CUDA 13.2, compute
capability 12.1. The same aarch64 Torch-TensorRT wheel ran on both boards in these runs (the branch wheel was tested on DGX Spark only).

• Torch-TensorRT: TensorRT 11.4 for aarch64 needs glibc 2.34 or newer, so build inside
  quay.io/pypa/manylinux_2_39_aarch64 instead of the manylinux 2.28 builder. That image has no CUDA: install the
  CUDA 13.2 toolkit for sbsa from NVIDIA's redistributable archives and bazel 8.4.2 (linux-arm64) into it, then run
  the same commands as on x86_64. If the link of torchtrtc stops on missing CUDA, cuDNN or NVPL libraries, add their
  pip folders with build --linkopt=-Wl,-rpath-link,<dirs> in a .bazelrc.user.
• ExecuTorch: build the wheel on the board, as above, with TORCH_CUDA_ARCH_LIST="11.0" on Thor (and
  CUDA_HOME=/usr/local/cuda, no pip toolkit needed) or TORCH_CUDA_ARCH_LIST="12.1" on Spark (with the pip
  toolkit when the board has no CUDA toolkit installed).
• DGX Spark's system Python can be newer than 3.12. Use a separate CPython 3.12 for the environment. If that CPython names
  clang as its compiler (standalone builds do) and the board has only gcc, set CC=gcc CXX=g++.

2. Install

python3.12 -m venv venv && . venv/bin/activate
pip install torch==2.15.0 \
  "dist/torch_tensorrt-2.15.0+praxis.edge-<platform>.whl[executorch]" \
  dist/torch_tensorrt_executorch_runtime-0.2.1-py3-none-<platform>.whl dist/executorch-1.6.0+praxis-*.whl
pip install "lerobot[praxis,export,dataset] @ git+https://github.com/shoumikhin/lerobot@c32f41a56159c7986cc7ee46088fb563bb5f4d47"

The export also needs a CUDA toolkit with nvcc on PATH: one small part of the model runs through ExecuTorch's
CUDA backend, which compiles its kernels with nvcc.

3. Export

The weights download is more than 50 GB. Point TMPDIR at a disk with room. On the boards the exports ran in a
memory-capped user scope (systemd-run --user --scope -p MemoryMax=118G ...) with OMP_NUM_THREADS=8.

export OMP_NUM_THREADS=8 TMPDIR=<a folder on a disk with room>
# bf16 (both boards)
lerobot-export --policy.path=runwayml/praxis-1-droid-lerobot --backend.type=executorch_tensorrt \
  --output_dir=praxis_bundle --task="pick up the red cube" --weights_as_inputs=true
# NVFP4 (Jetson AGX Thor)
lerobot-export --policy.path=runwayml/praxis-1-droid-lerobot --backend.type=executorch_tensorrt \
  --output_dir=praxis_bundle_nvfp4 --task="pick up the red cube" --weights_as_inputs=true \
  --quantize=nvfp4 --tolerance=1
# FP8 (DGX Spark)
lerobot-export --policy.path=runwayml/praxis-1-droid-lerobot --backend.type=executorch_tensorrt \
  --output_dir=praxis_bundle_fp8 --task="pick up the red cube" --weights_as_inputs=true \
  --quantize=fp8 --tolerance=1

What the exports took on the boards:

4. Run without PyTorch

lerobot-rollout --bundle.path=praxis_bundle --robot.type=<your robot> --fps=15

Accuracy method

32 real DROID chunks from a reset, the first 32 chunks of the same sequence, with the same noise for every run
(chunk k draws with seed k). The reference is LeRobot eager in float32 on an H100 (TF32 off). The error is
|difference| divided by the action scale ([0.167, 0.197, 0.201, 1, 1, 1, 0.5]; angles compared the short way round),
reported as the max and the mean over every action of every chunk. The bundles ran without PyTorch on the boards.
Runway's bf16 reference was measured on an H100; its NVFP4 accuracy was measured with its own code on Jetson AGX Thor.

What exactly ran

The two Spark rows were measured on the LeRobot branch tip itself. Both Thor rows were measured earlier, on the same
LeRobot changes before their last review fixes (refusal and guard checks, formatting, docs), and without the H100
recipe's last model change (1f0fdab24b, the action expert's cross attention without a dense mask). The Thor NVFP4
row also predates the branch's last commit (fused NVFP4 rescale and SwiGLU). A rerun on the tip did not fit in the
available Thor time.

For Torch-TensorRT and ExecuTorch, the edge runs did not install wheels of the branches. They used the H100 recipe's
builds with the same pull requests' changes laid over them, at earlier revisions: an earlier revision of #5043, an
earlier revision of #5044's bind-in-place runtime, an earlier form of the disk fix that became #25099, and a
backport of #25098. The branches carry the later heads, which add test, documentation and safety fixes. No speed
change is intended. What each edge row had on top of the H100 recipe's builds:

The Spark runs did not need the disk fix: that board had enough free disk for the extra copy.

The branch tip was also checked on an H100, in two ways. Nothing in these checks was timed.

• With the edge runs' runtime (the H100 recipe's builds with the same pull-request changes laid over them, without the
  disk fix, which only changes how the export moves files), all three exports with --weights_as_inputs=true (bf16, FP8, NVFP4)
  finished all 8 steps and passed their own checks (tolerance 1). Accuracy on the 32 real chunks, max / mean against
  float32: bf16 0.886 / 0.062, FP8 0.914 / 0.060, NVFP4 1.581 / 0.081.
• With wheels built from the three branches: the same three exports finished and passed their own checks (largest difference 0.202, 0.138 and 0.308,
  tolerance 1), and the bf16 bundle ran through lerobot-rollout with --fps=15 --duration=60 on a simulated robot with PyTorch
  blocked from loading, and sent 951 actions in 62.8 s of wall clock. After the last ExecuTorch update (the newest head of [PyBindings] Pass float8 tensors through DLPack #25098), the bf16 and FP8 exports and the
  rollout were run again on the new wheel and passed (largest difference 0.234 and 0.127).

On aarch64, the branch wheels were built and checked on a DGX Spark, but no full Praxis-1 export was run with them:

• Torch-TensorRT and its ExecuTorch runtime wheel were built inside manylinux_2_39_aarch64. ExecuTorch was built on
  the board: CPython 3.12, gcc 15.2, the CUDA 13.2.2 toolkit from pip, TORCH_CUDA_ARCH_LIST="11.0;12.1".
• In a new environment with these wheels and LeRobot from the branch, LeRobot's executorch_tensorrt tests passed
  (15 passed; 2 skipped because they need nvidia-modelopt). One of them exports a small model through TensorRT and
  ExecuTorch's CUDA backend, checks that both run it with no copies to or from the host, and runs it in a process
  without PyTorch, matching eager. LeRobot's torch-free tests and ExecuTorch's float8 tests passed too.
• No branch wheels were built or run on Jetson AGX Thor.

With the branch wheels, the full Praxis-1 exports and the rollout ran on x86_64 only. The Thor and Spark
rows in the table are full exports made on the boards themselves, with the runtime described above.

Time it side by side with Runway's code

The scripts in benchmarks/praxis_compare on the LeRobot branch run the comparison. They replay real DROID frames
through the bundle (PyTorch blocked) and through Runway's
praxis-1-oss at afef2e1 (configs/serve_droid.yaml, through
PraxisPipeline and Session.step), on the same GPU with the same timer: raw frames and the DROID state in, DROID
actions out, with the GPU synchronized at both ends. Its README has the details.

• Frames: real DROID frames from lerobot/droid_100, 766 control ticks, 33 chunks, task "pick up the red cube".
  Make the sequence on an x86_64 machine and copy it to the board: frames decoded on aarch64 differ.
• Passes: each pass is a fresh process for one side, with 3 warm-up and 30 timed chunks and OMP_NUM_THREADS=8.
  The sides alternate pass by pass, 3 rounds. The result is the median of the valid passes' medians.
• Validity rule: a pass is invalid when its median is more than 1.25 times the median of the same side's other
  passes, or when a warm-up chunk after the first takes more than 10 s. The 90th percentile is reported but never
  excludes a pass. On the Thor bf16 set, one Runway pass (4281.58 ms) was invalid, so a 4th round ran.
• Runway's "client resize" side resizes each frame once, outside the timer, as Runway's client does.
• Runway on Jetson AGX Thor runs as it is; its NVFP4 side (runway_nv) packs the bfloat16 checkpoint to NVFP4 when
  it loads. On DGX Spark, Runway's default paired expert GEMM needs more shared memory than GB10 has, so it runs with
  3 pipeline stages. Runway's NVFP4 refuses GB10, so the Spark FP8 row compares with Runway in bf16.

python make_sequence.py seq.npz       # on x86_64, then copy seq.npz to the board
# Jetson AGX Thor
SIDES="bundle runway runway_cr" bash run_ab.sh 0 praxis_bundle seq.npz out_thor_bf16 3
SIDES="bundle runway_nv" bash run_ab.sh 0 praxis_bundle_nvfp4 seq.npz out_thor_nvfp4 3
# DGX Spark
RUNWAY_PAIR_GEMM_STAGES=3 SIDES="bundle runway runway_cr" bash run_ab.sh 0 praxis_bundle seq.npz out_spark_bf16 3
RUNWAY_PAIR_GEMM_STAGES=3 SIDES="bundle runway runway_cr" bash run_ab.sh 0 praxis_bundle_fp8 seq.npz out_spark_fp8 3

Set BUNDLE_PY, RUNWAY_PY and RUNWAY_REPO as the README says. On the boards each side-pass also ran alone, with
the page cache dropped before it and at least 80 GB (Thor) or 100 GB (Spark) of free memory.

Questions and results on other GPUs are welcome here.