`UNet2DConditionModel` : `qk_norm` setting in `config.json` is ignored

[damian0815]

Describe the bug

Adding eg "qk_norm": "rms_norm" to config.json for a UNet2DConditionModel has no effect.

This is because the value is not propagated by the UNet2DContionalModel initialization logic through to Attention.__init__ in src/diffusers/models/attention_processor.py.

Reproduction

Default behaviour with empty config dict:

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

For supported models, QK norm modules show up as eg ... .attn1.norm_q and ... .attn1.norm_k. If we add "qk_norm" : "rms_norm" to the config then we should expect modules with these names to appear, but they don't:

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])
# expected output: ['down_blocks.0.attentions.0.transformer_blocks.0.attn1.norm_q', 'down_blocks.0.attentions.0.transformer_blocks.0.attn1.norm_k', ...]
# actual output: []

System Info

diffusers main branch commit 0c71189, macOS

Who can help?

@DN6 @yiyixuxu @sayakpaul

[damian0815]

PR to fix: #12051

[DN6]

Hi @damian0815 qk_norm isn't currently supported by the 2D UNet since it's typically applied in DiTs to help with scaling. What is your use case? We typically update the init arguments if there is a popular checkpoint on the Hub that requires it.

The UNet already has a fairly large number of init arguments so it might be better to define your own UNet with qk_norm support if you need it.

[damian0815]

i'm adapting Stable Diffusion 2.1 to flow matching. was having trouble with the attention calculations going NaN during training, and also at inference time when running the model at float16. numerous sources suggested this is a known problem and QK norm resolves it - which has indeed been my experience when finetuning from the branch linked to in the PR locally.

[damian0815]

sorry just to re-state - i'm seeing notably better performance with SD2.1 when I enable QK norm, and i fully expect that reverting back to regular SD2.1 training with the regular DDPM scheduler, i would also see improvements there.