[shardformer] vit support flash attention and jit operator

colossalai/shardformer/modeling/vit.py

@@ -0,0 +1,53 @@
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+from torch import nn
+
+
+def get_vit_flash_self_attention_forward():
+
+    from transformers.models.vit.modeling_vit import ViTSelfAttention
+
+    from colossalai.kernel.cuda_native.flash_attention import ColoAttention
+
+    def transpose_for_scores(x: torch.Tensor, num_attention_heads, attention_head_size) -> torch.Tensor:
+        new_x_shape = x.size()[:-1] + (num_attention_heads, attention_head_size)
+        x = x.view(new_x_shape)
+        return x
+
+    def forward(self: ViTSelfAttention,
+                hidden_states: torch.Tensor,
+                head_mask: Optional[torch.Tensor] = None,
+                output_attentions: bool = False) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
+        mixed_query_layer = self.query(hidden_states)
+
+        key_layer = transpose_for_scores(self.key(hidden_states), self.num_attention_heads, self.attention_head_size)
+        value_layer = transpose_for_scores(self.value(hidden_states), self.num_attention_heads,
+                                           self.attention_head_size)
+        query_layer = transpose_for_scores(mixed_query_layer, self.num_attention_heads, self.attention_head_size)
+
+        scale = 1.0 / math.sqrt(self.attention_head_size)
+        attention = ColoAttention(embed_dim=self.all_head_size,
+                                  num_heads=self.num_attention_heads,
+                                  dropout=self.dropout.p,
+                                  scale=scale)
+        context_layer = attention(query_layer, key_layer, value_layer)
+
+        outputs = (context_layer,)
+
+        return outputs
+
+    return forward
+
+
+def get_jit_fused_vit_output_forward():
+
+    from transformers.models.vit.modeling_vit import ViTOutput
+
+    def forward(self: ViTOutput, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout_add(hidden_states, input_tensor, self.dropout.p, self.dropout.training)
+        return hidden_states
+
+    return forward

colossalai/shardformer/policies/vit.py

@@ -10,6 +10,8 @@
     Linear1D_Row,
 )
 
+from ..modeling.jit import get_jit_fused_dropout_add_func
+from ..modeling.vit import get_jit_fused_vit_output_forward, get_vit_flash_self_attention_forward
 from .basepolicy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
 
 __all__ = ['ViTPolicy', 'ViTForImageClassificationPolicy', 'ViTForMaskedImageModelingPolicy']
@@ -24,7 +26,7 @@ def preprocess(self):
         return self.model
 
     def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
-        from transformers.models.vit.modeling_vit import ViTEmbeddings, ViTLayer, ViTModel
+        from transformers.models.vit.modeling_vit import ViTEmbeddings, ViTLayer, ViTModel, ViTOutput, ViTSelfAttention
 
         policy = {}
 
@@ -101,6 +103,19 @@ def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
                                                         policy=policy,
                                                         target_key=ViTLayer)
 
+        # use flash attention
+        if self.shard_config.enable_flash_attention:
+            policy[ViTSelfAttention] = ModulePolicyDescription(method_replacement={
+                'forward': get_vit_flash_self_attention_forward(),
+            })
+
+        # use jit fused operator
+        if self.shard_config.enable_jit_fused:
+            policy[ViTOutput] = ModulePolicyDescription(method_replacement={
+                'forward': get_jit_fused_vit_output_forward(),
+                'dropout_add': get_jit_fused_dropout_add_func(),
+            })
+
         return policy
 
     def new_model_class(self):

tests/test_shardformer/test_model/test_shard_vit.py

@@ -50,14 +50,17 @@ def check_forward_backward(org_model, sharded_model, data_gen_fn, output_transfo
         all_shard_grad = shard_grad
     assert torch.allclose(org_grad, all_shard_grad,
                           atol=1e-5), f"shard model grad is not equal to orgin model grad\n{org_grad}\n{shard_grad}"
-    
+
 
 @parameterize('enable_fused_normalization', [True, False])
 @parameterize('enable_tensor_parallelism', [True, False])
-def run_vit_test(enable_fused_normalization, enable_tensor_parallelism):
+@parameterize('enable_flash_attention', [True, False])
+@parameterize('enable_jit_fused', [True, False])
+def run_vit_test(enable_fused_normalization, enable_tensor_parallelism, enable_flash_attention, enable_jit_fused):
     sub_model_zoo = model_zoo.get_sub_registry('transformers_vit')
     for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items():
-        org_model, sharded_model = build_model(model_fn, enable_fused_normalization, enable_tensor_parallelism)
+        org_model, sharded_model = build_model(model_fn, enable_fused_normalization, enable_tensor_parallelism,
+                                               enable_flash_attention, enable_jit_fused)
         check_forward_backward(org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn)
     torch.cuda.empty_cache()