IPNDMScheduler and KDPM2DiscreteScheduler return all-NaN on MPS

[4ktLuffy]

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:

# 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

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

[4ktLuffy]

Closing the "not tested" item from the issue: both schedulers do reach this through a full pipeline run, and the result is all-NaN latents.

import torch
from diffusers import StableDiffusionPipeline, DDIMScheduler, IPNDMScheduler, KDPM2DiscreteScheduler

pipe = StableDiffusionPipeline.from_pretrained(
    "diffusers/tiny-stable-diffusion-torch", safety_checker=None, requires_safety_checker=False)
cfg = pipe.scheduler.config

def run(cls, device):
    pipe.scheduler = cls.from_config(cfg)
    return pipe.to(device)(prompt="a cat", num_inference_steps=4, guidance_scale=2.0,
                           generator=torch.Generator("cpu").manual_seed(0),
                           output_type="latent").images.float().cpu()

for cls in (DDIMScheduler, IPNDMScheduler, KDPM2DiscreteScheduler):
    c, m = run(cls, "cpu"), run(cls, "mps")
    print(f"{cls.__name__:24} cpu_nan={torch.isnan(c).sum().item():5d} mps_nan={torch.isnan(m).sum().item():5d}")

DDIMScheduler            cpu_nan=    0 mps_nan=    0    maxdiff=1.05e-05
IPNDMScheduler           cpu_nan=    0 mps_nan=16384    maxdiff=nan
KDPM2DiscreteScheduler   cpu_nan=    0 mps_nan=16384    maxdiff=nan

So on Apple Silicon this is a black image with no error, not just a scheduler-level numerical difference. DDIM through the same pipeline is fine, which isolates it to the two schedulers rather than to the pipeline or the models.

Two further data points from the same sweep:

• I ran a CPU-vs-MPS differential over UNet2DModel, UNet2DConditionModel, AutoencoderKL, AutoencoderTiny and VQModel with identical seeded inputs. All matched within float32 noise (maxdiff 2e-7 to 3e-6), so the models are not implicated.
• The upstream cause I filed for IPNDM (pytorch#191929) turned out to duplicate pytorch#187117, which was fixed on main by pytorch#187229 on 2026-06-16. That fix is not in torch 2.13.0 — the release branch diverged before it landed — so the CPU-coefficient path still misbehaves on every current release. Moving alphas/betas onto the device in set_timesteps fixes IPNDM regardless of torch version, and is what I verified earlier in this issue.

KDPM2DiscreteScheduler remains unexplained; the device move does not fix it and its NaN appears only at the final step.

[4ktLuffy]

Root cause for the KDPM2DiscreteScheduler half of this issue, which I had left unresolved above.

Where it comes from. scheduling_k_dpm_2_discrete.py:327:

sigmas_interpol = sigmas.log().lerp(sigmas.roll(1).log(), 0.5).exp()

The sigma schedule ends in 0.0, so sigmas.log() contains -inf, and the lerp is asked to interpolate from an infinity. That operation is not consistent across backends:

import torch
for dev in ("cpu", "mps"):
    a = torch.tensor([float("-inf")], device=dev)
    b = torch.tensor([-1.0], device=dev)
    print(dev, a.lerp(b, 0.5).item())
# cpu -inf
# mps nan

CPU and MPS use algebraically equivalent but numerically different formulas (b - (b-a)*(1-w) vs a + w*(b-a)), which agree for finite inputs and diverge at infinities. This is pytorch/pytorch#111374 — open since 2023-10-16, labelled triaged/module: mps, and its own description notes it "breaks several Diffusion schedulers". Since it has been open for over two years, waiting for it upstream does not seem practical.

The consequence in the scheduler. After .exp(), the trailing entries differ:

sigmas_interpol[-4:]   cpu: [0.142478, 0.0, 0.0, 0.0]
                       mps: [0.142478, nan, nan, nan]

The final step is 1st order, so its divisor sigma_hat = 0.0291675 is identical on both. What differs is dt = sigma_interpol - sigma_hat, which is nan on MPS, so prev_sample = sample + derivative * dt is nan and the whole latent is lost on the last step.

Verified as the sole cause. Replacing only the nan entries of sigmas_interpol with 0.0 — i.e. making MPS agree with CPU, and a no-op on CPU — makes MPS reproduce CPU exactly:

cpu  unpatched : nan=0
cpu  patched   : identical to unpatched   (confirms CPU already holds 0.0 there)
mps  unpatched : nan=128
mps  patched   : nan=0,  matches cpu, maxdiff=2.38e-07

Scope. Only this scheduler. KDPM2AncestralDiscreteScheduler lerps sigmas_down.log() rather than sigmas.roll(1).log() and produces no nan on either device; HeunDiscreteScheduler has no sigmas_interpol. So this is one site, not a pattern.

On a fix. The diagnosis above is verified; the patch is not, so treating it as a direction rather than a proposal. Computing the interpolated sigma without ever taking log(0) — for example a geometric mean guarded on the zero entries — gives identical, nan-free results on both backends in a standalone check. I did not wire that into set_timesteps and run the suite, so I would not want it taken as a tested change. Happy to prepare it properly if that is the direction you would prefer.

This is independent of the IPNDMScheduler half, which is the CPU-coefficient/storage_offset path and is fixed by moving alphas/betas onto the device.

[4ktLuffy]

Update — tested against torch nightly (2.14.0.dev20260802):

• IPNDMScheduler: clean on nightly (0/4096 NaN). Its root cause — the CPU-scalar storage_offset bug — is fixed upstream somewhere between 2.13.0 and current nightly, so this half resolves on its own with the next torch release.
• KDPM2DiscreteScheduler: still 4096/4096 NaN on nightly. The lerp-at--inf divergence (Tensor.lerp inconsistent when using -Infinity between MPS and CPU pytorch/pytorch#111374) is unchanged there, so this one still needs a scheduler-side fix.
• Both runs were verified with torch 2.13.0 side-by-side as the known-bad control (same scripts, NaNs reproduce on both schedulers there).

Historical note: this scheduler pair had a different MPS failure in 2022 — #1537, a loud float64 crash, fixed by #2241. The current one is silent, which is presumably why it survived.