Making attention mechanism stackable
Author: stas00
Pull request: no
Labels: name: Feature request

### Feature request

I'm porting [DeepSpeed Ulysses sequence parallelism](https://github.com/deepspeedai/DeepSpeed/blob/master/blogs/deepspeed-ulysses/README.md) from Megatron-Deepspeed to the HF transformers ecosphere so that everybody could use this SP implementation. I'm integrating it into [ArcticTraining](https://github.com/snowflakedb/ArcticTraining/pull/45) and then once everything is neat we can integrate it into HF Accelerate and make it available to many frameworks/users.

One of the nuances of this SP implementation is that we need to inject UlyssesAttention before the core attention, since it splits the long sequence across participating ranks and then the normal core attention is performed - which one is used fa2, sdpa, etc. - doesn't matter.

So I'm proposing that we create a new API that like `nn.Module` allows things to stack and each level will call next level's `__call__` or `forward` or whatever we call it.

Currently I have to perform this hack to inject the SP attn wrapper that splits projected qkv across shards:
```

            from transformers import AutoConfig
            from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
            mpu.initialize_model_parallel(sequence_parallel_size=self.config.sequence_parallel_size)

            # we don't have the model yet at this stage
            hf_model_config = AutoConfig.from_pretrained(self.config.model.name_or_path)

            core_attn_implementation = self.config.model.attn_implementation
            core_attn_function = ALL_ATTENTION_FUNCTIONS.get(core_attn_implementation, None)
            if core_attn_function is None:
                raise ValueError(f"{core_attn_implementation} is not a valid attn_implementation. The choices are {ALL_ATTENTION_FUNCTIONS.keys()}")

            uattn = UlyssesAttentionHF(
                attn=core_attn_function,
                local_seq_length=self.config.data.max_length // mpu.get_sequence_parallel_world_size(),
                global_seq_length=self.config.data.max_length,
                batch_size=self.config.micro_batch_size,
                attn_head_count=hf_model_config.num_attention_heads,
                attn_head_size=hf_model_config.hidden_size // hf_model_config.num_attention_heads,
                kv_head_count=hf_model_config.num_key_value_heads,
                process_group=mpu.get_sequence_parallel_group(),
                seq_length_is_variable=True,
            )

            def uattn_wrapper(
                module: torch.nn.Module,
                query: torch.Tensor,
                key: torch.Tensor,
                value: torch.Tensor,
                attention_mask: torch.Tensor,
                *args,
                **kwargs,
            ) -> Tuple[torch.Tensor, torch.Tensor]:

                attention_mask = None
                attn_output, attn_weights = uattn(
                    module,
                    query,
                    key,
                    value,
                    attention_mask,
                    *args,
                    **kwargs
                )
                return attn_output, attn_weights

            # we are overriding the original core attn implementation with `ulysses` and we have already passed the original core attn implementation to `UlyssesAttentionHF`
            self.config.model.attn_implementation = "ulysses"
            ALL_ATTENTION_FUNCTIONS["ulysses"] = uattn_wrapper
```

So this can break at any point if `transformers` modifies this private API. 

So let's think how to make this configurable.

I think 2 things are needed 
1. allow users to register a new attn implementation
2. create a concept of stackable attentions

for example if we envision just one level then say it can be called an `attn_preamble`

so if a user configures `attn_implementation = "sdpa"` and then enable sequence parallelism via `ulysses` or some other implementation, HF transformers will call:

instead of:
```
        attn_output, attn_weights = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            **kwargs,
        )
```
this:
```
    if self.config._attn_preamble_implementation != None:
        attention_preamble_interface = ALL_ATTENTION_PREAMBLE_FUNCTIONS[self.config._attn_preamble_implementation]
        attn_output, attn_weights = attention_preamble_interface(
            attention_interface,
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            **kwargs,
        )
    else:
        attn_output, attn_weights = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            **kwargs,
        )
```

And as I mentioned earlier we want an actual API to register new `ALL_ATTENTION_PREAMBLE_FUNCTIONS` and not use the private API.

I'm of course open to other ways of implementing this.

Eventually I think `ulysses` SP implementation might make into HF transformers core, but it's definitely going to end up in HF Accelerate I believe.


Comments:
This seems like quite a big change! Pinging @muellerzr and @sunmarc for accelerate and @arthurzucker @cyrilvallez for core maintenance / attention implementations

I think at the very least there should be an API for users to be able to register `ALL_ATTENTION_FUNCTIONS` so that the API is official and won't break down the road.

Agreed sorry that I missed this notification! 

The first part of this Issue has been solved in https://github.com/huggingface/transformers/pull/36889 - we now have a public API to get/set the attention functions - thank you @Cyrilvallez and @ArthurZucker 

Now, to part 2 - I'm fine with manually stealing the original core attention function, switching users config's `attn_implementation` to `ulysses` sp and passing the original attention implementation as an argument to it. It's a bit cumbersome but it works.

Do you want to discuss the possibility of using stackable attention mechanisms or park it for for the future if perhaps you have more similar requests emerging? If now doesn't feel right, we can close this Issue.

I think we can just do this, set the attn_implementation to `ulysses` as people using ctx parallel with ulysse would expect this to somewhat happen imo -> a warning or just in the repr of the attention layer just print the attn_implementation 

I want to make sure we have close to 0 changes to the modeling code, believe this is possible with the PR from Cyril now! 🤗 


My apologies, but I'm not sure what you're saying in the first para, Arthur.

To repeat ideally a user would configure:
```
from_pretrained(..., attn_implementation_wrapper="ulysses", attn_implementation="sdpa")
```
so they would know there is attn stacking happening `wrapper_fn(core_fn, *args, **kwargs)`

On the other hand to configure the wrapper function we need to pass many new args:
```
            uattn = UlyssesAttentionHF(
                attn=core_attn_function,
                local_seq_length=self.config.data.max_length // mpu.get_sequence_parallel_world_size(),
                global_seq_length=self.config.data.max_length,
                batch_size=self.config.micro_batch_size,
                attn_head_count=hf_model_config.num_attention_heads,
                attn_head_size=hf_model_config.hidden_size // hf_model_config.num_attention_heads,
                kv_head_count=hf_model_config.num_key_value_heads,
                process_group=mpu.get_sequence_parallel_group(),
                seq_length_is_variable=True,
            )
```

So probably at least for now, the manual injection is simpler as it requires no changes on the HF side.

Let's close this for now?


yep sounds good !