Skip to content

Commit 7d84f08

Browse files
authored
Fix CUDA low-memory external weights export (pytorch#21798)
## Summary - Stream AOTInductor external weights to disk instead of materializing a model-sized bytes object in host memory. - Keep the resulting weights blob file-backed during PTE serialization. - Avoid deep-copying immutable tensor storage for each delegated method. - Move compile-time clones of mutated buffers to CPU and temporarily release zero-initialized KV-cache storage. - Preserve the existing export path unless `low_memory_mode` is enabled. On MG 30B, peak host RSS decreased from **75.2 GiB to 48.8 GiB, a 26.3 GiB (35.1%) reduction**. A separate 131K-context export completed with approximately 17 GiB peak GPU. ## Testing - Added CUDA low-memory export and serialization unit tests. - Verified full MG 30B export and runtime successfully.
1 parent 46017a1 commit 7d84f08

2 files changed

Lines changed: 298 additions & 4 deletions

File tree

‎backends/cuda/cuda_backend.py‎

Lines changed: 168 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,10 @@
66

77

88
import contextlib
9+
import copy
10+
import ctypes
11+
import functools
12+
import gc
913
import logging
1014
import os
1115
import shutil
@@ -25,6 +29,7 @@
2529
from executorch.backends.cuda.triton.replacement_pass import (
2630
ReplaceEdgeOpWithTritonOpPass,
2731
)
32+
from executorch.exir._serialize._cord import FileBackedData
2833
from executorch.exir._warnings import experimental
2934
from executorch.exir.backend.backend_details import BackendDetails
3035
from executorch.exir.backend.compile_spec_schema import CompileSpec
@@ -60,6 +65,14 @@ def _is_cpu_clone_active() -> bool:
6065
return getattr(_CPU_CLONE_GUARD, "active", False)
6166

6267

68+
def _trim_host_memory() -> None:
69+
gc.collect()
70+
try:
71+
ctypes.CDLL(None).malloc_trim(0)
72+
except AttributeError:
73+
pass
74+
75+
6376
@contextlib.contextmanager
6477
def _keep_triton_reduction_loads_loop_scoped():
6578
"""Conservatively treating loads as reduction-masked keeps each definition in its
@@ -104,6 +117,12 @@ def _is_emptied(x) -> bool:
104117
)
105118

106119

120+
def _tensor_properties_for_low_memory(tensor, original):
121+
if _is_cpu_clone_active() and _is_emptied(tensor):
122+
return None
123+
return original(tensor)
124+
125+
107126
def _required_storage_nbytes(x: torch.Tensor) -> int:
108127
"""Return the backing storage required by ``x``'s logical view."""
109128
if x.numel() == 0:
@@ -156,10 +175,14 @@ def _rehydrate_emptied_tensors(tensors):
156175

157176

158177
@contextlib.contextmanager
159-
def _compile_time_cpu_clones(target_device: torch.device):
178+
def _compile_time_cpu_clones(target_device: torch.device): # noqa: C901
160179
"""Force AOTI's mutated-buffer clones onto CPU while preserving the
161180
serialized constants' target device."""
162-
from torch._inductor import compile_fx as _cfx, graph as _graph
181+
from torch._inductor import (
182+
codecache as _codecache,
183+
compile_fx as _cfx,
184+
graph as _graph,
185+
)
163186
from torch._inductor.codegen.cpp_wrapper_cpu import CppWrapperCpu as _Cpp
164187
from torch._inductor.graph import GraphLowering as _GL
165188
from torch._inductor.runtime.triton_heuristics import CachingAutotuner as _Autotuner
@@ -168,6 +191,17 @@ def _compile_time_cpu_clones(target_device: torch.device):
168191
orig_codegen_device = _Cpp.codegen_device
169192
orig_get_const = _GL.get_original_value_of_constant
170193
orig_is_same = _graph.is_same_tensor
194+
orig_tensor_properties = _codecache.TensorProperties
195+
orig_determine_aoti_mmap_flags = _codecache.determine_aoti_mmap_flags
196+
197+
def _force_external_weights_for_streaming(consts_size):
198+
# ``pickle_weights`` normally tells AOTI that no external binary blob
199+
# exists. We materialize that pickle output as a streamed blob below,
200+
# so the generated wrapper must use the matching external-weights ABI.
201+
if _is_cpu_clone_active():
202+
return True, False
203+
return orig_determine_aoti_mmap_flags(consts_size)
204+
171205
orig_autotuner_run = _Autotuner.run
172206

173207
def _is_same_skip_emptied(data, value):
@@ -241,6 +275,10 @@ def _codegen_device_target_aware(self, device):
241275
_Cpp.codegen_device = _codegen_device_target_aware
242276
_GL.get_original_value_of_constant = _get_const_synthesize_zeros
243277
_graph.is_same_tensor = _is_same_skip_emptied
278+
_codecache.TensorProperties = functools.partial(
279+
_tensor_properties_for_low_memory, original=orig_tensor_properties
280+
)
281+
_codecache.determine_aoti_mmap_flags = _force_external_weights_for_streaming
244282
_Autotuner.run = _autotuner_run_with_rehydrated_emptied_args
245283
prev_active = getattr(_CPU_CLONE_GUARD, "active", False)
246284
_CPU_CLONE_GUARD.active = True
@@ -252,6 +290,8 @@ def _codegen_device_target_aware(self, device):
252290
_Cpp.codegen_device = orig_codegen_device
253291
_GL.get_original_value_of_constant = orig_get_const
254292
_graph.is_same_tensor = orig_is_same
293+
_codecache.TensorProperties = orig_tensor_properties
294+
_codecache.determine_aoti_mmap_flags = orig_determine_aoti_mmap_flags
255295
_Autotuner.run = orig_autotuner_run
256296

257297

@@ -361,6 +401,41 @@ def _on_off_compile_spec_value(spec: CompileSpec) -> bool:
361401
return value == "ON"
362402

363403

404+
def _write_aoti_weights_blob(weights, blob_path: str) -> None:
405+
"""Stream AOTI tensor storages without creating a model-sized bytes object."""
406+
_trim_host_memory()
407+
tensors = [tensor for tensor, _ in weights.values()]
408+
all_cuda = all(tensor.is_cuda for tensor in tensors)
409+
chunk_size = 8 * 1024 * 1024
410+
411+
with open(blob_path, "wb") as output:
412+
for tensor in tensors:
413+
if tensor.is_mkldnn:
414+
raise RuntimeError("MKLDNN constants are not supported by CUDA AOTI")
415+
storage = tensor.untyped_storage()
416+
nbytes = storage.nbytes()
417+
if nbytes and tensor.is_cuda:
418+
byte_tensor = torch.empty(
419+
0, dtype=torch.uint8, device=tensor.device
420+
).set_(storage, 0, (nbytes,), (1,))
421+
for offset in range(0, nbytes, chunk_size):
422+
cpu_chunk = byte_tensor[offset : offset + chunk_size].cpu()
423+
output.write(memoryview(cpu_chunk.numpy()))
424+
del byte_tensor, cpu_chunk
425+
elif nbytes:
426+
raw_array = (ctypes.c_ubyte * nbytes).from_address(storage.data_ptr())
427+
raw_view = memoryview(raw_array).cast("B")
428+
for offset in range(0, nbytes, chunk_size):
429+
output.write(raw_view[offset : offset + chunk_size])
430+
del raw_view, raw_array
431+
# Match AOTInductor's binary_blob layout: CUDA-only constants are
432+
# packed, while CPU/mixed constants are aligned to 64 bytes.
433+
if not all_cuda and (padding := (-nbytes) % 64):
434+
output.write(bytes(padding))
435+
del storage
436+
_trim_host_memory()
437+
438+
364439
@final
365440
@experimental(
366441
"This API and all of cuda backend related functionality are experimental."
@@ -475,6 +550,81 @@ def save_data_externally(cls) -> bool:
475550
"""
476551
return True
477552

553+
@classmethod
554+
def load_weights_blob(
555+
cls, blob_path: str, compile_specs: List[CompileSpec]
556+
) -> tuple[Any, str]:
557+
"""Keep low-memory CUDA weights file-backed during PTE serialization.
558+
559+
The streamed file has the same layout as AOTInductor's ``binary_blob``.
560+
Keeping it file-backed avoids reading another model-sized copy into
561+
host memory without changing its bytes.
562+
"""
563+
if not cls._is_low_memory_mode(compile_specs):
564+
return super().load_weights_blob(blob_path, compile_specs)
565+
blob_data = FileBackedData.move_from(blob_path)
566+
return blob_data, blob_data.sha256().hex()
567+
568+
@classmethod
569+
def materialize_weights_blob(
570+
cls, paths: Any, compile_specs: List[CompileSpec]
571+
) -> Any:
572+
if not cls._is_low_memory_mode(compile_specs) or not isinstance(paths, list):
573+
return paths
574+
575+
from torch.export.pt2_archive._package_weights import Weights
576+
577+
weights = [path for path in paths if isinstance(path, Weights)]
578+
if not weights:
579+
return paths
580+
if len(weights) != 1:
581+
raise RuntimeError(
582+
f"Expected one CUDA AOTI weights output, got {len(weights)}"
583+
)
584+
585+
so_path = next(
586+
(
587+
path
588+
for path in paths
589+
if isinstance(path, str) and path.endswith(".wrapper.so")
590+
),
591+
None,
592+
)
593+
if so_path is None:
594+
raise RuntimeError(f"Expected a CUDA AOTI .wrapper.so output, got {paths}")
595+
blob_path = os.path.splitext(so_path)[0] + "_weights.blob"
596+
_write_aoti_weights_blob(weights[0], blob_path)
597+
598+
# Forcing the external-weights ABI makes Inductor emit an empty blob
599+
# path alongside the Weights object. Replace that file in place and do
600+
# not add a duplicate path to the returned package outputs.
601+
materialized = [path for path in paths if not isinstance(path, Weights)]
602+
if blob_path not in materialized:
603+
materialized.append(blob_path)
604+
return materialized
605+
606+
@classmethod
607+
def copy_exported_program_for_preprocess(
608+
cls, edge_program, compile_specs: List[CompileSpec]
609+
):
610+
"""Copy graph structure while sharing immutable tensor storage.
611+
612+
CUDA preprocessing replaces state-dict entries when moving them to the
613+
target device; it does not mutate the source tensors. Memoizing those
614+
tensors therefore avoids a model-sized host copy for every delegated
615+
method while preserving an independent graph and state-dict mapping.
616+
"""
617+
if not cls._is_low_memory_mode(compile_specs):
618+
return copy.deepcopy(edge_program)
619+
620+
tensor_memo = {
621+
id(tensor): tensor
622+
for values in (edge_program.state_dict, edge_program.constants)
623+
for tensor in values.values()
624+
if isinstance(tensor, torch.Tensor)
625+
}
626+
return copy.deepcopy(edge_program, tensor_memo)
627+
478628
@classmethod
479629
def get_supported_fallback_kernels(cls) -> Dict[str, Any]:
480630
# ROCm does not build the CUDA-only .cu fallback shims.
@@ -588,8 +738,13 @@ def get_aoti_compile_options(
588738
# Separate weight constants from the .so file
589739
"aot_inductor.package": True,
590740
"aot_inductor.package_constants_in_so": False,
591-
# Store weight constants on disk in a binary blob
592-
"aot_inductor.package_constants_on_disk_format": "binary_blob",
741+
# Store weight constants on disk in a binary blob. Low-memory mode
742+
# asks AOTI for a Weights object and streams the equivalent blob in
743+
# materialize_weights_blob; its context also forces the generated
744+
# wrapper to use the required external-weights ABI.
745+
"aot_inductor.package_constants_on_disk_format": cls._weights_format(
746+
compile_specs
747+
),
593748
# Enable maximum automatic tuning for optimal performance
594749
"max_autotune": True,
595750
# Use TRITON for GEMM (General Matrix Multiply) operations tuning only to avoid using operators in libtorch
@@ -696,6 +851,7 @@ def _combined():
696851
stack.enter_context(
697852
_compile_time_cpu_clones(torch.device(cls.get_device_name()))
698853
)
854+
_trim_host_memory()
699855
yield
700856

701857
return _combined()
@@ -708,6 +864,14 @@ def _is_low_memory_mode(compile_specs: List[CompileSpec]) -> bool:
708864
return spec.value.decode("utf-8").upper() == "ON"
709865
return False
710866

867+
@classmethod
868+
def _weights_format(cls, compile_specs: List[CompileSpec]) -> str:
869+
return (
870+
"pickle_weights"
871+
if cls._is_low_memory_mode(compile_specs)
872+
else "binary_blob"
873+
)
874+
711875
@classmethod
712876
def move_program_to_device(
713877
cls,

0 commit comments

Comments
 (0)