Loading WanTransformer3DModel using torch_dtype=torch.bfloat16 keeps some parameters as float32

[spezialspezial]

Describe the bug

Just checking if this is the expected behavior. Calling WanTransformer3DModel.from_pretrained with argument torch_dtype=torch.bfloat16 keeps some parameters as float32.

Reproduction

repo_id = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
transformer = WanTransformer3DModel.from_pretrained(repo_id, subfolder="transformer", torch_dtype=torch.bfloat16)
print(transformer.blocks[0].norm2.bias.dtype)  # torch.float32
print(transformer.blocks[0].scale_shift_table.dtype)  # torch.float32
print(transformer.blocks[0].attn1.norm_k.weight.dtype)  # torch.bfloat16
print(transformer.blocks[0].attn1.to_k.weight.dtype)  # torch.bfloat16

Logs

System Info

Diffusers #812b4e1eaa20fa8d88aa48b645b9d34ca45ecfde (2025-03-06), Linux, Python 3.10

Who can help?

Calling @DN6 @a-r-r-o-w

[a-r-r-o-w]

Yes, this is intended.

diffusers/src/diffusers/models/transformers/transformer_wan.py

Line 331 in 54ab475

_keep_in_fp32_modules = ["time_embedder", "scale_shift_table", "norm1", "norm2", "norm3"]

I did not test very extensively, but pure bf16 vs bf16 + fp32 norms and modulation produced slightly different videos. It's hard to understand the impact on quality from a few examples, and past experience tells that it's always better to just stick with fp32 norms and modulation, so we went with keeping these in fp32

For any reason, if you don't want these layers to be in fp32, you can do model.to(torch.bfloat16) and it should override the effect of _keep_in_fp32_modules

[spezialspezial]

Ah, okay, I suspected something like that. I actually just wanted to check different quantizations. Will cast manually for some tests.

[spezialspezial]

Thinking of it. Did you publish your tests somewhere? I would be interested in the quality degradation from bf16 or float32 to say Torchao int8 but can't test the bf16/float32 part.

[a-r-r-o-w]

Sadly I haven't published it anywhere and eyeballed it off ~15-20 generations on a simple personal dataset, which I delete after integrations are complete. It takes forever to compare at full resolution, but I can assure you that there was visible differences in generated video (especially high-motion videos, or ones with multiple subjects and objects) when comparing pure bf16 with bf16 + fp32 norms/modulation. I don't think the differences were bad in any way -- but just the fact that there are differences means we should do those operations in fp32 (because the original code does it in fp32 under autocast context as well. We don't use autocast in the codebase, so _keep_in_fp32_modules is the alternative way of doing it)

[spezialspezial]

I believe you :) just curious

I found another thing that is a problem while fighting for every MB of VRAM. Let me just quickly ask here before I log an issue.
CLIPVisionModelWithProjection is stored on disk as bf16 but the config says "torch_dtype": "float32". This causes an upcast and I'm leaning towards "not on purpose". Any idea?

Edit: "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
Edit2: While ignoring pipe = WanImageToVideoPipeline.from_pretrained(torch_dtype=torch.bfloat16)