[cehongwang]

Description

Documents two ways to lower compile-time memory that work today, in docsrc/tutorials/resource_memory/resource_management.rst:

• Compile on the CPU (GPU memory): export the model with CPU example inputs and compile the ExportedProgram with GPU inputs. TensorRT takes weights from host memory and the converters pass a CPU weight's own memory, so the model never has a GPU copy. Unlike offload_module_to_cpu, no host copy is added. Also covered:
  • the .cuda() the compiled module needs when some layers run in PyTorch, and the three reasons a layer can stay in PyTorch;
  • that moving the original model back to the CPU breaks the compiled module;
  • offload_module_to_cpu=True as an alternative to that .cuda();
  • that the engine is the same as a GPU compile's, so latency is unchanged; float32 folded constants can differ in the last bit;
  • that torch_tensorrt.compile(nn.Module) doesn't accept a model on the CPU, so the guide uses torch.export + torch_tensorrt.dynamo.compile;
  • the CUDA-only autocast pitfall, with the actual conversion error, and its fix torch.autocast("cpu", ...).
• Build several smaller engines (CPU memory): enable_resource_partitioning / cpu_memory_budget, and how the per-engine budget is derived. That includes the RSS caveat: on a CPU-compiled model, RSS already holds the weights. Also covered: the default budget, a link to the low_cpu_memory_compilation example, and the memory cost when a split cuts through weights reused across an unrolled loop.
• docsrc/contributors/partitioning.rst: "negligible runtime cost" now notes that exception.

Measured

pi0.5 (lerobot, bf16, about 3B parameters, 10-step unrolled denoising loop, Anthony Shoumikhin's ExecuTorch export), A40, TensorRT 11.0, load phase excluded:

• CPU compilation: host + GPU peak 41.8 → 34.9 GiB with one engine. GPU 13.6 → 5.5 GiB; host 28.2 → 29.4 GiB.
• 16 engines vs 1: engines 10.9 vs 5.3 GB, GPU at inference 11.5 vs 5.5 GB, latency +2% (93.8 vs 91.9 ms, C++ ExecuTorch runner).

The CPU-compilation snippet was run on main, both with a PyTorch fallback (default min_block_size, so the toy model stays in PyTorch) and fully converted (min_block_size=1). Outputs match eager. Without .cuda(), the fallback case fails with a device mismatch, which is why the guide includes it. With offload_module_to_cpu=True instead of .cuda(), the fallback case runs and matches eager.

Checked on main:

• Engines: GPU-compiled and CPU-compiled engines are byte-identical with a shared timing cache. The hash is over the 475 KB engine, and it changes when the weights do. A wider study found the same for ResNet-50, ViT, BERT, GPT-2 and a 4-layer Llama, and pi0.5's latency was unchanged (91.10 vs 91.06 ms).
• Front end: torch_tensorrt.compile(cpu_module, ir="dynamo") fails with FakeTensorDeviceMismatchError, with CPU or CUDA example inputs.
• Removed caveat: an earlier draft said a tensor created at run time (torch.arange) on a layer left to PyTorch fails with a device mismatch. I couldn't reproduce that: six variants pass, with static and dynamic batch, and with .cuda() or offload_module_to_cpu=True. So the caveat is gone.

Both pages parse. The two heading-level errors docutils reports in partitioning.rst are already on main.

Related settings in review, each documented in its own PR: #4853 (release_weights_after_build), #4854 (map_weights), #4855 (resource_partition_engine_cap).

Type of change

• Documentation update

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified