Merged PR 546 with this

abhishek-singh591 · abhishek-singh591 · commit 8f238cab2b70 · 2025-09-29T06:05:28.000Z
diff --git a/QEfficient/base/modeling_qeff.py b/QEfficient/base/modeling_qeff.py
@@ -47,11 +47,12 @@ class QEFFBaseModel(ABC):
     """
 
     _pytorch_transforms: List[PytorchTransform]
-    _onnx_transforms: List[OnnxTransform]
+    _onnx_transforms = ["FP16ClipTransform", "SplitTensorsTransform"]
 
     @classmethod
     def _transform_names(cls) -> List[str]:
-        return [x.__name__ for x in cls._pytorch_transforms + cls._onnx_transforms]
+        pytorch_names = [x.__name__ for x in cls._pytorch_transforms]
+        return pytorch_names + cls._onnx_transforms
 
     def __init__(self, model: torch.nn.Module, **kwargs) -> None:
         super().__init__()
@@ -321,9 +322,10 @@ def _export(
             }
             if onnx_transform_kwargs is not None:
                 transform_kwargs.update(onnx_transform_kwargs)
-
-            for transform in self._onnx_transforms:
-                model, transformed = transform.apply(model, **transform_kwargs)
+            # import pdb; pdb.set_trace()
+            transform_kwargs["transforms"] = self._onnx_transforms
+            # for transform in self._onnx_transforms:
+            model, transformed = OnnxTransform.apply(model, **transform_kwargs)
 
             model.metadata_props.append(
                 onnx.StringStringEntryProto(key="qeff_transforms", value=",".join(self._transform_names()))
diff --git a/QEfficient/base/onnx_transforms.py b/QEfficient/base/onnx_transforms.py
@@ -7,25 +7,29 @@
 
 import gc
 import logging
-from typing import Optional, Tuple
+import os
+import warnings
+from collections import namedtuple
+from concurrent.futures import ThreadPoolExecutor
+from typing import List, Optional, Tuple
 
 import numpy as np
-from onnx import ModelProto, external_data_helper, numpy_helper
+from onnx import ModelProto, TensorProto, external_data_helper, numpy_helper
 
 from QEfficient.utils.constants import ONNX_TRANSFORM_MEMORY_CLEANUP_INTERVAL
 
 logger = logging.getLogger(__name__)
 
 
-class OnnxTransform:
+class BaseOnnxTransform:
     """
     OnnxTransform is the base class for graph modifications on exported onnx.
     """
 
     _external_data_loaded_cache = {}  # Dict[int, bool]
 
     def __init__(self):
-        raise TypeError("Transform classes are not to be instantiated. Directly use the `apply` method.")
+        raise TypeError("Transform classes are not to be instantiated. Use the `apply` method directly.")
 
     @classmethod
     def apply(cls, model: ModelProto, **kwargs) -> Tuple[ModelProto, bool]:
@@ -47,15 +51,11 @@ def _check_external_data_loaded(cls, model: ModelProto) -> bool:
         :param model: The ONNX model to check
         :returns: True if external data is already loaded, False otherwise
         """
-        # Use object ID as key instead of the object itself
         model_id = id(model)
-        # Return cached result if available
         if model_id in cls._external_data_loaded_cache:
             return cls._external_data_loaded_cache[model_id]
 
-        # Load the model if not already loaded
         for tensor in external_data_helper._get_all_tensors(model):
-            # Check if tensor has external data but no raw data loaded
             if len(tensor.external_data) > 0 and not tensor.HasField("raw_data"):
                 cls._external_data_loaded_cache[model_id] = False
                 return False
