

TITLE
[RFC] Add `modeling_xxx_fusion.py` to support kernel fusion

COMMENTS
10

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 6
rocket: 0
eyes: 0

BODY
## Introduction
I am an engineer currently working on 3D model parallelism for transformers. When the tensor model parallelism (https://github.com/huggingface/transformers/pull/13726) is done, I am going to introduce [kernel fusion](https://stackoverflow.com/questions/53305830/cuda-how-does-kernel-fusion-improve-performance-on-memory-bound-applications-on) feature to transformers. 

![image](https://user-images.githubusercontent.com/38183241/135726581-ee305818-c78a-439f-90b4-30cd1edbc1fe.png)

For this, I want to create a new modeling file called `modeling_xxx_fusion.py`. This work is currently being discussed with @stas00 and @RezaYazdaniAminabadi (DeepSpeed team). 

## Kernel fusion API
```python
from transformers import BertForMaskedLM

# create model
model = BertForMaskedLM.from_pretrained("bert-base-cased")

# 1. fuse_modules 
# `fuse_modules` is function level fusion, It supports a wide variety of models.
# all arguments is `True` as default
model.fuse_modules()  

# fuse selective modules
model.fuse_modules(
    word_embedding=True,
    scale_mask_softmax=True,
    layer_norm=True,
    bias_act=True,
    bias_dropout_residual=False,
    cross_entropy=True,
)

# 2. fuse_layers 
# `fuse_layers` is block level (attention & mlp) fusion, only a few models are supported.
# argument (`inference`) is `None` -> `not self.training` of `torch.nn.Module` as default.
model.fuse_layers(inference=None)

# fuse layers for inference
model.fuse_layers(inference=True)

# fuse layers for training
model.fuse_layers(inference=False)
```

## Implementation
The internal module of each model will be re-implemented using kernel fusion method, and the existed module will be replaced with the fused module. The following example is an example of `BertOutput(nn.Module)`.

```python
# transformers/models/bert/modeling_bert.py

class BertOutput(nn.Module):
      def __init__(self, config):
            super().__init__()
            self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
            self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
            self.dropout = nn.Dropout(config.hidden_dropout_prob)

      def forward(self, hidden_states, input_tensor):
            hidden_states = self.dense(hidden_states)
            hidden_states = self.dropout(hidden_states)
            hidden_states = self.LayerNorm(hidden_states + input_tensor)
            return hidden_states
```
```python
# transformers/models/bert/modeling_bert_fusion.py

class FusedBertOutput(BertOutput):
      def forward(self, hidden_states, input_tensor):
            hidden_states = hidden_states @ self.dense.weight.t()
            hidden_states = FusedBiasDropoutResidual.apply(hidden_states, self.dense.bias, input_tensor)
            hidden_states = FusedLayerNorm.apply(hidden_states, self.LayerNorm.weight, self.LayerNorm.bias)
            return hidden_states
```
When the user calls the `fuse_modules()` method, the kernel fusion engine finds `BertOutput` and replaces it with `FusedBertOutput`. and user calls `fused_layers` method, engine finds `BertLayer` and replcases it with `FusedBertLayer`. This is the method that `parallelformers` parallelized transformers models flexibly, and the `deepspeed` also supports kernel fusion in this way. 

However, the current version of `deepspeed` fuses the entire transformer layer, so the supported models are very limited. For example, bigbird requires random attention mechanism. in this case random attention must be implemented in the custom cuda kernel. However, because the number of models is so large, it is impossible to implement them all. So I propose a flexible way to fuse the kernel on a per-function. This is a strategy of triage. The area that can be fused performs fusion, and the area that can not be fused uses the torch's default module.

```python
# kernel_fusion_utils.py

class KernelFusionMixin(object):
    
    def fuse_modules(...):
        assert self._is_able_to_fuse, "error message"
        ... implementation ...

    def fuse_layers(...)
        assert self._is_able_to_fuse, "error message"
        ... implementation ...
```
```python
# modeling_utils.py

class PreTrainedModel(..., KernelFusionMixin):
    _is_parallelizable = ...
    _is_able_to_fuse = False. # <--- Only models that can be fused have `True`.
```

This is a draft. The API can be changed at any time. I look forward to feedback. I'm going to show you this soon with a framework I'm making. (Like parallelformers, we will pre-open the repositories on our side and merge them later on transformers and deepspeed.)

cc. @Stas00 @RezaYazdaniAminabadi @Sylvain
 