Avoid rereading CUDA weights blob for SHA-256 computation (pytorch#21800)

## Summary
- Compute SHA-256 while streaming the CUDA AOTI external weights blob to
disk, avoiding a second full read of the approximately 19 GB weights
file during serialization.
- Preserve lazy SHA-256 computation for callers that do not provide a
digest.

## Performance
In an independent MG 30B cold export with an empty Inductor cache,
export time improved from **21:16.45 to 18:56.75**, a **10.9%
reduction**.

Author: Gasoonjia

backends/cuda/cuda_backend.py

@@ -10,6 +10,7 @@
 import ctypes
 import functools
 import gc
+import hashlib
 import logging
 import os
 import shutil
@@ -401,12 +402,17 @@ def _on_off_compile_spec_value(spec: CompileSpec) -> bool:
     return value == "ON"
 
 
-def _write_aoti_weights_blob(weights, blob_path: str) -> None:
-    """Stream AOTI tensor storages without creating a model-sized bytes object."""
+def _write_aoti_weights_blob(weights, blob_path: str) -> bytes:
+    """Stream AOTI tensor storages to disk and return their SHA-256 digest."""
     _trim_host_memory()
     tensors = [tensor for tensor, _ in weights.values()]
     all_cuda = all(tensor.is_cuda for tensor in tensors)
     chunk_size = 8 * 1024 * 1024
+    digest = hashlib.sha256()
+
+    def write_chunk(output, chunk) -> None:
+        digest.update(chunk)
+        output.write(chunk)
 
     with open(blob_path, "wb") as output:
         for tensor in tensors:
@@ -420,20 +426,21 @@ def _write_aoti_weights_blob(weights, blob_path: str) -> None:
                 ).set_(storage, 0, (nbytes,), (1,))
                 for offset in range(0, nbytes, chunk_size):
                     cpu_chunk = byte_tensor[offset : offset + chunk_size].cpu()
-                    output.write(memoryview(cpu_chunk.numpy()))
+                    write_chunk(output, memoryview(cpu_chunk.numpy()))
                 del byte_tensor, cpu_chunk
             elif nbytes:
                 raw_array = (ctypes.c_ubyte * nbytes).from_address(storage.data_ptr())
                 raw_view = memoryview(raw_array).cast("B")
                 for offset in range(0, nbytes, chunk_size):
-                    output.write(raw_view[offset : offset + chunk_size])
+                    write_chunk(output, raw_view[offset : offset + chunk_size])
                 del raw_view, raw_array
             # Match AOTInductor's binary_blob layout: CUDA-only constants are
             # packed, while CPU/mixed constants are aligned to 64 bytes.
             if not all_cuda and (padding := (-nbytes) % 64):
-                output.write(bytes(padding))
+                write_chunk(output, bytes(padding))
             del storage
     _trim_host_memory()
+    return digest.digest()
 
 
 @final
@@ -447,6 +454,11 @@ class CudaBackend(AotiBackend, BackendDetails):
     using the Executorch runtime.
     """
 
+    # AOTI calls materialize_weights_blob immediately before load_weights_blob
+    # for a given output path. A new materialization overwrites any digest left
+    # behind by an export that aborted before the consumer ran.
+    _materialized_blob_hashes: Dict[str, bytes] = {}
+
     @classmethod
     def get_device_name(cls) -> str:
         return "cuda"
@@ -562,8 +574,10 @@ def load_weights_blob(
         """
         if not cls._is_low_memory_mode(compile_specs):
             return super().load_weights_blob(blob_path, compile_specs)
-        blob_data = FileBackedData.move_from(blob_path)
-        return blob_data, blob_data.sha256().hex()
+        known_hash = cls._materialized_blob_hashes.pop(blob_path, None)
+        blob_data = FileBackedData.move_from(blob_path, sha256=known_hash)
+        weights_blob_hash = known_hash or blob_data.sha256()
+        return blob_data, weights_blob_hash.hex()
 
     @classmethod
     def materialize_weights_blob(
@@ -593,7 +607,9 @@ def materialize_weights_blob(
         if so_path is None:
             raise RuntimeError(f"Expected a CUDA AOTI .wrapper.so output, got {paths}")
         blob_path = os.path.splitext(so_path)[0] + "_weights.blob"
-        _write_aoti_weights_blob(weights[0], blob_path)
+        cls._materialized_blob_hashes[blob_path] = _write_aoti_weights_blob(
+            weights[0], blob_path
+        )
 
         # Forcing the external-weights ABI makes Inductor emit an empty blob
         # path alongside the Weights object. Replace that file in place and do

backends/cuda/tests/test_cuda_partitioner.py

@@ -79,6 +79,19 @@ def test_low_memory_weights_are_streamed_in_binary_blob_format(self) -> None:
             )
             self.assertEqual(expected, data)
 
+            # The streaming write computes the digest in the same pass. Loading
+            # the file-backed blob must not reread it solely for hashing.
+            with patch.object(
+                FileBackedData,
+                "sha256",
+                side_effect=AssertionError("unexpected blob reread"),
+            ):
+                blob, digest = CudaBackend.load_weights_blob(
+                    blob_path, [CompileSpec("low_memory_mode", b"ON")]
+                )
+            self.assertEqual(hashlib.sha256(expected).hexdigest(), digest)
+            self.assertEqual(expected, blob.to_bytes())
+
     def test_low_memory_weights_require_wrapper_so(self) -> None:
         tensor = torch.tensor([1], dtype=torch.int16)
         weights = Weights({"weight": (tensor, TensorProperties(tensor))})

exir/_serialize/_cord.py

@@ -18,10 +18,15 @@ class FileBackedData:
 
     _COPY_CHUNK_SIZE = 8 * 1024 * 1024
 
-    def __init__(self, path: str, cleanup: bool = False) -> None:
+    def __init__(
+        self,
+        path: str,
+        cleanup: bool = False,
+        sha256: Optional[bytes] = None,
+    ) -> None:
         self._path = path
         self._size = os.path.getsize(path)
-        self._sha256: Optional[bytes] = None
+        self._sha256 = sha256
         self._finalizer = (
             weakref.finalize(self, self._remove, path) if cleanup else None
         )
@@ -34,7 +39,7 @@ def _remove(path: str) -> None:
             pass
 
     @classmethod
-    def move_from(cls, path: str) -> "FileBackedData":
+    def move_from(cls, path: str, sha256: Optional[bytes] = None) -> "FileBackedData":
         """Take ownership of ``path`` without loading its contents."""
         directory = os.path.dirname(path) or "."
         fd, owned_path = tempfile.mkstemp(
@@ -46,7 +51,7 @@ def move_from(cls, path: str) -> "FileBackedData":
         except Exception:
             os.remove(owned_path)
             raise
-        return cls(owned_path, cleanup=True)
+        return cls(owned_path, cleanup=True, sha256=sha256)
 
     def __len__(self) -> int:
         return self._size