Minor fixes-1

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

Author: quic-amitraj

QEfficient/base/modeling_qeff.py

@@ -20,7 +20,7 @@
 import torch
 
 from QEfficient.base.onnx_transforms import OnnxTransform
-from QEfficient.base.pytorch_transforms import PytorchTransform
+from QEfficient.base.pytorch_transforms import PytorchTransform, append_tranform
 from QEfficient.compile.qnn_compiler import compile as qnn_compile
 from QEfficient.generation.cloud_infer import QAICInferenceSession
 from QEfficient.utils import constants, dump_qconfig
@@ -46,6 +46,7 @@ class QEFFBaseModel(ABC):
     def _transform_names(cls) -> List[str]:
         return [x.__name__ for x in cls._pytorch_transforms + cls._onnx_transforms]
 
+    @append_tranform
     def __init__(self, model: torch.nn.Module) -> None:
         super().__init__()
         self.model = model

QEfficient/base/pytorch_transforms.py

@@ -126,7 +126,9 @@ class SplitGateUpWeightsTransform(PytorchTransform):
     @classmethod
     def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]:
         transformed = False
-        model = model.language_model
+
+        model = model.language_model if hasattr(model, "language_model") else model
+
         num_layers = len(model.model.layers)
         delete_fused_key = True
         sd = model.state_dict()
@@ -158,3 +160,16 @@ def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]:
             print(f"[layer {layer_idx:02d}] loaded gate_proj & up_proj from fused tensor  (shape {fused.shape})")
             transformed = True
         return model, transformed
+
+
+VLM_SPLIT_GATE_UP_WEIGHTS = ["Llama4ForConditionalGeneration", "Llama4TextModel"]
+
+
+def append_tranform(func):
+    def wrapper(*args, **kwargs):
+        model_class = args[1].model.__class__.__name__ if hasattr(args[1], "model") else args[1].__class__.__name__
+        if model_class in VLM_SPLIT_GATE_UP_WEIGHTS:
+            args[0]._pytorch_transforms.append(SplitGateUpWeightsTransform)
+        return func(*args, **kwargs)
+
+    return wrapper

QEfficient/transformers/modeling_utils.py

@@ -384,6 +384,3 @@ def _create_causal_mask(
         attention_mask = attention_mask.unsqueeze(1)
 
     return attention_mask
-
-
-VLM_SPLIT_GATE_UP_WEIGHTS = ["Llama4ForConditionalGeneration"]

QEfficient/transformers/models/modeling_auto.py

@@ -27,15 +27,13 @@
 import QEfficient
 from QEfficient.base.modeling_qeff import QEFFBaseModel
 from QEfficient.base.onnx_transforms import FP16ClipTransform, SplitTensorsTransform
-from QEfficient.base.pytorch_transforms import SplitGateUpWeightsTransform
 from QEfficient.generation.cloud_infer import QAICInferenceSession
 from QEfficient.generation.text_generation_inference import (
     CloudAI100ExecInfoNew,
     PerfMetrics,
     calculate_latency,
     get_compilation_dims,
 )
-from QEfficient.transformers.modeling_utils import VLM_SPLIT_GATE_UP_WEIGHTS
 from QEfficient.transformers.models.pytorch_transforms import (
     CustomOpsTransform,
     KVCacheModuleMethodMapperTransform,
@@ -469,9 +467,6 @@ class QEffCausalLMForTextImageToTextModel(QEFFBaseModel):
     _onnx_transforms = [FP16ClipTransform, SplitTensorsTransform]
 
     def __init__(self, model):
-        if model.config.architectures[0] in VLM_SPLIT_GATE_UP_WEIGHTS:
-            self._pytorch_transforms.append(SplitGateUpWeightsTransform)
-
         super().__init__(model)
         self.model = model.get_qeff_language_decoder()
 

examples/llama4_lm_example.py

@@ -13,16 +13,12 @@
 from QEfficient.utils.constants import Constants
 from QEfficient.utils.run_utils import ApiRunner
 
-torch.manual_seed(42)
-
 model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct"
 model = Llama4ForCausalLM.from_pretrained(
     model_id, torch_dtype=torch.float32, use_cache=True, attn_implementation="eager"
 )
 model.eval()
 
-original_sd = model.state_dict()
-
 tokenizer = load_hf_tokenizer(pretrained_model_name_or_path=model_id)
 config = model.config
 batch_size = len(Constants.INPUT_STR)
@@ -37,7 +33,6 @@
 
 qeff_model = QEFFAutoModelForCausalLM(model)
 
-onnx_model_path = qeff_model.export()
 qpc_path = qeff_model.compile(
     prefill_seq_len=128,
     ctx_len=2048,