Enable vision encoder small op fusions (#3377)

Author: jianan-gu

intel_extension_for_pytorch/transformers/models/cpu/modules/decoder.py

@@ -6,6 +6,7 @@
     _IPEXlinearReluCPU,
     _IPEXlinearGeluCPU,
     _IPEXlinearMulCPU,
+    _IPEXlinearSiluCPU,
     _IPEXlinearSiluMulCPU,
 )
 
@@ -117,3 +118,32 @@ def __init__(self, module, config, tpp=False, woq=False):
                 )
         else:
             AssertionError(False, "Do not support the optimization of your model yet")
+
+
+class _IPEXEncoderLayerCPU(nn.Module):
+    def __init__(self, module, config, tpp=False, woq=False):
+        super().__init__()
+        for k, v in module.__dict__.items():
+            setattr(self, k, v)
+        for k, v in module.__class__.__dict__.items():
+            if k.startswith("__"):
+                continue
+            setattr(self.__class__, k, getattr(module.__class__, k))
+        if self.model_backbone in [
+            "MllamaForConditionalGeneration",
+        ]:
+            if not self.distributed:
+                if hasattr(module, "mlp_linear_add"):
+                    self.mlp_linear_add = _IPEXlinearAddCPU(
+                        module.mlp_linear_add.linear, tpp=tpp, woq=woq
+                    )
+                if hasattr(module, "mlp_linear_mul"):
+                    self.mlp_linear_mul = _IPEXlinearMulCPU(
+                        module.mlp_linear_mul.linear, tpp=tpp, woq=woq
+                    )
+            if hasattr(module, "linear_silu"):
+                self.linear_silu = _IPEXlinearSiluCPU(
+                    module.linear_silu.linear, tpp=tpp, woq=woq
+                )
+        else:
+            AssertionError(False, "Do not support the optimization of your model yet")

intel_extension_for_pytorch/transformers/models/reference/modules/decoder.py

@@ -8,8 +8,10 @@
     _IPEXlinearReluRef,
     _IPEXlinearGeluRef,
     _IPEXlinearMulRef,
+    _IPEXlinearSiluRef,
     _IPEXlinearSiluMulRef,
 )
+from .....llm.functional.fusions import add_layer_norm
 from torch.nn import functional as F
 from .....utils._logger import logger, WarningType
 
