Add file-backed serialization support

[Gasoonjia]

This PR adds file-backed serialization support to ExecuTorch, allowing large named-data blobs to be hashed, deduplicated, and written in chunks without loading them entirely into memory. It also introduces backend hooks for controlling program copying and AOTI weight materialization while preserving existing behavior by default.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 Unclassified Failures

As of commit a01dc1a with merge base 74f1ce8 ():

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

• pull / unittest / macos / macos-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  backends/xnnpack/test/recipes/test_xnnpack_recipes.py::TestXnnpackRecipes::test_all_models_with_recipes
• 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.

[metascroy]

@claude review this code

[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.

[claude]

Claude finished @metascroy's task in 3m 45s —— View job

---

Code Review

I reviewed the file-backed serialization changes across _cord.py, _named_data_store.py, the AOTI backend, and the backend API. The overall design is sound — chunked hashing/streaming to keep large blobs off the heap, with backward-compatible hooks. Below are the issues I found, roughly ordered by importance.

Correctness / lifecycle

1. FileBackedData cleanup relies solely on __del__ (_cord.py:65-74).
In the normal (non-duplicate) path, nothing ever calls close() — the temp file is only removed when the object is garbage-collected. For a low-memory export of large weights this means temp files can linger past the point they're needed, and __del__ firing at interpreter shutdown can raise/print (only FileNotFoundError is swallowed; a shutdown-time os.remove failure would surface). Relying on __del__ for resource cleanup is generally fragile. Consider making FileBackedData a context manager, or having the serializer explicitly close() buffers once write_to_file completes.

2. merge_named_data_store mutates the caller's other argument (_named_data_store.py:330-340).
_release_duplicate_file does other.buffers[data_entry.buffer_index] = retained and source.close() (which deletes the temp file). This has two consequences worth calling out:

• Merging the same NamedDataStoreOutput into two stores is no longer safe when the first merge dedups a file-backed buffer against the target — the underlying temp file is deleted, so the second merge would hit a deleted path in sha256()/prefix() (raising FileNotFoundError). test_file_backed_data_dedup actually asserts this mutation (assertIs(output2.buffers[0], file1)), so it's intentional, but the re-merge hazard should at minimum be documented on merge_named_data_store.
• After remapping, output2.buffers[0] points at a FileBackedData owned by store1 — ownership is now shared across stores. It's guarded by the _cleanup flag so you won't get a double-delete, but shared ownership via __del__ is easy to get wrong later.

3. Duplicate-key check silently upgraded from byte-equality to length+SHA-256 (_named_data_store.py:143-145).
This is a good change (avoids loading the file), and SHA-256 collision risk is negligible. Just noting the semantic shift from exact data != buffers[idx] comparison in case any caller depended on it. No action needed.

Minor

4. load_weights_blob default reads the whole blob into memory (aoti_backend.py:120-123).
The default defeats the file-backed goal, which is fine since subclasses override to return FileBackedData.move_from(...). But if add_named_data raises (e.g. duplicate key with different data) before merge, any file-backed source leaks until GC — see issue #1.

5. Cord.__bytes__ / FileBackedData.to_bytes() load the full file (_cord.py:57-59, 100-105).
Expected for the bytes(...) compatibility path (only used in tests here), but worth a one-line note that __bytes__ negates the streaming benefit, so callers on the memory-sensitive path should prefer write_to_file.

6. tuple[Any, str] return annotation (aoti_backend.py:117). Fine on py3.9+ (PEP 585); just flagging for consistency since the file otherwise imports from typing.

Looks good

• move_from correctly creates the temp file in the source's directory so os.replace stays on one filesystem, and cleans up the placeholder on failure.
• The backend_api.py refactor (hoisting backend_cls/method_to_compile_specs and routing the deep-copy through copy_exported_program_for_preprocess) is a clean, behavior-preserving change — the default hook still deepcopys. zip over call_submodule_nodes and the matching specs is length-safe since both derive from the same nodes.
• Chunked sha256() and shutil.copyfileobj streaming are the right primitives.

Suggestions

• Add a test covering the re-merge / double-consume scenario in issue Re-sync with internal repository #2 to pin down the intended semantics (or guard against it).
• Consider an explicit close()/context-manager cleanup path (issue Add support for quantized LeakyReLU #1) rather than depending on __del__.

Fix issues #1 and #2 →
• exir/file-backed-serialization

[Gasoonjia]

@claude check if current version solves your comment

[claude]

Claude finished @Gasoonjia's task in 1m 31s —— View job

---

Re-review: yes, the current version resolves my earlier comments

I re-read _cord.py, _named_data_store.py, aoti_backend.py, and the tests. Here's where things stand against my original review.

• Re-read changed files
• Verify issue Add support for quantized LeakyReLU #1 (__del__-only cleanup)
• Verify issue Re-sync with internal repository #2 (merge_named_data_store mutation / re-merge hazard)
• Report findings

✅ Issue #1 — cleanup no longer relies on __del__

FileBackedData now uses weakref.finalize for the GC-backstop and exposes an explicit close() plus __enter__/__exit__ context-manager support (_cord.py:25-27, 75-83). Resource lifetime is now caller-controlled — test_file_backed_data_dedup wraps buffers in a with block and asserts the temp dir is empty on exit (test_named_data_store.py:414-436). The shutdown-raise concern is also gone: _remove swallows OSError broadly rather than just FileNotFoundError (_cord.py:30-34). Resolved.

✅ Issue #2 — merge_named_data_store no longer mutates other

The _release_duplicate_file path is gone. merge_named_data_store now just replays add_named_data over other's entries (_named_data_store.py:296-324), so the source NamedDataStoreOutput and its backing temp files are left intact. The dedup happens naturally through the two-level fingerprint + SHA-256 check, and the duplicate buffer simply isn't appended to the target.

The re-merge hazard I flagged is now explicitly pinned by a test: after store1.merge_named_data_store(output2), output2.buffers[0] is still file2 (unchanged), both temp files still exist, and output2 can be merged again into a fresh store3 (test_named_data_store.py:425-434). This is exactly the double-consume scenario I asked for a test on. Resolved.

Remaining minor notes (non-blocking, unchanged from before)

• Rename _pt2e to pt2e #3 length + SHA-256 dedup for the duplicate-key check (_named_data_store.py:143-145) — good change, just a semantic shift from exact byte compare. No action.
• Add unlifting pass under private config #4 load_weights_blob default still reads the whole blob (aoti_backend.py:120-123) — intended; subclasses override for the file-backed path.
• [build] Add BUCK.oss #5 / Re-sync with internal repository #6 __bytes__/to_bytes() load the full file, and tuple[Any, str] annotation — cosmetic.

Overall: the two substantive concerns (explicit lifecycle + non-mutating, re-mergeable merge) are addressed cleanly and covered by tests. LGTM.
• exir/file-backed-serialization