diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index bf28077c62c..1f6a5df9ea2 100644 --- a/backends/cuda/cuda_backend.py +++ b/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,11 +95,21 @@ 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 @@ -94,6 +117,16 @@ def _compile_time_cpu_clones(target_device: torch.device): 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, diff --git a/backends/cuda/tests/test_cuda_partitioner.py b/backends/cuda/tests/test_cuda_partitioner.py index 89c1204ea00..38a65da1c6f 100644 --- a/backends/cuda/tests/test_cuda_partitioner.py +++ b/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.