huggingface / huggingface/diffusers

MLU backend utilities fail despite MLU device detection

Aperta
#14,779 5 commenti 0 reazioni 0 assegnatari Vedi su GitHub
bug models needs-env-info
Lingua principale
Python
Stelle
34.5k
Fork
7.3k
Merge medio
3g 3h
PR unite (30g)
91

Descrizione

### 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](https://github.com/huggingface/diffusers/blob/c5469b7ceb606edd7ba6570dcd17d38590a18db6/examples/controlnet/train_controlnet_sd3.py#L122-L124), 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

```text
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.

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

Start with the backend tables in src/diffusers/utils/torch_utils.py and reproduce the failing backend_empty_cache("mlu") call on an MLU machine. Then inspect free_memory() in src/diffusers/training_utils.py and device selection in tests/testing_utils.py. Done means MLU cache, synchronization, reset, count, and memory helpers work, free_memory() releases its cache, and TestMiniMaxH3TransformerTraining selects MLU when available.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Valutazione

Stack tecnologico
python, pytorch
Ambito
backend, machine-learning, testing-qa
Tipo di issue
Bug
Difficoltà
3/5
Tempo stimato
1-2 giorni
Stato di attività
Attiva
Chiarezza
Specificata chiaramente
Idoneità per principianti
74/100

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.