@@ -65,6 +67,54 @@ def LlamaDecoderLayer_forward(
     return outputs
 
 
+def MllamaVisionEncoderLayer_forward(
+    self,
+    hidden_state: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    output_attentions: bool = None,
+):
+    # Self Attention
+    residual = hidden_state
+    hidden_state = self.input_layernorm(hidden_state)
+    hidden_state, attn_weights = self.self_attn(
+        hidden_state, attention_mask=attention_mask
+    )
+    if self.is_gated:
+        hidden_state = self.gate_attn.tanh() * hidden_state
+
+    hidden_state = add_layer_norm(
+        residual,
+        hidden_state,
+        self.post_attention_layernorm.weight,
+        self.post_attention_layernorm.bias,
+        self.post_attention_layernorm.eps,
+        True,
+    )
+
+    hidden_states = self.self.linear_silu(hidden_states)
+
+    if self.is_gated:
+        if self.distributed:
+            hidden_states = self.mlp.fc2(hidden_states)
+            hidden_state = self.gate_ffn.tanh() * hidden_state
+        else:
+            hidden_state = self.mlp_linear_mul(hidden_state, self.gate_ffn.tanh())
+        hidden_state = residual + hidden_state
+    else:
+        if self.distributed:
+            hidden_states = self.mlp.fc2(hidden_states)
+            hidden_state = residual + hidden_state
+        else:
+            hidden_state = self.mlp_linear_add(hidden_state, residual)
+
+    outputs = (hidden_state,)
+
+    if output_attentions:
+        outputs += (attn_weights,)
+
+    return outputs
+
+
 def OPTDecoderLayer_forward(
     self,
     hidden_states: torch.Tensor,
@@ -2091,3 +2141,45 @@ def forward(
             )
         else:
             AssertionError(False, "Do not support the optimization of your model yet")
+
+
+class _IPEXEncoderLayerRef(nn.Module):
+    def __init__(self, module, config, distributed=False):
+        super().__init__()
+        for k, v in module.__dict__.items():
+            setattr(self, k, v)
+        for k, v in module.__class__.__dict__.items():
+            if k.startswith("__") or k.startswith("forward"):
+                continue
+            setattr(self.__class__, k, getattr(module.__class__, k))
+        self.distributed = distributed
+        self.model_backbone = config.architectures[0]
+        if self.model_backbone in [
+            "MllamaForConditionalGeneration",
+        ]:
+            if not self.distributed:
+                if self.is_gated:
+                    self.mlp_linear_mul = _IPEXlinearMulRef(module.mlp.fc2)
+                else:
+                    self.mlp_linear_add = _IPEXlinearAddRef(module.mlp.fc2)
+                del self.__dict__["_modules"]["mlp"].fc2
+            self.linear_silu = _IPEXlinearSiluRef(module.mlp.fc1)
+            del self.__dict__["_modules"]["mlp"].fc1
+        else:
+            AssertionError(False, "Do not support the optimization of your model yet")
+
+    def forward(
+        self,
+        hidden_state: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = None,
+    ):
+        if self.model_backbone == "MllamaForConditionalGeneration":
+            return MllamaVisionEncoderLayer_forward(
+                self,
+                hidden_state,
+                attention_mask,
+                output_attentions,
+            )
+        else:
+            AssertionError(False, "Do not support the optimization of your model yet")

intel_extension_for_pytorch/transformers/optimize.py

@@ -150,7 +150,10 @@ def model_convert_reference(_model):
     )
 
     # model wise optimization for Feedforward and Decoder layer modules
-    from .models.reference.modules.decoder import _IPEXDecoderLayerRef
+    from .models.reference.modules.decoder import (
+        _IPEXDecoderLayerRef,
+        _IPEXEncoderLayerRef,
+    )
 
     # generation length or model forward order
     from .models.reference.models import (
@@ -572,6 +575,16 @@ def model_convert_reference(_model):
             _model.config,
             distributed=distributed,
         )
+    for supported_encoder_class in [
+        transformers.models.mllama.modeling_mllama.MllamaVisionEncoderLayer
+    ]:
+        convert_class(
+            _model,
+            supported_encoder_class,
+            _IPEXEncoderLayerRef,
+            _model.config,
+            distributed=distributed,
+        )
     # special list that has not official transformers design
     if _model.config.architectures[0] == "BloomForCausalLM":
         convert_function(
@@ -1374,12 +1387,18 @@ def model_convert_lowering(
     cache_weight_for_large_batch=False,
 ):
     from .models.reference.modules.attentions import _IPEXAttentionRef
-    from .models.reference.modules.decoder import _IPEXDecoderLayerRef
+    from .models.reference.modules.decoder import (
+        _IPEXDecoderLayerRef,
+        _IPEXEncoderLayerRef,
+    )
 
     if device == "cpu":
         from .models.cpu.modules.attentions import _IPEXAttentionCPU
         from .models.cpu.fusions.mha_fusion import _IPEXRMSNormCPU
-        from .models.cpu.modules.decoder import _IPEXDecoderLayerCPU
+        from .models.cpu.modules.decoder import (
+            _IPEXDecoderLayerCPU,
+            _IPEXEncoderLayerCPU,
+        )
 
         _disable_tpp()
         if not is_quantization:
@@ -1479,7 +1498,15 @@ def model_convert_lowering(
                 tpp=True if _using_tpp() else False,
                 woq=woq,
             )
-
+        for supported_mlp_class in [_IPEXEncoderLayerRef]:
+            lowering_class_cpu(
+                _model,
+                supported_mlp_class,
+                _IPEXEncoderLayerCPU,
+                _model.config,
+                tpp=True if _using_tpp() else False,
+                woq=woq,
+            )
         for supported_mha_class in [_IPEXAttentionRef]:
             lowering_class_cpu(
                 _model,