From ef326f314bf7f6e54dda29f52086097223716965 Mon Sep 17 00:00:00 2001 From: DN6 Date: Wed, 16 Sep 2026 17:47:35 +0530 Subject: [PATCH] introduce TorchDeviceBackend --- docs/source/en/api/utilities.md | 4 + src/diffusers/hooks/group_offloading.py | 18 +- src/diffusers/loaders/lora_pipeline.py | 3 +- .../modular_pipelines/components_manager.py | 23 +- src/diffusers/pipelines/ltx2/__init__.py | 2 +- .../pipelines/pag/pipeline_pag_sana.py | 10 +- src/diffusers/pipelines/pipeline_utils.py | 2 +- src/diffusers/pipelines/sana/pipeline_sana.py | 10 +- .../sana/pipeline_sana_controlnet.py | 10 +- .../pipelines/sana/pipeline_sana_sprint.py | 10 +- .../sana/pipeline_sana_sprint_img2img.py | 10 +- .../sana_video/pipeline_sana_video.py | 10 +- .../sana_video/pipeline_sana_video_i2v.py | 10 +- .../quantizers/gguf/gguf_quantizer.py | 8 +- src/diffusers/training_utils.py | 11 +- src/diffusers/utils/torch_utils.py | 277 ++++++++---------- tests/hooks/test_group_offloading.py | 8 +- .../modular_pipelines/testing_utils/utils.py | 14 +- tests/others/test_utils.py | 91 ++++++ tests/testing_utils.py | 97 +----- 20 files changed, 273 insertions(+), 355 deletions(-) diff --git a/docs/source/en/api/utilities.md b/docs/source/en/api/utilities.md index 69e69742249f..91ffb1a41ae7 100644 --- a/docs/source/en/api/utilities.md +++ b/docs/source/en/api/utilities.md @@ -50,6 +50,10 @@ Utility and helper functions for working with 🤗 Diffusers. [[autodoc]] utils.torch_utils.randn_tensor +## TorchDeviceBackend + +[[autodoc]] utils.torch_utils.TorchDeviceBackend + ## apply_layerwise_casting [[autodoc]] hooks.layerwise_casting.apply_layerwise_casting diff --git a/src/diffusers/hooks/group_offloading.py b/src/diffusers/hooks/group_offloading.py index 10d3f0c245a1..662c49186a7b 100644 --- a/src/diffusers/hooks/group_offloading.py +++ b/src/diffusers/hooks/group_offloading.py @@ -23,6 +23,7 @@ import torch from ..utils import get_logger, is_accelerate_available, is_torchao_available +from ..utils.torch_utils import TorchDeviceBackend from ._common import _GO_LC_SUPPORTED_PYTORCH_LAYERS from .hooks import HookRegistry, ModelHook @@ -166,11 +167,7 @@ def __init__( else: self.cpu_param_dict = self._init_cpu_param_dict() - self._torch_accelerator_module = ( - getattr(torch, torch.accelerator.current_accelerator().type) - if hasattr(torch, "accelerator") - else torch.cuda - ) + self._torch_accelerator_module = TorchDeviceBackend(self.onload_device) @staticmethod def _to_cpu(tensor, low_cpu_mem_usage): @@ -664,14 +661,9 @@ def apply_group_offloading( offload_device = torch.device(offload_device) if isinstance(offload_device, str) else offload_device offload_type = GroupOffloadingType(offload_type) - stream = None - if use_stream: - if torch.cuda.is_available(): - stream = torch.cuda.Stream() - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - stream = torch.Stream() - else: - raise ValueError("Using streams for data transfer requires a CUDA device, or an Intel XPU device.") + if use_stream and onload_device.type == "cpu": + raise ValueError("Using streams for data transfer requires an accelerator onload device, got `cpu`.") + stream = TorchDeviceBackend(onload_device).Stream() if use_stream else None if not use_stream and record_stream: raise ValueError("`record_stream` cannot be True when `use_stream=False`.") diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index 739ff9d2b3b1..947f6877b3b7 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -30,6 +30,7 @@ logging, require_peft_backend, ) +from ..utils.torch_utils import get_device from .lora_base import ( # noqa LORA_WEIGHT_NAME, LORA_WEIGHT_NAME_SAFE, @@ -109,7 +110,7 @@ def _maybe_dequantize_weight_for_expanded_lora(model, module): if module.weight.device.type == "cpu": weight_on_cpu = True - device = torch.accelerator.current_accelerator().type if hasattr(torch, "accelerator") else "cuda" + device = get_device() if is_bnb_4bit_quantized or is_bnb_8bit_quantized: module_weight = dequantize_bnb_weight( module.weight.to(device) if weight_on_cpu else module.weight, diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index 87b43bc3a630..4d467e53ea45 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -28,7 +28,7 @@ is_accelerate_available, logging, ) -from ..utils.torch_utils import get_device +from ..utils.torch_utils import TorchDeviceBackend, empty_device_cache, get_device if is_accelerate_available(): @@ -179,13 +179,7 @@ def __call__(self, hooks, model_id, model, execution_device): except AttributeError: raise AttributeError(f"Do not know how to compute memory footprint of `{model.__class__.__name__}.") - device_type = execution_device.type - device_module = getattr(torch, device_type, torch.cuda) - try: - mem_on_device = device_module.mem_get_info(execution_device.index)[0] - except AttributeError: - raise AttributeError(f"Do not know how to obtain obtain memory info for {str(device_module)}.") - + mem_on_device = TorchDeviceBackend(execution_device).mem_get_info()[0] mem_on_device = mem_on_device - self.memory_reserve_margin if current_module_size < mem_on_device: return [] @@ -513,10 +507,7 @@ def remove(self, component_id: str = None): import gc gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - if torch.xpu.is_available(): - torch.xpu.empty_cache() + empty_device_cache() # YiYi TODO: rename to search_components for now, may remove this method def search_components( @@ -743,12 +734,8 @@ def enable_auto_cpu_offload( if not isinstance(device, torch.device): device = torch.device(device) - device_type = device.type - device_module = getattr(torch, device_type, torch.cuda) - if not hasattr(device_module, "mem_get_info"): - raise NotImplementedError( - f"`enable_auto_cpu_offload() relies on the `mem_get_info()` method. It's not implemented for {str(device.type)}." - ) + # Fail here rather than on the first forward: the strategy cannot run without a free-memory query. + TorchDeviceBackend(device).mem_get_info() if device.index is None: device = torch.device(f"{device.type}:{0}") diff --git a/src/diffusers/pipelines/ltx2/__init__.py b/src/diffusers/pipelines/ltx2/__init__.py index d4aa35127403..d48d890f4cb6 100644 --- a/src/diffusers/pipelines/ltx2/__init__.py +++ b/src/diffusers/pipelines/ltx2/__init__.py @@ -30,12 +30,12 @@ _import_structure["pipeline_ltx2_condition"] = ["LTX2ConditionPipeline", "LTX2VideoCondition"] _import_structure["pipeline_ltx2_dfr"] = ["LTX2DFRPipeline"] _import_structure["pipeline_ltx2_dfr_temporal_refine"] = ["LTX2DFRTemporalRefinePipeline"] - _import_structure["pipeline_output"] = ["LTX2DFRPipelineOutput", "LTX2PipelineOutput", "LTX2VideoDecodeOutput"] _import_structure["pipeline_ltx2_diffusion_decode"] = ["LTX2VideoDiffusionDecodePipeline"] _import_structure["pipeline_ltx2_hdr_lora"] = ["LTX2HDRPipeline", "LTX2HDRReferenceCondition"] _import_structure["pipeline_ltx2_ic_lora"] = ["LTX2InContextPipeline", "LTX2ReferenceCondition"] _import_structure["pipeline_ltx2_image2video"] = ["LTX2ImageToVideoPipeline"] _import_structure["pipeline_ltx2_latent_upsample"] = ["LTX2LatentUpsamplePipeline"] + _import_structure["pipeline_output"] = ["LTX2DFRPipelineOutput", "LTX2PipelineOutput", "LTX2VideoDecodeOutput"] _import_structure["vocoder"] = ["LTX2Vocoder", "LTX2VocoderWithBWE"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sana.py b/src/diffusers/pipelines/pag/pipeline_pag_sana.py index f621255325a0..2b8dd5f14d24 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sana.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sana.py @@ -35,7 +35,7 @@ logging, replace_example_docstring, ) -from ...utils.torch_utils import get_device, is_torch_version, randn_tensor +from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput from ..pixart_alpha.pipeline_pixart_alpha import ( ASPECT_RATIO_512_BIN, @@ -892,15 +892,9 @@ def __call__( image = latents else: latents = latents.to(self.vae.dtype) - torch_accelerator_module = getattr(torch, get_device(), torch.cuda) - oom_error = ( - torch.OutOfMemoryError - if is_torch_version(">=", "2.5.0") - else torch_accelerator_module.OutOfMemoryError - ) try: image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - except oom_error as e: + except torch.OutOfMemoryError as e: warnings.warn( f"{e}. \n" f"Try to use VAE tiling for large images. For example: \n" diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index 8986553eda3d..1510c154f6ed 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -1287,7 +1287,7 @@ def maybe_free_model_hooks(self): return # make sure the model is in the same state as before calling it - self.enable_model_cpu_offload(device=getattr(self, "_offload_device", "cuda")) + self.enable_model_cpu_offload(device=getattr(self, "_offload_device", get_device())) def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None): r""" diff --git a/src/diffusers/pipelines/sana/pipeline_sana.py b/src/diffusers/pipelines/sana/pipeline_sana.py index 553e45a628d9..f8a4fbe508e3 100644 --- a/src/diffusers/pipelines/sana/pipeline_sana.py +++ b/src/diffusers/pipelines/sana/pipeline_sana.py @@ -38,7 +38,7 @@ scale_lora_layers, unscale_lora_layers, ) -from ...utils.torch_utils import get_device, is_torch_version, randn_tensor +from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline from ..pixart_alpha.pipeline_pixart_alpha import ( ASPECT_RATIO_512_BIN, @@ -957,15 +957,9 @@ def __call__( image = latents else: latents = latents.to(self.vae.dtype) - torch_accelerator_module = getattr(torch, get_device(), torch.cuda) - oom_error = ( - torch.OutOfMemoryError - if is_torch_version(">=", "2.5.0") - else torch_accelerator_module.OutOfMemoryError - ) try: image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - except oom_error as e: + except torch.OutOfMemoryError as e: warnings.warn( f"{e}. \n" f"Try to use VAE tiling for large images. For example: \n" diff --git a/src/diffusers/pipelines/sana/pipeline_sana_controlnet.py b/src/diffusers/pipelines/sana/pipeline_sana_controlnet.py index de1910f68192..5205ede5904a 100644 --- a/src/diffusers/pipelines/sana/pipeline_sana_controlnet.py +++ b/src/diffusers/pipelines/sana/pipeline_sana_controlnet.py @@ -38,7 +38,7 @@ scale_lora_layers, unscale_lora_layers, ) -from ...utils.torch_utils import get_device, is_torch_version, randn_tensor +from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline from ..pixart_alpha.pipeline_pixart_alpha import ( ASPECT_RATIO_512_BIN, @@ -1053,15 +1053,9 @@ def __call__( image = latents else: latents = latents.to(self.vae.dtype) - torch_accelerator_module = getattr(torch, get_device(), torch.cuda) - oom_error = ( - torch.OutOfMemoryError - if is_torch_version(">=", "2.5.0") - else torch_accelerator_module.OutOfMemoryError - ) try: image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - except oom_error as e: + except torch.OutOfMemoryError as e: warnings.warn( f"{e}. \n" f"Try to use VAE tiling for large images. For example: \n" diff --git a/src/diffusers/pipelines/sana/pipeline_sana_sprint.py b/src/diffusers/pipelines/sana/pipeline_sana_sprint.py index 812441d8e462..9be1c57b07e7 100644 --- a/src/diffusers/pipelines/sana/pipeline_sana_sprint.py +++ b/src/diffusers/pipelines/sana/pipeline_sana_sprint.py @@ -38,7 +38,7 @@ scale_lora_layers, unscale_lora_layers, ) -from ...utils.torch_utils import get_device, is_torch_version, randn_tensor +from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline from ..pixart_alpha.pipeline_pixart_alpha import ASPECT_RATIO_1024_BIN from .pipeline_output import SanaPipelineOutput @@ -839,15 +839,9 @@ def __call__( image = latents else: latents = latents.to(self.vae.dtype) - torch_accelerator_module = getattr(torch, get_device(), torch.cuda) - oom_error = ( - torch.OutOfMemoryError - if is_torch_version(">=", "2.5.0") - else torch_accelerator_module.OutOfMemoryError - ) try: image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - except oom_error as e: + except torch.OutOfMemoryError as e: warnings.warn( f"{e}. \n" f"Try to use VAE tiling for large images. For example: \n" diff --git a/src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py b/src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py index e149e0c597b2..9f97ff835512 100644 --- a/src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py +++ b/src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py @@ -39,7 +39,7 @@ scale_lora_layers, unscale_lora_layers, ) -from ...utils.torch_utils import get_device, is_torch_version, randn_tensor +from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline from ..pixart_alpha.pipeline_pixart_alpha import ASPECT_RATIO_1024_BIN from .pipeline_output import SanaPipelineOutput @@ -929,15 +929,9 @@ def __call__( image = latents else: latents = latents.to(self.vae.dtype) - torch_accelerator_module = getattr(torch, get_device(), torch.cuda) - oom_error = ( - torch.OutOfMemoryError - if is_torch_version(">=", "2.5.0") - else torch_accelerator_module.OutOfMemoryError - ) try: image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - except oom_error as e: + except torch.OutOfMemoryError as e: warnings.warn( f"{e}. \n" f"Try to use VAE tiling for large images. For example: \n" diff --git a/src/diffusers/pipelines/sana_video/pipeline_sana_video.py b/src/diffusers/pipelines/sana_video/pipeline_sana_video.py index 7ae85639e358..8cf51d3bb0e8 100644 --- a/src/diffusers/pipelines/sana_video/pipeline_sana_video.py +++ b/src/diffusers/pipelines/sana_video/pipeline_sana_video.py @@ -37,7 +37,7 @@ scale_lora_layers, unscale_lora_layers, ) -from ...utils.torch_utils import get_device, is_torch_version, randn_tensor +from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline from .pipeline_output import SanaVideoPipelineOutput @@ -990,12 +990,6 @@ def __call__( video = latents else: latents = latents.to(self.vae.dtype) - torch_accelerator_module = getattr(torch, get_device(), torch.cuda) - oom_error = ( - torch.OutOfMemoryError - if is_torch_version(">=", "2.5.0") - else torch_accelerator_module.OutOfMemoryError - ) if isinstance(self.vae, AutoencoderKLLTX2Video): latents_mean = self.vae.latents_mean latents_std = self.vae.latents_std @@ -1014,7 +1008,7 @@ def __call__( latents = latents / latents_std + latents_mean try: video = self.vae.decode(latents, return_dict=False)[0] - except oom_error as e: + except torch.OutOfMemoryError as e: warnings.warn( f"{e}. \n" f"Try to use VAE tiling for large images. For example: \n" diff --git a/src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py b/src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py index 81df1d0759da..cc772af22665 100644 --- a/src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py +++ b/src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py @@ -39,7 +39,7 @@ scale_lora_layers, unscale_lora_layers, ) -from ...utils.torch_utils import get_device, is_torch_version, randn_tensor +from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline from .pipeline_output import SanaVideoPipelineOutput @@ -1043,12 +1043,6 @@ def __call__( video = latents else: latents = latents.to(self.vae.dtype) - torch_accelerator_module = getattr(torch, get_device(), torch.cuda) - oom_error = ( - torch.OutOfMemoryError - if is_torch_version(">=", "2.5.0") - else torch_accelerator_module.OutOfMemoryError - ) if isinstance(self.vae, AutoencoderKLLTX2Video): latents_mean = self.vae.latents_mean latents_std = self.vae.latents_std @@ -1067,7 +1061,7 @@ def __call__( latents = latents / latents_std + latents_mean try: video = self.vae.decode(latents, return_dict=False)[0] - except oom_error as e: + except torch.OutOfMemoryError as e: warnings.warn( f"{e}. \n" f"Try to use VAE tiling for large images. For example: \n" diff --git a/src/diffusers/quantizers/gguf/gguf_quantizer.py b/src/diffusers/quantizers/gguf/gguf_quantizer.py index 42ea3982c912..e2bee2e4023d 100644 --- a/src/diffusers/quantizers/gguf/gguf_quantizer.py +++ b/src/diffusers/quantizers/gguf/gguf_quantizer.py @@ -18,6 +18,7 @@ is_torch_available, logging, ) +from ...utils.torch_utils import get_device if is_torch_available() and is_gguf_available(): @@ -177,12 +178,7 @@ def _dequantize(self, model): logger.info( "Model was found to be on CPU (could happen as a result of `enable_model_cpu_offload()`). So, moving it to accelerator. After dequantization, will move the model back to CPU again to preserve the previous device." ) - device = ( - torch.accelerator.current_accelerator() - if hasattr(torch, "accelerator") - else torch.cuda.current_device() - ) - model.to(device) + model.to(get_device()) model = _dequantize_gguf_and_restore_linear(model, self.modules_to_not_convert) if is_model_on_cpu: diff --git a/src/diffusers/training_utils.py b/src/diffusers/training_utils.py index 68c0c5254234..384a73e51296 100644 --- a/src/diffusers/training_utils.py +++ b/src/diffusers/training_utils.py @@ -37,6 +37,7 @@ is_torchvision_available, is_transformers_available, ) +from .utils.torch_utils import empty_device_cache if is_transformers_available(): @@ -412,15 +413,7 @@ def free_memory(): Runs garbage collection. Then clears the cache of the available accelerator. """ gc.collect() - - if torch.cuda.is_available(): - torch.cuda.empty_cache() - elif torch.backends.mps.is_available(): - torch.mps.empty_cache() - elif is_torch_npu_available(): - torch_npu.npu.empty_cache() - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - torch.xpu.empty_cache() + empty_device_cache() @contextmanager diff --git a/src/diffusers/utils/torch_utils.py b/src/diffusers/utils/torch_utils.py index 9f0877eb1d13..d1982a922f00 100644 --- a/src/diffusers/utils/torch_utils.py +++ b/src/diffusers/utils/torch_utils.py @@ -24,9 +24,7 @@ from . import logging from .import_utils import ( is_torch_available, - is_torch_mlu_available, is_torch_neuronx_available, - is_torch_npu_available, is_torch_version, ) @@ -39,71 +37,6 @@ import torch from torch.fft import fftn, fftshift, ifftn, ifftshift - BACKEND_SUPPORTS_TRAINING = { - "cuda": True, - "xpu": True, - "cpu": True, - "mps": False, - "neuron": False, - "default": True, - } - BACKEND_EMPTY_CACHE = { - "cuda": torch.cuda.empty_cache, - "xpu": torch.xpu.empty_cache, - "cpu": None, - "mps": torch.mps.empty_cache, - "neuron": None, - "default": None, - } - BACKEND_DEVICE_COUNT = { - "cuda": torch.cuda.device_count, - "xpu": torch.xpu.device_count, - "cpu": lambda: 0, - "mps": lambda: 0, - "neuron": lambda: getattr(getattr(torch, "neuron", None), "device_count", lambda: 0)(), - "default": 0, - } - BACKEND_MANUAL_SEED = { - "cuda": torch.cuda.manual_seed, - "xpu": torch.xpu.manual_seed, - "cpu": torch.manual_seed, - "mps": torch.mps.manual_seed, - "neuron": torch.manual_seed, - "default": torch.manual_seed, - } - BACKEND_RESET_PEAK_MEMORY_STATS = { - "cuda": torch.cuda.reset_peak_memory_stats, - "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None), - "cpu": None, - "mps": None, - "neuron": None, - "default": None, - } - BACKEND_RESET_MAX_MEMORY_ALLOCATED = { - "cuda": torch.cuda.reset_max_memory_allocated, - "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None), - "cpu": None, - "mps": None, - "neuron": None, - "default": None, - } - BACKEND_MAX_MEMORY_ALLOCATED = { - "cuda": torch.cuda.max_memory_allocated, - "xpu": getattr(torch.xpu, "max_memory_allocated", None), - "cpu": 0, - "mps": 0, - "neuron": 0, - "default": 0, - } - BACKEND_SYNCHRONIZE = { - "cuda": torch.cuda.synchronize, - "xpu": getattr(torch.xpu, "synchronize", None), - "cpu": None, - "mps": None, - "neuron": getattr(getattr(torch, "neuron", None), "synchronize", None), - "default": None, - } - _FP64_UNSUPPORTED_DEVICES = frozenset({"mps", "npu", "neuron"}) _INT64_UNSUPPORTED_DEVICES = frozenset({"mps", "npu", "neuron"}) _DTYPE_DOWNCAST = {torch.float64: torch.float32, torch.int64: torch.int32} @@ -119,62 +52,6 @@ def maybe_allow_in_graph(cls): return cls -# This dispatches a defined function according to the accelerator from the function definitions. -def _device_agnostic_dispatch(device: str, dispatch_table: dict[str, callable], *args, **kwargs): - if device not in dispatch_table: - return dispatch_table["default"](*args, **kwargs) - - fn = dispatch_table[device] - - # Some device agnostic functions return values. Need to guard against 'None' instead at - # user level - if not callable(fn): - return fn - - return fn(*args, **kwargs) - - -# These are callables which automatically dispatch the function specific to the accelerator -def backend_manual_seed(device: str, seed: int): - return _device_agnostic_dispatch(device, BACKEND_MANUAL_SEED, seed) - - -def backend_synchronize(device: str): - return _device_agnostic_dispatch(device, BACKEND_SYNCHRONIZE) - - -def backend_empty_cache(device: str): - return _device_agnostic_dispatch(device, BACKEND_EMPTY_CACHE) - - -def backend_device_count(device: str): - return _device_agnostic_dispatch(device, BACKEND_DEVICE_COUNT) - - -def backend_reset_peak_memory_stats(device: str): - return _device_agnostic_dispatch(device, BACKEND_RESET_PEAK_MEMORY_STATS) - - -def backend_reset_max_memory_allocated(device: str): - return _device_agnostic_dispatch(device, BACKEND_RESET_MAX_MEMORY_ALLOCATED) - - -def backend_max_memory_allocated(device: str): - return _device_agnostic_dispatch(device, BACKEND_MAX_MEMORY_ALLOCATED) - - -# These are callables which return boolean behaviour flags and can be used to specify some -# device agnostic alternative where the feature is unsupported. -def backend_supports_training(device: str): - if not is_torch_available(): - return False - - if device not in BACKEND_SUPPORTS_TRAINING: - device = "default" - - return BACKEND_SUPPORTS_TRAINING[device] - - def maybe_adjust_dtype_for_device(dtype: "torch.dtype", device: "torch.device") -> "torch.dtype": unsupported = _DTYPE_UNSUPPORTED_DEVICES.get(dtype) return _DTYPE_DOWNCAST[dtype] if unsupported and device.type in unsupported else dtype @@ -319,38 +196,136 @@ def get_torch_cuda_device_capability(): return None -@functools.lru_cache -def get_device(): - if torch.cuda.is_available(): - return "cuda" - elif is_torch_npu_available(): - return "npu" - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - return "xpu" - elif torch.backends.mps.is_available(): - return "mps" - elif is_torch_mlu_available(): - return "mlu" - elif is_torch_neuronx_available() and hasattr(torch, "neuron") and torch.neuron.is_available(): - return "neuron" - else: +class TorchDeviceBackend: + """ + A proxy for the `torch.` namespace (`torch.cuda`, `torch.xpu`, `torch.mps`, ...) of one device. Attributes + the class does not define are the module's own (`synchronize`, `device_count`, `Stream`, `current_stream`, ...); + the methods defined here override the operations whose availability differs between backends and need a fallback: + cache clearing, seeding and memory queries. With no `device`, detects the host accelerator through + `torch.accelerator`. Raises if torch has no module for the backend rather than silently falling back to + `torch.cuda`. + """ + + def __init__(self, device: str | torch.device | None = None): + self.device = torch.device(self._detect_device_type() if device is None else device) + self.module = torch.get_device_module(self.device.type) + + @staticmethod + @functools.lru_cache + def _detect_device_type() -> str: + if torch.accelerator.is_available(): + return torch.accelerator.current_accelerator().type + + # Neuron is XLA-based and never registers as a torch accelerator. + if is_torch_neuronx_available() and hasattr(torch, "neuron") and torch.neuron.is_available(): + return "neuron" return "cpu" + def empty_cache(self) -> None: + # Backends without a caching allocator (cpu, neuron) have nothing to clear. + empty_cache = getattr(self.module, "empty_cache", None) + if empty_cache is not None: + empty_cache() + + def manual_seed(self, seed: int) -> None: + # `torch.manual_seed` seeds every device, so it is the correct fallback for backends without their own. + manual_seed = getattr(self.module, "manual_seed", None) + if manual_seed is None: + torch.manual_seed(seed) + return + manual_seed(seed) + + def __getattr__(self, name: str): + # Proxy: anything not overridden here is `torch.`'s own attribute. + if name == "module": + raise AttributeError(name) + return getattr(self.module, name) + + def _accelerator_serves(self, min_torch_version: str) -> bool: + # `torch.accelerator` only serves the process accelerator, and its memory API arrived in 2.9 (statistics) and + # 2.10 (`get_memory_info`). + current = torch.accelerator.current_accelerator() + return is_torch_version(">=", min_torch_version) and current is not None and current.type == self.device.type + + def mem_get_info(self) -> tuple[int, int]: + """Free and total device memory in bytes.""" + mem_get_info = getattr(self.module, "mem_get_info", None) + if mem_get_info is not None: + return mem_get_info(self.device.index) + if self._accelerator_serves("2.10"): + return torch.accelerator.get_memory_info(self.device) + raise NotImplementedError( + f"`torch.{self.device.type}` does not implement `mem_get_info()`, and `torch.accelerator.get_memory_info()` " + f"cannot serve `{self.device}` on torch {torch.__version__} (requires torch>=2.10 and the current accelerator)." + ) + + def max_memory_allocated(self) -> int: + """Peak memory allocated on the device in bytes since the last reset; 0 where the backend keeps no statistics.""" + max_memory_allocated = getattr(self.module, "max_memory_allocated", None) + if max_memory_allocated is not None: + return max_memory_allocated(self.device.index) + if self._accelerator_serves("2.9"): + return torch.accelerator.max_memory_allocated(self.device) + logger.warning( + f"`torch.{self.device.type}` keeps no memory statistics on torch {torch.__version__}; " + "`max_memory_allocated()` returns 0." + ) + return 0 + + def reset_peak_memory_stats(self) -> None: + reset_peak_memory_stats = getattr(self.module, "reset_peak_memory_stats", None) + if reset_peak_memory_stats is not None: + reset_peak_memory_stats(self.device.index) + return + if self._accelerator_serves("2.9"): + torch.accelerator.reset_peak_memory_stats(self.device) + return + logger.warning( + f"`torch.{self.device.type}` keeps no memory statistics on torch {torch.__version__}; " + "`reset_peak_memory_stats()` is a no-op." + ) + + +def get_device() -> str: + return TorchDeviceBackend._detect_device_type() + def empty_device_cache(device_type: str | None = None): - if device_type is None: - device_type = get_device() - if device_type in ["cpu"]: - return - device_mod = getattr(torch, device_type, torch.cuda) - device_mod.empty_cache() - - -def device_synchronize(device_type: str | None = None): - if device_type is None: - device_type = get_device() - device_mod = getattr(torch, device_type, torch.cuda) - device_mod.synchronize() + TorchDeviceBackend(device_type).empty_cache() + + +# Function-style spellings of `TorchDeviceBackend` for test code. +def backend_manual_seed(device: str, seed: int): + TorchDeviceBackend(device).manual_seed(seed) + + +def backend_synchronize(device: str): + TorchDeviceBackend(device).synchronize() + + +def backend_empty_cache(device: str): + TorchDeviceBackend(device).empty_cache() + + +def backend_device_count(device: str): + return TorchDeviceBackend(device).device_count() + + +def backend_reset_peak_memory_stats(device: str): + TorchDeviceBackend(device).reset_peak_memory_stats() + + +def backend_reset_max_memory_allocated(device: str): + # `reset_max_memory_allocated` is CUDA's deprecated alias of `reset_peak_memory_stats`. + TorchDeviceBackend(device).reset_peak_memory_stats() + + +def backend_max_memory_allocated(device: str): + return TorchDeviceBackend(device).max_memory_allocated() + + +def backend_supports_training(device: str): + return str(device).split(":")[0] not in ("mps", "neuron") def enable_full_determinism(): diff --git a/tests/hooks/test_group_offloading.py b/tests/hooks/test_group_offloading.py index a903186aa6b4..b48446d356e0 100644 --- a/tests/hooks/test_group_offloading.py +++ b/tests/hooks/test_group_offloading.py @@ -334,14 +334,10 @@ def test_warning_logged_if_group_offloaded_pipe_moved_to_accelerator(self, caplo assert f"The module '{self.model.__class__.__name__}' is group offloaded" in caplog.text def test_error_raised_if_streams_used_and_no_accelerator_device(self): - torch_accelerator_module = getattr(torch, torch_device, torch.cuda) - original_is_available = torch_accelerator_module.is_available - torch_accelerator_module.is_available = lambda: False - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="requires an accelerator onload device"): self.model.enable_group_offload( - onload_device=torch.device(torch_device), offload_type="leaf_level", use_stream=True + onload_device=torch.device("cpu"), offload_type="leaf_level", use_stream=True ) - torch_accelerator_module.is_available = original_is_available def test_error_raised_if_supports_group_offloading_false(self): self.model._supports_group_offloading = False diff --git a/tests/modular_pipelines/testing_utils/utils.py b/tests/modular_pipelines/testing_utils/utils.py index 32a69f3dd0d7..8b2f1e71ac26 100644 --- a/tests/modular_pipelines/testing_utils/utils.py +++ b/tests/modular_pipelines/testing_utils/utils.py @@ -21,7 +21,7 @@ import torch from huggingface_hub import hf_hub_download -from ...testing_utils import torch_device +from diffusers.utils.torch_utils import TorchDeviceBackend def backend_memory_allocated(device: str) -> int: @@ -37,15 +37,13 @@ def backend_memory_allocated(device: str) -> int: def patch_free_memory(free_bytes: int, total_bytes: int = 80 * 1024): """ - Simulate `free_bytes` of free device memory on whichever backend module (cuda/xpu/...) backs `torch_device`. + Simulate `free_bytes` of free device memory on whichever backend backs `torch_device`. - `mem_get_info` returns `(free, total)` and is the single point where `AutoOffloadStrategy` learns how much memory - is available, so patching it makes offloading decisions deterministic instead of dependent on the real free memory - of the test hardware (an 80GB GPU never runs low on a handful of KB-sized models). + `TorchDeviceBackend.mem_get_info` returns `(free, total)` and is the single point where `AutoOffloadStrategy` learns how + much memory is available, so patching it makes offloading decisions deterministic instead of dependent on the real + free memory of the test hardware (an 80GB GPU never runs low on a handful of KB-sized models). """ - device_type = torch.device(torch_device).type - device_module = getattr(torch, device_type, torch.cuda) - return mock.patch.object(device_module, "mem_get_info", return_value=(free_bytes, total_bytes)) + return mock.patch.object(TorchDeviceBackend, "mem_get_info", return_value=(free_bytes, total_bytes)) def get_specified_components(path_or_repo_id, cache_dir=None): diff --git a/tests/others/test_utils.py b/tests/others/test_utils.py index d1a59cec52f1..a2b59855657f 100755 --- a/tests/others/test_utils.py +++ b/tests/others/test_utils.py @@ -269,6 +269,97 @@ def _capture(target_device): assert "moved to" in cuda_out, f"Non-MPS target should still emit the CPU-fallback info log, got: {cuda_out}" +class TestTorchDeviceBackend: + """Tests for :class:`diffusers.utils.torch_utils.TorchDeviceBackend` on the CPU backend.""" + + def test_module_resolves_from_str_and_torch_device(self): + import torch + + from diffusers.utils.torch_utils import TorchDeviceBackend + + assert TorchDeviceBackend("cpu").module is torch.cpu, "str device type should resolve to torch.cpu" + assert TorchDeviceBackend(torch.device("cpu")).module is torch.cpu, "torch.device should resolve to torch.cpu" + assert TorchDeviceBackend("cuda:1").module is torch.cuda, "device index should be ignored for the module" + + def test_get_device_matches_torch_accelerator(self): + import torch + + from diffusers.utils.torch_utils import get_device + + expected = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu" + assert get_device() == expected, "get_device() should report what torch.accelerator reports" + + def test_default_device_is_the_detected_accelerator(self): + import torch + + from diffusers.utils.torch_utils import TorchDeviceBackend, get_device + + backend = TorchDeviceBackend() + assert backend.device == torch.device(get_device()), "no-arg backend should target get_device()" + assert backend.module is getattr(torch, get_device()), "module should be the torch namespace of that device" + + def test_unknown_backend_raises_instead_of_falling_back_to_cuda(self): + import pytest + + from diffusers.utils.torch_utils import TorchDeviceBackend + + with pytest.raises(RuntimeError, match="does not have a corresponding module"): + TorchDeviceBackend("privateuseone") + + def test_cpu_operations_without_a_caching_allocator(self): + import torch + + from diffusers.utils.torch_utils import TorchDeviceBackend, empty_device_cache + + empty_device_cache("cpu") + backend = TorchDeviceBackend("cpu") + backend.empty_cache() + backend.manual_seed(1234) + assert torch.initial_seed() == 1234, "cpu manual_seed should fall back to torch.manual_seed" + + def test_function_style_backend_helpers_on_cpu(self): + from diffusers.utils import torch_utils + + torch_utils.backend_manual_seed("cpu", 0) + torch_utils.backend_synchronize("cpu") + torch_utils.backend_empty_cache("cpu") + assert isinstance(torch_utils.backend_device_count("cpu"), int) + assert torch_utils.backend_supports_training("cpu") is True + assert torch_utils.backend_supports_training("mps") is False + + def test_unsupported_operations_raise_naming_the_backend(self): + import pytest + import torch + + from diffusers.utils import torch_utils + from diffusers.utils.torch_utils import TorchDeviceBackend + + from ..testing_utils import CaptureLogger + + backend = TorchDeviceBackend("cpu") + with pytest.raises(NotImplementedError, match="torch.cpu"): + backend.mem_get_info() + assert backend.Stream is torch.cpu.Stream, "undefined attributes forward to the backend module" + backend.synchronize() + with pytest.raises(AttributeError, match="Stream"): + TorchDeviceBackend("mps").Stream() + with CaptureLogger(torch_utils.logger) as cl: + backend.reset_peak_memory_stats() + assert backend.max_memory_allocated() == 0, "cpu keeps no memory statistics" + assert "no memory statistics" in cl.out, f"no-op memory calls should warn, got: {cl.out}" + + def test_memory_statistics_on_the_host_accelerator(self): + import pytest + + from diffusers.utils.torch_utils import TorchDeviceBackend, get_device + + if get_device() == "cpu": + pytest.skip("memory statistics need an accelerator") + backend = TorchDeviceBackend() + backend.reset_peak_memory_stats() + assert backend.max_memory_allocated() >= 0 + + # Copied from https://github.com/huggingface/transformers/blob/main/tests/utils/test_expectations.py class TestExpectations: def test_expectations(self): diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 9da89e626198..cce6f15325b4 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -52,6 +52,7 @@ is_transformers_available, ) from diffusers.utils.logging import get_logger +from diffusers.utils.torch_utils import TorchDeviceBackend if is_torch_available(): @@ -1505,107 +1506,38 @@ def _is_torch_fp64_available(device): else None ) - # Function definitions - BACKEND_EMPTY_CACHE = { - "cuda": torch.cuda.empty_cache, - "xpu": torch.xpu.empty_cache, - "cpu": None, - "mps": torch.mps.empty_cache, - "default": None, - } - BACKEND_DEVICE_COUNT = { - "cuda": torch.cuda.device_count, - "xpu": torch.xpu.device_count, - "cpu": lambda: 0, - "mps": lambda: 0, - "default": 0, - } - BACKEND_MANUAL_SEED = { - "cuda": torch.cuda.manual_seed, - "xpu": torch.xpu.manual_seed, - "cpu": torch.manual_seed, - "mps": torch.mps.manual_seed, - "default": torch.manual_seed, - } - BACKEND_RESET_PEAK_MEMORY_STATS = { - "cuda": torch.cuda.reset_peak_memory_stats, - "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None), - "cpu": None, - "mps": None, - "default": None, - } - BACKEND_RESET_MAX_MEMORY_ALLOCATED = { - "cuda": torch.cuda.reset_max_memory_allocated, - "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None), - "cpu": None, - "mps": None, - "default": None, - } - BACKEND_MAX_MEMORY_ALLOCATED = { - "cuda": torch.cuda.max_memory_allocated, - "xpu": getattr(torch.xpu, "max_memory_allocated", None), - "cpu": 0, - "mps": 0, - "default": 0, - } - BACKEND_SYNCHRONIZE = { - "cuda": torch.cuda.synchronize, - "xpu": getattr(torch.xpu, "synchronize", None), - "cpu": None, - "mps": None, - "default": None, - } - if _neuron_device is not None: - BACKEND_EMPTY_CACHE[_neuron_device] = None - BACKEND_DEVICE_COUNT[_neuron_device] = torch.neuron.device_count - BACKEND_MANUAL_SEED[_neuron_device] = torch.manual_seed - BACKEND_RESET_PEAK_MEMORY_STATS[_neuron_device] = None - BACKEND_RESET_MAX_MEMORY_ALLOCATED[_neuron_device] = None - BACKEND_MAX_MEMORY_ALLOCATED[_neuron_device] = 0 - BACKEND_SYNCHRONIZE[_neuron_device] = torch.neuron.synchronize BACKEND_SUPPORTS_TRAINING[_neuron_device] = False -# This dispatches a defined function according to the accelerator from the function definitions. -def _device_agnostic_dispatch(device: str, dispatch_table: dict[str, Callable], *args, **kwargs): - fn = dispatch_table[device] if device in dispatch_table else dispatch_table["default"] - - # Some device agnostic functions return values. Need to guard against 'None' instead at - # user level - if not callable(fn): - return fn - - return fn(*args, **kwargs) - - -# These are callables which automatically dispatch the function specific to the accelerator +# Device operations go through `TorchDeviceBackend`. def backend_manual_seed(device: str, seed: int): - return _device_agnostic_dispatch(device, BACKEND_MANUAL_SEED, seed) + TorchDeviceBackend(device).manual_seed(seed) def backend_synchronize(device: str): - return _device_agnostic_dispatch(device, BACKEND_SYNCHRONIZE) + TorchDeviceBackend(device).synchronize() def backend_empty_cache(device: str): - return _device_agnostic_dispatch(device, BACKEND_EMPTY_CACHE) + TorchDeviceBackend(device).empty_cache() def backend_device_count(device: str): - return _device_agnostic_dispatch(device, BACKEND_DEVICE_COUNT) + return TorchDeviceBackend(device).device_count() def backend_reset_peak_memory_stats(device: str): - return _device_agnostic_dispatch(device, BACKEND_RESET_PEAK_MEMORY_STATS) + TorchDeviceBackend(device).reset_peak_memory_stats() def backend_reset_max_memory_allocated(device: str): - return _device_agnostic_dispatch(device, BACKEND_RESET_MAX_MEMORY_ALLOCATED) + # `reset_max_memory_allocated` is CUDA's deprecated alias of `reset_peak_memory_stats`. + TorchDeviceBackend(device).reset_peak_memory_stats() def backend_max_memory_allocated(device: str): - return _device_agnostic_dispatch(device, BACKEND_MAX_MEMORY_ALLOCATED) + return TorchDeviceBackend(device).max_memory_allocated() # These are callables which return boolean behaviour flags and can be used to specify some @@ -1659,14 +1591,9 @@ def update_mapping_from_spec(device_fn_dict: dict[str, Callable], attribute_name torch_device = device_name - # Add one entry here for each `BACKEND_*` dictionary. - update_mapping_from_spec(BACKEND_MANUAL_SEED, "MANUAL_SEED_FN") - update_mapping_from_spec(BACKEND_EMPTY_CACHE, "EMPTY_CACHE_FN") - update_mapping_from_spec(BACKEND_DEVICE_COUNT, "DEVICE_COUNT_FN") + # `SUPPORTS_TRAINING` is the only per-device table left. Device operations come from the backend's own + # `torch.` module through `TorchDeviceBackend`, so a spec file no longer supplies them. update_mapping_from_spec(BACKEND_SUPPORTS_TRAINING, "SUPPORTS_TRAINING") - update_mapping_from_spec(BACKEND_RESET_PEAK_MEMORY_STATS, "RESET_PEAK_MEMORY_STATS_FN") - update_mapping_from_spec(BACKEND_RESET_MAX_MEMORY_ALLOCATED, "RESET_MAX_MEMORY_ALLOCATED_FN") - update_mapping_from_spec(BACKEND_MAX_MEMORY_ALLOCATED, "MAX_MEMORY_ALLOCATED_FN") # Modified from https://github.com/huggingface/transformers/blob/cdfb018d0300fef3b07d9220f3efe9c2a9974662/src/transformers..testing_utils.py#L3090