feat(dynamo): release model weights once their last engine is built

[cehongwang]

Description

Adds an opt-in compilation setting, release_weights_after_build, that frees each model weight once the last TensorRT engine that reads it is built.

TensorRT holds a network's weights only until its engine is built, and a built engine carries its own copy. Yet the model, the ExportedProgram and the partitioned graph modules keep every weight referenced until compilation ends. With several engines (resource partitioning, PyTorch fallbacks), compilation therefore holds all of a model's weights while it builds each engine. Torch-TensorRT already frees constant folding's _frozen_param tensors after each build (conversion/_conversion.py); the model's own parameters are not covered.

With the setting on, compile_module:

• counts, for each weight storage, the partitions that read it (get_attr), including partitions that run in PyTorch;
• before the first build, releases the lowered module's parameters that no partition reads (for example a model head the traced function never calls);
• after each TensorRT partition is built, releases the storages no remaining partition reads.

Storages read by PyTorch partitions or by the partitioned module's own graph are never released. Buffers are never released as unread: mutable buffers are lifted into engine inputs and inlined back into the compiled module. A storage that wasn't seen at partition time (for example, a weight rebound since) is kept.

The storage itself is released (resize_(0)), because dropping references can't free what the caller's model still points to. So the source module, and an ExportedProgram exported from it, must not be used after compilation, even if compilation fails; hence opt-in. Storages that can't be resized, such as views of an external buffer, are skipped. #4854, stacked on this PR, adds memory-mapped weights, which are released by dropping their pages instead.

Also:

• the single-engine fast path (require_full_compilation with every op supported) releases too;
• MutableTorchTensorRTModule rejects the setting, because it refits and recompiles from the source weights;
• with strip_engine_weights, a warning says to refit from a freshly loaded checkpoint;
• the setting has no effect under torch.compile, where parameters arrive as graph inputs; the docstrings say so.

Measurements

pi0.5 (lerobot, bf16, Anthony Shoumikhin's ExecuTorch export), model kept on the CPU through compilation, resource partitioner (9 engines), A40, TensorRT 11.0; host RSS + GPU memory, load phase excluded:

	Host + GPU peak
Without the setting	25.7 GiB
With it	18.9 GiB

4.85 GiB was released engine by engine, and 1.90 GiB of weights the graph never reads (mostly the two lm_heads, 1004 + 502 MiB) before the first engine. Engines and latency are unchanged.

Testing

tests/py/dynamo/models/test_release_weights_after_build.py (13 tests):

• Settings: off by default; old pickles restore to the default.
• CPU, release logic:
  • a weight is released after its last reader is built, and a shared one only after both;
  • one read by a PyTorch partition is never released;
  • a storage rebound after counting is kept;
  • unread parameters are released, but buffers are not;
  • a non-resizable storage is skipped.
• GPU, end to end: two TensorRT partitions split by ops forced to PyTorch, with a shared weight, a read-only buffer, a weight read only by the PyTorch partition and an unused head.
  • Run with both the fast and the global partitioner, the output matches eager on two calls, and exactly the expected storages are released (none with the setting off).
  • If the partitioned module's own graph isn't counted as a reader (mutation), the global-partitioner case uses a freed weight and the process aborts.
• GPU, other paths:
  • the single-engine fast path releases everything;
  • a buffer mutated in place (lifted into an engine input) keeps its storage while the weight is released;
  • MutableTorchTensorRTModule rejects the setting.

tests/py/dynamo/runtime/test_mutable_torchtrt_module.py passes. tests/py/dynamo/partitioning/ also passes, except test_000_fallback_reporting.py::test_fallback_is_reported_1_complex_dtype_refusal, which fails the same way on main in my environment (torch 2.13).

Type of change

• New feature (non-breaking change which adds functionality)

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

[narendasan]

Think this one looks mostly good, non invasive from a user level, do you think in a future version this should become the default?

[cehongwang]

This ruins the original model, so users cannot run the pytorch model anymore. Making it default could lead to some confusion

[cehongwang]

We can make it automatically when weights are mapped from safetensors though. That's safe