[Gasoonjia]

Summary

Currently we store weight in method granularity: if two methods do not share identical weight we have to keep two weight copies, even if they may just one tensor different.

This PR stores the weight in tensor granularity: we store the weight pickle format and use FQN as the key for each tensor. With this PR weight can be shared across different methods even if they are partially overlapped.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (2 Unrelated Failures)

As of commit fa7c98a with merge base d2f97d5 ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-qnn-wheel-packages-linux (3.10) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.

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

[Gasoonjia]

@claude review this PR plz

[claude]

Claude finished @Gasoonjia's task in 3m 54s —— View job

---

PR Review — share AOTI weights by FQN across methods

Nice work — this is a substantial and well-reasoned change. The manifest format, overflow-hardened view validation, content-addressed dedup, and the mutable/immutable sharing distinction are all thoughtfully done. Below are findings grouped by severity.

🟠 Correctness / robustness

1. Thread-local capture is fragile if AOTI materializes on a worker thread
preprocess sets _FQN_WEIGHTS_CAPTURE.current and relies on materialize_weights_blob reading it back on the same thread (cuda_backend.py:685-692, 785-789). If super().preprocess (AOTI packaging) ever invokes materialize_weights_blob from a worker/pool thread, capture is None and the export dies with "CUDA structured weights must be materialized inside preprocess". It works today because packaging is synchronous, but this couples correctness to an AOTI implementation detail. Worth a comment documenting the same-thread requirement, or threading the capture through explicitly rather than via TLS.

2. get_key() pointer identity as the mutable-scope key is a subtle assumption
acquire_fqn_weight_storage (cuda_backend.cpp:1069-1099) uses reinterpret_cast<uintptr_t>(key.get()) as the PTD-instance identity for scoping mutable weights. The comment explains the reasoning well, but this leans on get_key() returning a pointer owned by the underlying PTD map and stable for its lifetime. That's an implementation contract of NamedDataMap/MergedDataMap, not an interface guarantee — if an implementation ever returns a pointer into a temporary or a per-call buffer, mutable state would silently stop being shared (or worse, cross-contaminate after address reuse). Consider asserting/documenting this contract at the NamedDataMap interface, or deriving the scope from something with a documented lifetime.

3. Unsupported dtypes fail at runtime rather than export
_materialize_fqn_weights encodes scalar_type_enum(tensor.dtype) for any dtype, but the runtime whitelist is_supported_cuda_fqn_dtype / validate_fqn_weight_view only accepts a fixed set (no float64, complex, float8, etc.). A model with an unsupported constant dtype will export cleanly and only fail with InvalidProgram at load time. Validating against the supported set in _materialize_fqn_weights (or _encode_fqn_weight_manifest) would surface a clear error at export.

🟡 Minor / performance

4. O(entries × keys) scan for mutable weights
The get_num_keys()/get_key() linear scan in acquire_fqn_weight_storage runs once per non-shareable entry. For a model with many mutable buffers this is quadratic in named-data keys. Since the scan only maps storage_key → key pointer, you could build that map once per load_constants_from_fqn_manifest call and reuse it. Fix this →

5. Comment placement between if and else if
cuda_backend.cpp:433-439: the multi-line legacy comment now sits after the closing brace of the if (has_fqn_weights) block and before else if, so it reads as documentation for a branch it's separated from. Moving it directly above the else if (or into the block) would read more clearly.

6. Behavior change scope — confirm intended
_weights_format now unconditionally returns "pickle_weights" and materialize_weights_blob dropped its _is_low_memory_mode guard (cuda_backend.py:766, 1065-1071). So the FQN path now applies to all CUDA exports, not just low-memory mode. This matches the PR intent ("backend-wide"), but it's a broader behavior change than the title implies — please confirm non-low-memory paths were exercised end-to-end (the determine_aoti_mmap_flags patch still only forces the external-weights ABI under _is_cpu_clone_active(), so it'd be good to verify the pickle_weights + non-low-memory combination produces a correctly loadable artifact).

🟢 Things I liked / verified

