Optimize CUDA export host memory usage (#21617)

This PR reduces peak host memory usage during CUDA export by streaming
and file-backing serialized weight data instead of materializing large
in-memory copies. For Gemma 4 31B Q4_K_M, cold-start peak host memory
decreased from 110.7 GB to 54.9 GB (50.35%), with identical output and
no increase in GPU VRAM usage.

Author: Gasoonjia

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,29 @@ 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):
     """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
 
     def _is_same_skip_emptied(data, value):
         # KV buffers freed via resize_(0) all have data_ptr 0, so the stock
@@ -152,6 +176,9 @@ 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
+    )
     prev_active = getattr(_CPU_CLONE_GUARD, "active", False)
     _CPU_CLONE_GUARD.active = True
     try:
@@ -162,6 +189,7 @@ 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
 
 
 def _is_kv_buffer(name, v) -> bool:
@@ -270,6 +298,38 @@ def _on_off_compile_spec_value(spec: CompileSpec) -> bool:
     return value == "ON"
 
 
+def _write_aoti_weights_blob(weights, blob_path: str) -> None:
+    _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 +444,56 @@ def save_data_externally(cls) -> bool:
         """
         return True
 
+    @classmethod
+    def load_weights_blob(
+        cls, blob_path: str, compile_specs: List[CompileSpec]
+    ) -> tuple[Any, str]:
+        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)
+        return [path for path in paths if not isinstance(path, Weights)] + [blob_path]
+
+    @classmethod
+    def copy_exported_program_for_preprocess(
+        cls, edge_program, compile_specs: List[CompileSpec]
+    ):
+        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 {
@@ -459,7 +569,9 @@ def get_aoti_compile_options(
             "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",
+            "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 +706,7 @@ def _combined():
                     stack.enter_context(
                         _compile_time_cpu_clones(torch.device(cls.get_device_name()))
                     )
+                    _trim_host_memory()
                 yield
 
         return _combined()
@@ -606,6 +719,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_export.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import os
+import tempfile
 import unittest
 from typing import Tuple
 
@@ -130,9 +132,23 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         inputs = (torch.randn(3, 4), torch.randn(3, 4))
 
         # Test export
-        edge_program_manager = self._export_to_cuda_with_lower(module, inputs)
+        edge_program_manager = self._export_to_cuda_with_lower(
+            module,
+            inputs,
+            [
+                CudaBackend.generate_method_name_compile_spec("forward"),
+                CompileSpec("low_memory_mode", b"ON"),
+            ],
+        )
         self.assertIsNotNone(edge_program_manager, "Simple add operation export failed")
 
+        et_program = edge_program_manager.to_executorch()
+        with tempfile.TemporaryDirectory() as output_dir:
+            et_program.write_tensor_data_to_file(output_dir)
+            ptd_path = os.path.join(output_dir, "aoti_cuda_blob.ptd")
+            self.assertTrue(os.path.isfile(ptd_path))
+            self.assertGreater(os.path.getsize(ptd_path), 0)
+
     def test_conv2d(self):
         """Test CUDA export for 2D convolution."""
 

backends/cuda/tests/test_cuda_partitioner.py

@@ -4,19 +4,112 @@
 # 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
 
 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 TestCudaWeightsBlob(unittest.TestCase):
+    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 f:
+                f.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 f:
+                f.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_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")
+            paths = CudaBackend.materialize_weights_blob(
+                [so_path, weights], [CompileSpec("low_memory_mode", b"ON")]
+            )
+            blob_path = os.path.join(directory, "model.wrapper_weights.blob")
+
+            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_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.