diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 25684f8b32e..ede0cd591cd 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -273,9 +273,10 @@ jobs: export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH cmake --preset llm-release-cuda -DEXECUTORCH_BUILD_TESTS=ON - cmake --build cmake-out --target test_cuda_allocator test_cuda_mutable_state -j$(nproc) + cmake --build cmake-out --target test_cuda_allocator test_cuda_mutable_state test_cuda_weight_cache -j$(nproc) ctest --test-dir cmake-out -R test_cuda_allocator --output-on-failure -V ctest --test-dir cmake-out -R test_cuda_mutable_state --output-on-failure -V + ctest --test-dir cmake-out -R test_cuda_weight_cache --output-on-failure -V test-model-cuda-e2e: name: test-model-cuda-e2e-${{ matrix.model.name }}-${{ matrix.quant }} diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 1fbbf61be3c..2d0da4af2af 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -93,6 +93,7 @@ fbcode_target( name = "cuda_backend", srcs = [ "cuda_backend.py", + "cuda_weight_collector.py", ], visibility = ["PUBLIC"], deps = [ diff --git a/backends/cuda/CMakeLists.txt b/backends/cuda/CMakeLists.txt index 21f8f5e2914..c3e5d17809e 100644 --- a/backends/cuda/CMakeLists.txt +++ b/backends/cuda/CMakeLists.txt @@ -337,8 +337,9 @@ install( ) # CUDA backend implementation -set(_aoti_cuda_backend_sources runtime/cuda_backend.cpp - runtime/cuda_mutable_state.cpp +set(_aoti_cuda_backend_sources + runtime/cuda_backend.cpp runtime/cuda_mutable_state.cpp + runtime/cuda_weight_cache.cpp ) if(_cuda_is_msvc_toolchain) # MSVC links aoti_cuda_backend into portable_lib without relying on C++ @@ -448,4 +449,10 @@ if(BUILD_TESTING) EXTRA_LIBS aoti_cuda_backend ) target_compile_definitions(test_cuda_mutable_state PRIVATE CUDA_AVAILABLE=1) + + et_cxx_test( + test_cuda_weight_cache SOURCES runtime/test/test_cuda_weight_cache.cpp + EXTRA_LIBS aoti_cuda_backend + ) + target_compile_definitions(test_cuda_weight_cache PRIVATE CUDA_AVAILABLE=1) endif() diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 1d6e0b11370..bf4f66f6d32 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -7,10 +7,7 @@ import contextlib import copy -import ctypes import functools -import gc -import hashlib import logging import os import shutil @@ -21,6 +18,12 @@ import torch from executorch.backends.aoti.aoti_backend import AotiBackend +from executorch.backends.cuda.cuda_weight_collector import ( + AOTI_DEVICE_TYPE_CPU, + AOTI_DEVICE_TYPE_CUDA, + CudaWeightCollector, + trim_host_memory, +) from executorch.backends.cuda.passes.move_cond_predicate_to_cpu import ( MoveCondPredicateToCpuPass, ) @@ -32,7 +35,7 @@ ) 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.backend_details import BackendDetails, PreprocessResult from executorch.exir.backend.compile_spec_schema import CompileSpec from torch._inductor.decomposition import conv1d_to_conv2d from torch.nn.attention import SDPBackend @@ -66,12 +69,18 @@ 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 _aoti_device_type_for_weight(tensor: torch.Tensor) -> int: + # Low-memory compilation clones lifted CUDA buffers onto CPU, while the + # patched wrapper records them as CUDA constants. Mirror that target-device + # substitution in the metadata. Outside that scoped mode the serialized + # tensor's actual device is the AOTI constant's device. + if _is_cpu_clone_active() or tensor.device.type == "cuda": + return AOTI_DEVICE_TYPE_CUDA + if tensor.device.type == "cpu": + return AOTI_DEVICE_TYPE_CPU + raise RuntimeError( + f"Unsupported AOTI constant device for CUDA export: {tensor.device}" + ) @contextlib.contextmanager @@ -195,10 +204,11 @@ def _compile_time_cpu_clones(target_device: torch.device): # noqa: C901 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. + def _force_external_weights_for_fqn_binding(consts_size): + # Structured weights are serialized as independently named storages, + # but the generated AOTI wrapper must still use the external-weights + # ABI. That mode preserves the original constant view metadata when + # the runtime replaces the dense blob with user-managed FQN tensors. if _is_cpu_clone_active(): return True, False return orig_determine_aoti_mmap_flags(consts_size) @@ -279,7 +289,7 @@ def _codegen_device_target_aware(self, device): _codecache.TensorProperties = functools.partial( _tensor_properties_for_low_memory, original=orig_tensor_properties ) - _codecache.determine_aoti_mmap_flags = _force_external_weights_for_streaming + _codecache.determine_aoti_mmap_flags = _force_external_weights_for_fqn_binding _Autotuner.run = _autotuner_run_with_rehydrated_emptied_args prev_active = getattr(_CPU_CLONE_GUARD, "active", False) _CPU_CLONE_GUARD.active = True @@ -402,47 +412,6 @@ def _on_off_compile_spec_value(spec: CompileSpec) -> bool: return value == "ON" -def _write_aoti_weights_blob(weights, blob_path: str) -> bytes: - """Stream AOTI tensor storages to disk and return their SHA-256 digest.""" - _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 - digest = hashlib.sha256() - - def write_chunk(output, chunk) -> None: - digest.update(chunk) - output.write(chunk) - - 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() - write_chunk(output, 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): - write_chunk(output, raw_view[offset : offset + chunk_size]) - del raw_view, raw_array - # Match AOTInductor's binary_blob layout: CUDA-only constants are - # packed, while CPU/mixed constants are aligned to 64 bytes. - if not all_cuda and (padding := (-nbytes) % 64): - write_chunk(output, bytes(padding)) - del storage - _trim_host_memory() - return digest.digest() - - @final @experimental( "This API and all of cuda backend related functionality are experimental." @@ -454,11 +423,6 @@ class CudaBackend(AotiBackend, BackendDetails): using the Executorch runtime. """ - # AOTI calls materialize_weights_blob immediately before load_weights_blob - # for a given output path. A new materialization overwrites any digest left - # behind by an export that aborted before the consumer ran. - _materialized_blob_hashes: Dict[str, bytes] = {} - @classmethod def get_device_name(cls) -> str: return "cuda" @@ -557,33 +521,83 @@ def _setup_cuda_environment_for_fatbin() -> bool: @classmethod def save_data_externally(cls) -> bool: """ - CUDA backend saves SO blob and weights blob to an external .ptd file. + CUDA backend saves weight storages (and, when configured, SO blobs) in + external named data such as a .ptd file. This file must be provided at runtime via --data_path argument. """ return True + @classmethod + def _preprocess_with_weight_collector( + cls, + edge_program: Any, + compile_specs: List[CompileSpec], + collector: CudaWeightCollector, + ) -> PreprocessResult: + with collector.capture(_aoti_device_type_for_weight) as capture: + result = super().preprocess(edge_program, compile_specs) + if capture.artifact is None: + raise RuntimeError("CUDA AOTI did not return a structured Weights output") + collector.add_preprocess_result(result, capture.artifact, cls.get_device_name()) + return result + + @classmethod + def preprocess( + cls, edge_program: Any, compile_specs: List[CompileSpec] + ) -> PreprocessResult: + """Compile one CUDA method with its own weight collector.""" + collector = CudaWeightCollector() + result = cls._preprocess_with_weight_collector( + edge_program, compile_specs, collector + ) + collector.finish() + return result + + @classmethod + def preprocess_multimethod( + cls, + edge_programs: Dict[str, List[Any]], + compile_specs: Dict[str, List[List[CompileSpec]]], + ) -> Dict[str, List[PreprocessResult]]: + """Compile every method against one global ``(device, FQN)`` store.""" + collector = CudaWeightCollector() + results: Dict[str, List[PreprocessResult]] = {} + for method_name, programs in edge_programs.items(): + if method_name not in compile_specs: + raise ValueError(f"Missing CUDA compile specs for {method_name!r}") + method_specs = compile_specs[method_name] + if len(programs) != len(method_specs): + raise ValueError( + f"Method {method_name!r} has {len(programs)} partitions but " + f"{len(method_specs)} compile-spec lists" + ) + results[method_name] = [ + cls._preprocess_with_weight_collector(program, specs, collector) + for program, specs in zip(programs, method_specs) + ] + collector.finish() + return results + @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. + """Keep low-memory CUDA named data file-backed during 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. + New FQN artifacts use this path only for AotiBackend's empty + compatibility placeholder. Legacy binary-blob handling remains + unchanged for callers that still provide a real blob. """ if not cls._is_low_memory_mode(compile_specs): return super().load_weights_blob(blob_path, compile_specs) - known_hash = cls._materialized_blob_hashes.pop(blob_path, None) - blob_data = FileBackedData.move_from(blob_path, sha256=known_hash) - weights_blob_hash = known_hash or blob_data.sha256() - return blob_data, weights_blob_hash.hex() + 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): + if not isinstance(paths, list): return paths from torch.export.pt2_archive._package_weights import Weights @@ -607,13 +621,19 @@ def materialize_weights_blob( if so_path is None: raise RuntimeError(f"Expected a CUDA AOTI .wrapper.so output, got {paths}") blob_path = os.path.splitext(so_path)[0] + "_weights.blob" - cls._materialized_blob_hashes[blob_path] = _write_aoti_weights_blob( - weights[0], blob_path + capture = CudaWeightCollector.current_capture() + capture.artifact = capture.collector.materialize( + weights[0], os.path.dirname(blob_path), capture.device_type_for_weight ) - # 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. + # Keep AotiBackend's existing path contract intact. The compatibility + # blob is empty and ignored by the versioned CUDA runtime path; old + # artifacts continue to carry and load their original dense blob. + with open(blob_path, "wb"): + pass + + # Replace the structured Weights output with the compatibility path + # expected by AotiBackend's existing named-data packaging contract. materialized = [path for path in paths if not isinstance(path, Weights)] if blob_path not in materialized: materialized.append(blob_path) @@ -754,10 +774,8 @@ 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. 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. + # Ask AOTI for structured constants. CUDABackend collects them by + # (device, FQN) and emits the tensor metadata needed at runtime. "aot_inductor.package_constants_on_disk_format": cls._weights_format( compile_specs ), @@ -877,7 +895,7 @@ def _combined(): stack.enter_context( _compile_time_cpu_clones(torch.device(cls.get_device_name())) ) - _trim_host_memory() + trim_host_memory() yield return _combined() @@ -892,11 +910,10 @@ def _is_low_memory_mode(compile_specs: List[CompileSpec]) -> bool: @classmethod def _weights_format(cls, compile_specs: List[CompileSpec]) -> str: - return ( - "pickle_weights" - if cls._is_low_memory_mode(compile_specs) - else "binary_blob" - ) + # CUDA consumes the structured AOTI output directly and emits a + # versioned per-FQN metadata payload. This is backend-wide rather than a + # model/export-script option. + return "pickle_weights" @classmethod def move_program_to_device( diff --git a/backends/cuda/cuda_weight_collector.py b/backends/cuda/cuda_weight_collector.py new file mode 100644 index 00000000000..ac47f7c7e2c --- /dev/null +++ b/backends/cuda/cuda_weight_collector.py @@ -0,0 +1,355 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import contextlib +import ctypes +import gc +import hashlib +import os +import struct +import tempfile +import threading +from dataclasses import dataclass, replace +from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple + +import torch +from executorch.exir._serialize._cord import FileBackedData +from executorch.exir._serialize._named_data_store import NamedDataStore +from executorch.exir.backend.backend_details import PreprocessResult +from executorch.exir.tensor import scalar_type_enum + + +CUDA_WEIGHT_CACHE_MAGIC = b"ETCUDAFQN3" + +AOTI_DEVICE_TYPE_CPU = 0 +AOTI_DEVICE_TYPE_CUDA = 1 + + +@dataclass(frozen=True) +class CudaWeightEntry: + fqn: str + storage_key: str + storage_nbytes: int + dtype: int + device_type: int + storage_offset: int + sizes: Tuple[int, ...] + strides: Tuple[int, ...] + + +@dataclass +class CudaWeightArtifact: + entries: List[CudaWeightEntry] + storages: Dict[str, FileBackedData] + + +@dataclass +class _CudaWeightCapture: + collector: "CudaWeightCollector" + device_type_for_weight: Callable[[torch.Tensor], int] + artifact: Optional[CudaWeightArtifact] = None + + +def trim_host_memory() -> None: + gc.collect() + try: + ctypes.CDLL(None).malloc_trim(0) + except AttributeError: + pass + + +def _write_tensor_storage(tensor: torch.Tensor, path: str) -> bytes: + """Stream one AOTI storage to ``path`` and return its SHA-256 digest.""" + chunk_size = 8 * 1024 * 1024 + digest = hashlib.sha256() + + def write_chunk(output, chunk) -> None: + digest.update(chunk) + output.write(chunk) + + if tensor.is_mkldnn: + raise RuntimeError("MKLDNN constants are not supported by CUDA AOTI") + storage = tensor.untyped_storage() + nbytes = storage.nbytes() + with open(path, "wb") as output: + 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() + write_chunk(output, 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): + write_chunk(output, raw_view[offset : offset + chunk_size]) + del raw_view, raw_array + del storage + return digest.digest() + + +def _storage_key( + fqn: str, device_type: int, aoti_library_key: Optional[str] = None +) -> str: + if device_type == AOTI_DEVICE_TYPE_CPU: + device = "cpu" + elif device_type == AOTI_DEVICE_TYPE_CUDA: + device = "cuda" + else: + raise RuntimeError(f"Unsupported AOTI device type: {device_type}") + if aoti_library_key is not None: + fqn = f"{aoti_library_key}:{fqn}" + return f"cuda_fqn_weight:{device}:{fqn}" + + +def _is_aoti_library_local_fqn(fqn: str) -> bool: + # PyTorch assigns this prefix to TensorConstant entries that do not have a + # model-level FQN. The numbering restarts in every independently compiled + # AOTI library, so the library key is part of their global identity. + return fqn.startswith("_tensor_constant") + + +def encode_cuda_weight_metadata( + so_blob_key: str, entries: List[CudaWeightEntry] +) -> bytes: + """Encode the per-method FQN-to-tensor metadata consumed by CUDA runtime.""" + output = bytearray(CUDA_WEIGHT_CACHE_MAGIC) + + def write_string(value: str) -> None: + encoded = value.encode("utf-8") + output.extend(struct.pack(" value`` store for all methods.""" + + _active = threading.local() + + def __init__(self) -> None: + self._store = NamedDataStore() + self._results: List[PreprocessResult] = [] + + @contextlib.contextmanager + def capture( + self, device_type_for_weight: Callable[[torch.Tensor], int] + ) -> Iterator[_CudaWeightCapture]: + previous = getattr(self._active, "current", None) + capture = _CudaWeightCapture(self, device_type_for_weight) + self._active.current = capture + try: + yield capture + finally: + self._active.current = previous + + @classmethod + def current_capture(cls) -> _CudaWeightCapture: + capture = getattr(cls._active, "current", None) + if capture is None: + raise RuntimeError( + "CUDA structured weights must be materialized inside preprocess" + ) + return capture + + def materialize( + self, + weights: Any, + directory: str, + device_type_for_weight: Callable[[torch.Tensor], int], + ) -> CudaWeightArtifact: + """Turn AOTI ``Weights`` into one independently named blob per FQN.""" + trim_host_memory() + entries: List[CudaWeightEntry] = [] + storages: Dict[str, FileBackedData] = {} + + for fqn, (tensor, properties) in weights.items(): + storage = tensor.untyped_storage() + storage_nbytes = storage.nbytes() + del storage + device_type = device_type_for_weight(tensor) + expected_storage_nbytes = int( + getattr(properties, "storage_size", None) or 0 + ) + if storage_nbytes < expected_storage_nbytes: + raise RuntimeError( + "AOTI cloned storage is smaller than its TensorProperties " + f"({storage_nbytes} < {expected_storage_nbytes} bytes)" + ) + + fd, storage_path = tempfile.mkstemp( + prefix=".cuda_weight_", suffix=".storage", dir=directory + ) + os.close(fd) + try: + digest = _write_tensor_storage(tensor, storage_path) + data = FileBackedData.move_from(storage_path, sha256=digest) + except Exception: + try: + os.remove(storage_path) + except OSError: + pass + raise + + storage_key = _storage_key(fqn, device_type) + if storage_key in storages: + data.close() + raise RuntimeError(f"Duplicate CUDA FQN weight key for {fqn!r}") + storages[storage_key] = data + + sizes = tuple( + int(size) for size in getattr(properties, "shape", tensor.shape) + ) + strides = tuple( + int(stride) for stride in getattr(properties, "stride", tensor.stride()) + ) + storage_offset = int(getattr(properties, "offset", tensor.storage_offset())) + if ( + len(sizes) != len(strides) + or storage_offset < 0 + or any(size < 0 for size in sizes) + or any(stride < 0 for stride in strides) + ): + raise RuntimeError(f"AOTI view {fqn!r} has invalid tensor metadata") + required_nbytes = 0 + if all(size != 0 for size in sizes): + last_element = storage_offset + sum( + stride * (size - 1) for size, stride in zip(sizes, strides) + ) + required_nbytes = (last_element + 1) * tensor.element_size() + if required_nbytes > storage_nbytes: + raise RuntimeError( + f"AOTI view {fqn!r} requires {required_nbytes} bytes from a " + f"{storage_nbytes}-byte cloned storage" + ) + entries.append( + CudaWeightEntry( + fqn=fqn, + storage_key=storage_key, + storage_nbytes=storage_nbytes, + dtype=int(scalar_type_enum(tensor.dtype)), + device_type=device_type, + storage_offset=storage_offset, + sizes=sizes, + strides=strides, + ) + ) + + trim_host_memory() + return CudaWeightArtifact(entries=entries, storages=storages) + + def _merge_aoti_data( + self, + parent_store: Any, + compatibility_blob_key: Optional[str], + keep_compatibility_blob: bool, + ) -> None: + for key, entry in parent_store.pte_data.items(): + if key != compatibility_blob_key or keep_compatibility_blob: + self._store.add_named_data( + key, + parent_store.buffers[entry.buffer_index], + alignment=entry.alignment, + tensor_layout=entry.tensor_layout, + ) + for tag, entries in parent_store.external_data.items(): + for key, entry in entries.items(): + if key != compatibility_blob_key or keep_compatibility_blob: + self._store.add_named_data( + key, + parent_store.buffers[entry.buffer_index], + alignment=entry.alignment, + external_tag=tag, + tensor_layout=entry.tensor_layout, + ) + + def _add_weight( + self, + entry: CudaWeightEntry, + data: FileBackedData, + external_tag: str, + ) -> None: + is_duplicate = entry.storage_key in self._store.key_to_buffer_idx + try: + self._store.add_named_data( + entry.storage_key, + data, + alignment=1, + external_tag=external_tag, + ) + except Exception: + data.close() + raise + if is_duplicate: + data.close() + + def add_preprocess_result( + self, + result: PreprocessResult, + artifact: CudaWeightArtifact, + device_name: str, + ) -> None: + if result.data_store_output is None: + raise RuntimeError("CUDA AOTI preprocess returned no named data") + try: + parent_keys = result.processed_bytes.decode("utf-8").splitlines() + except UnicodeDecodeError as error: + raise RuntimeError("Malformed CUDA AOTI named-data payload") from error + if not parent_keys or not parent_keys[0]: + raise RuntimeError("CUDA AOTI payload is missing its shared-object key") + so_blob_key = parent_keys[0] + compatibility_blob_key = parent_keys[1] if len(parent_keys) > 1 else None + + parent_store = result.data_store_output + self._merge_aoti_data( + parent_store, + compatibility_blob_key, + keep_compatibility_blob=not artifact.storages, + ) + + external_tag = f"aoti_{device_name}_blob" + serialized_entries = [] + for entry in artifact.entries: + data = artifact.storages[entry.storage_key] + if _is_aoti_library_local_fqn(entry.fqn): + entry = replace( + entry, + storage_key=_storage_key( + entry.fqn, entry.device_type, aoti_library_key=so_blob_key + ), + ) + self._add_weight(entry, data, external_tag) + serialized_entries.append(entry) + + result.processed_bytes = encode_cuda_weight_metadata( + so_blob_key, serialized_entries + ) + self._results.append(result) + + def finish(self) -> None: + shared_output = self._store.get_named_data_store_output() + for result in self._results: + result.data_store_output = shared_output diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 349082ad690..29cba8b5ada 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -46,6 +46,7 @@ #include #include #include +#include #include #include #include @@ -165,8 +166,8 @@ class ET_EXPERIMENTAL CudaBackend final return shared_cuda_stream_ != nullptr; } - // Enable cross-method per-FQN weight caching. Set via the - // kWeightSharingAcrossMethods runtime backend option. + // Enable the legacy dense-blob per-FQN cache. New FQN artifacts use + // their FQN-addressed data keys automatically. void set_weight_sharing_across_methods(bool enabled) { weight_sharing_across_methods_.store(enabled, std::memory_order_relaxed); } @@ -177,7 +178,7 @@ class ET_EXPERIMENTAL CudaBackend final Error load_function_pointers_into_handle( void* so_handle, - AOTIDelegateHandle* handle) const { + cuda::CudaDelegateHandle* handle) const { #define LOAD_SYMBOL(member, name) \ do { \ auto symbol_res = get_function(so_handle, #name); \ @@ -226,6 +227,8 @@ class ET_EXPERIMENTAL CudaBackend final LOAD_OPTIONAL_SYMBOL( get_constant_original_fqn, AOTInductorModelContainerGetConstantOriginalFQN); + LOAD_OPTIONAL_SYMBOL( + get_constant_dtype, AOTInductorModelContainerGetConstantDtype); LOAD_OPTIONAL_SYMBOL( extract_constants_map, AOTInductorModelContainerExtractConstantsMap); LOAD_OPTIONAL_SYMBOL( @@ -327,10 +330,21 @@ class ET_EXPERIMENTAL CudaBackend final std::string so_blob_key; std::string weights_blob_key; - ET_CHECK_OK_OR_RETURN_ERROR( - executorch::backends::aoti::resolve_blob_keys( - processed, method_name, so_blob_key, weights_blob_key), - "Malformed named-data key payload"); + CudaWeightCache::Metadata fqn_weights; + const bool has_fqn_weights = + CudaWeightCache::is_serialized(processed->data(), processed->size()); + if (has_fqn_weights) { + ET_CHECK_OK_OR_RETURN_ERROR( + CudaWeightCache::parse( + processed->data(), processed->size(), fqn_weights), + "Malformed CUDA FQN weight metadata"); + so_blob_key = fqn_weights.so_blob_key; + } else { + ET_CHECK_OK_OR_RETURN_ERROR( + executorch::backends::aoti::resolve_blob_keys( + processed, method_name, so_blob_key, weights_blob_key), + "Malformed named-data key payload"); + } const NamedDataMap* named_data_map = context.get_named_data_map(); auto aoti_dso_buffer = named_data_map->get_data(so_blob_key.c_str()); @@ -404,14 +418,13 @@ class ET_EXPERIMENTAL CudaBackend final handle->container_handle = container_handle; - // Load constants. When weight_sharing_across_methods is enabled (opt-in - // via the kWeightSharingAcrossMethods runtime backend option set by the - // runner), use the per-weight FQN cache so methods that share weights - // (e.g. prefill/decode) avoid duplicate GPU allocations. Otherwise fall - // back to the legacy per-method blob load — required for models whose - // methods are independent sub-graphs that may have FQN collisions - // (e.g. parakeet). - if (is_weight_sharing_across_methods_enabled()) { + // Versioned artifacts load each (device, FQN) through the same process-wide + // cross-method cache model used by the legacy path. The payload only adds + // the tensor metadata needed to reconstruct independently named PTD blobs. + if (has_fqn_weights) { + ET_CHECK_OK_OR_RETURN_ERROR( + fqn_weight_cache_.load(handle, named_data_map, fqn_weights)); + } else if (is_weight_sharing_across_methods_enabled()) { ET_CHECK_OK_OR_RETURN_ERROR(load_constants_with_cache( handle, named_data_map, method_name, weights_blob_key)); } else { @@ -859,20 +872,14 @@ class ET_EXPERIMENTAL CudaBackend final // NOTE: AOTInductorModelContainerDelete does not work correctly with // multiple .so files. Deleting one container frees shared resources, // which causes segmentation faults when attempting to delete other - // containers. As a workaround, we skip explicit container deletion - // and defer cleanup to the OS. + // containers. As a workaround, we skip explicit container deletion and + // defer cleanup to the OS. The corresponding shared library must remain + // loaded as well: the leaked container still owns objects whose code and + // process-wide state live in that library, so dlclose/FreeLibrary can + // invalidate them and crash during multi-method teardown. // TODO(gasoonjia): Find a proper solution for safe container deletion. // AOTInductorModelContainerDelete(handle->container_handle); - // Now close the shared library - if (handle->so_handle != nullptr) { - Error err = close_library(handle->so_handle); - ET_CHECK_OR_LOG_ERROR( - err == Error::Ok, - "Failed to close shared library for %s", - handle->so_path.c_str()); - } - // Remove the temporary shared library file if (!handle->so_path.empty()) { std::error_code remove_error; @@ -899,9 +906,9 @@ class ET_EXPERIMENTAL CudaBackend final mutable std::mutex cuda_stream_mutex_; std::shared_ptr shared_cuda_stream_ = nullptr; - // Whether to enable cross-method per-FQN weight caching at init time. + // Whether to enable cross-method caching for legacy dense-blob artifacts. // Toggled by the kWeightSharingAcrossMethods runtime backend option. Default - // OFF — see set_weight_sharing_across_methods() for safety constraints. + // OFF; versioned FQN artifacts do not consult this option. std::atomic weight_sharing_across_methods_{false}; // --------------------------------------------------------------- @@ -1266,6 +1273,8 @@ class ET_EXPERIMENTAL CudaBackend final // explicitly deleted — see destroy() comment). mutable std::unordered_map shared_constant_tensors_; + + mutable CudaWeightCache fqn_weight_cache_; }; } // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index 83d88b65c5a..32144ce139e 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -9,7 +9,10 @@ #pragma once #include +#include #include +#include +#include #include #include @@ -17,6 +20,50 @@ namespace executorch { namespace backends { namespace cuda { +using AOTInductorModelContainerGetConstantDtypeFunc = + aoti::AOTIRuntimeError (*)( + aoti::AOTInductorModelContainerHandle container_handle, + size_t idx, + int32_t* dtype); +struct CudaWeightStorage { + void* data{nullptr}; + size_t nbytes{0}; + aoti::slim::c10::DeviceType device_type{aoti::slim::c10::DeviceType::CUDA}; + int device_index{0}; + + CudaWeightStorage( + void* data_, + size_t nbytes_, + aoti::slim::c10::DeviceType device_type_, + int device_index_) + : data(data_), + nbytes(nbytes_), + device_type(device_type_), + device_index(device_index_) {} + + ~CudaWeightStorage() { + if (data == nullptr) { + return; + } + if (device_type == aoti::slim::c10::DeviceType::CPU) { + std::free(data); + return; + } + int previous_device = 0; + const cudaError_t get_device_error = cudaGetDevice(&previous_device); + if (get_device_error == cudaSuccess && previous_device != device_index) { + (void)cudaSetDevice(device_index); + } + (void)cudaFree(data); + if (get_device_error == cudaSuccess && previous_device != device_index) { + (void)cudaSetDevice(previous_device); + } + } + + CudaWeightStorage(const CudaWeightStorage&) = delete; + CudaWeightStorage& operator=(const CudaWeightStorage&) = delete; +}; + // Shared CUDA stream wrapper with proper RAII cleanup. // This ensures the stream is destroyed when all handles using it are destroyed. struct CudaStreamDeleter { @@ -148,6 +195,9 @@ struct CudaGraphState { // CUDA-specific delegate handle that extends AOTIDelegateHandle. // This consolidates CUDA stream management into a single location. struct CudaDelegateHandle : public aoti::AOTIDelegateHandle { + // Extra AOTI metadata used to validate per-FQN weights before binding. + AOTInductorModelContainerGetConstantDtypeFunc get_constant_dtype{nullptr}; + // CUDA stream for this handle, support both shared mode and single mode. // In shared mode, all cuda delegate handles share the same stream (e.g., for // skip-copy optimization), they will all hold a reference to the same @@ -168,6 +218,11 @@ struct CudaDelegateHandle : public aoti::AOTIDelegateHandle { // CUDA graph state (warmup, capture, replay, static buffers) CudaGraphState cuda_graph_state; + + // Per-FQN weight artifacts keep the allocations and their + // SlimTensor handles alive for as long as AOTI may reference their views. + std::vector> fqn_weight_storages; + std::vector> fqn_weight_tensors; }; } // namespace cuda diff --git a/backends/cuda/runtime/cuda_weight_cache.cpp b/backends/cuda/runtime/cuda_weight_cache.cpp new file mode 100644 index 00000000000..09b4bd76312 --- /dev/null +++ b/backends/cuda/runtime/cuda_weight_cache.cpp @@ -0,0 +1,566 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace executorch::backends::cuda { + +using aoti::AOTInductorConstantMapEntry; +using aoti::AtenTensorHandle; +using aoti::slim::SlimTensor; +using aoti::slim::c10::Device; +using runtime::Error; +using runtime::NamedDataMap; + +namespace { + +class MetadataReader final { + public: + MetadataReader(const void* data, size_t size) + : cursor_(static_cast(data)), end_(cursor_ + size) {} + + bool skip(size_t size) { + if (remaining() < size) { + return false; + } + cursor_ += size; + return true; + } + + bool read_u32(uint32_t& value) { + uint64_t wide = 0; + if (!read_unsigned(wide, 4)) { + return false; + } + value = static_cast(wide); + return true; + } + + bool read_i32(int32_t& value) { + uint32_t raw = 0; + if (!read_u32(raw)) { + return false; + } + std::memcpy(&value, &raw, sizeof(value)); + return true; + } + + bool read_u64(uint64_t& value) { + return read_unsigned(value, 8); + } + + bool read_i64(int64_t& value) { + uint64_t raw = 0; + if (!read_u64(raw)) { + return false; + } + std::memcpy(&value, &raw, sizeof(value)); + return true; + } + + bool read_string(std::string& value) { + uint32_t size = 0; + if (!read_u32(size) || remaining() < size) { + return false; + } + value.assign(reinterpret_cast(cursor_), size); + cursor_ += size; + return true; + } + + bool empty() const { + return cursor_ == end_; + } + + private: + size_t remaining() const { + return static_cast(end_ - cursor_); + } + + bool read_unsigned(uint64_t& value, size_t width) { + if (remaining() < width) { + return false; + } + value = 0; + for (size_t index = 0; index < width; ++index) { + value |= static_cast(cursor_[index]) << (index * 8); + } + cursor_ += width; + return true; + } + + const uint8_t* cursor_; + const uint8_t* end_; +}; + +bool is_supported_dtype(int32_t dtype) { + switch (dtype) { + case 0: // Byte + case 1: // Char + case 2: // Short + case 3: // Int + case 4: // Long + case 5: // Half + case 6: // Float + case 11: // Bool + case 15: // BFloat16 + return true; + default: + return false; + } +} + +bool is_supported_device_type(int32_t device_type) { + return device_type == 0 || device_type == 1; // CPU or CUDA +} + +} // namespace + +bool CudaWeightCache::is_serialized(const void* data, size_t size) { + return data != nullptr && size >= kFormatMagicSize && + std::memcmp(data, kFormatMagic, kFormatMagicSize) == 0; +} + +Error CudaWeightCache::parse( + const void* data, + size_t size, + Metadata& metadata) { + if (!is_serialized(data, size)) { + return Error::InvalidProgram; + } + + MetadataReader reader(data, size); + if (!reader.skip(kFormatMagicSize) || + !reader.read_string(metadata.so_blob_key) || + metadata.so_blob_key.empty()) { + return Error::InvalidProgram; + } + + uint32_t num_entries = 0; + constexpr uint32_t kMaxEntries = 1U << 20; + if (!reader.read_u32(num_entries) || num_entries > kMaxEntries) { + return Error::InvalidProgram; + } + metadata.entries.clear(); + metadata.entries.reserve(num_entries); + + constexpr uint32_t kMaxTensorDimensions = 64; + for (uint32_t index = 0; index < num_entries; ++index) { + Entry entry; + uint32_t ndim = 0; + if (!reader.read_string(entry.fqn) || entry.fqn.empty() || + !reader.read_string(entry.storage_key) || entry.storage_key.empty() || + !reader.read_u64(entry.storage_nbytes) || + !reader.read_i32(entry.dtype) || !is_supported_dtype(entry.dtype) || + !reader.read_i32(entry.device_type) || + !is_supported_device_type(entry.device_type) || + !reader.read_i64(entry.storage_offset) || !reader.read_u32(ndim) || + ndim > kMaxTensorDimensions) { + return Error::InvalidProgram; + } + + entry.sizes.resize(ndim); + entry.strides.resize(ndim); + for (uint32_t dim = 0; dim < ndim; ++dim) { + if (!reader.read_i64(entry.sizes[dim]) || entry.sizes[dim] < 0) { + return Error::InvalidProgram; + } + } + for (uint32_t dim = 0; dim < ndim; ++dim) { + if (!reader.read_i64(entry.strides[dim]) || entry.strides[dim] < 0) { + return Error::InvalidProgram; + } + } + if (entry.storage_offset < 0) { + return Error::InvalidProgram; + } + metadata.entries.push_back(std::move(entry)); + } + + return reader.empty() ? Error::Ok : Error::InvalidProgram; +} + +Error CudaWeightCache::validate_view(const Entry& entry) { + uint64_t item_size = 0; + switch (static_cast(entry.dtype)) { + case aoti::slim::c10::ScalarType::Byte: + case aoti::slim::c10::ScalarType::Char: + case aoti::slim::c10::ScalarType::Bool: + item_size = 1; + break; + case aoti::slim::c10::ScalarType::Short: + case aoti::slim::c10::ScalarType::Half: + case aoti::slim::c10::ScalarType::BFloat16: + item_size = 2; + break; + case aoti::slim::c10::ScalarType::Int: + case aoti::slim::c10::ScalarType::Float: + item_size = 4; + break; + case aoti::slim::c10::ScalarType::Long: + item_size = 8; + break; + default: + return Error::InvalidProgram; + } + + ET_CHECK_OR_RETURN_ERROR( + entry.storage_nbytes <= std::numeric_limits::max(), + InvalidProgram, + "CUDA FQN storage '%s' is too large for this platform", + entry.storage_key.c_str()); + + bool empty = false; + uint64_t last_element = static_cast(entry.storage_offset); + for (size_t dim = 0; dim < entry.sizes.size(); ++dim) { + const uint64_t size = static_cast(entry.sizes[dim]); + const uint64_t stride = static_cast(entry.strides[dim]); + if (size == 0) { + empty = true; + break; + } + const uint64_t extent = size - 1; + ET_CHECK_OR_RETURN_ERROR( + extent == 0 || stride <= std::numeric_limits::max() / extent, + InvalidProgram, + "CUDA FQN weight '%s' has overflowing shape/stride metadata", + entry.fqn.c_str()); + const uint64_t span = stride * extent; + ET_CHECK_OR_RETURN_ERROR( + last_element <= std::numeric_limits::max() - span, + InvalidProgram, + "CUDA FQN weight '%s' has overflowing storage metadata", + entry.fqn.c_str()); + last_element += span; + } + + uint64_t required_nbytes = 0; + if (!empty) { + ET_CHECK_OR_RETURN_ERROR( + last_element < std::numeric_limits::max() && + last_element + 1 <= + std::numeric_limits::max() / item_size, + InvalidProgram, + "CUDA FQN weight '%s' has overflowing storage size", + entry.fqn.c_str()); + required_nbytes = (last_element + 1) * item_size; + } + ET_CHECK_OR_RETURN_ERROR( + required_nbytes <= entry.storage_nbytes, + InvalidProgram, + "CUDA FQN weight '%s' requires %llu bytes from a %llu-byte storage", + entry.fqn.c_str(), + static_cast(required_nbytes), + static_cast(entry.storage_nbytes)); + return Error::Ok; +} + +Error CudaWeightCache::acquire_storage( + const NamedDataMap* named_data_map, + const Entry& entry, + uintptr_t logical_scope, + int device_index, + std::shared_ptr& storage, + bool& reused) const { + reused = false; + ET_CHECK_OR_RETURN_ERROR( + named_data_map != nullptr, + InvalidArgument, + "CUDA FQN weights require a named data map"); + + const auto device_type = + static_cast(entry.device_type); + const bool is_cuda_storage = device_type == aoti::slim::c10::DeviceType::CUDA; + const int storage_device_index = is_cuda_storage ? device_index : 0; + const std::string cache_key = entry.storage_key + "@" + + std::to_string(logical_scope) + + (is_cuda_storage ? "@cuda:" + std::to_string(device_index) : "@cpu"); + + std::unique_lock lock(mutex_); + auto cached = storages_.find(cache_key); + if (cached != storages_.end()) { + storage = cached->second.lock(); + } + if (storage != nullptr) { + ET_CHECK_OR_RETURN_ERROR( + storage->nbytes == entry.storage_nbytes && + storage->device_type == device_type && + storage->device_index == storage_device_index, + InvalidProgram, + "CUDA FQN weight '%s' has inconsistent allocation metadata", + entry.fqn.c_str()); + reused = true; + return Error::Ok; + } + + auto host_data = named_data_map->get_data(entry.storage_key.c_str()); + ET_CHECK_OR_RETURN_ERROR( + host_data.ok(), + NotFound, + "CUDA FQN storage '%s' is missing from named data", + entry.storage_key.c_str()); + if (host_data->size() != entry.storage_nbytes) { + const size_t actual_size = host_data->size(); + host_data->Free(); + ET_LOG( + Error, + "CUDA FQN storage '%s' has size %zu, expected %llu", + entry.storage_key.c_str(), + actual_size, + static_cast(entry.storage_nbytes)); + return Error::InvalidProgram; + } + + void* storage_data = nullptr; + const size_t allocation_size = + std::max(1, static_cast(entry.storage_nbytes)); + if (is_cuda_storage) { + const cudaError_t allocation_error = + cudaMalloc(&storage_data, allocation_size); + if (allocation_error != cudaSuccess) { + host_data->Free(); + ET_LOG( + Error, + "cudaMalloc failed for FQN storage '%s': %s", + entry.storage_key.c_str(), + cudaGetErrorString(allocation_error)); + return Error::MemoryAllocationFailed; + } + } else { + storage_data = std::malloc(allocation_size); + if (storage_data == nullptr) { + host_data->Free(); + ET_LOG( + Error, + "malloc failed for CPU FQN storage '%s'", + entry.storage_key.c_str()); + return Error::MemoryAllocationFailed; + } + } + + const auto free_storage_data = [&]() { + if (is_cuda_storage) { + (void)cudaFree(storage_data); + } else { + std::free(storage_data); + } + }; + + cudaError_t copy_error = cudaSuccess; + if (entry.storage_nbytes > 0) { + if (is_cuda_storage) { + copy_error = cudaMemcpy( + storage_data, + host_data->data(), + static_cast(entry.storage_nbytes), + cudaMemcpyHostToDevice); + } else { + std::memcpy( + storage_data, + host_data->data(), + static_cast(entry.storage_nbytes)); + } + } + host_data->Free(); + if (copy_error != cudaSuccess) { + free_storage_data(); + ET_LOG( + Error, + "cudaMemcpy failed for FQN storage '%s': %s", + entry.storage_key.c_str(), + cudaGetErrorString(copy_error)); + return Error::Internal; + } + + storage = std::make_shared( + storage_data, + static_cast(entry.storage_nbytes), + device_type, + storage_device_index); + storages_[cache_key] = storage; + return Error::Ok; +} + +Error CudaWeightCache::load( + CudaDelegateHandle* handle, + const NamedDataMap* named_data_map, + const Metadata& metadata) const { + ET_CHECK_OR_RETURN_ERROR( + named_data_map != nullptr, + InvalidArgument, + "CUDA FQN weights require a named data map"); + ET_CHECK_OR_RETURN_ERROR( + handle->get_num_constants && handle->get_constant_name && + handle->get_constant_original_fqn && handle->get_constant_dtype && + handle->update_user_managed_constant_buffer_pairs, + NotSupported, + "AOTI library does not expose the APIs required by CUDA FQN weights"); + + size_t num_constants = 0; + ET_CHECK_OK_OR_RETURN_ERROR( + handle->get_num_constants(handle->container_handle, &num_constants), + "Failed to enumerate CUDA AOTI constants"); + std::unordered_map> + fqn_to_internal_names; + std::unordered_map fqn_to_aoti_dtype; + for (size_t index = 0; index < num_constants; ++index) { + const char* internal_name = nullptr; + const char* fqn = nullptr; + int32_t dtype = 0; + ET_CHECK_OK_OR_RETURN_ERROR( + handle->get_constant_name( + handle->container_handle, index, &internal_name), + "Failed to read CUDA AOTI constant name at index %zu", + index); + ET_CHECK_OK_OR_RETURN_ERROR( + handle->get_constant_original_fqn( + handle->container_handle, index, &fqn), + "Failed to read CUDA AOTI constant FQN at index %zu", + index); + ET_CHECK_OK_OR_RETURN_ERROR( + handle->get_constant_dtype(handle->container_handle, index, &dtype), + "Failed to read CUDA AOTI constant dtype at index %zu", + index); + if (internal_name != nullptr && fqn != nullptr && fqn[0] != '\0') { + fqn_to_internal_names[fqn].emplace_back(internal_name); + auto [metadata, inserted] = fqn_to_aoti_dtype.emplace(fqn, dtype); + ET_CHECK_OR_RETURN_ERROR( + inserted || metadata->second == dtype, + InvalidProgram, + "CUDA AOTI constant FQN '%s' has inconsistent metadata", + fqn); + } + } + + int device_index = 0; + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaGetDevice(&device_index)); + // Each method receives a different MergedDataMap wrapper, but get_key() + // forwards the stable key owned by the shared PTD map. Use that identity to + // scope the process-wide FQN cache to one loaded model. + std::unordered_map key_scopes; + auto num_keys = named_data_map->get_num_keys(); + ET_CHECK_OR_RETURN_ERROR( + num_keys.ok(), + InvalidProgram, + "Failed to enumerate CUDA named data while loading FQN weights"); + key_scopes.reserve(num_keys.get()); + for (uint32_t index = 0; index < num_keys.get(); ++index) { + auto key = named_data_map->get_key(index); + ET_CHECK_OR_RETURN_ERROR( + key.ok() && key.get() != nullptr, + InvalidProgram, + "Failed to read CUDA named data key %u", + index); + key_scopes.emplace(key.get(), reinterpret_cast(key.get())); + } + std::vector pairs; + pairs.reserve(metadata.entries.size()); + std::unordered_set bound_fqns; + size_t reused_storages = 0; + handle->fqn_weight_tensors.reserve(metadata.entries.size()); + + for (const Entry& entry : metadata.entries) { + ET_CHECK_OK_OR_RETURN_ERROR( + validate_view(entry), "Invalid CUDA FQN view '%s'", entry.fqn.c_str()); + auto internal_names = fqn_to_internal_names.find(entry.fqn); + ET_CHECK_OR_RETURN_ERROR( + internal_names != fqn_to_internal_names.end(), + InvalidProgram, + "CUDA FQN weight '%s' is not present in its AOTI library", + entry.fqn.c_str()); + const auto aoti_dtype = fqn_to_aoti_dtype.find(entry.fqn); + ET_CHECK_OR_RETURN_ERROR( + aoti_dtype != fqn_to_aoti_dtype.end() && + aoti_dtype->second == entry.dtype, + InvalidProgram, + "CUDA FQN weight '%s' dtype does not match its AOTI library " + "(serialized=%d, AOTI=%d)", + entry.fqn.c_str(), + entry.dtype, + aoti_dtype == fqn_to_aoti_dtype.end() ? -1 : aoti_dtype->second); + ET_CHECK_OR_RETURN_ERROR( + bound_fqns.emplace(entry.fqn).second, + InvalidProgram, + "CUDA FQN weight '%s' appears more than once in serialized metadata", + entry.fqn.c_str()); + + std::shared_ptr storage; + bool reused = false; + const auto logical_scope = key_scopes.find(entry.storage_key); + ET_CHECK_OR_RETURN_ERROR( + logical_scope != key_scopes.end(), + NotFound, + "CUDA FQN storage '%s' is missing from named data", + entry.storage_key.c_str()); + ET_CHECK_OK_OR_RETURN_ERROR( + acquire_storage( + named_data_map, + entry, + logical_scope->second, + device_index, + storage, + reused), + "Failed to load CUDA FQN storage '%s'", + entry.storage_key.c_str()); + reused_storages += reused ? 1 : 0; + handle->fqn_weight_storages.push_back(storage); + + auto tensor = std::make_unique(aoti::slim::from_blob( + storage->data, + aoti::slim::makeArrayRef(entry.sizes), + aoti::slim::makeArrayRef(entry.strides), + static_cast(entry.dtype), + Device( + static_cast(entry.device_type), + entry.device_type == + static_cast(aoti::slim::c10::DeviceType::CUDA) + ? device_index + : 0), + entry.storage_offset)); + AtenTensorHandle tensor_handle = + reinterpret_cast(tensor.get()); + handle->fqn_weight_tensors.push_back(std::move(tensor)); + for (const std::string& internal_name : internal_names->second) { + pairs.push_back({internal_name.c_str(), tensor_handle}); + } + } + + ET_CHECK_OK_OR_RETURN_ERROR( + handle->update_user_managed_constant_buffer_pairs( + handle->container_handle, + pairs.data(), + pairs.size(), + /*use_inactive=*/false, + /*validate_full_update=*/true), + "Failed to bind CUDA FQN weights"); + ET_LOG( + Info, + "Loaded %zu CUDA FQN weights (%zu reused across methods)", + metadata.entries.size(), + reused_storages); + return Error::Ok; +} + +} // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/cuda_weight_cache.h b/backends/cuda/runtime/cuda_weight_cache.h new file mode 100644 index 00000000000..eb58ebc8fb1 --- /dev/null +++ b/backends/cuda/runtime/cuda_weight_cache.h @@ -0,0 +1,72 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace executorch::backends::cuda { + +class CudaWeightCache final { + public: + static constexpr char kFormatMagic[] = "ETCUDAFQN3"; + static constexpr size_t kFormatMagicSize = sizeof(kFormatMagic) - 1; + + struct Entry { + std::string fqn; + std::string storage_key; + uint64_t storage_nbytes{0}; + int32_t dtype{0}; + int32_t device_type{0}; + int64_t storage_offset{0}; + std::vector sizes; + std::vector strides; + }; + + struct Metadata { + std::string so_blob_key; + std::vector entries; + }; + + static bool is_serialized(const void* data, size_t size); + + static runtime::Error + parse(const void* data, size_t size, Metadata& metadata); + + runtime::Error load( + CudaDelegateHandle* handle, + const runtime::NamedDataMap* named_data_map, + const Metadata& metadata) const; + + private: + static runtime::Error validate_view(const Entry& entry); + + runtime::Error acquire_storage( + const runtime::NamedDataMap* named_data_map, + const Entry& entry, + uintptr_t logical_scope, + int device_index, + std::shared_ptr& storage, + bool& reused) const; + + mutable std::mutex mutex_; + mutable std::unordered_map> + storages_; +}; + +} // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/targets.bzl b/backends/cuda/runtime/targets.bzl index cef4988536f..e8372a33462 100644 --- a/backends/cuda/runtime/targets.bzl +++ b/backends/cuda/runtime/targets.bzl @@ -124,10 +124,12 @@ def define_common_targets(is_fbcode = False): srcs = [ "cuda_backend.cpp", "cuda_mutable_state.cpp", + "cuda_weight_cache.cpp", ], headers = [ "cuda_delegate_handle.h", "cuda_mutable_state.h", + "cuda_weight_cache.h", ], # @lint-ignore BUCKLINT: Avoid `link_whole=True` (https://fburl.com/avoid-link-whole) link_whole = True, @@ -180,6 +182,25 @@ def define_common_targets(is_fbcode = False): ), ) + cpp_unittest( + name = "test_cuda_weight_cache", + srcs = [ + "test/test_cuda_weight_cache.cpp", + ], + deps = [ + ":cuda_backend", + "//executorch/runtime/core:core", + ], + external_deps = [ + ("cuda", None, "cuda-lazy"), + ], + preprocessor_flags = ["-DCUDA_AVAILABLE=1"], + keep_gpu_sections = True, + remote_execution = re_test_utils.remote_execution( + platform = "gpu-remote-execution", + ), + ) + cpp_unittest( name = "test_cuda_allocator", srcs = ["test/test_cuda_allocator.cpp"], diff --git a/backends/cuda/runtime/test/test_cuda_weight_cache.cpp b/backends/cuda/runtime/test/test_cuda_weight_cache.cpp new file mode 100644 index 00000000000..d16e7143058 --- /dev/null +++ b/backends/cuda/runtime/test/test_cuda_weight_cache.cpp @@ -0,0 +1,116 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include + +#include +#include +#include + +namespace cuda = ::executorch::backends::cuda; +using ::executorch::runtime::Error; + +namespace { + +void append_u32(std::vector& output, uint32_t value) { + for (size_t index = 0; index < 4; ++index) { + output.push_back(static_cast(value >> (index * 8))); + } +} + +void append_u64(std::vector& output, uint64_t value) { + for (size_t index = 0; index < 8; ++index) { + output.push_back(static_cast(value >> (index * 8))); + } +} + +void append_string(std::vector& output, const std::string& value) { + append_u32(output, static_cast(value.size())); + output.insert(output.end(), value.begin(), value.end()); +} + +std::vector serialized_metadata( + uint32_t dtype = 6, + uint32_t device_type = 1) { + std::vector output( + cuda::CudaWeightCache::kFormatMagic, + cuda::CudaWeightCache::kFormatMagic + + cuda::CudaWeightCache::kFormatMagicSize); + append_string(output, "so-key"); + append_u32(output, 1); // entries + append_string(output, "model.weight"); + append_string(output, "storage-key"); + append_u64(output, 24); // storage bytes + append_u32(output, dtype); // dtype + append_u32(output, device_type); // device type (CUDA) + append_u64(output, 0); // storage offset + append_u32(output, 2); // ndim + append_u64(output, 2); + append_u64(output, 3); + append_u64(output, 3); + append_u64(output, 1); + return output; +} + +} // namespace + +TEST(CudaWeightCacheTest, LegacyPayloadIsNotMisdetected) { + const std::string legacy = "so-key\nweights-key"; + EXPECT_FALSE( + cuda::CudaWeightCache::is_serialized(legacy.data(), legacy.size())); +} + +TEST(CudaWeightCacheTest, ParsesSerializedMetadata) { + const std::vector bytes = serialized_metadata(); + cuda::CudaWeightCache::Metadata metadata; + ASSERT_EQ( + cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), + Error::Ok); + ASSERT_EQ(metadata.so_blob_key, "so-key"); + ASSERT_EQ(metadata.entries.size(), 1u); + const auto& entry = metadata.entries[0]; + EXPECT_EQ(entry.fqn, "model.weight"); + EXPECT_EQ(entry.storage_key, "storage-key"); + EXPECT_EQ(entry.storage_nbytes, 24u); + EXPECT_EQ(entry.dtype, 6); + EXPECT_EQ(entry.device_type, 1); + EXPECT_EQ(entry.sizes, (std::vector{2, 3})); + EXPECT_EQ(entry.strides, (std::vector{3, 1})); +} + +TEST(CudaWeightCacheTest, RejectsTruncationAndTrailingData) { + std::vector bytes = serialized_metadata(); + cuda::CudaWeightCache::Metadata metadata; + ASSERT_GT(bytes.size(), 1u); + EXPECT_EQ( + cuda::CudaWeightCache::parse(bytes.data(), bytes.size() - 1, metadata), + Error::InvalidProgram); + bytes.push_back(0); + EXPECT_EQ( + cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), + Error::InvalidProgram); +} + +TEST(CudaWeightCacheTest, RejectsUnsupportedDtype) { + const std::vector bytes = + serialized_metadata(7); // Double is unsupported. + cuda::CudaWeightCache::Metadata metadata; + EXPECT_EQ( + cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), + Error::InvalidProgram); +} + +TEST(CudaWeightCacheTest, RejectsUnsupportedDeviceType) { + const std::vector bytes = serialized_metadata(6, 2); + cuda::CudaWeightCache::Metadata metadata; + EXPECT_EQ( + cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), + Error::InvalidProgram); +} diff --git a/backends/cuda/tests/test_cuda_partitioner.py b/backends/cuda/tests/test_cuda_partitioner.py index a3e8b6bbee1..153828b2cb0 100644 --- a/backends/cuda/tests/test_cuda_partitioner.py +++ b/backends/cuda/tests/test_cuda_partitioner.py @@ -13,9 +13,21 @@ from unittest.mock import patch import torch -from executorch.backends.cuda.cuda_backend import CudaBackend +from executorch.backends.cuda.cuda_backend import ( + _aoti_device_type_for_weight, + CudaBackend, +) from executorch.backends.cuda.cuda_partitioner import CudaPartitioner +from executorch.backends.cuda.cuda_weight_collector import ( + AOTI_DEVICE_TYPE_CPU, + AOTI_DEVICE_TYPE_CUDA, + CUDA_WEIGHT_CACHE_MAGIC, + CudaWeightCollector, + encode_cuda_weight_metadata, +) from executorch.exir._serialize._cord import FileBackedData +from executorch.exir._serialize._named_data_store import NamedDataStore +from executorch.exir.backend.backend_details import PreprocessResult from executorch.exir.backend.compile_spec_schema import CompileSpec from executorch.exir.backend.partitioner import PartitionResult from executorch.exir.delegate import executorch_call_delegate @@ -26,10 +38,24 @@ 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 + @staticmethod + def _materialize(weights, directory, device_type=AOTI_DEVICE_TYPE_CPU): + return CudaWeightCollector().materialize( + weights, directory, lambda _: device_type + ) + + @staticmethod + def _parent_result(so_key: str) -> PreprocessResult: + store = NamedDataStore() + store.add_named_data(so_key, so_key.encode()) + store.add_named_data("empty_weights", b"", external_tag="aoti_cuda_blob") + return PreprocessResult( + processed_bytes=f"{so_key}\nempty_weights".encode(), + data_store_output=store.get_named_data_store_output(), + ) + @patch.object(CudaBackend, "_setup_cuda_environment_for_fatbin", return_value=True) + def test_all_cuda_exports_request_structured_weights(self, _) -> None: options = CudaBackend.get_aoti_compile_options( [CompileSpec("low_memory_mode", b"ON")] ) @@ -37,15 +63,14 @@ def test_low_memory_streaming_keeps_external_weights_abi(self, _) -> None: options["aot_inductor.package_constants_on_disk_format"], "pickle_weights", ) + self.assertEqual( + CudaBackend.get_aoti_compile_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: + def test_weights_are_materialized_as_independent_storages(self) -> None: first = torch.tensor([1, 2, 3], dtype=torch.int16) second = torch.tensor([4, 5], dtype=torch.int32) weights = Weights( @@ -56,41 +81,254 @@ def test_low_memory_weights_are_streamed_in_binary_blob_format(self) -> None: ) 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")], + artifact = self._materialize(weights, directory) + self.assertEqual(2, len(artifact.entries)) + self.assertEqual(2, len(artifact.storages)) + self.assertEqual(artifact.entries[0].device_type, 0) + self.assertEqual(artifact.entries[1].device_type, 0) + self.assertEqual( + bytes(first.untyped_storage()), + artifact.storages[artifact.entries[0].storage_key].to_bytes(), + ) + self.assertEqual( + bytes(second.untyped_storage()), + artifact.storages[artifact.entries[1].storage_key].to_bytes(), + ) + + metadata = encode_cuda_weight_metadata("so-key", artifact.entries) + self.assertTrue(metadata.startswith(CUDA_WEIGHT_CACHE_MAGIC)) + self.assertIn(b"first", metadata) + self.assertIn(b"second", metadata) + for storage in artifact.storages.values(): + storage.close() + + @patch( + "executorch.backends.cuda.cuda_backend._is_cpu_clone_active", + return_value=True, + ) + def test_low_memory_cpu_clones_keep_cuda_device_type(self, _) -> None: + tensor = torch.tensor([1, 2], dtype=torch.int16) + weights = Weights({"weight": (tensor, TensorProperties(tensor))}) + + with tempfile.TemporaryDirectory() as directory: + artifact = CudaWeightCollector().materialize( + weights, directory, _aoti_device_type_for_weight + ) + self.assertEqual(artifact.entries[0].device_type, AOTI_DEVICE_TYPE_CUDA) + for storage in artifact.storages.values(): + storage.close() + + def test_different_fqn_views_have_distinct_logical_storage(self) -> None: + base = torch.arange(12, dtype=torch.float32).reshape(3, 4) + view = base[:, 1:] + weights = Weights( + { + "base": (base, TensorProperties(base)), + # AOTI may return a cloned value tensor; TensorProperties is + # the source of truth for reconstructing the original view. + "view": (base, TensorProperties(view)), + } + ) + + with tempfile.TemporaryDirectory() as directory: + artifact = self._materialize(weights, directory) + self.assertEqual(2, len(artifact.storages)) + self.assertNotEqual( + artifact.entries[0].storage_key, + artifact.entries[1].storage_key, + ) + self.assertEqual(1, artifact.entries[1].storage_offset) + self.assertEqual((3, 3), artifact.entries[1].sizes) + self.assertEqual((4, 1), artifact.entries[1].strides) + for storage in artifact.storages.values(): + storage.close() + + def test_identical_values_keep_distinct_fqn_keys(self) -> None: + first = torch.zeros(4) + second = torch.zeros(4) + weights = Weights( + { + "first": (first, TensorProperties(first)), + "second": (second, TensorProperties(second)), + } + ) + + with tempfile.TemporaryDirectory() as directory: + artifact = self._materialize(weights, directory) + self.assertEqual(2, len(artifact.storages)) + self.assertNotEqual(artifact.entries[0].fqn, artifact.entries[1].fqn) + self.assertNotEqual( + artifact.entries[0].storage_key, + artifact.entries[1].storage_key, + ) + store = NamedDataStore() + for entry in artifact.entries: + store.add_named_data( + entry.storage_key, artifact.storages[entry.storage_key] + ) + self.assertEqual(1, len(store.buffers)) + for storage in artifact.storages.values(): + storage.close() + + def test_same_fqn_with_different_data_is_rejected_when_merged(self) -> None: + first = torch.zeros(4) + second = torch.ones(4) + + with ( + tempfile.TemporaryDirectory() as first_dir, + tempfile.TemporaryDirectory() as second_dir, + ): + first_artifact = self._materialize( + Weights({"weight": (first, TensorProperties(first))}), first_dir + ) + second_artifact = self._materialize( + Weights({"weight": (second, TensorProperties(second))}), second_dir ) + first_key = first_artifact.entries[0].storage_key + second_key = second_artifact.entries[0].storage_key + self.assertEqual(first_key, second_key) - 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) + collector = CudaWeightCollector() + collector.add_preprocess_result( + self._parent_result("first_so"), first_artifact, "cuda" ) - self.assertEqual(expected, data) - - # The streaming write computes the digest in the same pass. Loading - # the file-backed blob must not reread it solely for hashing. - with patch.object( - FileBackedData, - "sha256", - side_effect=AssertionError("unexpected blob reread"), - ): - blob, digest = CudaBackend.load_weights_blob( - blob_path, [CompileSpec("low_memory_mode", b"ON")] + with self.assertRaisesRegex(ValueError, "different data"): + collector.add_preprocess_result( + self._parent_result("second_so"), second_artifact, "cuda" ) - self.assertEqual(hashlib.sha256(expected).hexdigest(), digest) - self.assertEqual(expected, blob.to_bytes()) + + for artifact in (first_artifact, second_artifact): + for storage in artifact.storages.values(): + storage.close() + + def test_same_fqn_with_different_metadata_shares_storage(self) -> None: + tensor = torch.arange(4) + view = tensor.reshape(2, 2) + with ( + tempfile.TemporaryDirectory() as first_dir, + tempfile.TemporaryDirectory() as second_dir, + ): + first_artifact = self._materialize( + Weights({"weight": (tensor, TensorProperties(tensor))}), first_dir + ) + second_artifact = self._materialize( + Weights({"weight": (tensor, TensorProperties(view))}), second_dir + ) + collector = CudaWeightCollector() + first_result = self._parent_result("first_so") + second_result = self._parent_result("second_so") + collector.add_preprocess_result(first_result, first_artifact, "cuda") + collector.add_preprocess_result(second_result, second_artifact, "cuda") + collector.finish() + + self.assertIs( + first_result.data_store_output, second_result.data_store_output + ) + self.assertNotEqual( + first_result.processed_bytes, second_result.processed_bytes + ) + weight_key = first_artifact.entries[0].storage_key + self.assertIn( + weight_key, + first_result.data_store_output.external_data["aoti_cuda_blob"], + ) + + for artifact in (first_artifact, second_artifact): + for storage in artifact.storages.values(): + storage.close() + + def test_aoti_local_fqns_are_scoped_by_library(self) -> None: + first = torch.zeros(4) + second = torch.ones(5) + + with ( + tempfile.TemporaryDirectory() as first_dir, + tempfile.TemporaryDirectory() as second_dir, + ): + first_artifact = self._materialize( + Weights( + { + "_tensor_constant0": ( + first, + TensorProperties(first), + ) + } + ), + first_dir, + ) + second_artifact = self._materialize( + Weights( + { + "_tensor_constant0": ( + second, + TensorProperties(second), + ) + } + ), + second_dir, + ) + first_result = self._parent_result("first_so") + second_result = self._parent_result("second_so") + collector = CudaWeightCollector() + collector.add_preprocess_result(first_result, first_artifact, "cuda") + collector.add_preprocess_result(second_result, second_artifact, "cuda") + collector.finish() + + external_data = first_result.data_store_output.external_data[ + "aoti_cuda_blob" + ] + first_key = "cuda_fqn_weight:cpu:first_so:_tensor_constant0" + second_key = "cuda_fqn_weight:cpu:second_so:_tensor_constant0" + self.assertIn(first_key, external_data) + self.assertIn(second_key, external_data) + self.assertIn(first_key.encode(), first_result.processed_bytes) + self.assertIn(second_key.encode(), second_result.processed_bytes) + + for artifact in (first_artifact, second_artifact): + for storage in artifact.storages.values(): + storage.close() + + def test_methods_share_one_collected_named_data_store(self) -> None: + tensor = torch.arange(4) + weights = Weights({"weight": (tensor, TensorProperties(tensor))}) + with ( + tempfile.TemporaryDirectory() as first_dir, + tempfile.TemporaryDirectory() as second_dir, + ): + first_artifact = self._materialize(weights, first_dir) + second_artifact = self._materialize(weights, second_dir) + first_result = self._parent_result("first_so") + second_result = self._parent_result("second_so") + collector = CudaWeightCollector() + collector.add_preprocess_result(first_result, first_artifact, "cuda") + collector.add_preprocess_result(second_result, second_artifact, "cuda") + collector.finish() + + self.assertIs( + first_result.data_store_output, second_result.data_store_output + ) + weight_key = first_artifact.entries[0].storage_key + self.assertIn( + weight_key, + first_result.data_store_output.external_data["aoti_cuda_blob"], + ) + for artifact in (first_artifact, second_artifact): + for storage in artifact.storages.values(): + storage.close() + + def test_device_is_part_of_the_fqn_key(self) -> None: + tensor = torch.zeros(4) + weights = Weights({"weight": (tensor, TensorProperties(tensor))}) + with ( + tempfile.TemporaryDirectory() as cpu_dir, + tempfile.TemporaryDirectory() as cuda_dir, + ): + cpu = self._materialize(weights, cpu_dir, AOTI_DEVICE_TYPE_CPU) + cuda = self._materialize(weights, cuda_dir, AOTI_DEVICE_TYPE_CUDA) + self.assertNotEqual(cpu.entries[0].storage_key, cuda.entries[0].storage_key) + for artifact in (cpu, cuda): + for storage in artifact.storages.values(): + storage.close() def test_low_memory_weights_require_wrapper_so(self) -> None: tensor = torch.tensor([1], dtype=torch.int16) diff --git a/runtime/core/named_data_map.h b/runtime/core/named_data_map.h index dbd5b21a66f..e6a53b7d46e 100644 --- a/runtime/core/named_data_map.h +++ b/runtime/core/named_data_map.h @@ -65,8 +65,9 @@ class NamedDataMap { * Get the key at the given index. * * @param index The index of the key to retrieve. - * @return Result containing the key at the given index. Note: the returned - * pointer is only valid for the lifetime of the DataMap. + * @return Result containing the key at the given index. The returned pointer + * is owned by the DataMap, remains stable across calls, and is only valid for + * the lifetime of the DataMap. */ ET_NODISCARD virtual Result get_key(uint32_t index) const = 0; };