• Overflow-safe span computation in validate_fqn_weight_view (cuda_backend.cpp:1063-1097) mirrors the Python-side required_nbytes check — good defense in depth.
• Content-addressed storage_key (sha256 + "_cuda_weight_storage") correctly dedups identical immutable storages while keeping identical mutable buffers in distinct groups — nicely covered by test_identical_mutable_storages_remain_distinct_groups.
• CudaWeightStorage's destructor restores the prior CUDA device around cudaFree — correct for multi-GPU.
• weak_ptr cache with size re-validation on hit avoids stale-alias bugs; expired entries fall through to a fresh allocation.
• Error paths in acquire_fqn_weight_storage consistently cudaFree before returning. host_data->Free() is called after the H2D copy.
• Manifest parser is strict (magic check, entry cap, dim cap, trailing-byte rejection) and the new test_cuda_weight_manifest.cpp covers legacy-misdetection, truncation, trailing data, and unsupported dtype.

Suggestions

• Add a runtime/integration test that actually loads a two-method model sharing a partial weight set and asserts the "reused across methods" path (the log line at cuda_backend.cpp:1454), so the core value prop is regression-covered beyond the parser unit tests.
• Consider centralizing the dtype↔size table: is_supported_cuda_fqn_dtype, the item_size switch in validate_fqn_weight_view, and the Python whitelist are three places that must stay in sync.

Overall the design is sound and the safety checks are thorough. The main things I'd want resolved before merge are #1 (TLS coupling) and #3 (late dtype failure); #2 is worth at least a documented contract.
· branch codex/cuda-fqn-weight-sharing

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: Gasoonjia / name: Songhao Jia (27ae1e8, 29939a0, 29c08b5, 38b0c99, 89de71f, 905edad, 915db81, fa7c98a)
• ✅ login: Gasoonjia / name: gasoonjia (02304be, 1d6744e, 6cf890d, fb5f24b)

[Gasoonjia]

/easycla

[digantdesai]

what about this, how do we handle this if the fqns are same but content is indeed different?

[Gasoonjia]

Right now we just raised error: from eager model perspective different methods are different "forward" functions in a same model, which means they shouldn't contain weights sharing same fqn.

[digantdesai]

@claude review this PR.

[digantdesai]

what were the memory footprint gains?

[claude]

Claude finished @digantdesai's task in 3m 48s —— View job

---

PR Review — share AOTI weights by FQN across methods

Focused on @digantdesai's two open questions, plus what changed since my last review.

---

❓ "How do we handle same FQN but different content?" (cuda_backend.cpp:413)

There are two distinct cases, handled in different places:

1. Anonymous library-local constants (_tensor_constant*). These are the dangerous ones — PyTorch restarts the numbering in every independently-compiled AOTI library, so _tensor_constant0 in method A and _tensor_constant0 in method B are unrelated tensors that just happen to share a name. The PR namespaces them by the owning library so they can never collide (cuda_weight_collector.py:110-114, 337-343):

def _is_aoti_library_local_fqn(fqn: str) -> bool:
    return fqn.startswith("_tensor_constant")
...
if _is_aoti_library_local_fqn(entry.fqn):
    entry = replace(entry, storage_key=_storage_key(
        entry.fqn, entry.device_type, aoti_library_key=so_blob_key))

The so_blob_key is a content hash of the .so, so each library gets its own storage key → distinct-content-under-same-name stays distinct. This is exactly what the old parakeet fallback protected against, now handled structurally instead of by disabling sharing.

2. Genuine model-level FQNs (e.g. encoder.weight). The storage key is cuda_fqn_weight:{device}:{fqn}, and the bytes are written through NamedDataStore.add_named_data. That store is content-verified (_named_data_store.py:141-150): same key + identical bytes → deduped to one buffer; same key + different bytes → hard ValueError("Duplicate key … with different data") at export time. So a real FQN collision with divergent content fails loudly at export rather than silently binding the wrong tensor. That's the correct default — a shared model-level FQN that legitimately holds different tensors across methods is a genuine conflict, and content-addressed keying surfaces it instead of guessing.

