Grok-1Modelling changes and On device sampling (#447)

The PR include on device sampling, Grok1, llama3.2 TF bug fix and bug
fix on model hashing for multi device

---------

Signed-off-by: Amit Raj <[email protected]>
Signed-off-by: Onkar Chougule <[email protected]>
Signed-off-by: Rishin Raj <[email protected]>
Signed-off-by: Abukhoyer Shaik <[email protected]>
Signed-off-by: Dipankar Sarkar <[email protected]>
Signed-off-by: Dhiraj Kumar Sah <[email protected]>
Signed-off-by: quic-sanising <[email protected]>
Co-authored-by: Amit Raj <[email protected]>
Co-authored-by: Onkar Chougule <[email protected]>
Co-authored-by: Abukhoyer Shaik <[email protected]>
Co-authored-by: Dipankar Sarkar <[email protected]>
Co-authored-by: Dhiraj Kumar Sah <[email protected]>
Co-authored-by: Sanidhya Singal <[email protected]>

Author: 7 people

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,12 @@ 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, *args, **kwargs)
 
-        class_name = MODEL_CLASS_MAPPING.get(config.__class__.__name__, None)
+        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)

QEfficient/base/modeling_qeff.py

@@ -299,6 +299,8 @@ def _compile(
 
         if num_speculative_tokens:
             compile_hash.update(to_hashable({"num_speculative_tokens": num_speculative_tokens}))
+        # Hash num_devices too, since default value would always be 1.
+        compile_hash.update(to_hashable(mdp_ts_num_devices))
 
         # Check if already compiled
         compile_hash = compile_hash.hexdigest()[:16]

QEfficient/base/pytorch_transforms.py

@@ -92,7 +92,7 @@ def mutate(cls, original_module: nn.Module, parent_module: nn.Module):
         raise NotImplementedError("Please implement your own method by inheriting this class")
 
 
-class ModuleMethodMapperTransform(PytorchTransform):
+class ExternalModuleMapperTransform(PytorchTransform):
     """
     Serves as base class for any transform that want to map a particular method of a class to a new method implementation.
     """
@@ -109,6 +109,10 @@ def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]:
             ):
                 for orig_method_name, mapped_method in repl_method_map.items():
                     setattr(module, orig_method_name, MethodType(mapped_method, module))
+
+                    if hasattr(module, "__qeff_init__"):
+                        module.__qeff_init__()
+
                     transformed = True
 
         return model, transformed

QEfficient/cloud/infer.py

@@ -111,6 +111,7 @@ def main(
     allow_mxint8_mdp_io: bool = False,
     enable_qnn: Optional[bool] = False,
     qnn_config: Optional[str] = None,
+    trust_remote_code: Optional[bool] = False,
     **kwargs,
 ) -> None:
     """
@@ -140,6 +141,7 @@ def main(
         :allow_mxint8_mdp_io (bool): Allows MXINT8 compression of MDP IO traffic. ``Defaults to False.``
         :enable_qnn (bool): Enables QNN Compilation. ``Defaults to False.``
         :qnn_config (str): Path of QNN Config parameters file. ``Defaults to None.``
+        :trust_remote_code (bool): Trust remote code execution. ``Defaults to False.``
         :kwargs: Pass any compiler option as input. Any flag that is supported by `qaic-exec` can be passed. Params are converted to flags as below:
                 -allocator_dealloc_delay=1 -> -allocator-dealloc-delay=1
                 -qpc_crc=True -> -qpc-crc
@@ -164,6 +166,7 @@ def main(
         hf_token=hf_token,
         full_batch_size=full_batch_size,
         local_model_dir=local_model_dir,
+        trust_remote_code=trust_remote_code,
     )
 
     image_path = kwargs.pop("image_path", None)
@@ -264,6 +267,12 @@ def main(
         action="store_true",
         help="Compress constant MatMul weights to MXFP6 E2M3, default is no compression",
     )
+    parser.add_argument(
+        "--trust_remote_code",
+        action="store_true",
+        default=False,
+        help="Enable trusting remote code when loading models. Default is False; set to True by passing this flag.",
+    )
     parser.add_argument(
         "--mxint8",
         "--mxint8_kv_cache",

QEfficient/transformers/modeling_utils.py

@@ -283,6 +283,8 @@ 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/__init__.py

@@ -0,0 +1,7 @@
+# -----------------------------------------------------------------------------
+#
+# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# ----------------------------------------------------------------------------
+