MLU backend utilities fail despite MLU device detection

[SeptPonts]

Describe the bug

MLU device detection was added in #12629, but the backend utility tables, training_utils.free_memory(), and test-device selection still omit MLU.

On an MLU machine, get_device() returns "mlu", while backend_empty_cache("mlu") raises TypeError. The same problem affects synchronization, both peak-memory reset helpers, device count, and peak allocated-memory queries. An existing caller is the SD3 ControlNet training example, which passes accelerator.device.type to backend_empty_cache after prompt encoding.

I would like to submit a small follow-up to #12629 covering three files:

• src/diffusers/utils/torch_utils.py: register MLU operations in the existing backend tables.
• src/diffusers/training_utils.py: release unused MLU allocator cache in free_memory().
• tests/testing_utils.py: select an available MLU automatically and register its test backend operations.

The existing H3 training test runs successfully on MLU when the device is selected explicitly. The missing test-device selection currently makes it run on CPU by default on this machine.

Reproduction

Run with Diffusers and a working torch_mlu installation on an MLU machine:

import torch
import torch_mlu

from diffusers.utils.torch_utils import backend_empty_cache, get_device

print(torch.mlu.is_available())
print(get_device())
backend_empty_cache("mlu")

Logs

True
mlu
TypeError: 'NoneType' object is not callable

I compared the baseline with a local patch on the same device:

Check	Baseline	Local patch
Six production cache/sync/reset/count/memory helpers with "mlu"	All raise TypeError	All succeed; count and allocated-memory results match direct MLU calls
free_memory() after allocating and deleting a 64 MiB MLU tensor	Allocated: 0; reserved: 67,108,864 bytes	Allocated: 0; reserved: 0 bytes
Existing TestMiniMaxH3TransformerTraining::test_training, default device	Passes on CPU	Passes on mlu:0
Same baseline test with DIFFUSERS_TEST_DEVICE=mlu	Passes on mlu:0	—

For the cache check, calling torch.mlu.empty_cache() directly also releases the reserved memory on the baseline. The SD3 ControlNet example above is a source-level caller of the reproduced failing helper; I have not run that complete training example.

System Info

• Device: Cambricon MLU590-M9DK, one visible device.
• PyTorch: 2.12.1+cpu; torch_mlu: 1.34.1+torch2.12.1.
• Accelerate: 1.14.0; Transformers: 5.14.1; PEFT: 0.18.1.
• Tested Diffusers baseline: c5469b7ceb606edd7ba6570dcd17d38590a18db6.
• Also checked current main, 759164b7ad116e091e9d3e222211c9aa27d835f6: these three files are unchanged from the tested baseline.

Who can help?

@sayakpaul, would this scope be welcome as a follow-up PR? This is AI-assisted work, and I am opening the issue first to coordinate as requested in CONTRIBUTING.md. I checked related open and closed PRs; #13824 addresses Z-Image RoPE compatibility and #14039 adds TPU support, neither covers these MLU utility gaps.

[L4XB]

I have a fix ready for this and, per CONTRIBUTING, I am asking for a maintainer ack on the issue before opening the PR.

Confirmed the diagnosis and found three more places with the same gap that the report does not name.

Why it raises rather than falling back. _device_agnostic_dispatch in utils/torch_utils.py reaches for the "default" entry when a device is missing from a table and calls it without the callable check it applies to known devices. So BACKEND_EMPTY_CACHE["default"] is None → TypeError: 'NoneType' object is not callable, and BACKEND_DEVICE_COUNT["default"] is the literal 0 → TypeError: 'int' object is not callable. Every table there has keys cuda / xpu / cpu / mps / neuron / default, and get_device() can return mlu.

BACKEND_MANUAL_SEED is the quiet one: it does not raise, because "default" is torch.manual_seed. It just seeds the CPU generator only, so an MLU run is silently unseeded. BACKEND_SUPPORTS_TRAINING already resolves True via the default; registering it explicitly changes nothing.

Three more, beyond the tables:

1. training_utils.free_memory() walks cuda → mps → npu → xpu and stops, so the MLU allocator cache is never released — which matches your measurement of reserved staying at 67,108,864 bytes.
2. empty_device_cache() and device_synchronize() resolve the module with getattr(torch, device_type, torch.cuda). torch.mlu only exists once torch_mlu has been imported, and diffusers never imports it — so on a process that had not imported it by hand, these two silently ran on CUDA. Importing torch_mlu under the availability guard (the same thing training_utils.py already does for torch_npu) closes that.
3. tests/testing_utils.py keeps its own copy of the seven tables and its own _device_agnostic_dispatch, which does apply the callable check to the default entry — so there the helpers do not raise, they quietly return None/0. Its device selection goes cuda → xpu → neuron → cpu, which is why your MLU box runs the suite on CPU.