At runtime the same guarantee holds from a different angle: acquire_storage keys its GPU-allocation cache on storage_key @ logical_scope @ device (cuda_weight_cache.cpp:297-299), and re-validates nbytes/device on a cache hit (308-313) — so a stale or mismatched alias can never be reused.

One thing worth confirming: the _tensor_constant detection is a prefix string match. If a real model ever has a user-named parameter starting with _tensor_constant, it would be incorrectly library-scoped. Very unlikely, but a comment noting this is a PyTorch-internal naming contract (not user-facing) would be worth it.

---

❓ "What were the memory footprint gains?"

I can't run the benchmark from here, but the mechanism gives a clear bound. Previously weights were stored at method granularity — two methods sharing all-but-one tensor still kept two full copies, both on disk (.ptd) and in GPU memory. Now storage is tensor granularity, deduped in two places:

• On disk: NamedDataStore dedups identical storages to a single buffer (content-addressed).
• In GPU memory: the runtime weak_ptr cache (storages_) hands the same CudaWeightStorage to every method that references a given (storage_key, scope, device).

So for a prefill/decode pair sharing S bytes of the T-byte weight set, the saving is roughly (num_methods − 1) × S on both axes; a partial overlap now saves the overlapping subset instead of nothing. It'd strengthen the PR to drop the measured before/after .ptd size and peak GPU memory for the model that motivated this into the description.

---

🔎 Re-review — status of prior findings

• ✅ Container/library leak (cuda_backend.cpp:872-885): good catch keeping the .so loaded — the leaked container still owns code/state in that library, so dlclose during multi-method teardown was a real crash risk. Comment explains it well.
• ✅ Overflow-hardened view validation (cuda_weight_cache.cpp:205-277) and strict manifest parser (magic, entry cap, dim cap, trailing-byte rejection) are retained and solid.
• ✅ Scope identity via get_key() pointer is now hoisted into a key_scopes map built once per load() (cuda_weight_cache.cpp:462-477) — this also resolves my prior perf note (Add unlifting pass under private config #4) about the O(entries×keys) scan. Nice.
• 🟠 Unsupported dtypes still fail at runtime, not export. materialize encodes scalar_type_enum(tensor.dtype) for any dtype (cuda_weight_collector.py:253), but the runtime whitelist is_supported_dtype (cuda_weight_cache.cpp:118-133) rejects float64/complex/float8/etc. A model with such a constant exports cleanly and only fails at load. Validating against the supported set in materialize would surface a clear export-time error. Fix this →
• 🟡 Three-place dtype↔size table. is_supported_dtype (enum list), the item_size switch in validate_view, and the (proposed) Python whitelist must stay in sync — worth a shared source of truth or at least a cross-referencing comment.
• 🟡 TLS same-thread coupling remains. materialize_weights_blob reads CudaWeightCollector.current_capture() from thread-local state set in preprocess (cuda_backend.py:531-540, cuda_weight_collector.py:169-176). Correct as long as AOTI packaging stays synchronous on the calling thread; current_capture() raises a clear error otherwise, so it fails safe. A one-line comment documenting the same-thread contract would help.
• 🟡 Behavior is now backend-wide. _weights_format unconditionally returns "pickle_weights" and materialize_weights_blob dropped its low-memory guard, so the FQN path applies to all CUDA exports. This matches the PR intent, but note _force_external_weights_for_fqn_binding only forces the external-weights ABI under _is_cpu_clone_active() (cuda_backend.py:207-217) — good to confirm the pickle_weights + non-low-memory combo produces a loadable artifact end-to-end.

Overall: the design is sound and the two hardest correctness concerns (anonymous-constant collisions, stale GPU aliasing) are handled structurally. Nothing here is blocking given the approval — the dtype-at-export item is the main one I'd still land.
· branch codex/cuda-fqn-weight-sharing

[Gasoonjia]

what were the memory footprint gains?

no obvious memory footprint change, but ptd size reduced a lot