Addressed comments and added support of the external model using Infer

Signed-off-by: Amit Raj <[email protected]>

Author: quic-amitraj

QEfficient/base/common.py

@@ -18,7 +18,7 @@
 from transformers import AutoConfig
 
 from QEfficient.base.modeling_qeff import QEFFBaseModel
-from QEfficient.transformers.modeling_utils import MODEL_CLASS_MAPPING
+from QEfficient.transformers.modeling_utils import EXTERNAL_MODEL_CLASS_MAPPING, MODEL_CLASS_MAPPING
 from QEfficient.utils import login_and_download_hf_lm
 
 
@@ -40,9 +40,13 @@ def from_pretrained(cls, pretrained_model_name_or_path: str, *args, **kwargs) ->
         """
         Downloads HuggingFace model if already doesn't exist locally, returns QEFFAutoModel object based on type of model.
         """
-        config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True)
 
-        class_name = MODEL_CLASS_MAPPING.get(config.__class__.__name__, None)
+        # class_name = MODEL_CLASS_MAPPING.get(config.__class__.__name__, None) OR MODEL_EXTERNAL_CLASS_MAPPING(config.__class__.__name__)
+        class_name = (
+            MODEL_CLASS_MAPPING.get(config.__class__.__name__, None)
+            or EXTERNAL_MODEL_CLASS_MAPPING[config.__class__.__name__]
+        )
         if class_name:
             module = __import__("QEfficient.transformers.models.modeling_auto")
             model_class = getattr(module, class_name)
@@ -61,6 +65,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str, *args, **kwargs) ->
             pretrained_model_name_or_path=(local_model_dir if local_model_dir else pretrained_model_name_or_path),
             token=hf_token,
             continuous_batching=continuous_batching,
+            trust_remote_code=True,
             **kwargs,
         )
         return qeff_model

QEfficient/base/modeling_qeff.py

@@ -52,7 +52,7 @@ def __init__(self, model: torch.nn.Module) -> None:
         self.onnx_path: Optional[str] = None
         self.qpc_path: Optional[str] = None
         self.qpc_session: Optional[QAICInferenceSession] = None
-
+        model = model.to(torch.float32)
         # Apply the transformations
         any_transformed = False
         for transform in self._pytorch_transforms:

QEfficient/transformers/modeling_utils.py

@@ -283,6 +283,7 @@ def build_model_class_mapping(auto_model_class, qeff_class_name):
     }
 
 
