SD3 model generates noise on the bfloat16 datatype on CPU. about diffusers

[asomoza]

As a reference, if you use bfloat16 on CPU with other models like SDXL or SD15 it works?

[0x6b64]

Thanks for the response. Confirmed that SDXL with bfloat16 produces expected images. Ran it on the same execution environment.

from diffusers import AutoPipelineForText2Image
import torch

pipeline_text2image = AutoPipelineForText2Image.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, use_safetensors=True
)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipeline_text2image(prompt=prompt).images[0]
image.save(f"cpu_img.png")

[asomoza]

I see, I'll try to test it later, sadly I have an AMD 5900x and testing SD3 with all the text encoders probably will take all day to generate one image for me, so I can't help that much here.

Probably will help to know what kind of environment you're working on, specially the CPU.

cc @sayakpaul

[sayakpaul]

I don't have a capable AMD machine either. Would also help to know where do the activations start to diverge in the model.

Ccing @AmericanPresidentJimmyCarter @bghira if you have faced this.

[0x6b64]

Thanks Sayak for your comment, in my spare time I can compare the activations as well.
I'm running on Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz, Ubuntu20.

[0x6b64]

Looks like bfloat16 on CPU does generate image, I have two environments (pytorch/numpy/diffusers/transformers version are the same; CPU's are similar-ish)

diffusers                          0.29.0
torch                              2.1.2
transformers                       4.41.2
numpy                              1.26.4

Here are the activations on the working CPU environment
Model name: Intel(R) Xeon(R) Platinum 8275CL CPU @ 3.00GHz
https://gist.github.com/0x6b64/fa34cb1d7b26ac7a73265590aebfda03

Here are the activation on the failing CPU environment
Model name: Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz
https://gist.github.com/0x6b64/9fd1c11f57688720d1919759e6ebcaed

This hook was used to inspect the i/o values of the modules during inference (executed only 2 steps)

import torch
from diffusers import StableDiffusion3Pipeline
import torch.nn.functional as F
import numpy as np
torch.manual_seed(1)
np.random.seed(1)
generator = torch.Generator().manual_seed(1)
pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.bfloat16, token=token)

def hook(mod, inputs, outputs):
  for i, input in enumerate(inputs):
    if isinstance(input, torch.Tensor):
      print(f"{mod.__class__.__name__} input={i} dtype={input.dtype} min={input.min()} max={input.max()}"), 

  for i, output in enumerate(outputs):
    if isinstance(output, torch.Tensor):
      print(f"{mod.__class__.__name__} output={i} dtype={output.dtype} min={output.min()} max={output.max()}"), 

torch.nn.modules.module.register_module_forward_hook(hook, always_call=True)

image = pipe(
  "A cat holding a sign that says hello world",
  negative_prompt="",
  num_inference_steps=2,
  guidance_scale=7.0,
  generator=generator,
).images[0]

image.save("image.png")

A couple of things to point out are that

1. CLIPEncoderLayer input=2 dtype=torch.bfloat16 min=-3.3895313892515355e+38 max=0.0 This input looks wrong. However, the same behavior is on the GPU activations as well, so even though this looks wrong, it would not affect the diffusion process.
2. In the failing CPU activations, the JointTransformerBlock module seems to get more inputs. How can this be the case?
3. I did a line by line comparison of the I/O values fed into the forward module hook && there are differences in the values (surprisingly a lot of them are in the Linear module). The random seeds were set identically on both environments.

Bottom line is that the fact the the CPU bfloat16 code works in one of my environments, suggests that there is some delta between the two. Is there any obvious dependency that I should check for?

[0x6b64]

Resolving this as this is something in my workspace;