Revert "Optimize CUDA export host memory usage" (pytorch#21698)

Reverts pytorch#21617

Author: Gasoonjia

backends/cuda/cuda_backend.py

@@ -6,10 +6,6 @@
 
 
 import contextlib
-import copy
-import ctypes
-import functools
-import gc
 import logging
 import os
 import shutil
@@ -29,7 +25,6 @@
 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
@@ -65,14 +60,6 @@ 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``.
 
@@ -95,29 +82,18 @@ 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 (
-        codecache as _codecache,
-        compile_fx as _cfx,
-        graph as _graph,
-    )
+    from torch._inductor import 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
@@ -176,9 +152,6 @@ 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:
@@ -189,7 +162,6 @@ 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:
@@ -298,38 +270,6 @@ 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."
@@ -444,56 +384,6 @@ 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 {
@@ -569,9 +459,7 @@ 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": cls._weights_format(
-                compile_specs
-            ),
+            "aot_inductor.package_constants_on_disk_format": "binary_blob",
             # 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
@@ -706,7 +594,6 @@ def _combined():
                     stack.enter_context(
                         _compile_time_cpu_clones(torch.device(cls.get_device_name()))
                     )
-                    _trim_host_memory()
                 yield
 
         return _combined()
@@ -719,14 +606,6 @@ 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,8 +4,6 @@
 # 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
 
@@ -132,23 +130,9 @@ 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,
-            [
-                CudaBackend.generate_method_name_compile_spec("forward"),
-                CompileSpec("low_memory_mode", b"ON"),
-            ],
-        )
+        edge_program_manager = self._export_to_cuda_with_lower(module, inputs)
         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,112 +4,19 @@
 # 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.