The registration block would mirror the existing if _neuron_device is not None: block in tests/testing_utils.py, and matches how transformers registers the same backend (if is_torch_mlu_available(): BACKEND_EMPTY_CACHE["mlu"] = torch.mlu.empty_cache, …).

Three things I deliberately would not touch, and would like your call on:

• src/diffusers/utils/testing_utils.py carries a third copy of these tables. It emits a FutureWarning on import, is scheduled for removal in 1.0.0, and its deprecation message points at diffusers.utils.torch_utils. It is already missing the neuron rows that tests/testing_utils.py has, so adding MLU there deepens a copy on its way out. Say the word if you would rather keep the three in step.
• training_utils.set_seed() has an NPU branch and otherwise calls torch.cuda.manual_seed_all, a no-op without CUDA — so on MLU the accelerator RNG is left unseeded. One line to fix, but I have no MLU to confirm torch.mlu.manual_seed_all exists, and guessing wrong turns a silent no-op into an AttributeError.
• NPU has the identical gap in utils/torch_utils.py: get_device() can return "npu" and no table has an npu key. Different vendor API (torch_npu.npu.*, not torch.npu.*, as free_memory shows) and I cannot test it, so it belongs in a separate change by someone who can.
• _FP64_UNSUPPORTED_DEVICES / _INT64_UNSUPPORTED_DEVICES list mps, npu, neuron. Whether MLU belongs there is a hardware question I cannot answer.

On evidence: I do not have MLU hardware, so I cannot paste a real before/after from one. What I can show is the dispatch path and a torch.mlu-shaped stub proving each helper resolves to the right callable after registration, plus the unchanged behaviour on CPU/CUDA. If you would rather have this confirmed on real hardware first, the reporter is in a position to run it — happy to wait for that before the PR.

Let me know if the shape is right and I will open it.

[L4XB]

Correcting myself: I said three of those were places "the report does not name", and two of them are
in fact in your own three-file plan — training_utils.free_memory() and tests/testing_utils.py are
both listed in the issue. Your scope was right; I misread it. Sorry.

What is actually beyond it is one thing, plus one detail inside the tables:

• empty_device_cache() and device_synchronize() in utils/torch_utils.py resolve the module with
  getattr(torch, device_type, torch.cuda). Since diffusers never imports torch_mlu, torch.mlu
  is absent unless the caller imported it by hand — so those two silently ran on CUDA rather than
  raising. Importing torch_mlu under the availability guard (what training_utils.py already does
  for torch_npu) fixes that as a side effect of your fix.
• BACKEND_MANUAL_SEED is the quiet one in the tables: it does not raise, because "default" is
  torch.manual_seed, so an MLU run just leaves the accelerator generator unseeded.

Everything else in my earlier comment stands, including the two I would not touch without hardware
(set_seed()'s missing torch.mlu.manual_seed_all, and the identical NPU gap) and the third,
deprecated copy of the tables in src/diffusers/utils/testing_utils.py.

One thing I should have said plainly the first time: I do not have an MLU, so four of the eight
torch.mlu.* names in the branch — synchronize, reset_peak_memory_stats,
reset_max_memory_allocated, max_memory_allocated — rest entirely on your measurement. The other
four match how transformers registers the same backend. If a maintainer would rather have those four
written as getattr(torch.mlu, ..., None), which is what the tables already do for the XPU
equivalents, that is a one-line change.

[sayakpaul]

@SeptPonts the changes look reasonable to me. Feel free to open the PR. @L4XB thanks for offering to help, however, since @SeptPonts already laid down a solid approach, I think they deserve the chance better. I hope we will have future opportunities for contributions.

[L4XB]

Completely fair, and the right call — @SeptPonts did the measurement on real hardware, which is the part I could not do. Looking forward to the PR.

Everything I found is in the thread above for them to take or leave, in particular the two beyond the three-file scope: empty_device_cache() / device_synchronize() resolving through getattr(torch, device_type, torch.cuda) and therefore silently running on CUDA unless something has already imported torch_mlu, and BACKEND_MANUAL_SEED falling through to torch.manual_seed so the accelerator generator is never seeded. Neither raises, which is why they are easy to miss next to the six that do.

@SeptPonts — happy to answer anything about those two if it helps, and no need to credit me.

[SeptPonts]

Thanks @sayakpaul! I'll open a PR with the scope outlined above and include the before/after validation results from MLU hardware.
Thanks @L4XB for taking a look and sharing your findings.