Fixed torch patch for subfunction with VLMs (quic#750)

Signed-off-by: abhishek-singh591 <sabhis@qti.qualcomm.com>

Author: abhishek-singh591

QEfficient/peft/auto.py

@@ -289,8 +289,8 @@ def export(self, export_dir: Optional[str] = None, **kwargs) -> str:
 
         return self._export(
             example_inputs,
-            output_names,
-            dynamic_axes,
+            output_names=output_names,
+            dynamic_axes=dynamic_axes,
             do_constant_folding=False,  # To avoid merging adapter weights with base weights
             onnx_transform_kwargs={"adapter_name": self.model.active_adapter},
             export_dir=export_dir,

QEfficient/peft/lora/auto.py

@@ -384,8 +384,8 @@ def export(self, export_dir: Optional[str] = None, **kwargs) -> str:
 
         return self._export(
             example_inputs,
-            output_names,
-            dynamic_axes,
+            output_names=output_names,
+            dynamic_axes=dynamic_axes,
             export_dir=export_dir,
             **kwargs,
         )

QEfficient/transformers/models/modeling_auto.py

@@ -344,8 +344,8 @@ def export(self, export_dir: Optional[str] = None, **kwargs) -> str:
 
         return self._export(
             example_inputs,
-            output_names,
-            dynamic_axes,
+            output_names=output_names,
+            dynamic_axes=dynamic_axes,
             export_dir=export_dir,
             use_onnx_subfunctions=kwargs.get("use_onnx_subfunctions", False),
         )
@@ -623,8 +623,8 @@ def export(self, inputs, output_names, dynamic_axes, export_dir=None, offload_pt
         """
         return self._export(
             inputs,
-            output_names,
-            dynamic_axes,
+            output_names=output_names,
+            dynamic_axes=dynamic_axes,
             export_dir=export_dir,
             offload_pt_weights=offload_pt_weights,
             use_onnx_subfunctions=kwargs.get("use_onnx_subfunctions", False),
@@ -768,8 +768,8 @@ def export(self, inputs, output_names, dynamic_axes, export_dir=None, offload_pt
         """
         return self._export(
             inputs,
-            output_names,
-            dynamic_axes,
+            output_names=output_names,
+            dynamic_axes=dynamic_axes,
             export_dir=export_dir,
             offload_pt_weights=offload_pt_weights,
             use_onnx_subfunctions=kwargs.get("use_onnx_subfunctions", False),
@@ -1708,8 +1708,8 @@ def export(
         output_names = self.model.get_output_names()
         return self._export(
             inputs,
-            output_names,
-            dynamic_axes,
+            output_names=output_names,
+            dynamic_axes=dynamic_axes,
             export_dir=export_dir,
             use_onnx_subfunctions=use_onnx_subfunctions,
         )
@@ -2706,8 +2706,8 @@ def export(
             )
         return self._export(
             example_inputs,
-            output_names,
-            dynamic_axes,
+            output_names=output_names,
+            dynamic_axes=dynamic_axes,
             export_dir=export_dir,
             use_onnx_subfunctions=kwargs.get("use_onnx_subfunctions", False),
             offload_pt_weights=kwargs.get("offload_pt_weights", True),
@@ -3300,8 +3300,8 @@ def export(self, export_dir: Optional[str] = None, **kwargs) -> str:
         output_names = self.model.get_output_names()
         return self._export(
             inputs,
-            output_names,
-            dynamic_axes,
+            output_names=output_names,
+            dynamic_axes=dynamic_axes,
             export_dir=export_dir,
             use_onnx_subfunctions=kwargs.get("use_onnx_subfunctions", False),
         )
@@ -3676,8 +3676,8 @@ def export(self, export_dir: Optional[str] = None, **kwargs) -> str:
 
         return self._export(
             example_inputs,
-            output_names,
-            dynamic_axes,
+            output_names=output_names,
+            dynamic_axes=dynamic_axes,
             export_dir=export_dir,
             use_onnx_subfunctions=kwargs.get("use_onnx_subfunctions", False),
         )

QEfficient/utils/export_utils.py

@@ -161,15 +161,18 @@ def _setup_onnx_subfunctions(qeff_model, args, kwargs):
     # Apply torch patches for subfunction support
     apply_torch_patches()
     InvalidIndexProvider.SUBFUNC_ENABLED = True
+
     # Transform output names for subfunction compatibility
     if "output_names" in kwargs:
         kwargs["output_names"] = [
             re.sub("_RetainedState", "_InternalRetainedState", name) for name in kwargs["output_names"]
         ]
     else:
-        args = list(args)
-        args[1] = [re.sub("_RetainedState", "_InternalRetainedState", name) for name in args[1]]
-        args = tuple(args)
+        warnings.warn(
+            "ONNX subfunctions are enabled, but no retained-state output names were found to rewrite. "
+            "Ensure `output_names` includes key/value retained states if subfunction compatibility is required."
+        )
+
     # Add subfunction-specific ONNX transforms
     qeff_model._onnx_transforms.append(RenameFunctionOutputsTransform)
     qeff_model._onnx_transforms.append(CustomOpTransform)

QEfficient/utils/torch_patches.py

@@ -11,6 +11,8 @@
 import torch.onnx.utils as onnx_utils
 from torch import _C
 
+from QEfficient.utils.logging_utils import logger
+
 # Store original references before patching
 _original_setup_trace_module_map = onnx_utils._setup_trace_module_map
 _original_get_module_attributes = getattr(onnx_utils, "_get_module_attributes", None)
@@ -37,9 +39,10 @@ def _track_module_attributes_forward_hook(module, input, output):
             if hasattr(module, attr_name):
                 onnx_attrs = getattr(module, attr_name)
                 delattr(module, attr_name)
-            # FIX: use empty dict to avoid type mismatch
-            onnx_attrs = {}
-            _C._jit_pass_onnx_track_scope_attributes(graph, onnx_attrs)
+            try:
+                _C._jit_pass_onnx_track_scope_attributes(graph, onnx_attrs)
+            except Exception as e:
+                logger.warning(f"Failed to track ONNX scope attributes: {e}. Skipping this step.")
 
         for m in model.modules():
             m.register_forward_hook(_track_module_attributes_forward_hook)