66
77
88import contextlib
9+ import copy
10+ import ctypes
11+ import functools
12+ import gc
913import logging
1014import os
1115import shutil
2529from executorch .backends .cuda .triton .replacement_pass import (
2630 ReplaceEdgeOpWithTritonOpPass ,
2731)
32+ from executorch .exir ._serialize ._cord import FileBackedData
2833from executorch .exir ._warnings import experimental
2934from executorch .exir .backend .backend_details import BackendDetails
3035from 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
6477def _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+
107126def _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