Optimize CUDA export host memory usage

[Gasoonjia]

This PR reduces peak host memory usage during CUDA export by streaming and file-backing serialized weight data instead of materializing large in-memory copies. For Gemma 4 31B Q4_K_M, cold-start peak host memory decreased from 110.7 GB to 54.9 GB (50.35%), with identical output and no increase in GPU VRAM usage.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 12 Pending, 1 Unclassified Failure

As of commit 49c2980 with merge base 74f1ce8 ():

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

• Test CUDA Builds / export-model-cuda-artifact (SocialLocalMobile, Qwen3.5-35B-A3B-HQQ-INT4, quantized-int4-tile-packed) / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

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.

[digantdesai]

nit: do you want to split in two PRs? EXIR and CUDA?

[digantdesai]

@claude review this PR.

[claude]

Claude finished @digantdesai's task in 4m 23s —— View job

---

PR Review: Optimize CUDA export host memory usage

I reviewed the full diff against origin/main. This is a well-structured change — the file-backing abstraction is clean, dedup semantics are preserved via size+SHA-256, and the default (non-low-memory) path is untouched. Below are findings, ordered by importance.

Correctness / lifetime

1. FileBackedData ownership is shared but cleanup is __del__-driven — worth a sanity check. In merge_named_data_store, when a buffer is not a duplicate, self.buffers and other.buffers end up holding the same FileBackedData object (_release_duplicate_file only swaps on the dedup path). Cleanup happens in FileBackedData.__del__ → close() → os.remove. This is safe only because __del__ fires on refcount-zero, so the file survives as long as either store references it. But relying on __del__ for multi-GB temp files is fragile: it may run late, out of order at interpreter shutdown, or not at all if a reference is captured in a traceback/cycle. Consider making Cord/the serializer the explicit owner, or documenting that the temp file lifetime is tied to all referencing stores. (exir/_serialize/_cord.py:47-59, _named_data_store.py:330-341)

2. The all-CUDA streaming path — the one that actually runs on GPU export — is not unit-tested. In _write_aoti_weights_blob, alignment padding is written only when not all_cuda:

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

test_low_memory_weights_are_streamed_in_binary_blob_format uses CPU tensors (all_cuda=False), so it only exercises the padded branch. The real export path (all_cuda=True, streaming each CUDA storage to disk in chunks via .cpu()) is only covered by test_cuda_export.py, which requires a GPU. Two things to confirm:

• That omitting 64-byte padding in the all-CUDA case byte-matches what AOTI's pickle_weights writer / the compiled .so expects when reading constants back by offset. A mismatch here would silently corrupt weights. The PR says output is identical, so presumably it matches — but this is the highest-risk line and deserves a comment explaining why padding is conditional on not all_cuda.
• Whether a small CPU-tensor test of the raw byte-copy path (the elif nbytes: branch using ctypes.from_address) is worth adding for non-GPU CI coverage. Fix this →

3. materialize_weights_blob raises an opaque StopIteration. The next(...) that locates the .wrapper.so will raise a bare StopIteration if no matching path exists, unlike the adjacent explicit RuntimeError for the weights-count check. A next(..., None) + explicit error message would be more debuggable. (backends/cuda/cuda_backend.py:469-474)

Minor

4. Trailing whitespace on backends/cuda/cuda_backend.py:324 (blank line after _write_aoti_weights_blob). lintrunner -a will fix it.

5. del cpu_chunk after the loop (cuda_backend.py:313) assumes the loop ran at least once. It does, since that branch is guarded by nbytes and tensor.is_cuda so range(0, nbytes, ...) is non-empty — correct, but a touch subtle. Not a bug.

6. Cord.__bytes__ now materializes every FileBackedData via to_bytes(). That's the correct fallback, but any caller doing bytes(cord) silently defeats the memory optimization. Worth confirming the hot serialization path uses write_to_file (streaming) rather than __bytes__ — it appears to, given FlatTensorSerializer is tested through the streaming path.

On @digantdesai's split suggestion

I agree splitting is reasonable. The _cord.py / _named_data_store.py / data_serializer.py / backend_api.py / backend_details.py changes form a self-contained, backend-agnostic capability (file-backed buffers + a copy_exported_program_for_preprocess hook) that's reviewable and testable on its own. The CUDA-specific streaming (cuda_backend.py, aoti_backend.py) then layers on top. Splitting would isolate the risk in finding #2 to the CUDA PR and let the generic EXIR plumbing land independently. Not blocking, but it would make the blame/bisect story cleaner.

Things done well

• Two-level dedup (size + 32-byte prefix fingerprint, then SHA-256) is correctly generalized to FileBackedData without loading full contents.
• move_from uses mkstemp + os.replace in the same directory — atomic, avoids a cross-device copy.
• The backend_api.py refactor (hoisting backend_name_to_subclass and method_to_compile_specs, then delegating the deep-copy to copy_exported_program_for_preprocess) is clean and the zip alignment is correct.
• Default paths are genuinely unchanged; _is_low_memory_mode gates every new behavior and the test_default_* tests assert the fallback.
  · branch cuda/gemma4-31b-host-memory