[cehongwang]

Description

Adds a weight-mapping tool, torch_tensorrt.dynamo.tools (py/torch_tensorrt/dynamo/tools/weight_mapping.py), so a model can be compiled from a safetensors checkpoint's weights in place, without loading them into memory first.

The problem. To compile, every weight is normally loaded into host memory before the first engine is built. CPU compilation already avoids a GPU copy, since TensorRT takes weights as host pointers, but all the weights still sit in host memory for the whole compile.

The fix. map_weights(module, path) maps the file with mmap and points each parameter and buffer at a zero-copy view of its bytes. Nothing is read until something touches a weight:

• torch.export reads only shapes, so after mapping and exporting, none of the file is resident.
• TensorRT reads a weight while it builds the engine that uses it, so only that partition's pages become resident.
• With release_weights_after_build=True (feat(dynamo): release model weights once their last engine is built #4853; this PR is stacked on it), those pages are dropped once the weight's last engine is built. The tensor stays valid and is re-read from the file if used again.

API:

Function	What it does
map_weights(module, path, *, strict=False, prefix="") -> MappedWeights	Points the module's tensors at views of the file. With prefix, each name is looked up as prefix + name, as in load_state_dict, so a submodule maps straight from a checkpoint of the whole model. A module built on meta never allocates its weights, and tied weights stay tied. A real tensor keeps its object and gets .data = view. Each tensor must have the model's dtype and shape, or it raises. Names missing from the file go in MappedWeights.missing; strict=True raises for any missing state_dict name, but not for non-persistent buffers. On any error, the module is unchanged.
release_mapped_storage(storage) -> bool	Drops a mapped weight's pages with madvise(MADV_DONTNEED). Only whole pages inside the storage are dropped, so a neighbor's writes survive.
save_weights_for_mapping(module, path, *, source=None)	Writes a module's weights in its dtypes, with a tied weight once, for checkpoints whose loader casts (e.g. stored in fp32, run in bf16). It's a plain safetensors writer, with no new dependency. It writes one tensor at a time, to a temporary name renamed when complete. With source, it converts a checkpoint without loading it: each tensor of a model built on meta comes from source(name, tensor), whole or in chunks along dimension 0, and is cast to the model's dtype as it is written.
checkpoint_source(path, key=lambda name: name, *, chunk_bytes=64 << 20)	A source that reads a safetensors checkpoint under key(name) in chunks of at most chunk_bytes, so no whole checkpoint tensor is in memory.

Change to #4853's code: 5 lines in _weight_release.py. If release_mapped_storage recognizes the storage, its pages are dropped instead of calling resize_(0), which fails on a view of a file.

Docs:

• the API reference entries (Tools section): map_weights, save_weights_for_mapping, checkpoint_source, release_mapped_storage and MappedWeights;
• a "Compiling from Memory-Mapped Weights" section in the resource management guide, including converting a checkpoint without loading it;
• a runnable example, examples/dynamo/compile_from_mapped_weights.py.

Measured

pi0.5 (lerobot, bf16): 6.74 GiB of weights mapped, nothing missing.

• Before compilation: host RSS 8.7 → 2.9 GiB.

• The whole rolled-up ExecuTorch export (load, compile, save, with feat(dynamo): release model weights once their last engine is built #4853, 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): 6.60 GiB.

• Without releasing the mapped pages (no feat(dynamo): release model weights once their last engine is built #4853), the save peaks at 11.21 GiB.

• Writing the mapping file from the 14.5 GB float32 checkpoint (source=checkpoint_source(...) on a meta-built policy, cast to bf16): 0.91 GiB of process memory and 25 s, against 9.3 GiB for loading the policy on the CPU and writing it. The output is the same 820 tensors, and is byte-identical across runs. This makes the one-time conversion possible on the 8 GB Jetson itself.

• On a Jetson AGX Thor with the whole stack (pytorch/TensorRT:cehongw/pi05-low-memory-repro): the export peaks at 4,983 MiB in 233 s, against 11,752 MiB and 394 s for the chain of one process per part.

The example: 192 MiB mapped. After compiling, 0 MiB of the file is resident (192 MiB, the whole file, without release_weights_after_build), and the output matches the original model within 6.4e-06.

Testing

tests/py/dynamo/models/test_weight_mapping.py (14 tests):

• Mapping:
  • a meta module, in fp32 and bf16: every tensor points into the mapping, values match, tied weights stay tied;
  • a real module keeps its Parameter objects.
• Names and errors:
  • missing names are listed, or raise with strict=True without changing the module;
  • a file from the safetensors package's save_file(state_dict) maps with strict=True without its non-persistent buffer;
  • a dtype or shape mismatch raises without changing the module;
  • malformed headers raise;
  • a submodule maps from the whole model's file with prefix.
• Converting:
  • a bf16 model built on meta, written from a float32 checkpoint stored under other names, in 64-byte chunks: no chunk is a whole tensor, and the file maps back to the checkpoint's values with ties kept;
  • a source of the wrong size raises and leaves no file, and a meta tensor without source raises.
• Release:
  • one buffer's release keeps its neighbor's writes;
  • mappings are freed with their views;
  • feat(dynamo): release model weights once their last engine is built #4853's release_storage drops a mapped weight's pages, and still resizes an ordinary storage.
• GPU: a meta-built model, mapped, exported and compiled on the CPU with release_weights_after_build=True, matches eager; every released weight was a mapped one, and the weights read correctly afterwards.

Mutation checks:

• If map_weights copies instead of mapping, 7 tests fail.
• If the released page range is rounded outward, the neighbor test fails.

#4853's tests pass on this branch (27 together), and mypy is clean.

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]

This is quite cool but I think we need to work on the APIs a bit

[narendasan]

It would be cool if we could simplify this boilerplat

[narendasan]

Maybe like some of the other PRs where we might want to make a weight optimized API, we can fold a lot of these things in to a default recipe

[cehongwang]

RFC (draft): an API for low-memory compilation

Status (2026-10-15): proposal, not implemented. Builds on
#4854 (map_weights),
#4853 (release_weights_after_build),
#4855 (engine cap) and
#4863 (engines on disk). Tracking:
#4867.

Two changes:

1. A context manager, torch_tensorrt.dynamo.low_memory, that builds the model on meta, maps its
   weights, and sets defaults for every dynamo.compile call inside the block.
2. release_weights_after_build releases mapped weights by default.

The existing primitives (map_weights, save_weights_for_mapping, release_mapped_storage, the compile
settings) stay as they are.

Part 1: low_memory

class low_memory:
    def __init__(
        self,
        wm_model_init_fn: Callable[[], torch.nn.Module] | None = None,  # builds the model to weight-map
        weights: str | os.PathLike | None = None,     # safetensors file to map; required with wm_model_init_fn
        strict: bool = True,                          # every weight must be in the file
        engine_cap: int | Callable[[ExportedProgram], int] | None = None,  # bytes per engine (#4855)
        engine_dir: str | os.PathLike | None = None,  # engine_output_dir (#4863)
        release_weights: bool | None = None,          # release_weights_after_build (part 2)
        **settings: Any,                              # any other dynamo.compile setting, as a default
    ) -> None: ...
    def __enter__(self) -> torch.nn.Module | None: ...  # the weight-mapped model with wm_model_init_fn, else None
    def __exit__(self, *exc) -> None: ...
    report: MappedWeights | None                      # mapped_bytes, missing; set on entry when mapping

Your own script:

with torch_tensorrt.dynamo.low_memory(
    lambda: MyModel(config), weights="weights.safetensors", engine_cap=512 << 20, engine_dir="engines/"
) as model:
    exported = torch.export.export(model, (example,))           # example on the CPU
    trt_model = torch_tensorrt.dynamo.compile(exported, inputs=[example.cuda()])

A script you don't own (no wm_model_init_fn; the block only sets compile defaults):

with torch_tensorrt.dynamo.low_memory(engine_cap=512 << 20, engine_dir="engines/", optimization_level=1):
    run_export_script()

Entry

• With wm_model_init_fn: runs it under torch.device("meta"), maps the weights with
  map_weights(module, weights, strict=strict), and returns the mapped module, on the CPU. Without it:
  returns None.
• wm_model_init_fn and weights come together; passing one without the other raises.
• wm_model_init_fn must not move the model to a real device (a move out of meta raises).
• Raises if any parameter or buffer is still on meta after mapping, listing them.
• The mapping result is on the context object: ctx = low_memory(...), with ctx as model: ...,
  ctx.report.
• Compile with torch.export on CPU inputs, then dynamo.compile; torch_tensorrt.compile(module)
  doesn't accept a CPU module today.

Compile defaults

• compile is decorated once in _compiler.py; entry only sets a contextvars.ContextVar:

  _low_memory_settings: ContextVar[dict[str, Any]] = ContextVar("low_memory_settings", default={})
  
  def _with_low_memory_defaults(fn):
      @functools.wraps(fn)
      def wrapper(exported_program, *args, **kwargs):
          settings = dict(_low_memory_settings.get())
          if callable(settings.get("resource_partition_engine_cap")):
              settings["resource_partition_engine_cap"] = settings["resource_partition_engine_cap"](exported_program)
          return fn(exported_program, *args, **{**settings, **kwargs})
      return wrapper
  
  @_with_low_memory_defaults
  def compile(exported_program, inputs=None, *, ...): ...

• An explicit argument to compile wins over the context.

• engine_cap sets resource_partition_engine_cap and enable_resource_partitioning=True; engine_dir
  sets engine_output_dir; release_weights sets release_weights_after_build. **settings passes any
  other compile setting through; unknown names raise on entry.

• engine_cap may be a callable that takes the ExportedProgram being compiled and returns its cap in
  bytes, so one block can give different programs different caps.

• Contexts nest: the innermost wins, and exit restores the previous defaults.

• Covers dynamo.compile and its callers (torch_tensorrt.compile(ir="dynamo"),
  MutableTorchTensorRTModule), including a compile imported before entry. A worker thread starts
  with an empty context and doesn't see the defaults.

Exit

• Drops the pages of every mapped weight (release_mapped_storage), also on an exception.
• Restores the previous compile defaults.
• Never unmaps the file, and leaves engine_dir alone: the model and the compiled module still read
  both. The file must stay in place, unchanged, until they are gone.

Part 2: release mapped weights by default

release_weights_after_build: Optional[bool] = None   # was: bool = False
#   None  -> release mapped weights only (the new default)
#   True  -> also free ordinary weights with resize_(0) (destroys the source model; opt-in, as today)
#   False -> release nothing

• Mapped weights are released at two points: at the start of compile (after lowering, before
  partitioning), and after each weight's last engine is built.
• Releasing a mapped weight is madvise(MADV_DONTNEED): the tensor stays valid and reloads from the file,
  so the source model keeps working.
• A write to a mapped weight is lost when its pages drop. Training or in-place fine-tuning on a mapped
  model should pass False.
• Old pickled settings keep loading; a missing field defaults to None.