SDXL: enable_attention_slicing() + enable_model_cpu_offload() produces all-black images on MPS (Apple Silicon)

[pupa3066]

SDXL: enable_attention_slicing() + enable_model_cpu_offload() produces all-black images on MPS

Describe the bug

When using enable_model_cpu_offload() combined with enable_attention_slicing() on SDXL with Apple Silicon (MPS backend), the pipeline produces all-black output images due to NaN in the UNet attention layers. Removing enable_attention_slicing() produces valid output.

enable_sequential_cpu_offload() with attention slicing works correctly.

Root Cause

The SlicedAttnProcessor produces NaN when the UNet is loaded via model_cpu_offload hooks on MPS. Attention slices appear to operate on uninitialized/stale memory regions. The NaN propagates through the UNet latents, and the VAE faithfully decodes corrupted latents to zeros (black).

The VAE is not the failure point — the NaN originates in the UNet during denoising.

Isolation Results

Configuration	Result
model_cpu_offload() without enable_attention_slicing()	✅ Valid image (pixel mean=140.73)
model_cpu_offload() with enable_attention_slicing()	❌ All black (pixel mean=0.0, NaN)
sequential_cpu_offload() with enable_attention_slicing()	✅ Valid image (pixel mean=140.72)
sequential_cpu_offload() without enable_attention_slicing()	✅ Valid image

The bug is:

• Prompt-agnostic (anime, photorealistic, landscape, abstract — all fail)
• Resolution-agnostic (512×512 and 768×768 both fail)
• Seed-agnostic (tested seeds 42, 123, 7777)
• Scheduler-agnostic (Euler Ancestral, DPM++ 2M both fail)
• VAE-agnostic (both madebyollin/sdxl-vae-fp16-fix and default SDXL VAE)
• Hardware-agnostic (reproduced on M1 8GB and M5 Pro 24GB — not memory-pressure-related)

Reproduction

import torch
from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler, AutoencoderKL

vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    vae=vae,
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True,
)
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
pipe.enable_model_cpu_offload()
pipe.enable_attention_slicing()  # ← THIS is the trigger. Remove this line and output is valid.

result = pipe(
    prompt="1girl, anime, walking through cherry blossom park, high quality, masterpiece",
    negative_prompt="ugly, deformed, blurry, low quality, realistic",
    num_inference_steps=20,
    guidance_scale=7.5,
    width=512,
    height=512,
    generator=torch.Generator("cpu").manual_seed(42),
)

import numpy as np
img_array = np.array(result.images[0])
print(f"Pixel mean: {img_array.mean()}")  # 0.0
print(f"Pixel max: {img_array.max()}")    # 0

Expected: Valid image with pixel values in [0, 255].
Actual: All-zero (black) image.

Workaround

pipe.enable_model_cpu_offload()
# pipe.enable_attention_slicing()  # Do not use with model_cpu_offload on MPS

Environment

Component	Version
OS	macOS (Apple Silicon)
Hardware	Apple M1 8GB / M5 Pro 24GB (both reproduce)
Python	3.13.7
PyTorch	2.13.0
diffusers	0.39.0
transformers	5.14.1
accelerate	1.14.0

Who can help?

@yiyixuxu @sayakpaul @DN6

[RudraMantri123]

Hi @pupa3066 reproduced, and I can answer your unconfirmed-hardware question: this is not limited to 8 GB devices.

Apple M5 Pro, 24 GB unified memory, macOS 26.4.1, Python 3.14, torch 2.13.0 (MPS available), diffusers 0.39.0 — running your script verbatim: pixel mean 0.0, pixel max 0, all-black PNG, and the same RuntimeWarning: invalid value encountered in cast at image_processor.py:142. All 20 denoising steps complete normally; only the VAE decode is affected, matching your description.

Since it reproduces with 3× your RAM, memory pressure may not be the (only) trigger — this looks more like a deterministic issue in the fp16 VAE decode on MPS under component-level offload.

If nobody has picked this up, I'd like to work on it. Plan:

