[Gasoonjia]

Summary

• Stream AOTInductor external weights to disk instead of materializing a model-sized bytes object in host memory.
• Keep the resulting weights blob file-backed during PTE serialization.
• Avoid deep-copying immutable tensor storage for each delegated method.
• Move compile-time clones of mutated buffers to CPU and temporarily release zero-initialized KV-cache storage.
• Preserve the existing export path unless low_memory_mode is enabled.

On MG 30B, peak host RSS decreased from 75.2 GiB to 48.8 GiB, a 26.3 GiB (35.1%) reduction. A separate 131K-context export completed with approximately 17 GiB peak GPU.

Testing

• Added CUDA low-memory export and serialization unit tests.
• Verified full MG 30B export and runtime successfully.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21798

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[Gasoonjia]

@claude review this PR plz

[claude]

Claude finished @Gasoonjia's task in 8m 19s —— View job

---

PR Review: Fix CUDA low-memory external weights export

• Read the diff (cuda_backend.py, test_cuda_partitioner.py)
• Cross-check base-class contracts (aoti_backend.py) and runtime consumer (cuda_backend.cpp)
• Verify FileBackedData / NamedDataStore integration
• Assess correctness, edge cases, and test coverage

---

Overall this is a well-structured change. Every new behavior is gated behind low_memory_mode="ON" and the default export path is untouched, the streaming writer genuinely avoids the model-sized bytes object, and the FileBackedData handoff into NamedDataStore.add_named_data is type-compatible (_named_data_store.py accepts bytes | FileBackedData | Tensor). Nice work, and the docstrings carry the non-obvious rationale well.

A few things worth confirming or tightening before merge.

1. Blob alignment: the production (all-CUDA) path is packed, but only the padded path is unit-tested

In _write_aoti_weights_blob (backends/cuda/cuda_backend.py:339) the 64-byte inter-tensor padding is written only when not all_cuda:

if not all_cuda and (padding := (-nbytes) % 64):
    output.write(bytes(padding))

So the real MG 30B export (all constants on CUDA) writes tensors back-to-back with no alignment padding, while the CPU/mixed path pads each tensor to a 64-byte boundary. The blob's constant offsets are baked into the AOTI-generated wrapper at compile time (update_constants_from_blob in cuda_backend.cpp:1105), so the streamed layout must match exactly what the wrapper expects — a mismatch produces silently garbage weights, not a crash.

The docstring claims "same layout as AOTInductor's binary_blob", and the PR says the full MG 30B export + runtime were verified, which is the real evidence here. But test_low_memory_weights_are_streamed_in_binary_blob_format uses CPU tensors, so it only exercises the padded branch — the packed all-CUDA branch that production actually takes has no unit coverage. Two suggestions:

• Add a short comment explaining why CUDA constants are packed while CPU/mixed are 64-byte aligned (this asymmetry is surprising — layout is normally a compile-time property independent of where the source tensor currently lives).
• If feasible in CI, add a test asserting the packed (no-padding) layout for the CUDA case, or at least document that it's validated only via the end-to-end MG 30B export.

2. next(...) raises an opaque StopIteration if no .so is present

backends/cuda/cuda_backend.py:491:

so_path = next(
    path for path in paths
    if isinstance(path, str) and path.endswith(".wrapper.so")
)

If Inductor's outputs ever lack a .wrapper.so, this raises a bare StopIteration rather than a descriptive error — inconsistent with the explicit RuntimeError you raise two lines up for the weights count. Consider a next(..., None) + explicit raise RuntimeError(f"...got {paths}"). Fix this →

3. _codecache.TensorProperties is patched process-globally, not gated by the thread-local guard

In _compile_time_cpu_clones (:188), determine_aoti_mmap_flags is wrapped so it only diverges when _is_cpu_clone_active() (thread-local) is true, but TensorProperties is replaced with an unconditional functools.partial. The blast radius is small — _tensor_properties_for_low_memory only diverges (returns None) for emptied tensors and otherwise delegates — so a concurrent export in another thread is unlikely to be affected. Still, for consistency with the other patches, consider gating it on _is_cpu_clone_active() too. Minor.

4. Minor / nits

• chunk_size = 8 * 1024 * 1024 (:317) duplicates FileBackedData._COPY_CHUNK_SIZE; could reference the shared constant.
• all_cuda = all(...) is True for an empty tensors list (:316) → no padding on an empty weights set. Harmless, but the len(weights) != 1 guard already makes empty weights impossible to reach here, so it's moot.
• The shared-storage deepcopy in copy_exported_program_for_preprocess (:521) is safe because move_program_to_device and release_moved_tensors only replace state-dict entries (via .to() / _empty_strided_on_device) and resize the new device tensors — they never resize_(0) the shared source storage in place. That invariant is load-bearing for correctness across multi-method exports; worth a one-line note at the call site so a future change to release_moved_tensors doesn't silently corrupt the shared source tensors.

None of these are blockers. #1 is the one I'd most like confirmed, since it's the highest-consequence correctness point and the only one not covered by the unit tests.
• branch low-host-mem-cuda-export-fix

[digantdesai]

any noticeable increase in export time because of all these writes?

[Gasoonjia]

no the export time remains the same