UNet2DConditionModel: add support for QK Normalization by propagating qk_norm value from config through to child attention modules

[damian0815]

What does this PR do?

Fixes #12050

QK Normalization was already implemented in Attention __init__, but adding eg "qk_norm": "rms_norm" to the config.json for a UNet2DConfitionModel had no effect.

This PR makes config qk_norm have an effect by propagating its value through the various UNet2DConfitionModel block initialization logic.

Without this PR:

from diffusers import UNet2DConditionModel
config_minimal = {"qk_norm": "rms_norm"}
model = UNet2DConditionModel.from_config(config_minimal)
print([n for n, _ in model.named_modules()
       if 'attn1.norm_' in n])
# output: []

With this PR:

from diffusers import UNet2DConditionModel
config_minimal = {"qk_norm": "rms_norm"}
model = UNet2DConditionModel.from_config(config_minimal)
print([n for n, _ in model.named_modules()
       if 'attn1.norm_' in n])
# output: ['down_blocks.0.attentions.0.transformer_blocks.0.attn1.norm_q', 'down_blocks.0.attentions.0.transformer_blocks.0.attn1.norm_k', ...]

(I have successfully finetuned a model after making this change)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Pings based on git blame:
@yiyixuxu @gnobitab @sayakpaul

[damian0815]

For the tests I find myself writing code like this:

for down_block_type in ["DownBlock2D", "ResnetDownsampleBlock2D",
                     "AttnDownBlock2D", "CrossAttnDownBlock2D",
                     "SimpleCrossAttnDownBlock2D", "SkipDownBlock2D",
                     "AttnSkipDownBlock2D", "DownEncoderBlock2D",
                     "AttnDownEncoderBlock2D", "KDownBlock2D",
                     "KCrossAttnDownBlock2D"]:
    block = get_down_block(
        down_block_type=down_block_type,
        ...

Is there a canonical way of obtaining the list of block types so it doesn't have to be hardcoded in the test?

[damian0815]

I wasn't able to run the full test suite - significant components seem to be broken on macOS/mps

[ppbrown]

@damian0815 you mention what you changed in the code.
Might you also give some specific examples of what takes advantage of it?
okay, the value is propagated.. but what is looking for the value?

[damian0815]

the Attention modules get q_norm and k_norm (known as "QK Normalization" in the literature). In the SD3 paper they state that they were having problems with attention calculations getting larger and larger with training, which I also saw when using flow matching. QK Normalization brings that back under control.

[ppbrown]

@damian0815 Is this an SD3 only kind of thing?
Otherwise, I was hoping of a more "diffusers class XYZ can take advantage of this", so we can understand the more general category of things that will benefit from the change.