huggingface / huggingface/diffusers

IPNDMScheduler and KDPM2DiscreteScheduler return all-NaN on MPS

Aperta
#14,368 3 commenti 0 reazioni 0 assegnatari Vedi su GitHub
bug schedulers
Lingua principale
Python
Stelle
34.5k
Fork
7.3k
Merge medio
3g 3h
PR unite (30g)
91

Descrizione

### Describe the bug

A CPU-vs-MPS differential sweep over the schedulers (same seed, same inputs, 36 classes) found two that return entirely NaN output on MPS while being correct on CPU. There is no error or warning — a user on Apple Silicon gets a black image.

```
IPNDMScheduler cpu: finite, mean -0.0586 mps: 4096/4096 NaN
KDPM2DiscreteScheduler cpu: finite, mean -1.2485 mps: 4096/4096 NaN
```

They have **different causes** — I have fully diagnosed the first and not the second.

### 1. IPNDMScheduler — diagnosed, fix verified

`set_timesteps` moves `timesteps` to the device but leaves the coefficient tables on the CPU:

```python
# scheduling_ipndm.py
self.betas = torch.sin(steps * math.pi / 2) ** 2 # CPU
self.alphas = (1.0 - self.betas**2) ** 0.5 # CPU
self.timesteps = timesteps.to(device) # only this moves
```

`step()` then does `sample * self.betas[i]` — an MPS tensor times a 0-dim CPU tensor with a non-zero storage offset. That combination silently returns the wrong value on MPS; I filed the upstream bug as pytorch/pytorch#191929. From step 1 onward the scheduler multiplies by the wrong coefficient and diverges into NaN.

Moving the two tables onto the device makes MPS match CPU exactly:

```
mps as-is : nan=4096
mps + moving alphas/betas to device : nan=0, matches cpu, maxdiff=0.00e+00
```

### 2. KDPM2DiscreteScheduler — confirmed, cause not established

Same symptom, but the NaN appears **only at the final step** (`t=0.0`), and moving `alphas`, `betas` and `sigmas` onto the device does **not** fix it. I have not isolated the cause and I am not assuming it is the same one — filing it here because it was found by the same sweep and is the same user-visible failure.

Incidentally `sigmas_interpol[0]` differs between backends (`nan` on CPU, `0.0` on MPS). #14213 notes that element is never read by this scheduler, so it is probably incidental rather than the cause.

### Reproduction

```python
import torch
from diffusers import IPNDMScheduler, KDPM2DiscreteScheduler

def run(cls, device):
torch.manual_seed(0)
s = cls(); s.set_timesteps(6, device=device)
out = torch.randn(1, 4, 32, 32, generator=torch.Generator().manual_seed(1)).to(device)
for t in s.timesteps:
out = s.step(torch.full_like(out, 0.05), t, out).prev_sample
return out.float().cpu()

for cls in (IPNDMScheduler, KDPM2DiscreteScheduler):
cpu, mps = run(cls, "cpu"), run(cls, "mps")
print(f"{cls.__name__:24} cpu nan={torch.isnan(cpu).sum().item():5d} mps nan={torch.isnan(mps).sum().item():5d}")
```

### Scope

The sweep covered 36 schedulers: 28 matched between CPU and MPS, 5 could not run on CPU with my synthetic inputs and were skipped, and 3 diverged — the two above plus `DPMSolverSDEScheduler`, which stays finite and differs only in magnitude. That one is stochastic and uses a Brownian noise sampler, so a backend RNG difference is expected; I have not investigated it and do not think it belongs here.

Not tested: whether either scheduler is reachable on MPS through a full pipeline run, and whether `npu`/`neuron` behave the same way.

### System Info

- diffusers 0.40.0.dev0 (`main`, `a8345366e`), torch 2.13.0
- macOS 26.0.1, arm64 (Apple Silicon), Python 3.11

### Who can help?

@yiyixuxu @dg845

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

For IPNDMScheduler, start in scheduling_ipndm.py and inspect how set_timesteps and step use the coefficient tables and device. Run the supplied CPU-versus-MPS reproduction to verify finite, matching output after the diagnosed issue is addressed. Then trace KDPM2DiscreteScheduler's final t=0.0 step separately and confirm both schedulers no longer return NaNs on MPS.

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

Valutazione

Stack tecnologico
python, pytorch
Ambito
machine-learning
Tipo di issue
Bug
Difficoltà
4/5
Tempo stimato
3-5 giorni
Stato di attività
Tranquilla
Chiarezza
Abbastanza chiara
Idoneità per principianti
50/100

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.