Localize where the first NaN appears in the VAE decode under enable_model_cpu_offload() (forward hooks on the decoder blocks), comparing against the working enable_sequential_cpu_offload() path.
Test whether upcasting the VAE decode to fp32 under model offload on MPS resolves it.
Turn the finding into a fix + regression test (or a documented guard if it's an upstream PyTorch/MPS issue).
@yiyixuxu @sayakpaul @DN6 — happy to open a PR if this scope looks right.

AI disclosure per the contribution policy: I use AI assistance for code navigation and drafting; all debugging and verification are run and reviewed by me on real hardware.

[pupa3066]

Update: Corrected Root Cause

After further isolation testing, the root cause is more specific than initially reported:

The bug is triggered by enable_attention_slicing() combined with enable_model_cpu_offload() on MPS. The VAE is not the failure point — the NaN originates in the UNet's SlicedAttnProcessor during denoising.

Isolation results (Apple M1 8GB, same script otherwise):

Configuration	Result
model_cpu_offload() without enable_attention_slicing()	✅ Valid image (pixel mean=140.73)
model_cpu_offload() with enable_attention_slicing()	❌ All black (pixel mean=0.0)
sequential_cpu_offload() with enable_attention_slicing()	✅ Valid image (pixel mean=140.72)

Why:

SlicedAttnProcessor splits attention into chunks. Under model_cpu_offload on MPS, attention slices appear to operate on uninitialized/stale memory regions left by previous component offloads via accelerate's hooks. The resulting NaN propagates through the UNet latents and the VAE faithfully decodes the corrupted latents to zeros (black).

With sequential_cpu_offload, each layer is freshly loaded before computation, so slices always see valid data. Without attention slicing, the full attention path avoids the issue entirely.

Additional evidence:

• @RudraMantri123's reproduction on M5 Pro 24GB confirms this is not memory-pressure-related — it is deterministic regardless of available RAM
• Bug is prompt-agnostic, seed-agnostic, scheduler-agnostic, and VAE-agnostic (tested with both madebyollin/sdxl-vae-fp16-fix and default SDXL VAE)

Immediate workaround:

pipe.enable_model_cpu_offload()
# pipe.enable_attention_slicing()  # Do not use with model_cpu_offload on MPS

Note on my original reproduction script:

The reproduction code in the issue body includes pipe.enable_attention_slicing() — which is indeed the trigger. Removing that single line produces valid output. Apologies for the initial misdirection toward the VAE — the controlled isolation came later.

I am working on a proper fix (either auto-disabling attention slicing on MPS + model_cpu_offload, or adding sync barriers in SlicedAttnProcessor) and will open a PR once the edge cases are covered.

[RudraMantri123]

Great isolation, @pupa3066 — your table matches my config matrix exactly (seeds 42/7/2024 across five configurations). I can add the missing piece: the precise mechanism, traced to a specific PyTorch MPS defect — which affects which fix will work.

Root cause: Attention.get_attention_scores computes scores as torch.baddbmm(input=torch.empty(...), beta=0, ...). Per the PyTorch docs, beta=0 means the input is ignored and NaN/Inf in it must not propagate. On MPS, that contract is violated — garbage in the uninitialized buffer leaks into the output. Only the sliced processors call this code (the default path uses SDPA), which is why slicing is the trigger. Minimal repro, no diffusers involved:

import torch
B, T, D = 10, 4096, 64
junk = torch.full((B, T, T), float("nan"), device="mps", dtype=torch.float16)
del junk  # allocator will recycle these NaN pages
q = torch.randn(B, T, D, device="mps", dtype=torch.float16)
k = torch.randn(B, T, D, device="mps", dtype=torch.float16)
buf = torch.empty(B, T, T, device="mps", dtype=torch.float16)
print(torch.isnan(torch.baddbmm(buf, q, k.mT, beta=0, alpha=0.125)).any())  # True — bug
print(torch.isnan(torch.bmm(q, k.mT) * 0.125).any())                        # False — control

This refines two conclusions from your update:

1. It isn't offload-specific — plain pipe.to("mps") + enable_attention_slicing("max") produces the same NaNs with no offload involved; offload just churns the allocator so torch.empty recycles dirty pages far more often.
2. Sync barriers won't fix it — I tested single-sync-point variants (before UNet call 0, per step-end) and they still NaN; it's recycled uninitialized memory, not a timing race. And auto-disabling slicing would remove the memory-saving feature from exactly the 8GB configurations that need it — it can be made safe instead.

Fix: opened #14459 — on MPS with no mask, compute scores with a buffer-free scaled bmm. No reliance on the broken beta=0 contract, one less scores-sized allocation (lower peak memory), and ~35% faster than the current path at SDXL slice dimensions on Apple Silicon. fp32 matches the CPU reference exactly. Includes a regression test that fails deterministically on current main, plus a CPU-parity test. Your original script with slicing enabled now renders correctly across seeds on my M5 Pro.

@pupa3066 — since you were heading toward a PR yourself: would you co-verify #14459 on your M1 8GB? That's the memory-constrained case I can't cover, and your isolation work already shaped the fix. I'll also file the underlying baddbmm contract violation upstream against PyTorch.

AI disclosure per the contribution policy: AI-assisted debugging and drafting; all experiments run and verified by me on real hardware.