[MPS] SDXL pipeline fails inference in fp16 mode

[bghira]

Describe the bug

When running a pipeline under 🤗 Accelerate, my M3 system crashes with a tensor dtype mismatch during validation inference.

This issue does not occur without 🤗 Accelerate.

Reproduction

Unfortunately, the way the reproducer is executed within a training script makes it exceedingly difficult/time-consuming to put together a small reproducer.

I added print statements to the pipeline so that I could determine where the broadcast compatibility error originates, and discovered this line:

                print("compute the previous noisy sample x_t -> x_t-1")
                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]

which seems to bump the dtype from float16 to float32 on the 2nd step. the 1st step is correct and fine.

just on a whim, i modified the code to the following:

                print("compute the previous noisy sample x_t -> x_t-1")
                old_dtype = latents.dtype
                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
                if latents.dtype is not old_dtype:
                    latents = latents.to(old_dtype)

which then allows validations to complete and finish.

i'm sure there's a deeper issue and this is not a valid solution, so i have not opened up a PR with this solution.

it seems like the DDIM scheduler simply should not change the dtype.

Logs

Before resolution:


2024-03-21 10:14:22,290 [DEBUG] (validation) Generating validation image: 
 Device allocations:
 -> unet on mps:0
 -> text_encoder on None
 -> vae on mps:0
 -> current_validation_prompt_embeds on mps:0
 -> current_validation_pooled_embeds on mps:0
 -> validation_negative_prompt_embeds on mps:0
 -> validation_negative_pooled_embeds on mps:0
2024-03-21 10:14:22,361 [DEBUG] (validation) Generating validation image: 
 Weight dtypes:
 -> unet: torch.float16
 -> text_encoder: None
 -> vae: torch.float16
 -> current_validation_prompt_embeds: torch.float16
 -> current_validation_pooled_embeds: torch.float16
 -> validation_negative_prompt_embeds: torch.float16
 -> validation_negative_pooled_embeds: torch.float16
2024-03-21 10:14:22,361 [DEBUG] (validation) Generating validation image: 
 -> Number of images: 1
 -> Number of inference steps: 30
 -> Guidance scale: 7.5
 -> Guidance rescale: 0.7
 -> Resolution: 1024.0
 -> Extra validation kwargs: {'generator': <torch._C.Generator object at 0x6cb88de10>}
1. Check inputs. Raise error if not correct
Prompt embeds shape: torch.Size([1, 77, 2048])
Negative prompt embeds shape: torch.Size([1, 77, 2048])
2. Define call parameters
3. Encode input prompt
4. Prepare timesteps
5. Prepare latent variables
6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
7. Prepare added time ids & embeddings
8. Denoising loop
8.1 Apply denoising_end
9. Optionally get Guidance Scale Embedding
expand the latents if we are doing classifier free guidance
predict the noise residual
 -> dtype: torch.float16
 -> dtype (latents): torch.float16
 -> dtype (prompt_embeds): torch.float16
 -> dtype (add_text_embeds): torch.float16
perform guidance
Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
compute the previous noisy sample x_t -> x_t-1
call the callback, if provided
expand the latents if we are doing classifier free guidance
predict the noise residual
 -> dtype: torch.float32
 -> dtype (latents): torch.float32
 -> dtype (prompt_embeds): torch.float16
 -> dtype (add_text_embeds): torch.float16
loc("mps_add"("(mpsFileLoc): /AppleInternal/Library/BuildRoots/ce725a5f-c761-11ee-a4ec-b6ef2fd8d87b/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphUtilities.mm":233:0)): error: input types 'tensor<2x1280xf32>' and 'tensor<1280xf16>' are not broadcast compatible
LLVM ERROR: Failed to infer result type(s).

After resolution:

 Weight dtypes:
 -> unet: torch.float16
 -> text_encoder: None
 -> vae: torch.float16
 -> current_validation_prompt_embeds: torch.float16
 -> current_validation_pooled_embeds: torch.float16
 -> validation_negative_prompt_embeds: torch.float16
 -> validation_negative_pooled_embeds: torch.float16
2024-03-21 10:12:17,921 [DEBUG] (validation) Generating validation image:
 -> Number of images: 1
 -> Number of inference steps: 30
 -> Guidance scale: 7.5
 -> Guidance rescale: 0.7
 -> Resolution: 1024.0
 -> Extra validation kwargs: {'generator': <torch._C.Generator object at 0x35ae4be50>}
