Repository navigation
[Story] Memory-efficient compilation of pi0.5 (ExecuTorch + TensorRT) for 8 GB devices #4867
Description
Goal
Compile pi0.5 (lerobot, bf16) to ExecuTorch + TensorRT within the memory of a Jetson-class 8 GB device, where CPU and GPU share one pool. Anthony Shoumikhin reported about 50 GB for the export on Jetson Thor, then about 29 GB after #4774/#4775 and ExecuTorch #23339/#23340. This story tracks the PRs by @cehongwang and @narendasan that bring it down further.
Where it stands
Measured on one A40, peak host RSS + GPU memory, for Anthony's ExecuTorch export with the denoising loop rolled up (prefix program in 8 engines, one step engine with the timestep conditioning precomputed as a table). The run is load, compile, and save of the three .pte files in one process:
| Peak | |
|---|---|
| His script as-is, one engine (load excluded) | 41.8 GiB |
| One engine, every compile technique below | 26.0 GiB (the TensorRT build is the floor) |
Rolled-up, main + the PRs below, whole export (load included) |
6.59 GiB: load 2.12, compile 5.64, save 6.59 |
- Accuracy: checking the saved
.ptefiles in a separate process gives a max action error of 0.0338 (export tolerance 5). - Latency: with GPU-resident I/O on a matched ExecuTorch 1.6 runtime, a chunk takes 93.2 ms, against 91.9 ms for one whole engine.
- Run-time memory: 5.99 GiB steady.
The accuracy check isn't run in the export process: it would load all 9 engines and peak at 9.6 GiB.
PRs
Engines and plans
| PR | Author | What | Measured on pi0.5 | State |
|---|---|---|---|---|
| #4863 | @cehongwang | engine_output_dir: write each engine to disk at build time and load it lazily. The ExecuTorch save copies it from the file and drops the file's pages afterwards. |
compile 14.17 → 5.59 GiB; file-backed RSS during the save stays at 0.86 GiB instead of growing 4.3–5.7 | Open |
| #4794 → #4795 → #4796 | @narendasan | Build engines from a view of the plan, drop the host plan after setup, and hand the plan over without a copy (stacked) | #4795: −8.7 GiB at compile with partitioning | Open |
| #4821 | @cehongwang | ExecuTorch save stores each engine once, as named data (on #4796) | export peak 17.8 → 11.6 GiB (13 engines) | Open |
Weights
| PR | Author | What | Measured on pi0.5 | State |
|---|---|---|---|---|
| #4853 | @cehongwang | release_weights_after_build: free a weight once its last engine is built, and never-read weights before the first build |
25.7 → 18.9 GiB (9 engines) | Open |
| #4854 | @cehongwang | map_weights: compile from memory-mapped checkpoint weights, dropping their pages once built (on #4853) |
host RSS before compile 8.7 → 2.9 GiB | Open |
| #4793 | @narendasan | Offload modules to the CPU without holding both copies | not measured alone | Merged |
Partitioning
| PR | Author | What | Measured on pi0.5 | State |
|---|---|---|---|---|
| #4783 | @cehongwang | Keep a getitem with its multi-output producer when splitting by size |
makes small caps possible (prefix split at 0.5 GiB) | Open |
| #4855 | @cehongwang | resource_partition_engine_cap: a fixed per-engine size |
same split on every run (27 vs 13 engines without it); prefix build 17.6 GiB unsplit → 5.3 | Draft |
| #4820 | @cehongwang | Free memory before the resource partitioner reads RSS (on #4783) | −2.8 GiB RSS at partitioning after a CPU load | Draft |
Build and save
| PR | Author | What | Measured on pi0.5 | State |
|---|---|---|---|---|
| #4792 | @narendasan | Return the builder heap to the OS after each engine build (malloc_trim) |
−14.9 GiB | Merged |
| #4864 | @cehongwang | The exporter shares the compiled module's tensors instead of deep-copying them | −6.74 GiB in the actions program's save |
Open |
| #4865 | @cehongwang | Free each saved program at the end of save(output_format="executorch"). ExecuTorch's is_same_node cache and FakeTensorMode's dispatch cache kept every lowered program, engines included, alive. |
whole export 7.62 → 6.59 GiB (7.73 → 6.59 on ExecuTorch 1.6) | Open |
Docs
| PR | Author | What | State |
|---|---|---|---|
| #4858 | @cehongwang | CPU compilation and resource partitioning in the resource management guide. CPU compilation alone: 41.8 → 34.9 GiB, same engine and latency. | Draft |
| #4797 | @narendasan | Compile memory guidance | Merged |
Stacks: #4794 → #4795 → #4796 → #4821; #4783 → #4820; #4853 → #4854. #4863, #4864 and #4865 are each off main and merge cleanly with each other.
Related work by others
- fix(executorch): copy the engine once when building the delegate blob #4774, fix(executorch): share tensors when copying a program for preprocess #4775 and ExecuTorch #23339 / #23340 (@shoumikhin, merged): one engine copy in the save, −6.6 GiB.
Still open
- Land the PRs above.
- ExecuTorch upstream: scope
is_same_node'slru_cache(maxsize=None)to oneis_identical_graphcall, soto_backendstops pinning lowered programs; fix(executorch): free each saved program after save(output_format="executorch") #4865 works around it. - The save's remaining peak is the prefix program's 4.3 GiB of delegate blobs held as
bytes. perf(executorch): store each engine once, as named data, when saving a .pte #4821, plus an ExecuTorch serializer that writes file-backed blobs in chunks, would bring it down to about one engine. - lerobot: build the policy on the
metadevice and map its weights (load 11.1 → 2.12 GiB, done in our harness), and recompute non-persistent buffers instead of reading them from a dump. - Validate on a Jetson Orin Nano 8 GB: about 7.4 GiB in total, so the margin is a few hundred MiB headless.
Activity