Skip to content

Fix MLU device utilities and test backend support #14791

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
SeptPonts wants to merge 2 commits into
base: main
Choose a base branch
from
+25 −1
Open

Fix MLU device utilities and test backend support #14791

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
be1ba28
Add MLU device utilities and test backend support
SeptPonts Sep 8, 2026
21e7313
Fix MLU availability import ordering
SeptPonts Sep 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions src/diffusers/training_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
deprecate,
is_accelerate_available,
is_peft_available,
is_torch_mlu_available,
is_torch_npu_available,
is_torchvision_available,
is_transformers_available,
Expand All @@ -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):
"""
Expand Down Expand Up @@ -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()

Expand Down
8 changes: 8 additions & 0 deletions src/diffusers/utils/torch_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@

BACKEND_SUPPORTS_TRAINING = {
"cuda": True,
"mlu": True,
"xpu": True,
"cpu": True,
"mps": False,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down
12 changes: 11 additions & 1 deletion tests/testing_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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
Expand All @@ -1508,48 +1511,55 @@ 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),

@sayakpaul sayakpaul Sep 17, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this the case? is empty_cache not available generally in Torch MLU?

@SeptPonts SeptPonts Sep 17, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, empty_cache() is available in torch_mlu, and I verified it on hardware.

The guard is for the optional backend namespace: unlike the built-in torch.cuda and torch.xpu modules, torch.mlu is registered by the external torch_mlu package through PyTorch’s PrivateUse1 integration. Without the extension loaded, accessing torch.mlu.empty_cache while constructing this module-level dictionary raises AttributeError: module 'torch' has no attribute 'mlu'. This prevents the shared utilities from being imported even for CPU/CUDA tests, before any MLU operation is requested.

I followed the guarded-access pattern already used for Neuron’s device_count and synchronize entries in src/diffusers/utils/torch_utils.py. The same reasoning applies to the device_count comment below.

"xpu": torch.xpu.empty_cache,
"cpu": None,
"mps": torch.mps.empty_cache,
"default": None,
}
BACKEND_DEVICE_COUNT = {
"cuda": torch.cuda.device_count,
"mlu": lambda: getattr(getattr(torch, "mlu", None), "device_count", lambda: 0)(),

@sayakpaul sayakpaul Sep 17, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same as above.

"xpu": torch.xpu.device_count,
"cpu": lambda: 0,
"mps": lambda: 0,
"default": 0,
}
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,
"default": torch.manual_seed,
}
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,
"default": None,
}
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,
"default": None,
}
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,
"default": 0,
}
BACKEND_SYNCHRONIZE = {
"cuda": torch.cuda.synchronize,
"mlu": getattr(getattr(torch, "mlu", None), "synchronize", None),
"xpu": getattr(torch.xpu, "synchronize", None),
"cpu": None,
"mps": None,
Expand Down
Loading