#14779·diffusers

MLU backend utilities fail despite MLU device detection

Author: SeptPontsCreated Sep 15, 2026Updated Sep 16, 2026
Labelsbugmodelsneeds-env-info

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:

python
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.