1. Check inputs. Raise error if not correct
Prompt embeds shape: torch.Size([1, 77, 2048])
Negative prompt embeds shape: torch.Size([1, 77, 2048])
2. Define call parameters
3. Encode input prompt
4. Prepare timesteps
5. Prepare latent variables
6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
7. Prepare added time ids & embeddings
8. Denoising loop
8.1 Apply denoising_end
9. Optionally get Guidance Scale Embedding
expand the latents if we are doing classifier free guidance
predict the noise residual
 -> dtype: torch.float16
 -> dtype (latents): torch.float16
 -> dtype (prompt_embeds): torch.float16
 -> dtype (add_text_embeds): torch.float16
perform guidance
Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
compute the previous noisy sample x_t -> x_t-1
call the callback, if provided
expand the latents if we are doing classifier free guidance
predict the noise residual
 -> dtype: torch.float16
 -> dtype (latents): torch.float16
 -> dtype (prompt_embeds): torch.float16
 -> dtype (add_text_embeds): torch.float16

System Info

• diffusers version: 0.26.3
• Platform: macOS-14.4-arm64-arm-64bit
• Python version: 3.10.13
• PyTorch version (GPU?): 2.2.1 (False)
• Huggingface_hub version: 0.21.4
• Transformers version: 4.40.0.dev0
• Accelerate version: 0.26.1
• xFormers version: not installed
• Using GPU in script?: M3 Max
• Using distributed or parallel set-up in script?: False

Who can help?

@pcuenca

[sayakpaul]

Are you running this under autocast? Also, when you say you're running validation, is it the final validation run (i.e., after training's fully done and you're loading the weights) or an intermediate one?

[bghira]

this is intermediate validations where the model is unwrapped before inference.

here is my autocast config for validations:

            with torch.autocast(
                str(accelerator.device).replace(":0", ""),
                enabled=(
                    accelerator.mixed_precision == "fp16"
                    or accelerator.mixed_precision == "bf16"
                ),
            ):

it seems to work like this for NVIDIA systems. maybe it should be disabled on MPS?

[sayakpaul]

Interesting. Thanks for providing the information. Could you confirm if disabling the autocast when using MPS doesn't lead to anything unexpected? If so, happy to open a PR.

[bghira]

will test and get back to you

[bghira]

expand the latents if we are doing classifier free guidance
predict the noise residual
 -> dtype: torch.float16
 -> dtype (latents): torch.float16
 -> dtype (prompt_embeds): torch.float16
 -> dtype (add_text_embeds): torch.float16
perform guidance
Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
compute the previous noisy sample x_t -> x_t-1
call the callback, if provided
expand the latents if we are doing classifier free guidance
predict the noise residual
 -> dtype: torch.float32
 -> dtype (latents): torch.float32
 -> dtype (prompt_embeds): torch.float16
 -> dtype (add_text_embeds): torch.float16
loc("mps_add"("(mpsFileLoc): /AppleInternal/Library/BuildRoots/ce725a5f-c761-11ee-a4ec-b6ef2fd8d87b/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphUtilities.mm":233:0)): error: input types 'tensor<2x1280xf32>' and 'tensor<1280xf16>' are not broadcast compatible

unfortunately it fails right away without my fix, when autocast is disabled.

i will try autocast disabled with the fix now, as there is a VAE dtype issue sometimes, if i don't force it to fp16

[bghira]

note: forced VAE upcasting is disabled, as i'm using the fp16 fixed VAE:

make sure the VAE is in float32 mode, as it overflows in float16? False

• with autocast disabled, and no change to the pipeline code, it fails on the 2nd diffusion step
• with autocast disabled, and the proposed change to the pipeline code, it makes it all the way to the VAE decode, where it fails as:

