Split Q,K,V vs Fused QKV Projections in diffusers Discussion

[dg845]

diffusers generally splits fused QKV projections to separate Q, K, V projections in its checkpoint format (for example, for diffusion transformer models). I am interested in collecting feedback on this design choice and any problems which may arise from it. To be honest upfront, I think it is unlikely that we will change this design in the near future, but we would appreciate your feedback nevertheless.

EDIT: the reasoning for using split Q, K, V projections is described in more detail in #14003 (comment).

[dg845]

@bghira, with respect to #13861 (comment), would something like having an inplace flag on AttentionMixin.fuse_qkv_projections, where inplace=True fuses the weights and then deletes the old ones, mitigate your training concerns? I think doing the fusion in-place should make the fused QKV projection more friendly for training.

[dxqb]

a good starting point for this discussion would be to understand, why diffusers is doing this. I don't understand it.

[dg845]

My understanding is that the main reason is code readability and consistency across the codebase. Having a single convention makes it easier to understand the attention code across different models. I believe the reason the split convention is used in particular is partly historical (older models like Stable Diffusion (pre-SD3) used the split convention natively) and partly because it is more expressive modeling-wise (for example, fused projections have a harder time expressing per-projection techniques like IP-Adapter).