[WIP] Refactor Attention Modules #11685
Conversation
| @@ -0,0 +1,1251 @@ | |||
| # Copyright 2024 The HuggingFace Team. All rights reserved. | |||
There was a problem hiding this comment.
I have no strong feelings about the idea of having a common modeling file, but I think if we could minify the implementation a bit and specialize per-model, it will be really clean. Basically, the entire model implementation is top-to-bottom present in a single without requiring imports from diffusers.*. I can help complete this if we're okay with the idea
| logger = logging.get_logger(__name__) | ||
|
|
||
|
|
||
| def _chunked_feed_forward(ff: nn.Module, hidden_states: torch.Tensor, chunk_dim: int, chunk_size: int): |
There was a problem hiding this comment.
Ideally, I would like to try and deprecate this method. We should, in theory, be able to chunk across all the dimensions on which a model does not operate on. We discussed a bit in the past about a SplitInferenceHook but I never quite got around to it. I can revive the idea if it sounds like something we would be okay with doing.
| return hidden_states | ||
|
|
||
|
|
||
| class LuminaFeedForward(nn.Module): |
| value = torch.cat([encoder_value, value], dim=2) | ||
|
|
||
| if image_rotary_emb is not None: | ||
| from ..embeddings import apply_rotary_emb |
There was a problem hiding this comment.
For the attention refactor, I think having per-model implementation of RoPE preparation would be nice. This aligns with having one-file-per-model and the full implementation being available at once place
| return hidden_states | ||
|
|
||
|
|
||
| class FluxIPAdapterAttnProcessor(torch.nn.Module): |
There was a problem hiding this comment.
I don't know the context behind making an attention processor a torch.nn.Module and why this was done in the past. Is there a way we could decouple the module part from the attention processor by creating a new "AttentionIPAdapter" class and setting its processor to this?
|
|
||
| return processors | ||
|
|
||
| def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): |
There was a problem hiding this comment.
I think this implementation could be simplified a bit too to not use recursion
| if "Added" in str(attn_processor.__class__.__name__): | ||
| raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") |
There was a problem hiding this comment.
IMO, we should support fusing the KV projections into single linear. Anything that's commonly used for inference optimization would benefit from having a reference implementation. Not really a priority though, so just making a note
| backend = AttentionBackendName(backend.lower()) | ||
| self.processor._attention_backend = backend | ||
|
|
||
| def set_use_npu_flash_attention(self, use_npu_flash_attention: bool) -> None: |
There was a problem hiding this comment.
Do we want to keep the npu/xla/xformers-specific functions here? I think attention backends should have it covered already, no?
| self.set_attention_backend("xformers") | ||
|
|
||
| @torch.no_grad() | ||
| def fuse_projections(self): |
There was a problem hiding this comment.
Having a fuse/unfuse method per "Attention" class (which may exist at per-model level if we are aligned with my above comments) would be nice and super clean IMO
|
|
||
| self.fused_projections = False | ||
|
|
||
| def set_attention_slice(self, slice_size: int) -> None: |
There was a problem hiding this comment.
Ideally, I would like to completely get rid of the 4-5 functions that follow and do model-specific implementations. For batch/head slicing, re: SplitInferenceHook.
What does this PR do?
Draft PR to address some issues with the attention modules.
Moves a number of transformer related blocks/objects in the attention module that would have a better home under models/transformers
Moves the definition of the Attention module into
attention.py, rather than have it live inattention_prcessors.py.We have a very large number of processors, but with Attention Dispatcher #11368 we should no longer need a good chunk of them and they can be deprecated. I think with these changes we would end up with ~3 processors per model (Attn, IPAdapter, PAG)
Make it so that we can bump up our minimum supported Torch version to
>=2.0and use theF.sdpaAPI for all processors.There was some discussion around naming of the processors:
https://huggingface.slack.com/archives/C065E480NN9/p1737130514639479
We landed on calling the processors something like AttnProcessorsSDPA, but with Attention Dispatcher #11368 we no longer need to use a dedicated processor per backend, so I think it's okay to just have the class be named AttnProcessor.
Move processor definitions into the model files, so we don't end up with very large files containing all processors.
Introduce AttentionModuleMixin that contains all common methods related to attention operations. New attention modules would inherit from this class.
Introduce an
AttentionMixinfor models so that methods likeset_processorare not duplicated across models. Although we can probably just add the methods of this Mixin toModelMixinUsing Flux as an example here to show how we can define a single Processor to support both fused/unfused qkv projections.
Fixes # (issue)
Before submitting
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.