Traceback (most recent call last):                                                                                           
  File "/Users/bghira/src/SimpleTuner/train_sdxl.py", line 1604, in <module>
    main()
  File "/Users/bghira/src/SimpleTuner/train_sdxl.py", line 1416, in main
    log_validations(
  File "/Users/bghira/src/SimpleTuner/helpers/legacy/validation.py", line 342, in log_validations
    pipeline(
  File "/Users/bghira/src/SimpleTuner/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/Users/bghira/src/SimpleTuner/.venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py", line 1317, in __call__
    image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
  File "/Users/bghira/src/SimpleTuner/.venv/lib/python3.10/site-packages/diffusers/utils/accelerate_utils.py", line 46, in wrapper
    return method(self, *args, **kwargs)
  File "/Users/bghira/src/SimpleTuner/.venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl.py", line 302, in decode
    decoded = self._decode(z).sample
  File "/Users/bghira/src/SimpleTuner/.venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl.py", line 272, in _decode
    z = self.post_quant_conv(z)
  File "/Users/bghira/src/SimpleTuner/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/Users/bghira/src/SimpleTuner/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/bghira/src/SimpleTuner/.venv/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 460, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/Users/bghira/src/SimpleTuner/.venv/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 456, in _conv_forward
    return F.conv2d(input, weight, bias, self.stride,
RuntimeError: Input type (c10::Half) and bias type (float) should be the same

[sayakpaul]

This is starting to feel like a PyTorch thingy on MPS. Any possibility to try to reproduce this more minimally?

[BuildBackBuehler]

note: forced VAE upcasting is disabled, as i'm using the fp16 fixed VAE:

make sure the VAE is in float32 mode, as it overflows in float16? False

• with autocast disabled, and no change to the pipeline code, it fails on the 2nd diffusion step
• with autocast disabled, and the proposed change to the pipeline code, it makes it all the way to the VAE decode, where it fails as:

Traceback (most recent call last):                                                                                           
  File "/Users/bghira/src/SimpleTuner/train_sdxl.py", line 1604, in <module>
    main()
  File "/Users/bghira/src/SimpleTuner/train_sdxl.py", line 1416, in main
    log_validations(
  File "/Users/bghira/src/SimpleTuner/helpers/legacy/validation.py", line 342, in log_validations
    pipeline(
  File "/Users/bghira/src/SimpleTuner/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/Users/bghira/src/SimpleTuner/.venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py", line 1317, in __call__
    image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
  File "/Users/bghira/src/SimpleTuner/.venv/lib/python3.10/site-packages/diffusers/utils/accelerate_utils.py", line 46, in wrapper
    return method(self, *args, **kwargs)
  File "/Users/bghira/src/SimpleTuner/.venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl.py", line 302, in decode
    decoded = self._decode(z).sample
  File "/Users/bghira/src/SimpleTuner/.venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl.py", line 272, in _decode
    z = self.post_quant_conv(z)
  File "/Users/bghira/src/SimpleTuner/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/Users/bghira/src/SimpleTuner/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/bghira/src/SimpleTuner/.venv/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 460, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/Users/bghira/src/SimpleTuner/.venv/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 456, in _conv_forward
    return F.conv2d(input, weight, bias, self.stride,
RuntimeError: Input type (c10::Half) and bias type (float) should be the same

This has also been interfering with training a LoRA. I actually don't know how to turn off autocast except in the script. Is there a way to do this via CLI?

And yeah, this seems like a mounting issue/desired fix when it comes to pytorch. I'm personally going to install the PR and give 'er a whirl!

pytorch/pytorch#99272

[bghira]

        if not output_type == "latent":
            needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast

            if needs_upcasting:
                self.upcast_vae()
                latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
            elif torch.backends.mps.is_available() and latents.dtype != self.vae.dtype:
                print(f"Updating VAE precision to {latents.dtype}")
                self.vae = self.vae.to(latents.dtype)

this works for me, combined with the initial proposed fix.

i agree, it's a pytorch issue. but it works this way, and the model trains, which is what is important

[sayakpaul]

Thanks much! Should PR this. Would you want me to do that and add you as a co-author?

[bghira]

@sayakpaul sure, it would be nice if @pcuenca the apple SME could take a look at this too, but getting it working imo is more important

[pcuenca]

Possibly related: pytorch/pytorch#99272
Edit: already mentioned above, initially missed it.

[bghira]

that bug is definitely the cause of this, however:

• the main downside is no autocast support
• the secondary downside is no bfloat16 support
  • but for SDXL/SD1.x/SD2.x anyway, bf16 isn't really needed if you aren't training CLIP

i have not yet encountered any other impact from the reduced mps support in pytorch other than some dtype shifting, which we can easily work around for now