feat(dynamo): map_weights - compile from memory-mapped checkpoint weights

[cehongwang]

Description

Adds torch_tensorrt.dynamo.map_weights, which points a model's weights at zero-copy views of a safetensors checkpoint, and save_weights_for_mapping, which writes a model's weights in a form map_weights can map.

Compiling on the CPU (model and inputs kept on the CPU through torch_tensorrt.dynamo.compile) already avoids a GPU copy of the weights: TensorRT takes weights as host pointers, and the converters pass a CPU tensor's own data_ptr(). But the weights still have to be loaded into host memory first, all of them, before the first engine is built. With map_weights the checkpoint is mapped copy-on-write (mmap.ACCESS_COPY, so the file is never written) and each parameter and buffer becomes a torch.frombuffer view of it. torch.export reads only shapes, so a weight's pages become resident only when TensorRT reads them, while building the engine of the partition that uses the weight.

• map_weights(model_or_exported_program, path, *, key_map=None, strict=False) -> MappedWeights
  • A module's tensors may be on the meta device: they are replaced wherever they are registered, so tied weights stay tied. Real tensors keep their Parameter objects and get .data = view; on any device they end up on the CPU, ready for CPU compilation.
  • An ExportedProgram's state_dict and constants are remapped, tied weights included.
  • key_map (a callable or a dict) maps model names to checkpoint names.
  • strict=True raises unless every tensor in the model's state_dict is mapped without a copy: no name missing, as with load_state_dict(strict=True), and no dtype mismatch. Non-persistent buffers, such as rotary frequencies, are not required.
  • Otherwise a dtype mismatch is cast into new memory and reported in copied, and missing names in missing (still_meta for those on meta). A shape mismatch always raises.
  • Every tensor is resolved before any is changed, so an error leaves the model untouched.
  • The header is validated: length, JSON object, and data_offsets in bounds and matching dtype × shape.
• MappedWeights is a context manager. close() drops only the handle's reference: the weights stay mapped and releasable while they're in use, so mapping in a with block and compiling afterwards still releases pages.
• release_mapped_storage(storage) -> bool drops a mapped storage's resident pages with madvise(MADV_DONTNEED). The tensor stays valid and is re-read from the file if used again. Only pages entirely inside the storage are dropped, so a neighbor's in-place writes survive.
• With release_weights_after_build (feat(dynamo): release model weights once their last engine is built #4853; this PR is stacked on it, merge feat(dynamo): release model weights once their last engine is built #4853 first):
  • feat(dynamo): release model weights once their last engine is built #4853 frees a weight after its last engine is built by calling resize_(0) on its storage. That fails on a mapped weight, which is a view of a file.
  • So this PR adds one call at the top of feat(dynamo): release model weights once their last engine is built #4853's release_storage: if release_mapped_storage recognizes the storage, its pages are dropped and resize_(0) is skipped.
  • That's the whole change to feat(dynamo): release model weights once their last engine is built #4853's code: 5 lines in _weight_release.py.
  • The registry holds weak references, so a mapping lives exactly as long as its views or its handle.
• save_weights_for_mapping(model_or_exported_program, path) writes one entry per tensor object (a tied weight once), in the model's current dtypes. That's for checkpoints cast while loading, e.g. stored in fp32 and run in bf16.
  • It's a plain safetensors writer with no new dependency. Offsets are ordered by element size, so every view is aligned.
  • Files written by the safetensors package map too, and vice versa.
• The four names are added to docsrc/py_api/dynamo.rst.

Compared to safetensors.torch.load_file + load_state_dict(assign=True), which also return file-backed views:

• it needs no dependency;
• it handles ties and ExportedPrograms;
• it knows the mapped address ranges, which is what lets compilation drop each weight's pages after its last engine.

Measurements

pi0.5 (lerobot, bf16, Anthony Shoumikhin's ExecuTorch export), CPU compilation, resource partitioner with a fixed per-engine cap (16 engines), weights released after each engine, A40, TensorRT 11.0; host RSS + GPU memory, load phase excluded. The checkpoint is fp32, so save_weights_for_mapping wrote the bf16 weights once (6.74 GiB); the run maps them with this PR's map_weights(policy, path, strict=True) in place of my earlier prototype:

	Host at partitioning	Host + GPU peak
Prototype mapping	3.21 GiB	7.20 GiB
map_weights (this PR)	2.23 GiB	7.07 GiB

6.74 GiB mapped, nothing missing or copied. For comparison, loading the weights into host memory instead of mapping them peaks at 10.1 GiB (27-engine split, otherwise the same).

Testing

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

• With feat(dynamo): release model weights once their last engine is built #4853:
  • release_storage drops a mapped weight's pages without resizing it, and still resizes an ordinary storage;
  • compiling a mapped model with release_weights_after_build=True releases every weight through release_mapped_storage, the output matches eager, and the weights still read correctly afterwards.
• Mapping: a meta module, a real module, a bf16 checkpoint, a real ExportedProgram and a meta ExportedProgram with tied weights.
  • Every tensor's data_ptr() is inside the mapping, values match, tied weights stay tied, and nothing is reported missing.
  • ExportedProgram.module() reproduces the original output.
• Names and errors:
  • key_map and missing names (still_meta);
  • strict=True raises on missing persistent names and on dtype mismatches, leaving the model unchanged, but accepts a missing non-persistent buffer (a safetensors state_dict() checkpoint); without strict, a dtype mismatch is cast and reported;
  • a shape mismatch raises, and an error partway through leaves the model unchanged;
  • malformed headers raise (too-short, negative and out-of-bounds offsets; an oversized header length).
• Interop: a file written by safetensors.torch.save_file maps.
• Release:
  • releasing one buffer keeps its neighbor's in-place writes, and the released buffer is re-read from the file; the buffers are sized from mmap.PAGESIZE, so this holds on 64 KiB-page kernels too;
  • other storages are not touched, and a closed handle doesn't stop release;
  • unclosed mappings are freed with their views.
• GPU: a meta-built model mapped and compiled from its CPU weights matches eager.

Mutation checks: if map_weights copies instead of mapping, 6 tests fail; with the earlier outward page rounding, the neighbor test fails. mypy (strict, with torch) is clean. #4853's tests pass on this branch.

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