@@ -77,6 +77,13 @@ def _load_external_data(cls, model: ModelProto, onnx_base_dir: Optional[str] = N
         else:
             logger.info("External data already loaded (or cached). Skipping bulk load.")
 
+    @classmethod
+    def _cleanup_memory(cls):
+        """
+        Force garbage collection to free up memory after tensor processing.
+        """
+        gc.collect()
+
     @classmethod
     def _cleanup_external_data_and_cache(cls, model: ModelProto):
         """
@@ -94,108 +101,99 @@ def _cleanup_external_data_and_cache(cls, model: ModelProto):
 
         logger.info("External data and cache cleaned up.")
 
-    @classmethod
-    def _cleanup_memory(cls):
-        """
-        Force garbage collection to free up memory after tensor processing.
-        """
-        gc.collect()
-
-
-class FP16ClipTransform(OnnxTransform):
-    """
-    Clips the tensor values to be in FP16 range, but preserves -inf values.
-    """
 
+class OnnxTransform(BaseOnnxTransform):
     @classmethod
-    def apply(cls, model: ModelProto, *, onnx_base_dir: Optional[str] = None, **kwargs) -> Tuple[ModelProto, bool]:
-        """
-        :param onnx_base_dir: Base directory to load tensors
-        """
+    def apply(
+        cls,
+        model: ModelProto,
+        *,
+        transforms: List[str],
+        model_name: str = "",
+        onnx_base_dir: Optional[str] = None,
+        file_chunk_size: int = 10 * 2**30,
+        size_threshold: int = 1024,
+        **kwargs,
+    ) -> Tuple[ModelProto, bool]:
+        if len(transforms) == 0:
+            warnings.warn("Transform is empty. Skipping transformation.")
+            return model, False
+
         try:
-            # --- FIX: Ensure external data is loaded efficiently BEFORE processing ---
             cls._load_external_data(model, onnx_base_dir)
+            tensors = external_data_helper._get_all_tensors(model)
 
-            finfo = np.finfo(np.float16)
-            fp16_max = finfo.max
-            fp16_min = finfo.min
-            transformed = False
+            TensorInfo = namedtuple("TensorInfo", ["tensor", "tsize"])
+            tensor_infos = [
+                TensorInfo(tensor, len(tensor.raw_data) if tensor.HasField("raw_data") else 0) for tensor in tensors
+            ]
 
-            processed_count = 0
-            for tensor in external_data_helper._get_all_tensors(model):
-                nptensor = numpy_helper.to_array(tensor)  # Removed onnx_base_dir as data is already loaded
-                if nptensor.dtype == np.float32 and (np.any(nptensor > fp16_max) or np.any(nptensor < fp16_min)):
-                    neg_inf_mask = np.isinf(nptensor) & (nptensor < 0)
-                    clipped_tensor = np.clip(nptensor, fp16_min, fp16_max)
+            fp16_min, fp16_max = np.finfo(np.float16).min, np.finfo(np.float16).max
+            file_num_tracker = {"num": 0, "size": 0}
 
-                    # Restore -inf values
-                    if neg_inf_mask.any():
-                        clipped_tensor = np.where(neg_inf_mask, np.float32("-inf"), clipped_tensor)
+            # Track which transforms were requested and which were actually applied
+            requested_transforms = set(transforms)
+            applied_transforms = {name: False for name in requested_transforms}
 
-                    new_tensor = numpy_helper.from_array(clipped_tensor, tensor.name)
-                    tensor.CopyFrom(new_tensor)
-                    transformed = True
+            def process_tensor(index_info: Tuple[int, TensorInfo]) -> List[str]:
+                idx, info = index_info
+                tensor, tsize = info
 
-                    del neg_inf_mask, clipped_tensor, new_tensor
+                local_applied = []
 
-                del nptensor
-                processed_count += 1
+                if "FP16ClipTransform" in requested_transforms:
+                    if cls._clip_tensor(tensor, fp16_min, fp16_max):
+                        local_applied.append("FP16ClipTransform")
 
-                if processed_count % ONNX_TRANSFORM_MEMORY_CLEANUP_INTERVAL == 0:
-                    cls._cleanup_memory()
+                if "SplitTensorsTransform" in requested_transforms and tsize > size_threshold:
+                    if file_num_tracker["size"] + tsize > file_chunk_size:
+                        file_num_tracker["num"] += 1
+                        file_num_tracker["size"] = tsize
+                    else:
+                        file_num_tracker["size"] += tsize
 
-            return model, transformed
-        finally:
-            # Ensure cleanup happens even if an exception occurs
-            cls._cleanup_memory()
+                    cls._split_tensor(tensor, model_name, file_num_tracker["num"])
+                    local_applied.append("SplitTensorsTransform")
 
+                if (idx + 1) % ONNX_TRANSFORM_MEMORY_CLEANUP_INTERVAL == 0:
+                    cls._cleanup_memory()
 
-class SplitTensorsTransform(OnnxTransform):
-    """
-    Split external tensors file
-    """
+                return local_applied
 
-    @classmethod
-    def apply(
-        cls,
-        model: ModelProto,
-        *,
-        model_name: str,
-        onnx_base_dir: Optional[str] = None,
-        file_chunk_size: int = 10 * 2**30,  # 10 GiB
-        size_threshold: int = 1024,
-        **kwargs,
-    ) -> Tuple[ModelProto, bool]:
-        """
-        :param model_name: Used for naming external files. i.e. {model_name}_0.onnx.data
-        :param onnx_base_dir: Base directory to load tensors (if not already loaded).
-        :param file_chunk_size: Chunk size to split external files into.
-        :param size_threshold: Only tensors greater than this threshold (in bytes) will be saved externally.
-        """
-        try:
-            file_num = 0
-            current_file_size = 0
-            transformed = False
+            with ThreadPoolExecutor(max_workers=os.cpu_count() * 4) as executor:
+                results = list(executor.map(process_tensor, enumerate(tensor_infos)))
 
-            # --- Adjustment: The initial check and load will now use the new bulk loader ---
-            # This will either use the cache (if FP16ClipTransform loaded it) or perform the bulk load itself.
-            cls._load_external_data(model, onnx_base_dir)
+            for result in results:
+                for transform_name in result:
+                    applied_transforms[transform_name] = True
 
-            processed_count = 0
-            for tensor in external_data_helper._get_all_tensors(model):
-                if tensor.HasField("raw_data") and ((tsize := len(tensor.raw_data)) > size_threshold):
-                    transformed = True
-                    current_file_size += tsize
-                    if current_file_size > file_chunk_size:
-                        file_num += 1
-                        current_file_size = tsize
-                    external_data_helper.set_external_data(tensor, f"{model_name}_{file_num}.onnx.data")
-
-                processed_count += 1
-                if processed_count % ONNX_TRANSFORM_MEMORY_CLEANUP_INTERVAL == 0:
-                    cls._cleanup_memory()
+            for name in requested_transforms:
+                if applied_transforms[name]:
+                    logger.info(f"Transform '{name}' was applied.")
+                else:
+                    logger.warning(f"Transform '{name}' was requested but not applied.")
+
+            return model, any(applied_transforms.values())
 
-            return model, transformed
         finally:
-            # Ensure cleanup happens even if an exception occurs
             cls._cleanup_memory()
+
+    @staticmethod
+    def _clip_tensor(tensor, fp16_min, fp16_max) -> bool:
+        if tensor.data_type != TensorProto.FLOAT:
+            return False
+
+        nptensor = numpy_helper.to_array(tensor)
+        if np.any(nptensor > fp16_max) or np.any(nptensor < fp16_min):
+            neg_inf_mask = np.isinf(nptensor) & (nptensor < 0)
+            clipped_tensor = np.clip(nptensor, fp16_min, fp16_max)
+            if neg_inf_mask.any():
+                clipped_tensor = np.where(neg_inf_mask, np.float32("-inf"), clipped_tensor)
+            new_tensor = numpy_helper.from_array(clipped_tensor, tensor.name)
+            tensor.CopyFrom(new_tensor)
+            return True
+        return False
+
+    @staticmethod
+    def _split_tensor(tensor, model_name: str, file_num: int) -> None:
+        external_data_helper.set_external_data(tensor, f"{model_name}_{file_num}.onnx.data")
diff --git a/QEfficient/exporter/export_utils.py b/QEfficient/exporter/export_utils.py
@@ -17,7 +17,7 @@
 import torch
 from onnx import external_data_helper
 
-from QEfficient.base.onnx_transforms import FP16ClipTransform
+from QEfficient.base.onnx_transforms import OnnxTransform
 
 
 def export_onnx(
@@ -218,7 +218,13 @@ def fix_onnx_fp16(
         :str: Updated base name of exported ONNX model.
     """
     model = onnx.load(os.path.join(gen_models_path, f"{model_base_name}.onnx"))
-    model, fp16_fix = FP16ClipTransform.apply(model, onnx_base_dir=gen_models_path)
+    if "model" in locals():
+        OnnxTransform._cleanup_external_data_and_cache(gen_models_path)
+        OnnxTransform._cleanup_memory()
+
+    model, fp16_fix = OnnxTransform.apply(
+        model, model_name="", onnx_base_dir=gen_models_path, transforms=["FP16ClipTransform"]
+    )
 
     if fp16_fix:
         # Save FP16 model
diff --git a/QEfficient/peft/auto.py b/QEfficient/peft/auto.py
@@ -18,7 +18,7 @@
 from transformers.generation.streamers import BaseStreamer
 
 from QEfficient.base.modeling_qeff import QEFFBaseModel
-from QEfficient.base.onnx_transforms import FP16ClipTransform, OnnxTransform, SplitTensorsTransform
+from QEfficient.base.onnx_transforms import BaseOnnxTransform, OnnxTransform
 from QEfficient.base.pytorch_transforms import PytorchTransform
 from QEfficient.generation.cloud_infer import QAICInferenceSession
 from QEfficient.peft.lora import QEffAutoLoraModelForCausalLM
@@ -58,7 +58,7 @@ class QEffAutoPeftModelForCausalLM(QEFFBaseModel):
     """
 
     _pytorch_transforms: List[PytorchTransform] = [CustomOpsTransform, KVCacheTransform, PeftModelInputsTransform]
-    _onnx_transforms: List[OnnxTransform] = [FP16ClipTransform, AdapterWeightsToInputsTransform, SplitTensorsTransform]
+    _onnx_transforms: List[BaseOnnxTransform] = [OnnxTransform, AdapterWeightsToInputsTransform]
     _hf_auto_class = AutoPeftModelForCausalLM
 
     def __init__(self, model: nn.Module):
diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py
diff --git a/tests/base/test_onnx_transforms.py b/tests/base/test_onnx_transforms.py
