[Story] Memory-efficient compilation of pi0.5 (ExecuTorch + TensorRT) for 8 GB devices #4867

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 (2026-10-30)

Measured on a Jetson AGX Thor. Peak is system used memory plus GPU memory above idle; GPU memory on this board comes from a hugepage pool, which is counted. Latency is the median of 3 × 100 chunks through ExecuTorch's Python runtime with the Torch-TensorRT delegate. Stack: torch 2.15.0.dev20260930, ExecuTorch 1.6.0.dev20261020, TensorRT 11.4, Torch-TensorRT Python-only.

• On a Jetson Orin Nano 8 GB (JetPack 7.2, run by @shoumikhin with the reproduction branches below), the export peaks at about 3.3 GB above idle. A chunk takes about 615 ms, and the largest action error is 0.044. It also works on the Torch-TensorRT 2.15 release stack with these PRs.
• Earlier on one A40 (load, compile and save; peak host RSS + GPU memory), the same flow peaked at 6.59 GiB.
• The one-time weights dump maps from a bf16 file. It's now written straight from the float32 checkpoint, without loading it, in 0.91 GiB (feat(dynamo): map_weights - compile from memory-mapped weights, and convert a checkpoint without loading it #4854), so it runs on the 8 GB board too.

To reproduce: pytorch/TensorRT:cehongw/pi05-low-memory-repro holds every Torch-TensorRT PR below combined on main. cehongwang/lerobot:pi05-6gb-repro, examples/export/pi05_low_memory/, holds the export script, the versions and the commands.

PRs

Each State badge links to the PR and updates on its own.

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 streams it from the file into the .pte, as FileBackedData, without reading it into memory.	compile 14.17 → 5.59 GiB
#4794	@narendasan	Build C++ engines from a view of the plan instead of copies
#4795	@narendasan	Don't keep a host copy of the plan once the engine is set up	−8.7 GiB at compile with partitioning
#4796	@narendasan	Hand the built engine's plan to the runtime without copying it
#4821	@cehongwang	ExecuTorch save stores each engine once, as named data	export peak 17.8 → 11.6 GiB (13 engines)

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)
#4854	@cehongwang	Weight mapping tools (on #4853): • map_weights: compile from memory-mapped checkpoint weights, dropping their pages once built; prefix= maps a submodule from a whole-model file. • save_weights_for_mapping(source=checkpoint_source(...)): write the mapping file from a checkpoint without loading it.	host RSS before compile 8.7 → 2.9 GiB; writing the mapping file 9.3 → 0.91 GiB
#4942	@cehongwang	low_memory, the RFC on #4854 (on #4854): • one block builds the model on meta, maps its weights and sets compile defaults, also for code you don't own; • release_weights_after_build defaults to releasing mapped weights only.	the API on top of #4854's measured savings; no new memory measurement
#4793	@narendasan	Offload modules to the CPU without holding both copies	not measured alone

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)
#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
#4820	@cehongwang	Free memory before the resource partitioner reads RSS (on #4783)	−2.8 GiB RSS at partitioning after a CPU load
#4961	@cehongwang	The resource partitioner recomputes large input-only values (attention mask, rotary cos/sin) in each engine; values of 64 KiB or less cross between engines. The graph walks are iterative, so chains deeper than Python's recursion limit work.	Without it, every Gemma attention engine fails to build on sm_110 (TensorRT 11.0–11.4): "uses too much shared data (0x39174 bytes, 0x38c00 max)"

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
#4864	@cehongwang	The exporter shares the compiled module's tensors instead of deep-copying them	−6.74 GiB in the actions program's save
#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)

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.
#4797	@narendasan	Compile memory guidance

Stacks:

• feat(dynamo): release model weights once their last engine is built #4853 → feat(dynamo): map_weights - compile from memory-mapped weights, and convert a checkpoint without loading it #4854 → feat(dynamo): low_memory - compile from a weights file in one block, and release mapped weights by default #4942;
• fix(partitioning): keep getitem with its multi-output producer when splitting by size #4783 → perf(partitioning): free memory before the resource partitioner reads RSS #4820.

#4863, #4864, #4865 and #4961 are each off main, and all are up to date with it (2026-10-30). Each also applies cleanly to release/2.15 on its own. Applied one after another, #4853 and #4863 conflict only where both add a setting at the same spot, and #4863 and #4865 where both add a test; both sides are kept.

Related work by others

PR	Author	What	State
#4774	@shoumikhin	Copy the engine once when building the delegate blob
#4775	@shoumikhin	Share tensors when copying a program for preprocess
ExecuTorch #23339	@shoumikhin	Use the backend's copy hook in single-backend to_backend
ExecuTorch #23340	@shoumikhin	Restore program segments from a view of the input

Together, #4774, #4775 and the two ExecuTorch PRs leave one engine copy in the save: −6.6 GiB.

Still open

• Land the PRs above.
• Torch-TensorRT 2.15: cherry-pick feat(dynamo): release model weights once their last engine is built #4853 and feat(dynamo): map_weights - compile from memory-mapped weights, and convert a checkpoint without loading it #4854, feat(dynamo): engine_output_dir - keep each built engine on disk, not in memory #4863, perf(export): share the compiled module's tensors instead of deep-copying them #4864 and fix(executorch): free each saved program after save(output_format="executorch") #4865 (the Orin Nano runs out of memory without them), and fix(partitioning): recompute large input-only values in each resource-split engine #4961 (the Thor needs it).
• File the partitioner fix as a PR: fix(partitioning): recompute large input-only values in each resource-split engine #4961.
• Bring feat(dynamo): engine_output_dir - keep each built engine on disk, not in memory #4863 up to date with main and perf(executorch): store each engine once, as named data, when saving a .pte #4821: engine files stream into the .pte as FileBackedData.
• ExecuTorch upstream: scope is_same_node's lru_cache(maxsize=None) to one is_identical_graph call, so to_backend stops pinning lowered programs; fix(executorch): free each saved program after save(output_format="executorch") #4865 works around it.
• lerobot: build the policy on the meta device and map its weights (load 11.1 → 2.12 GiB), as the export in cehongwang/lerobot:pi05-6gb-repro does.
• Validate on a Jetson Orin Nano 8 GB: about 3.3 GB above idle, 615 ms per chunk.