Fix CUDA low-memory external weights export

backends/cuda/cuda_backend.py

@@ -6,6 +6,10 @@
 
 
 import contextlib
+import copy
+import ctypes
+import functools
+import gc
 import logging
 import os
 import shutil
@@ -25,6 +29,7 @@
 from executorch.backends.cuda.triton.replacement_pass import (
     ReplaceEdgeOpWithTritonOpPass,
 )
+from executorch.exir._serialize._cord import FileBackedData
 from executorch.exir._warnings import experimental
 from executorch.exir.backend.backend_details import BackendDetails
 from executorch.exir.backend.compile_spec_schema import CompileSpec
@@ -60,6 +65,14 @@ def _is_cpu_clone_active() -> bool:
     return getattr(_CPU_CLONE_GUARD, "active", False)
 
 
+def _trim_host_memory() -> None:
+    gc.collect()
+    try:
+        ctypes.CDLL(None).malloc_trim(0)
+    except AttributeError:
+        pass
+
+
 def _full_zeros_preserving_strides(x: torch.Tensor, device) -> torch.Tensor:
     """Allocate a zero-filled tensor matching ``x``'s size/stride/dtype on ``device``.
 
@@ -82,18 +95,38 @@ def _is_emptied(x) -> bool:
     )
 
 
+def _tensor_properties_for_low_memory(tensor, original):
+    if _is_emptied(tensor):
+        return None
+    return original(tensor)
+
+
 @contextlib.contextmanager
-def _compile_time_cpu_clones(target_device: torch.device):
+def _compile_time_cpu_clones(target_device: torch.device):  # noqa: C901
     """Force AOTI's mutated-buffer clones onto CPU while preserving the
     serialized constants' target device."""
-    from torch._inductor import compile_fx as _cfx, graph as _graph
+    from torch._inductor import (
+        codecache as _codecache,
+        compile_fx as _cfx,
+        graph as _graph,
+    )
     from torch._inductor.codegen.cpp_wrapper_cpu import CppWrapperCpu as _Cpp
     from torch._inductor.graph import GraphLowering as _GL
 
     orig_clone = _cfx.clone_preserve_strides
     orig_codegen_device = _Cpp.codegen_device
     orig_get_const = _GL.get_original_value_of_constant
     orig_is_same = _graph.is_same_tensor
+    orig_tensor_properties = _codecache.TensorProperties
+    orig_determine_aoti_mmap_flags = _codecache.determine_aoti_mmap_flags
+
+    def _force_external_weights_for_streaming(consts_size):
+        # ``pickle_weights`` normally tells AOTI that no external binary blob
+        # exists. We materialize that pickle output as a streamed blob below,
+        # so the generated wrapper must use the matching external-weights ABI.
+        if _is_cpu_clone_active():
+            return True, False
+        return orig_determine_aoti_mmap_flags(consts_size)
 
     def _is_same_skip_emptied(data, value):
         # KV buffers freed via resize_(0) all have data_ptr 0, so the stock
@@ -152,6 +185,10 @@ def _codegen_device_target_aware(self, device):
     _Cpp.codegen_device = _codegen_device_target_aware
     _GL.get_original_value_of_constant = _get_const_synthesize_zeros
     _graph.is_same_tensor = _is_same_skip_emptied
+    _codecache.TensorProperties = functools.partial(
+        _tensor_properties_for_low_memory, original=orig_tensor_properties
+    )
+    _codecache.determine_aoti_mmap_flags = _force_external_weights_for_streaming
     prev_active = getattr(_CPU_CLONE_GUARD, "active", False)
     _CPU_CLONE_GUARD.active = True
     try:
@@ -162,6 +199,8 @@ def _codegen_device_target_aware(self, device):
         _Cpp.codegen_device = orig_codegen_device
         _GL.get_original_value_of_constant = orig_get_const
         _graph.is_same_tensor = orig_is_same
+        _codecache.TensorProperties = orig_tensor_properties
+        _codecache.determine_aoti_mmap_flags = orig_determine_aoti_mmap_flags
 
 
 def _is_kv_buffer(name, v) -> bool:
@@ -270,6 +309,39 @@ 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."""
+    _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
+
+    with open(blob_path, "wb") as output:
+        for tensor in tensors:
+            if tensor.is_mkldnn:
+                raise RuntimeError("MKLDNN constants are not supported by CUDA AOTI")
+            storage = tensor.untyped_storage()
+            nbytes = storage.nbytes()
+            if nbytes and tensor.is_cuda:
+                byte_tensor = torch.empty(
+                    0, dtype=torch.uint8, device=tensor.device
+                ).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()))
+                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])
+                del raw_view, raw_array
+            if not all_cuda and (padding := (-nbytes) % 64):
+                output.write(bytes(padding))
+            del storage
+    _trim_host_memory()
+
+
 @final
 @experimental(
     "This API and all of cuda backend related functionality are experimental."
@@ -384,6 +456,76 @@ def save_data_externally(cls) -> bool:
         """
         return True
 
