support gelu_bias_fused for gpt2

fix

fix

fix

Author: flybird11111

colossalai/shardformer/modeling/bert.py

@@ -1287,3 +1287,16 @@ def forward(
         )
 
     return forward
+
+
+def get_jit_fused_bert_intermediate_forward():
+    from transformers.models.bert.modeling_bert import BertIntermediate
+
+    from colossalai.kernel.jit.bias_gelu import GeLUFunction as JitGeLUFunction
+
+    def forward(self: BertIntermediate, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states, bias = self.dense(hidden_states)
+        hidden_states = JitGeLUFunction.apply(hidden_states, bias)
+        return hidden_states
+
+    return forward

colossalai/shardformer/modeling/blip2.py

@@ -129,3 +129,17 @@ def forward(
         return hidden_states
 
     return forward
+
+
+def get_jit_fused_blip2_mlp_forward():
+    from transformers.models.blip_2.modeling_blip_2 import Blip2MLP
+
+    from colossalai.kernel.jit.bias_gelu import GeLUFunction as JitGeLUFunction
+
+    def forward(self: Blip2MLP, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states, bias = self.fc1(hidden_states)
+        hidden_states = JitGeLUFunction.apply(hidden_states, bias)
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
+
+    return forward

colossalai/shardformer/modeling/gpt2.py

@@ -1310,3 +1310,18 @@ def forward(
         )
 
     return forward
+
+
+def get_jit_fused_gpt2_mlp_forward():
+    from transformers.models.gpt2.modeling_gpt2 import GPT2MLP
+
+    from colossalai.kernel.jit.bias_gelu import GeLUFunction as JitGeLUFunction
+
+    def forward(self: GPT2MLP, hidden_states: Optional[Tuple[torch.FloatTensor]]) -> torch.FloatTensor:
+        hidden_states, bias = self.c_fc(hidden_states)
+        hidden_states = JitGeLUFunction.apply(hidden_states, bias)
+        hidden_states = self.c_proj(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        return hidden_states
+
+    return forward

colossalai/shardformer/modeling/vit.py

@@ -372,3 +372,15 @@ def forward(self: ViTOutput, hidden_states: torch.Tensor, input_tensor: torch.Te
         return hidden_states
 
     return forward
+
+
+def get_jit_fused_vit_intermediate_forward():
+    from colossalai.kernel.jit.bias_gelu import GeLUFunction as JitGeLUFunction
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states, bias = self.dense(hidden_states)
+        hidden_states = JitGeLUFunction.apply(hidden_states, bias)
+
+        return hidden_states
+
+    return forward

colossalai/shardformer/policies/bert.py

@@ -12,6 +12,7 @@
     BertPipelineForwards,
     bert_sequence_parallel_forward_fn,
     get_bert_flash_attention_forward,
+    get_jit_fused_bert_intermediate_forward,
     get_jit_fused_bert_output_forward,
     get_jit_fused_bert_self_output_forward,
 )
@@ -38,11 +39,13 @@ def config_sanity_check(self):
 
     def preprocess(self):
         self.tie_weight = self.tie_weight_check()
+        self.enable_bias_gelu_fused = self.model.config.hidden_act == "gelu"
         return self.model
 
     def module_policy(self):
         from transformers.models.bert.modeling_bert import (
             BertEmbeddings,
+            BertIntermediate,
             BertLayer,
             BertModel,
             BertOutput,
@@ -131,6 +134,7 @@ def module_policy(self):
                         kwargs={
                             "seq_parallel_mode": sp_mode,
                             "overlap": overlap,
+                            "skip_bias_add": self.enable_bias_gelu_fused,
                         },
                     ),
                     SubModuleReplacementDescription(
@@ -231,6 +235,14 @@ def module_policy(self):
                 policy=policy,
                 target_key=BertOutput,
             )
+        if self.enable_bias_gelu_fused:
+            self.append_or_create_method_replacement(
+                description={
+                    "forward": get_jit_fused_bert_intermediate_forward(),
+                },
+                policy=policy,
+                target_key=BertIntermediate,
+            )
 
         return policy
 

colossalai/shardformer/policies/blip2.py

@@ -3,6 +3,7 @@
 from ..modeling.blip2 import (
     forward_fn,
     get_blip2_flash_attention_forward,
+    get_jit_fused_blip2_mlp_forward,
     get_jit_fused_blip2_QFormer_output_forward,
     get_jit_fused_blip2_QFormer_self_output_forward,
 )
@@ -18,12 +19,14 @@ def config_sanity_check(self):
 
     def preprocess(self):
         self.tie_weight = self.tie_weight_check()
+        self.enable_bias_gelu_fused = self.model.config.hidden_act == "gelu"
         return self.model
 
     def module_policy(self):
         from transformers.models.blip_2.modeling_blip_2 import (
             Blip2Attention,
             Blip2EncoderLayer,
+            Blip2MLP,
             Blip2QFormerLayer,
             Blip2QFormerModel,
             Blip2QFormerOutput,
@@ -73,6 +76,7 @@ def module_policy(self):
                     SubModuleReplacementDescription(
                         suffix="mlp.fc1",
                         target_module=col_nn.Linear1D_Col,
+                        kwargs={"skip_bias_add": self.enable_bias_gelu_fused},
                     ),
                     SubModuleReplacementDescription(
                         suffix="mlp.fc2",
@@ -359,6 +363,14 @@ def module_policy(self):
                 policy=policy,
                 target_key=Blip2QFormerOutput,
             )
+        if self.enable_bias_gelu_fused:
+            self.append_or_create_method_replacement(
+                description={
+                    "forward": get_jit_fused_blip2_mlp_forward(),
+                },
+                policy=policy,
+                target_key=Blip2MLP,
+            )
 
         return policy
 

colossalai/shardformer/policies/gpt2.py

@@ -6,11 +6,11 @@
 
 import colossalai.shardformer.layer as col_nn
 
-from ..layer.fused_ops import Bias_Gelu
 from ..modeling.gpt2 import (
     GPT2PipelineForwards,
     get_gpt2_flash_attention_forward,
     get_gpt_model_forward_for_flash_attn,
+    get_jit_fused_gpt2_mlp_forward,
     get_lm_forward_with_dist_cross_entropy,
     gpt2_sequence_parallel_forward_fn,
 )
@@ -37,10 +37,11 @@ def preprocess(self):
         """
         self.tie_weight = self.tie_weight_check()
         self.origin_attn_implement = self.model.config._attn_implementation
+        self.enable_bias_gelu_fused = self.model.config.activation_function == "gelu"
         return self.model
 
     def module_policy(self):
-        from transformers.models.gpt2.modeling_gpt2 import GPT2Attention, GPT2Block, GPT2Model
+        from transformers.models.gpt2.modeling_gpt2 import GPT2MLP, GPT2Attention, GPT2Block, GPT2Model
 
         ATTN_IMPLEMENTATION = {
             "eager": GPT2Attention,
@@ -120,6 +121,7 @@ def module_policy(self):
                             "n_fused": 1,
                             "seq_parallel_mode": sp_mode,
                             "overlap": overlap,
+                            "skip_bias_add": self.enable_bias_gelu_fused,
                         },
                     ),
                     SubModuleReplacementDescription(
@@ -200,39 +202,21 @@ def module_policy(self):
                 policy[GPT2Model].method_replacement = {
                     "forward": get_gpt_model_forward_for_flash_attn(self.shard_config)
                 }
+        if self.enable_bias_gelu_fused:
+            self.append_or_create_method_replacement(
+                description={
+                    "forward": get_jit_fused_gpt2_mlp_forward(),
+                },
+                policy=policy,
+                target_key=GPT2MLP,
+            )
 
         if sp_mode is not None:
             policy[GPT2Model].method_replacement = {"forward": gpt2_sequence_parallel_forward_fn(self.shard_config)}
 
         return policy
 
     def postprocess(self):
-        import torch
-
-        from colossalai.shardformer._utils import setattr_
-
-        def bias_gelu_substitute_gpt2(module):
-            target_linear = None
-            for name, child in module.named_children():
-                bias_gelu_substitute_gpt2(child)
-                if name == "c_fc" and isinstance(child, col_nn.GPT2FusedLinearConv1D_Col):
-                    target_linear = child
-                elif target_linear is not None:
-                    if name == "act":
-                        replace_sub_module = Bias_Gelu(target_linear.bias)
-                        target_linear.bias = None
-                        setattr_(module, "act", replace_sub_module)
-
-                        target_linear = None
-
-        def trial(module):
-            if torch.distributed.get_rank() == 0:
-                print(module.__class__.__name__)
-            for name, child in module.named_children():
-                trial(child)
-
-        bias_gelu_substitute_gpt2(self.model)
-        trial(self.model)
         return self.model
 
     def get_held_layers(self) -> List[nn.Module]:

colossalai/shardformer/policies/vit.py

@@ -11,6 +11,7 @@
     ViTForImageClassification_pipeline_forward,
     ViTForMaskedImageModeling_pipeline_forward,
     ViTModel_pipeline_forward,
+    get_jit_fused_vit_intermediate_forward,
     get_jit_fused_vit_output_forward,
     get_vit_flash_self_attention_forward,
 )
@@ -24,10 +25,17 @@ def config_sanity_check(self):
         pass
 
     def preprocess(self):
+        self.enable_bias_gelu_fused = self.model.config.hidden_act == "gelu"
         return self.model
 
     def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
-        from transformers.models.vit.modeling_vit import ViTEmbeddings, ViTLayer, ViTOutput, ViTSelfAttention
+        from transformers.models.vit.modeling_vit import (
+            ViTEmbeddings,
+            ViTIntermediate,
+            ViTLayer,
+            ViTOutput,
+            ViTSelfAttention,
+        )
 
         policy = {}
 
@@ -83,6 +91,9 @@ def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
                     SubModuleReplacementDescription(
                         suffix="intermediate.dense",
                         target_module=col_nn.Linear1D_Col,
+                        kwargs={
+                            "skip_bias_add": self.enable_bias_gelu_fused,
+                        },
                     ),
                     SubModuleReplacementDescription(
                         suffix="output.dense",
@@ -115,6 +126,14 @@ def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
                 policy=policy,
                 target_key=ViTOutput,
             )
+        if self.enable_bias_gelu_fused:
+            self.append_or_create_method_replacement(
+                description={
+                    "forward": get_jit_fused_vit_intermediate_forward(),
+                },
+                policy=policy,
+                target_key=ViTIntermediate,
+            )
         return policy
 
     def new_model_class(self):