[executorch][cuda] Write a weight storage the methods share only once

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom):
* #23787
* #23740
* #23739
* __->__ #23738
* #23737
* #23736

A multi-method CUDA export compiles every method against one `CudaWeightCollector`, so a storage the methods share (all of the model's weights, for a batching export) is stored under one `cuda_fqn_weight:<device>:<FQN>` key. But every method's `materialize` still copied each storage from the GPU to the host in 8 MB chunks, SHA-256'd it on one core, and wrote it to a temp file; only then did the `NamedDataStore` drop the duplicate. For an 18 GB model that is ~90 s per method (measured on 2 GiB: copy 1.9 s, SHA-256 5.9 s, write 1.3 s), all of it thrown away after the first method.

The collector now remembers a fingerprint of each CUDA storage it has written, by storage key. A later method's storage under the same key is fingerprinted on the GPU first; if it matches, its write is skipped and the method's manifest points at the stored entry. Otherwise it is written and added exactly as before, so a changed storage still fails with the store's "Duplicate key ... with different data".

The fingerprint is the storage's size, two independent random linear hashes of its 8-byte words (int64, chunk order kept significant) and its trailing bytes: ~0.8 s for 18 GB on an A100, and a single flipped bit changes it. Library-local constants (`_tensor_constant*`, keyed per shared library), off-graph KV entries and CPU storages keep the existing path.

Full Muse Glimmer 30B batching export (`export_solo_batching.py`, 10 methods; A100; empty Triton and Inductor caches), stacked on D124496888:

| | before | this diff |
|---|---|---|
| total | 135.1 min | 125.7 min |
| main thread in `cuda_weight_collector` | ~21.4 min | ~7.7 min |
| per method after the first (AOTI), min | 13.2 / 13.3 / 9.7 / 12.3 / 14.5 / 19.0 / 10.6 / 2.6 | 13.6 / 12.1 / 8.3 / 11.3 / 13.2 / 17.0 / 8.9 / 1.0 |

(forward_2 .. forward_64, forward_others, sample. forward_1 writes everything once, as before.) Most of what is left in later methods is the collector's `trim_host_memory` (`gc.collect` + `malloc_trim`).

The weight file is byte-identical with and without this diff (`aoti_cuda_blob.ptd`, 17,856,903,936 bytes, same SHA-256), and so is every method's weight manifest; only the methods' compiled libraries differ, as between any two cold exports. The export generates the same tokens (greedy, a 512-token prompt and 4 x 2k prompts) at the same speed (62.5 vs 62.7 tok/s; 110.0 vs 111.3).

Differential Revision: D124713681

Pull Request resolved: #23738

Author: Gasoonjia

backends/cuda/cuda_weight_collector.py

@@ -12,8 +12,8 @@
 import struct
 import tempfile
 import threading
-from dataclasses import dataclass, replace
-from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
+from dataclasses import dataclass, field, replace
+from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Tuple
 
 import torch
 from executorch.backends.cuda.passes.lower_offgraph_kv import OFFGRAPH_KV_FQN_PREFIX
@@ -45,6 +45,9 @@ class CudaWeightEntry:
 class CudaWeightArtifact:
     entries: List[CudaWeightEntry]
     storages: Dict[str, FileBackedData]
+    # Storage keys an earlier method of this collector already stored with the
+    # same bytes; they are not written again.
+    stored: Set[str] = field(default_factory=set)
 
 
 @dataclass(frozen=True)
@@ -126,6 +129,22 @@ def _materialize_storage(tensor: torch.Tensor, directory: str) -> FileBackedData
         raise
 
 
+def _storage_bytes(tensor: torch.Tensor) -> torch.Tensor:
+    storage = tensor.untyped_storage()
+    return torch.empty(0, dtype=torch.uint8, device=tensor.device).set_(
+        storage, 0, (storage.nbytes(),), (1,)
+    )
+
+
+def _same_cuda_storage(first: torch.Tensor, second: torch.Tensor) -> bool:
+    """Whether two CUDA tensors' whole storages hold the same bytes, compared on the device."""
+    if first.device != second.device:
+        return False
+    if first.untyped_storage().nbytes() != second.untyped_storage().nbytes():
+        return False
+    return torch.equal(_storage_bytes(first), _storage_bytes(second))
+
+
 def _required_view_nbytes(
     fqn: str,
     sizes: Tuple[int, ...],
@@ -365,6 +384,11 @@ class CudaWeightCollector:
     def __init__(self) -> None:
         self._store = NamedDataStore()
         self._results: List[PreprocessResult] = []
+        # The CUDA storages already written to the store, by storage key. Methods
+        # share their weights, so a later method's identical storage is compared
+        # with the first one on the device and not written again. The weights
+        # must not change until finish() releases them.
+        self._stored_storages: Dict[str, torch.Tensor] = {}
 
     @contextlib.contextmanager
     def capture(
@@ -397,6 +421,7 @@ def materialize(
         trim_host_memory()
         entries: List[CudaWeightEntry] = []
         storages: Dict[str, FileBackedData] = {}
+        stored: Set[str] = set()
 
         for fqn, (tensor, properties) in weights.items():
             is_offgraph_kv = _is_offgraph_kv_fqn(fqn)
@@ -429,11 +454,22 @@ def materialize(
 
             storage_key = _storage_key(fqn, device_type)
             if not is_offgraph_kv:
-                data = _materialize_storage(tensor, directory)
-                if storage_key in storages:
-                    data.close()
+                if storage_key in storages or storage_key in stored:
                     raise RuntimeError(f"Duplicate CUDA FQN weight key for {fqn!r}")
-                storages[storage_key] = data
+                # Library-local constants are keyed per shared library, so
+                # only FQN weights can match an earlier method's storage.
+                shareable = tensor.is_cuda and not _is_aoti_library_local_fqn(fqn)
+                first = self._stored_storages.get(storage_key)
+                if (
+                    shareable
+                    and first is not None
+                    and _same_cuda_storage(first, tensor)
+                ):
+                    stored.add(storage_key)
+                else:
+                    storages[storage_key] = _materialize_storage(tensor, directory)
+                    if shareable:
+                        self._stored_storages.setdefault(storage_key, tensor)
 
             sizes = tuple(
                 int(size) for size in getattr(properties, "shape", tensor.shape)
@@ -478,7 +514,7 @@ def materialize(
             )
 
         trim_host_memory()
-        return CudaWeightArtifact(entries=entries, storages=storages)
+        return CudaWeightArtifact(entries=entries, storages=storages, stored=stored)
 
     def _merge_aoti_data(
         self,
@@ -564,6 +600,11 @@ def add_preprocess_result(
                 )
             if data is not None:
                 self._add_weight(entry, data, external_tag)
+            elif entry.storage_key in artifact.stored:
+                if entry.storage_key not in self._store.key_to_buffer_idx:
+                    raise RuntimeError(
+                        f"CUDA FQN weight {entry.fqn!r} was reused but never stored"
+                    )
             elif not _is_offgraph_kv_fqn(entry.fqn):
                 raise RuntimeError(f"Missing CUDA FQN weight storage for {entry.fqn!r}")
             serialized_entries.append(entry)
@@ -578,3 +619,4 @@ def finish(self) -> None:
         shared_output = self._store.get_named_data_store_output()
         for result in self._results:
             result.data_store_output = shared_output
+        self._stored_storages.clear()

backends/cuda/tests/targets.bzl

@@ -306,6 +306,26 @@ def define_common_targets(is_fbcode = False):
         ],
     )
 
+    python_unittest_remote_gpu(
+        name = "test_cuda_weight_reuse",
+        srcs = [
+            "test_cuda_weight_reuse.py",
+        ],
+        visibility = [
+            "//executorch/...",
+        ],
+        deps = [
+            "//caffe2:torch",
+            "//executorch/backends/cuda:cuda_backend",
+            "//executorch/exir:lib",
+        ],
+        keep_gpu_sections = True,
+        remote_execution = re_test_utils.remote_execution(
+            platform = "gpu-remote-execution",
+            subplatform = "A100-exclusive",
+        ),
+    )
+
     python_unittest(
         name = "test_cuda_partitioner",
         srcs = [

backends/cuda/tests/test_cuda_weight_reuse.py

@@ -0,0 +1,171 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""CudaWeightCollector writes a weight storage the methods share only once."""
+
+import tempfile
+import unittest
+from unittest import mock
+
+import torch
+from executorch.backends.cuda import cuda_weight_collector
+from executorch.backends.cuda.cuda_weight_collector import (
+    AOTI_DEVICE_TYPE_CUDA,
+    CudaWeightCollector,
+    decode_cuda_aoti_metadata,
+)
+from executorch.exir._serialize._named_data_store import NamedDataStore
+from executorch.exir.backend.backend_details import PreprocessResult
+from torch.export.pt2_archive._package_weights import TensorProperties, Weights
+
+
+def _weights(**tensors: torch.Tensor) -> Weights:
+    return Weights({fqn: (t, TensorProperties(t)) for fqn, t in tensors.items()})
+
+
+def _parent_result(so_key: str) -> PreprocessResult:
+    store = NamedDataStore()
+    store.add_named_data(so_key, so_key.encode())
+    store.add_named_data("empty_weights", b"", external_tag="aoti_cuda_blob")
+    return PreprocessResult(
+        processed_bytes=f"{so_key}\nempty_weights".encode(),
+        data_store_output=store.get_named_data_store_output(),
+    )
+
+
+def _as_bytes(buffer) -> bytes:
+    return buffer if isinstance(buffer, bytes) else buffer.to_bytes()
+
+
+class CudaWeightReuseTest(unittest.TestCase):
+    def setUp(self) -> None:
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA required")
+        self.directory = tempfile.TemporaryDirectory()
+        self.addCleanup(self.directory.cleanup)
+
+    def _method(self, collector, so_key, weights):
+        """One method's preprocess: materialize, then add to the shared store."""
+        artifact = collector.materialize(
+            weights, self.directory.name, lambda _: AOTI_DEVICE_TYPE_CUDA
+        )
+        result = _parent_result(so_key)
+        collector.add_preprocess_result(result, artifact, "cuda")
+        return artifact, result
+
+    def test_a_storage_shared_by_methods_is_written_once(self) -> None:
+        weight = torch.randn(1024, 33, device="cuda")
+        bias = torch.arange(7, dtype=torch.int8, device="cuda")
+        collector = CudaWeightCollector()
+        writes = []
+        materialize_storage = cuda_weight_collector._materialize_storage
+
+        def counting(tensor, directory):
+            writes.append(tensor.untyped_storage().nbytes())
+            return materialize_storage(tensor, directory)
+
+        with mock.patch.object(cuda_weight_collector, "_materialize_storage", counting):
+            first, first_result = self._method(
+                collector, "so-0", _weights(w=weight, b=bias)
+            )
+            # Another method's copies: same bytes, other tensors.
+            second, second_result = self._method(
+                collector, "so-1", _weights(w=weight.clone(), b=bias.clone())
+            )
+        collector.finish()
+        self.assertEqual(len(writes), 2)
+        self.assertEqual(set(first.storages), {e.storage_key for e in first.entries})
+        self.assertEqual(second.storages, {})
+        self.assertEqual(second.stored, {e.storage_key for e in second.entries})
+        for entry in first.entries:
+            self.assertIn(entry.storage_key, collector._store.key_to_buffer_idx)
+        first_manifest = decode_cuda_aoti_metadata(first_result.processed_bytes)
+        second_manifest = decode_cuda_aoti_metadata(second_result.processed_bytes)
+        self.assertEqual(second_manifest.entries, first_manifest.entries)
+
+    def test_the_stored_bytes_match_writing_every_method(self) -> None:
+        tensors = {
+            "w": torch.randn(4096, 129, device="cuda"),
+            "odd": torch.arange(13, dtype=torch.uint8, device="cuda"),
+        }
+
+        def stored_data(reuse: bool):
+            collector = CudaWeightCollector()
+            same_storage = (
+                cuda_weight_collector._same_cuda_storage
+                if reuse
+                # Never equal, so every method writes its storages.
+                else (lambda first, second: False)
+            )
+            with mock.patch.object(
+                cuda_weight_collector, "_same_cuda_storage", same_storage
+            ):
+                for index in range(3):
+                    self._method(
+                        collector,
+                        f"so-{index}",
+                        _weights(**{k: v.clone() for k, v in tensors.items()}),
+                    )
+            collector.finish()
+            store = collector._store
+            return {
+                key: _as_bytes(store.buffers[index])
+                for key, index in store.key_to_buffer_idx.items()
+            }
+
+        self.assertEqual(stored_data(reuse=True), stored_data(reuse=False))
+
+    def test_changed_bytes_are_written_and_rejected_as_before(self) -> None:
+        weight = torch.randn(256, 8, device="cuda")
+        changed = weight.clone()
+        changed[3, 5] += 1
+        collector = CudaWeightCollector()
+        self._method(collector, "so-0", _weights(w=weight))
+        with self.assertRaisesRegex(ValueError, "with different data"):
+            self._method(collector, "so-1", _weights(w=changed))
+
+    def test_library_local_constants_are_written_per_library(self) -> None:
+        constant = torch.arange(5, dtype=torch.float32, device="cuda")
+        collector = CudaWeightCollector()
+        first, first_result = self._method(
+            collector, "so-0", _weights(_tensor_constant0=constant)
+        )
+        second, second_result = self._method(
+            collector, "so-1", _weights(_tensor_constant0=constant.clone())
+        )
+        self.assertEqual(len(first.storages), 1)
+        self.assertEqual(len(second.storages), 1)
+        self.assertEqual(second.stored, set())
+        first_key = decode_cuda_aoti_metadata(first_result.processed_bytes).entries[0]
+        second_key = decode_cuda_aoti_metadata(second_result.processed_bytes).entries[0]
+        self.assertNotEqual(first_key.storage_key, second_key.storage_key)
+
+    def test_storage_differing_in_two_high_bits_is_rejected(self) -> None:
+        weight = torch.arange(3 << 20, dtype=torch.int64, device="cuda")
+        changed = weight.clone()
+        view = changed.view(torch.uint8)
+        view[7] ^= 0x80
+        view[-1] ^= 0x80
+        collector = CudaWeightCollector()
+        self._method(collector, "so-0", _weights(w=weight))
+        with self.assertRaisesRegex(ValueError, "with different data"):
+            self._method(collector, "so-1", _weights(w=changed))
+
+    def test_same_cuda_storage_tells_one_flipped_bit_apart(self) -> None:
+        base = torch.randint(0, 255, (3 << 20,), dtype=torch.uint8, device="cuda")
+        same = cuda_weight_collector._same_cuda_storage
+        self.assertTrue(same(base, base.clone()))
+        for index in (0, 12345, base.numel() - 1):
+            flipped = base.clone()
+            flipped[index] ^= 1
+            self.assertFalse(same(base, flipped), index)
+
+    def test_same_cuda_storage_rejects_other_sizes_and_empty_storages(self) -> None:
+        same = cuda_weight_collector._same_cuda_storage
+        base = torch.zeros(16, dtype=torch.uint8, device="cuda")
+        self.assertFalse(same(base, base[:8].clone()))
+        empty = torch.empty(0, dtype=torch.uint8, device="cuda")
+        self.assertTrue(same(empty, empty.clone()))