+    @classmethod
+    def load_weights_blob(
+        cls, blob_path: str, compile_specs: List[CompileSpec]
+    ) -> tuple[Any, str]:
+        """Keep low-memory CUDA weights file-backed during PTE serialization.
+
+        The streamed file has the same layout as AOTInductor's ``binary_blob``.
+        Keeping it file-backed avoids reading another model-sized copy into
+        host memory without changing its bytes.
+        """
+        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()
+
+    @classmethod
+    def materialize_weights_blob(
+        cls, paths: Any, compile_specs: List[CompileSpec]
+    ) -> Any:
+        if not cls._is_low_memory_mode(compile_specs) or not isinstance(paths, list):
+            return paths
+
+        from torch.export.pt2_archive._package_weights import Weights
+
+        weights = [path for path in paths if isinstance(path, Weights)]
+        if not weights:
+            return paths
+        if len(weights) != 1:
+            raise RuntimeError(
+                f"Expected one CUDA AOTI weights output, got {len(weights)}"
+            )
+
+        so_path = next(
+            path
+            for path in paths
+            if isinstance(path, str) and path.endswith(".wrapper.so")
+        )
+        blob_path = os.path.splitext(so_path)[0] + "_weights.blob"
+        _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
+        # not add a duplicate path to the returned package outputs.
+        materialized = [path for path in paths if not isinstance(path, Weights)]
+        if blob_path not in materialized:
+            materialized.append(blob_path)
+        return materialized
+
+    @classmethod
+    def copy_exported_program_for_preprocess(
+        cls, edge_program, compile_specs: List[CompileSpec]
+    ):
+        """Copy graph structure while sharing immutable tensor storage.
+
+        CUDA preprocessing replaces state-dict entries when moving them to the
+        target device; it does not mutate the source tensors.  Memoizing those
+        tensors therefore avoids a model-sized host copy for every delegated
+        method while preserving an independent graph and state-dict mapping.
+        """
+        if not cls._is_low_memory_mode(compile_specs):
+            return copy.deepcopy(edge_program)
+
+        tensor_memo = {
+            id(tensor): tensor
+            for values in (edge_program.state_dict, edge_program.constants)
+            for tensor in values.values()
+            if isinstance(tensor, torch.Tensor)
+        }
+        return copy.deepcopy(edge_program, tensor_memo)
+
     @classmethod
     def get_supported_fallback_kernels(cls) -> Dict[str, Any]:
         return {
@@ -458,8 +600,13 @@ def get_aoti_compile_options(
             # Separate weight constants from the .so file
             "aot_inductor.package": True,
             "aot_inductor.package_constants_in_so": False,
-            # Store weight constants on disk in a binary blob
-            "aot_inductor.package_constants_on_disk_format": "binary_blob",
+            # Store weight constants on disk in a binary blob. Low-memory mode
+            # asks AOTI for a Weights object and streams the equivalent blob in
+            # materialize_weights_blob; its context also forces the generated
+            # wrapper to use the required external-weights ABI.
+            "aot_inductor.package_constants_on_disk_format": cls._weights_format(
+                compile_specs
+            ),
             # Enable maximum automatic tuning for optimal performance
             "max_autotune": True,
             # Use TRITON for GEMM (General Matrix Multiply) operations tuning only to avoid using operators in libtorch
@@ -594,6 +741,7 @@ def _combined():
                     stack.enter_context(
                         _compile_time_cpu_clones(torch.device(cls.get_device_name()))
                     )
+                    _trim_host_memory()
                 yield
 
         return _combined()
@@ -606,6 +754,14 @@ def _is_low_memory_mode(compile_specs: List[CompileSpec]) -> bool:
                 return spec.value.decode("utf-8").upper() == "ON"
         return False
 
+    @classmethod
+    def _weights_format(cls, compile_specs: List[CompileSpec]) -> str:
+        return (
+            "pickle_weights"
+            if cls._is_low_memory_mode(compile_specs)
+            else "binary_blob"
+        )
+
     @classmethod
     def move_program_to_device(
         cls,

backends/cuda/tests/test_cuda_partitioner.py

@@ -4,19 +4,138 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import hashlib
 import operator
+import os
+import tempfile
 import unittest
 from typing import Tuple
+from unittest.mock import patch
 
 import torch
+from executorch.backends.cuda.cuda_backend import CudaBackend
 from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
+from executorch.exir._serialize._cord import FileBackedData
+from executorch.exir.backend.compile_spec_schema import CompileSpec
 from executorch.exir.backend.partitioner import PartitionResult
 from executorch.exir.delegate import executorch_call_delegate
 from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
 from torch.export import export
+from torch.export.pt2_archive._package_weights import TensorProperties, Weights
 from torch.fx.passes.utils.fuser_utils import validate_partition
 
 
+class TestCudaLowMemoryExport(unittest.TestCase):
+    @patch.object(CudaBackend, "_setup_cuda_environment_for_fatbin", return_value=True)
+    def test_low_memory_streaming_keeps_external_weights_abi(self, _) -> None:
+        from torch._inductor import codecache
+
+        options = CudaBackend.get_aoti_compile_options(
+            [CompileSpec("low_memory_mode", b"ON")]
+        )
+        self.assertEqual(
+            options["aot_inductor.package_constants_on_disk_format"],
+            "pickle_weights",
+        )
+
+        original = codecache.determine_aoti_mmap_flags
+        with CudaBackend.get_extra_aoti_compile_context_manager(
+            [CompileSpec("low_memory_mode", b"ON")]
+        ):
+            self.assertEqual(codecache.determine_aoti_mmap_flags(0), (True, False))
+        self.assertIs(codecache.determine_aoti_mmap_flags, original)
+
+    def test_low_memory_weights_are_streamed_in_binary_blob_format(self) -> None:
+        first = torch.tensor([1, 2, 3], dtype=torch.int16)
+        second = torch.tensor([4, 5], dtype=torch.int32)
+        weights = Weights(
+            {
+                "first": (first, TensorProperties(first)),
+                "second": (second, TensorProperties(second)),
+            }
+        )
+
+        with tempfile.TemporaryDirectory() as directory:
+            so_path = os.path.join(directory, "model.wrapper.so")
+            blob_path = os.path.join(directory, "model.wrapper_weights.blob")
+            # AOTI emits this empty placeholder when the wrapper is compiled
+            # with the external-weights ABI and the tensor values are pickled.
+            with open(blob_path, "wb"):
+                pass
+
+            paths = CudaBackend.materialize_weights_blob(
+                [so_path, blob_path, weights],
+                [CompileSpec("low_memory_mode", b"ON")],
+            )
+
+            self.assertEqual([so_path, blob_path], paths)
+            with open(blob_path, "rb") as blob:
+                data = blob.read()
+            expected = (
+                bytes(first.untyped_storage())
+                + bytes(58)
+                + bytes(second.untyped_storage())
+                + bytes(56)
+            )
+            self.assertEqual(expected, data)
+
+    def test_low_memory_blob_stays_file_backed(self) -> None:
+        data = b"cuda weights"
+        with tempfile.TemporaryDirectory() as directory:
+            path = os.path.join(directory, "weights.blob")
+            with open(path, "wb") as output:
+                output.write(data)
+
+            blob, digest = CudaBackend.load_weights_blob(
+                path, [CompileSpec("low_memory_mode", b"ON")]
+            )
+
+            self.assertIsInstance(blob, FileBackedData)
+            self.assertEqual(hashlib.sha256(data).hexdigest(), digest)
+            self.assertEqual(data, blob.to_bytes())
+            self.assertFalse(os.path.exists(path))
+
+    def test_default_blob_behavior_is_unchanged(self) -> None:
+        data = b"cuda weights"
+        with tempfile.TemporaryDirectory() as directory:
+            path = os.path.join(directory, "weights.blob")
+            with open(path, "wb") as output:
+                output.write(data)
+
+            blob, digest = CudaBackend.load_weights_blob(path, [])
+
+            self.assertIsInstance(blob, bytes)
+            self.assertEqual(data, blob)
+            self.assertEqual(hashlib.sha256(data).hexdigest(), digest)
+            self.assertFalse(os.path.exists(path))
+
+    def test_low_memory_program_copy_shares_tensor_storage(self) -> None:
+        module = torch.nn.Linear(4, 3)
+        program = export(module, (torch.randn(2, 4),), strict=True)
+
+        copied = CudaBackend.copy_exported_program_for_preprocess(
+            program, [CompileSpec("low_memory_mode", b"ON")]
+        )
+
+        self.assertIsNot(program, copied)
+        self.assertIsNot(program.graph_module, copied.graph_module)
+        self.assertIs(program.state_dict["weight"], copied.state_dict["weight"])
+        copied._state_dict["weight"] = torch.nn.Parameter(torch.zeros(3, 4))
+        self.assertFalse(torch.count_nonzero(program.state_dict["weight"]) == 0)
+
+    def test_default_program_copy_has_independent_tensor_storage(self) -> None:
+        module = torch.nn.Linear(4, 3)
+        program = export(module, (torch.randn(2, 4),), strict=True)
+
+        copied = CudaBackend.copy_exported_program_for_preprocess(program, [])
+
+        self.assertIsNot(program.state_dict["weight"], copied.state_dict["weight"])
+        self.assertNotEqual(
+            program.state_dict["weight"].data_ptr(),
+            copied.state_dict["weight"].data_ptr(),
+        )
+
+
 class TestCudaPartitioner(unittest.TestCase):
     """
     Test CUDA partitioner functionality.