From be1ba28b7f9ddf99d111a3fbcc6ab8e63b35b040 Mon Sep 17 00:00:00 2001 From: Qitong Li Date: Tue, 8 Sep 2026 12:40:38 +0800 Subject: [PATCH 1/2] Add MLU device utilities and test backend support --- src/diffusers/training_utils.py | 6 ++++++ src/diffusers/utils/torch_utils.py | 8 ++++++++ tests/testing_utils.py | 12 +++++++++++- 3 files changed, 25 insertions(+), 1 deletion(-) diff --git a/src/diffusers/training_utils.py b/src/diffusers/training_utils.py index 68c0c5254234..f22191750e85 100644 --- a/src/diffusers/training_utils.py +++ b/src/diffusers/training_utils.py @@ -34,6 +34,7 @@ is_accelerate_available, is_peft_available, is_torch_npu_available, + is_torch_mlu_available, is_torchvision_available, is_transformers_available, ) @@ -57,6 +58,9 @@ if is_torch_npu_available(): import torch_npu # noqa: F401 +if is_torch_mlu_available(): + import torch_mlu # noqa: F401 + def set_seed(seed: int): """ @@ -419,6 +423,8 @@ def free_memory(): torch.mps.empty_cache() elif is_torch_npu_available(): torch_npu.npu.empty_cache() + elif is_torch_mlu_available(): + torch.mlu.empty_cache() elif hasattr(torch, "xpu") and torch.xpu.is_available(): torch.xpu.empty_cache() diff --git a/src/diffusers/utils/torch_utils.py b/src/diffusers/utils/torch_utils.py index 9f0877eb1d13..8aec71a8a7f9 100644 --- a/src/diffusers/utils/torch_utils.py +++ b/src/diffusers/utils/torch_utils.py @@ -41,6 +41,7 @@ BACKEND_SUPPORTS_TRAINING = { "cuda": True, + "mlu": True, "xpu": True, "cpu": True, "mps": False, @@ -49,6 +50,7 @@ } BACKEND_EMPTY_CACHE = { "cuda": torch.cuda.empty_cache, + "mlu": getattr(getattr(torch, "mlu", None), "empty_cache", None), "xpu": torch.xpu.empty_cache, "cpu": None, "mps": torch.mps.empty_cache, @@ -57,6 +59,7 @@ } BACKEND_DEVICE_COUNT = { "cuda": torch.cuda.device_count, + "mlu": lambda: getattr(getattr(torch, "mlu", None), "device_count", lambda: 0)(), "xpu": torch.xpu.device_count, "cpu": lambda: 0, "mps": lambda: 0, @@ -65,6 +68,7 @@ } BACKEND_MANUAL_SEED = { "cuda": torch.cuda.manual_seed, + "mlu": getattr(getattr(torch, "mlu", None), "manual_seed", torch.manual_seed), "xpu": torch.xpu.manual_seed, "cpu": torch.manual_seed, "mps": torch.mps.manual_seed, @@ -73,6 +77,7 @@ } BACKEND_RESET_PEAK_MEMORY_STATS = { "cuda": torch.cuda.reset_peak_memory_stats, + "mlu": getattr(getattr(torch, "mlu", None), "reset_peak_memory_stats", None), "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None), "cpu": None, "mps": None, @@ -81,6 +86,7 @@ } BACKEND_RESET_MAX_MEMORY_ALLOCATED = { "cuda": torch.cuda.reset_max_memory_allocated, + "mlu": getattr(getattr(torch, "mlu", None), "reset_peak_memory_stats", None), "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None), "cpu": None, "mps": None, @@ -89,6 +95,7 @@ } BACKEND_MAX_MEMORY_ALLOCATED = { "cuda": torch.cuda.max_memory_allocated, + "mlu": getattr(getattr(torch, "mlu", None), "max_memory_allocated", 0), "xpu": getattr(torch.xpu, "max_memory_allocated", None), "cpu": 0, "mps": 0, @@ -97,6 +104,7 @@ } BACKEND_SYNCHRONIZE = { "cuda": torch.cuda.synchronize, + "mlu": getattr(getattr(torch, "mlu", None), "synchronize", None), "xpu": getattr(torch.xpu, "synchronize", None), "cpu": None, "mps": None, diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 9da89e626198..ef9a9c30eaa6 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -45,6 +45,7 @@ is_sdnq_available, is_timm_available, is_torch_available, + is_torch_mlu_available, is_torch_neuronx_available, is_torch_version, is_torchao_available, @@ -100,6 +101,8 @@ else: if torch.cuda.is_available(): torch_device = "cuda" + elif is_torch_mlu_available() and hasattr(torch, "mlu") and torch.mlu.is_available(): + torch_device = "mlu" elif torch.xpu.is_available(): torch_device = "xpu" elif is_torch_neuronx_available() and hasattr(torch, "neuron") and torch.neuron.is_available(): @@ -1494,7 +1497,7 @@ def _is_torch_fp64_available(device): # Guard these lookups for when Torch is not used - alternative accelerator support is for PyTorch if is_torch_available(): # Behaviour flags - BACKEND_SUPPORTS_TRAINING = {"cuda": True, "xpu": True, "cpu": True, "mps": False, "default": True} + BACKEND_SUPPORTS_TRAINING = {"cuda": True, "mlu": True, "xpu": True, "cpu": True, "mps": False, "default": True} # Neuron device key: torch.neuron.current_device() returns an int (e.g. 0). # We capture it once at import time if torch_neuronx is available so we can add it @@ -1508,6 +1511,7 @@ def _is_torch_fp64_available(device): # Function definitions BACKEND_EMPTY_CACHE = { "cuda": torch.cuda.empty_cache, + "mlu": getattr(getattr(torch, "mlu", None), "empty_cache", None), "xpu": torch.xpu.empty_cache, "cpu": None, "mps": torch.mps.empty_cache, @@ -1515,6 +1519,7 @@ def _is_torch_fp64_available(device): } BACKEND_DEVICE_COUNT = { "cuda": torch.cuda.device_count, + "mlu": lambda: getattr(getattr(torch, "mlu", None), "device_count", lambda: 0)(), "xpu": torch.xpu.device_count, "cpu": lambda: 0, "mps": lambda: 0, @@ -1522,6 +1527,7 @@ def _is_torch_fp64_available(device): } BACKEND_MANUAL_SEED = { "cuda": torch.cuda.manual_seed, + "mlu": getattr(getattr(torch, "mlu", None), "manual_seed", torch.manual_seed), "xpu": torch.xpu.manual_seed, "cpu": torch.manual_seed, "mps": torch.mps.manual_seed, @@ -1529,6 +1535,7 @@ def _is_torch_fp64_available(device): } BACKEND_RESET_PEAK_MEMORY_STATS = { "cuda": torch.cuda.reset_peak_memory_stats, + "mlu": getattr(getattr(torch, "mlu", None), "reset_peak_memory_stats", None), "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None), "cpu": None, "mps": None, @@ -1536,6 +1543,7 @@ def _is_torch_fp64_available(device): } BACKEND_RESET_MAX_MEMORY_ALLOCATED = { "cuda": torch.cuda.reset_max_memory_allocated, + "mlu": getattr(getattr(torch, "mlu", None), "reset_peak_memory_stats", None), "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None), "cpu": None, "mps": None, @@ -1543,6 +1551,7 @@ def _is_torch_fp64_available(device): } BACKEND_MAX_MEMORY_ALLOCATED = { "cuda": torch.cuda.max_memory_allocated, + "mlu": getattr(getattr(torch, "mlu", None), "max_memory_allocated", 0), "xpu": getattr(torch.xpu, "max_memory_allocated", None), "cpu": 0, "mps": 0, @@ -1550,6 +1559,7 @@ def _is_torch_fp64_available(device): } BACKEND_SYNCHRONIZE = { "cuda": torch.cuda.synchronize, + "mlu": getattr(getattr(torch, "mlu", None), "synchronize", None), "xpu": getattr(torch.xpu, "synchronize", None), "cpu": None, "mps": None, From 21e73135c1a2ce24f32783fecf34c687c1dec39e Mon Sep 17 00:00:00 2001 From: Qitong Li Date: Wed, 16 Sep 2026 16:57:04 +0800 Subject: [PATCH 2/2] Fix MLU availability import ordering --- src/diffusers/training_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/training_utils.py b/src/diffusers/training_utils.py index f22191750e85..4610a14e3341 100644 --- a/src/diffusers/training_utils.py +++ b/src/diffusers/training_utils.py @@ -33,8 +33,8 @@ deprecate, is_accelerate_available, is_peft_available, - is_torch_npu_available, is_torch_mlu_available, + is_torch_npu_available, is_torchvision_available, is_transformers_available, )