Added support of subfunction for VLMs (quic#699)

Signed-off-by: Abhishek Kumar Singh <sabhis@qti.qualcomm.com>
Signed-off-by: abhishek-singh591 <sabhis@qti.qualcomm.com>
Signed-off-by: Abhishek kumar singh <sabhis@qti.qualcomm.com>

Author: abhishek-singh591

QEfficient/transformers/models/codegen/modeling_codegen.py

@@ -7,7 +7,7 @@
 
 """PyTorch Codegen model."""
 
-from typing import Optional, Tuple, Union
+from typing import Optional, Tuple, Type, Union
 
 import torch
 from torch import nn
@@ -296,6 +296,15 @@ class QEffCodeGenForCausalLM(CodeGenForCausalLM):
     - update the hidden_states, and fix for onnx model
     """
 
+    def get_submodules_for_export(self) -> Type[nn.Module]:
+        """
+        Return the set of class used as the repeated layer across the model for subfunction extraction.
+        Notes:
+            This method should return the *class object* (not an instance).
+            Downstream code can use this to find/build subfunctions for repeated blocks.
+        """
+        return {QEffCodeGenBlock}
+
     def forward(
         self,
         input_ids: Optional[torch.LongTensor] = None,

QEfficient/transformers/models/falcon/modeling_falcon.py

@@ -8,9 +8,10 @@
 """PyTorch Falcon model."""
 
 import math
-from typing import Optional, Tuple, Union
+from typing import Optional, Tuple, Type, Union
 
 import torch
+import torch.nn as nn
 import torch.utils.checkpoint
 from torch.nn import functional as F
 from transformers.cache_utils import Cache
@@ -353,6 +354,15 @@ class QEffFalconForCausalLM(FalconForCausalLM):
     - update the hidden_states, and fix for onnx model
     """
 
+    def get_submodules_for_export(self) -> Type[nn.Module]:
+        """
+        Return the set of class used as the repeated layer across the model for subfunction extraction.
+        Notes:
+            This method should return the *class object* (not an instance).
+            Downstream code can use this to find/build subfunctions for repeated blocks.
+        """
+        return {QEffFalconDecoderLayer}
+
     def forward(
         self,
         input_ids: torch.LongTensor = None,

QEfficient/transformers/models/gemma/modeling_gemma.py

@@ -5,7 +5,7 @@
 #
 # -----------------------------------------------------------------------------
 
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Type, Union
 
 import torch
 from torch import nn
@@ -336,6 +336,15 @@ class QEffGemmaForCausalLM(GemmaForCausalLM):
     - add new args cache idx for the kv retention
     """
 
+    def get_submodules_for_export(self) -> Type[nn.Module]:
+        """
+        Return the set of class used as the repeated layer across the model for subfunction extraction.
+        Notes:
+            This method should return the *class object* (not an instance).
+            Downstream code can use this to find/build subfunctions for repeated blocks.
+        """
+        return {QEffGemmaDecoderLayer}
+
     def forward(
         self,
         input_ids: torch.LongTensor = None,

QEfficient/transformers/models/gemma2/modeling_gemma2.py

@@ -5,7 +5,7 @@
 #
 # -----------------------------------------------------------------------------
 
-from typing import Callable, List, Optional, Tuple, Union
+from typing import Callable, List, Optional, Tuple, Type, Union
 
 import torch
 from torch import nn
@@ -388,6 +388,15 @@ class QEffGemma2ForCausalLM(Gemma2ForCausalLM, GenerationMixin):
     - add new args cache idx for the kv retention
     """
 
+    def get_submodules_for_export(self) -> Type[nn.Module]:
+        """
+        Return the set of class used as the repeated layer across the model for subfunction extraction.
+        Notes:
+            This method should return the *class object* (not an instance).
+            Downstream code can use this to find/build subfunctions for repeated blocks.
+        """
+        return {QEffGemma2DecoderLayer}
+
     def forward(
         self,
         input_ids: torch.LongTensor = None,

QEfficient/transformers/models/gemma3/modeling_gemma3.py

@@ -6,7 +6,7 @@
 # -----------------------------------------------------------------------------
 
 import copy
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Type, Union
 
 import torch
 from torch import nn
@@ -589,6 +589,15 @@ def __init__(self, model):
         self.model = model
         self.model.vision_model = self.model.vision_tower
 
+    def get_submodules_for_export(self) -> Type[nn.Module]:
+        """
+        Return the set of class used as the repeated layer across the model for subfunction extraction.
+        Notes:
+            This method should return the *class object* (not an instance).
+            Downstream code can use this to find/build subfunctions for repeated blocks.
+        """
+        return {self.model.vision_tower.vision_model.encoder.layers[0].__class__}
+
     def forward(self, pixel_values):
         image_features = self.model.get_image_features(pixel_values=pixel_values)
         return image_features
@@ -602,6 +611,15 @@ def __init__(self, model):
         self.config = self.model.config
         self.lm_head = self.model.lm_head
 
+    def get_submodules_for_export(self) -> Type[nn.Module]:
+        """
+        Return the set of class used as the repeated layer across the model for subfunction extraction.
+        Notes:
+            This method should return the *class object* (not an instance).
+            Downstream code can use this to find/build subfunctions for repeated blocks.
+        """
+        return {QEffGemma3DecoderLayer}
+
     def forward(
         self,
         input_ids,

QEfficient/transformers/models/gpt2/modeling_gpt2.py

@@ -5,7 +5,7 @@
 #
 # -----------------------------------------------------------------------------
 
-from typing import Callable, Optional, Tuple, Union
+from typing import Callable, Optional, Tuple, Type, Union
 
 import torch
 from torch import nn
@@ -397,6 +397,15 @@ class QEffGPT2LMHeadModel(GPT2LMHeadModel):
     - add new args position idx for the cache_kwargs for kv retention
     """
 
+    def get_submodules_for_export(self) -> Type[nn.Module]:
+        """
+        Return the set of class used as the repeated layer across the model for subfunction extraction.
+        Notes:
+            This method should return the *class object* (not an instance).
+            Downstream code can use this to find/build subfunctions for repeated blocks.
+        """
+        return {QEffGPT2Block}
+
     def forward(
         self,
         input_ids: Optional[torch.LongTensor] = None,

QEfficient/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py

@@ -7,7 +7,7 @@
 
 """PyTorch GPTBigCode model."""
 
-from typing import Optional, Tuple, Union
+from typing import Optional, Tuple, Type, Union
 
 import torch
 import torch.utils.checkpoint
@@ -378,6 +378,15 @@ def forward(
 
 
 class QEffGPTBigCodeForCausalLM(GPTBigCodeForCausalLM):
+    def get_submodules_for_export(self) -> Type[nn.Module]:
+        """
+        Return the set of class used as the repeated layer across the model for subfunction extraction.
+        Notes:
+            This method should return the *class object* (not an instance).
+            Downstream code can use this to find/build subfunctions for repeated blocks.
+        """
+        return {QEffGPTBigCodeBlock}
+
     def forward(
         self,
         input_ids: Optional[torch.Tensor] = None,

QEfficient/transformers/models/gpt_oss/modeling_gpt_oss.py

@@ -6,7 +6,7 @@
 # -----------------------------------------------------------------------------
 import math
 import os
-from typing import Callable, Optional, Union
+from typing import Callable, Optional, Type, Union
 
 import torch
 from torch import nn
@@ -1205,6 +1205,16 @@ def forward(
 
 
 class QEffGptOssForCausalLM(GptOssForCausalLM):
+    def get_submodules_for_export(self) -> Type[nn.Module]:
+        """
+        Return the set of class used as the repeated layer across the model for subfunction extraction.
+
+        Notes:
+            This method should return the *class object* (not an instance).
+            Downstream code can use this to find/build subfunctions for repeated blocks.
+        """
+        return {QEffGptOssDecoderLayer}
+
     def forward(
         self,
         input_ids: Optional[torch.LongTensor] = None,

QEfficient/transformers/models/gptj/modeling_gptj.py

@@ -7,7 +7,7 @@
 
 """PyTorch GPT-J model."""
 
-from typing import Optional, Tuple, Union
+from typing import Optional, Tuple, Type, Union
 
 import torch
 from torch import nn
@@ -318,6 +318,15 @@ class QEffGPTJForCausalLM(GPTJForCausalLM):
     - update the hidden_states, and fix for onnx model
     """
 
+    def get_submodules_for_export(self) -> Type[nn.Module]:
+        """
+        Return the set of class used as the repeated layer across the model for subfunction extraction.
+        Notes:
+            This method should return the *class object* (not an instance).
+            Downstream code can use this to find/build subfunctions for repeated blocks.
+        """
+        return {QEffGPTJBlock}
+
     def forward(
         self,
         input_ids: Optional[torch.LongTensor] = None,

QEfficient/transformers/models/granite/modeling_granite.py

@@ -5,7 +5,7 @@
 #
 # -----------------------------------------------------------------------------
 
-from typing import Callable, List, Optional, Tuple, Union
+from typing import Callable, List, Optional, Tuple, Type, Union
 
 import torch
 from torch import nn
@@ -347,6 +347,15 @@ class QEffGraniteForCausalLM(GraniteForCausalLM):
     Copied from GraniteForCausalLM: https://github.com/huggingface/transformers/blob/main/src/transformers/models/granite/modeling_granite.py
     """
 
+    def get_submodules_for_export(self) -> Type[nn.Module]:
+        """
+        Return the set of class used as the repeated layer across the model for subfunction extraction.
+        Notes:
+            This method should return the *class object* (not an instance).
+            Downstream code can use this to find/build subfunctions for repeated blocks.
+        """
+        return {QEffGraniteDecoderLayer}
+
     def forward(
         self,
         input_ids: torch.LongTensor = None,