Fused Softmax Kernels
Author: Sanger2000
Pull request: no

### Feature request

Optional Fused Softmax Cuda kernels for transformer implementations.

Megatron-LM has implemented these [here](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/fused_kernel), and they offer massive speedups for models under 10B params when training at 2048 sequence lengths. In my experience, this amounts to 2x improvements in throughput. As you can see from [this example](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/transformer.py#L203-L209), it's relatively straightforward to add fused kernels to a model.


### Motivation

From profiling the transformers models, it seems like they achieve at best 20% of peak hardware utilization on V100s and A100s for 2048 token contexts. With just the addition of fused kernels from the Megatron codebase, I see around 40% utilization. This is supported by the findings from [YaLM-100B](https://medium.com/yandex/yandex-publishes-yalm-100b-its-the-largest-gpt-like-neural-network-in-open-source-d1df53d0e9a6).

For massive models, the performance improvements are less substantial (175B-500B params) but NVIDIA notes 10-20% speedups in section 5.8 of [this paper](https://cs.stanford.edu/~matei/papers/2021/sc_megatron_lm.pdf).

### Your contribution

I have my own GPT-2 implementation that uses Megatron's kernels and I would be happy to contribute. I don't have the time to implement the full feature request - which would be providing the ability to use these fused kernels for most of the hugging face models - but I think this would be very valuable for the ecosystem.

A 2x improvement in throughput at medium to large scales (around 100M-1B params) would be a substantial cost improvement for users. 

Comments:
@Sanger2000 could you add a link to the kernels from Megatron-LM? I'm curious if it could also be easily combined with a fused kernel for attention-dot-product, like FLASH attention.

Raw Cuda and C++ code: https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/fused_kernels

This can then be easily added to a model like here: https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/transformer.py#L203-L209

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

@Sanger2000 Are you still interested in this?

Not anymore. But I very much agree with Abhi about using FlashAttention instead of the Megatron kernels, since it is a decent bit faster and consumes far less memory (esp for longer sequences).

It would massively speed up all transformer implementations if they had the option of using flash attention for their attention computation. The only downside is it won't be possible for all models since Flash Attention limits the head dimensions that can be used (I don't believe it supports anything larger than 128 last time I checked).

I see, thanks for the feedback. There is an integration of the nn.TransformerEncoderLayer fastpath in [Optimum](https://huggingface.co/docs/optimum/bettertransformer/overview), but it is only for inference for now - the training support + flash attention will come in a next pytorch release.

I've been thinking about massively supporting xformers or HazyResearch/flash-attention for transformers since some people may be interested in already benefiting from memory efficient attention / flash attention for training, and don't want to wait 2 months. I'm not just not aware if other solutions as deepspeed or others already allow to use it or not, in which case I'd rather avoid doing double work.

Hi @fxmarty, just want to flag I am also very interested in this. I've been digging around and haven't seen anything active regarding deepspeed integration for training (only inference too).

It does look like there's been some activity over the past few days in pytorch on this though. 

https://triton-lang.org/main/getting-started/tutorials/02-fused-softmax.html