+EXTERNAL_MODEL_CLASS_MAPPING = {"Grok1Config": "QEFFAutoModelForCausalLM"}
 MODEL_CLASS_MAPPING = {
     **build_model_class_mapping(mapping.AutoModelForCausalLM, "QEFFAutoModelForCausalLM"),
     **build_model_class_mapping(mapping.AutoModelForImageTextToText, "QEFFAutoModelForImageTextToText"),

QEfficient/transformers/models/grok_1/modeling_grok1.py

@@ -28,12 +28,25 @@ class QEFFGrok1CustomRMSNormAIC(nn.Module):
     """
 
     def forward(self, hidden_states):
+        """
+        Forward pass of the RMSNorm module.
+
+        Args:
+            hidden_states (torch.Tensor): Input tensor to be normalized.
+
+        Returns:
+            torch.Tensor: Normalized tensor.
+        """
         return CustomRMSNormFunc.apply(
             hidden_states, self.scale, self.variance_epsilon if hasattr(self, "variance_epsilon") else self.eps
         )
 
 
 class QEffGrok1MultiHeadAttention(nn.Module):
+    """
+    Multi-head attention module.
+    """
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -46,6 +59,22 @@ def forward(
         use_cache: bool = False,
         **kwargs,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        """
+        Forward pass of the multi-head attention module.
+
+        Args:
+            hidden_states (torch.Tensor): Input tensor.
+            layer_idx (int): Layer index.
+            attention_mask (Optional[torch.Tensor], optional): Attention mask. Defaults to None.
+            position_ids (Optional[torch.LongTensor], optional): Position ids. Defaults to None.
+            past_key_value (Optional[Tuple[torch.Tensor]], optional): Past key value. Defaults to None.
+            batch_index (Optional[torch.LongTensor], optional): Batch index. Defaults to None.
+            output_attentions (bool, optional): Whether to output attentions. Defaults to False.
+            use_cache (bool, optional): Whether to use cache. Defaults to False.
+
+        Returns:
+            Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: Attention output, attention weights, and past key value.
+        """
         bsz, q_len, _ = hidden_states.size()
 
         query_states = self.q_proj(hidden_states)
@@ -101,7 +130,20 @@ def forward(
 
 
 class QEffGrok1MoeBlock(nn.Module):
+    """
+    Mixture of experts (MoE) block.
+    """
+
     def forward(self, hidden_states: torch.Tensor):
+        """
+        Forward pass of the MoE block.
+
+        Args:
+            hidden_states (torch.Tensor): Input tensor.
+
+        Returns:
+            torch.Tensor: MoE output.
+        """
         batch_size, sequence_length, hidden_dim = hidden_states.shape
         hidden_states = hidden_states.view(-1, hidden_dim)
         router_logits = self.gate(hidden_states)
@@ -116,8 +158,8 @@ def forward(self, hidden_states: torch.Tensor):
             torch.nn.functional.one_hot(selected_experts[:, 1], num_classes=self.num_experts).bool().T.unsqueeze(-1)
         )
 
-        gateupout1 = torch.zeros(hidden_states.shape[0], 32768)  # T, hs
-        gateupout2 = torch.zeros(hidden_states.shape[0], 32768)  # T, hs
+        gateupout1 = torch.zeros(hidden_states.shape[0], self.ffn_dim)  # T, hs
+        gateupout2 = torch.zeros(hidden_states.shape[0], self.ffn_dim)  # T, hs
         for expert_idx in range(self.num_experts):
             expert_layer = self.experts[expert_idx]
             current_expert_output = expert_layer.act_fn(expert_layer.linear(hidden_states)) * expert_layer.linear_v(
@@ -150,6 +192,16 @@ def forward(self, hidden_states: torch.Tensor):
 
 
 class QEffGrok1DecoderLayer(nn.Module):
+    """
+    Decoder block of Grok1 model.
+    """
+
+    def __qeff_init__(self):
+        """
+        Assigning extra args to Moe block of decoder.
+        """
+        self.moe_block.ffn_dim = self.config.intermediate_size
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -162,6 +214,22 @@ def forward(
         use_cache: Optional[bool] = False,
         **kwargs,
     ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        """
+        Initialize the decoder layer.
+
+        Args:
+            hidden_states (torch.Tensor): Input tensor.
+            attention_mask (Optional[torch.Tensor], optional): Attention mask. Defaults to None.
+            position_ids (Optional[torch.LongTensor], optional): Position ids. Defaults to None.
+            past_key_value (Optional[Tuple[torch.Tensor]], optional): Past key value. Defaults to None.
+            batch_index (Optional[torch.LongTensor], optional): Batch index. Defaults to None.
+            output_attentions (Optional[bool], optional): Whether to output attentions. Defaults to False.
+            output_router_logits (Optional[bool], optional): Whether to output router logits. Defaults to False.
+            use_cache (Optional[bool], optional): Whether to use cache. Defaults to False.
+
+        Returns:
+            Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: Decoder output, attention weights, and past key value.
+        """
         residual = hidden_states
         hidden_states = self.pre_attn_norm(hidden_states)
         hidden_states, attention_weights, present_key_value = self.attn(
@@ -194,9 +262,17 @@ def forward(
 
 
 class QEffGrok1Model(nn.Module):
+    """
+    Grok1 model
+    """
+
     def __qeff_init__(self):
+        """
+        Initialize the extra args to model.
+        """
         for idx, layer in enumerate(self.layers):
             layer.layer_idx = idx
+            layer.config = self.config
 
     def forward(
         self,
@@ -212,6 +288,24 @@ def forward(
         output_router_logits: Optional[bool] = None,
         return_dict: Optional[bool] = None,
     ) -> Union[Tuple, MoeModelOutputWithPast]:
+        """
+        Forward pass of the Grok1 model.
+        Args:
+            input_ids (torch.LongTensor, optional): Input ids. Defaults to None.
+            attention_mask (Optional[torch.Tensor], optional): Attention mask. Defaults to None.
+            position_ids (Optional[torch.LongTensor], optional): Position ids. Defaults to None.
+            past_key_values (Optional[List[torch.FloatTensor]], optional): Past key values. Defaults to None.
+            batch_index (Optional[torch.LongTensor], optional): Batch index. Defaults to None.
+            inputs_embeds (Optional[torch.FloatTensor], optional): Input embeddings. Defaults to None.
+            use_cache (Optional[bool], optional): Whether to use cache. Defaults to None.
+            output_attentions (Optional[bool], optional): Whether to output attentions. Defaults to None.
+            output_hidden_states (Optional[bool], optional): Whether to output hidden states. Defaults to None.
+            output_router_logits (Optional[bool], optional): Whether to output router logits. Defaults to None.
+            return_dict (Optional[bool], optional): Whether to return a dictionary. Defaults to None.
+
+        Returns:
+            Union[Tuple, MoeModelOutputWithPast]: Model output.
+        """
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
@@ -294,6 +388,10 @@ def forward(
 
 
 class QEffGrok1ModelForCausalLM(nn.Module):
+    """
+    Grok model for causal language modeling.
+    """
+
     def forward(
         self,
         input_ids: torch.LongTensor = None,
@@ -310,6 +408,26 @@ def forward(
         return_dict: Optional[bool] = None,
         **kwargs,
     ):
+        """
+        Forward pass for Grok model for causal language modeling
+
+        Args:
+            input_ids (torch.LongTensor, optional): Input ids. Defaults to None.
+            attention_mask (Optional[torch.Tensor], optional): Attention mask. Defaults to None.
+            position_ids (Optional[torch.LongTensor], optional): Position ids. Defaults to None.
+            past_key_values (Optional[List[torch.FloatTensor]], optional): Past key values. Defaults to None.
+            batch_index (Optional[torch.LongTensor], optional): Batch index. Defaults to None.
+            inputs_embeds (Optional[torch.FloatTensor], optional): Input embeddings. Defaults to None.
+            labels (Optional[torch.LongTensor], optional): Labels. Defaults to None.
+            use_cache (Optional[bool], optional): Whether to use cache. Defaults to None.
+            output_attentions (Optional[bool], optional): Whether to output attentions. Defaults to None.
+            output_hidden_states (Optional[bool], optional): Whether to output hidden states. Defaults to None.
+            output_router_logits (Optional[bool], optional): Whether to output router logits. Defaults to None.
+            return_dict (Optional[bool], optional): Whether to return a dictionary. Defaults to None.
+
+        Returns:
+            MoeCausalLMOutputWithPast: Model output.
+        """
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_router_logits = (
             output_router_logits if output_router_logits is not None else self.config.output_router_logits

QEfficient/transformers/models/pytorch_transforms.py

@@ -499,7 +499,10 @@ class KVCacheExternalModuleMapperTransform(ExternalModuleMapperTransform):
             "forward": QEffGrok1Model.forward,
             "__qeff_init__": QEffGrok1Model.__qeff_init__,
         },
-        "DecoderLayer": {"forward": QEffGrok1DecoderLayer.forward},
+        "DecoderLayer": {
+            "forward": QEffGrok1DecoderLayer.forward,
+            "__qeff_init__": QEffGrok1DecoderLayer.__qeff_init__,
+        },
         "MoeBlock": {"forward": QEffGrok1MoeBlock.forward},
         "MultiHeadAttention": {
             "forward": QEffGrok1MultiHeadAttention.forward,

README.md

@@ -6,6 +6,7 @@
 ---
 
 *Latest news* :fire: <br>
+- [06/2025] Added support of model `hpcai-tech/grok-1` [hpcai-tech/grok-1](https://huggingface.co/hpcai-tech/grok-1)
 - [03/2025] Added support for swiftkv model [Snowflake/Llama-3.1-SwiftKV-8B-Instruct](https://huggingface.co/Snowflake/Llama-3.1-SwiftKV-8B-Instruct)
 - [02/2025] [VLMs support](https://github.com/quic/efficient-transformers/pull/267) added for the models [InternVL-1B](https://huggingface.co/OpenGVLab/InternVL2_5-1B), [Llava](https://huggingface.co/llava-hf/llava-1.5-7b-hf) and [Mllama](https://huggingface.co/meta-llama/Llama-3.2-11B-Vision-Instruct)
 - [01/2025] [FP8 models support](https://huggingface.co/collections/neuralmagic/fp8-llms-for-vllm-666742ed2b78b7ac8df13127) Added support for inference of FP8 models.

docs/source/validate.md

@@ -34,7 +34,7 @@
 | **QwenForCausalLM**     | DeepSeek-R1-Distill-Qwen | [DeepSeek-R1-Distill-Qwen-32B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B)                                                   | ✔️          |
 |                         | Qwen2, Qwen2.5     | [Qwen/Qwen2-1.5B-Instruct](https://huggingface.co/Qwen/Qwen2-1.5B-Instruct)                                                            | ✔️          |
 | **LlamaSwiftKVForCausalLM**  | swiftkv            | [Snowflake/Llama-3.1-SwiftKV-8B-Instruct](https://huggingface.co/Snowflake/Llama-3.1-SwiftKV-8B-Instruct)                                                  | ✔️          |
-
+| **Grok1ModelForCausalLM**  |  grok-1          | [hpcai-tech/grok-1](https://huggingface.co/hpcai-tech/grok-1)                                                  | ✔️          |
 ## Embedding Models
 
 ### Text Embedding Task