feat: remove fullgraph and get generate patch pipe back

Author: llcnt

src/pruna/algorithms/compilation/utils.py

@@ -556,6 +556,17 @@ def generate(self, *args, **kwargs) -> torch.Tensor:
         self.top_k = kwargs.pop("top_k", self.top_k)
         self.use_cache = kwargs.pop("use_cache", self.use_cache)
 
+        # Handle generation when the user does not provide max_new_tokens, but a generation_config.
+        # This also fixes the evaluation test.
+        generation_config = kwargs.pop("generation_config", None)
+        if (
+            generation_config is not None
+            and "max_new_tokens" not in kwargs
+            and hasattr(generation_config, "max_new_tokens")
+            and getattr(generation_config, "max_new_tokens") is not None
+        ):
+            kwargs["max_new_tokens"] = int(generation_config.max_new_tokens)
+
         # Log any kwargs that are not explicitly handled
         unhandled_kwargs = {
             k: v

src/pruna/engine/handler/handler_pipeline.py

@@ -43,99 +43,43 @@ def __init__(self, pipeline: Any = None, model_args: Optional[Dict[str, Any]] =
         self.pipeline = pipeline
         self.tokenizer = getattr(pipeline, "tokenizer", None) if pipeline else None
 
-        # Patch the pipeline's model generate method if it exists to handle generation_config properly
-        if pipeline and hasattr(pipeline, "model") and hasattr(pipeline.model, "generate"):
-            self._patch_generate_method(pipeline.model)
-
-        # Also patch the pipeline's __call__ method to handle evaluation contexts
-        if pipeline:
+        # Patch the pipeline's __call__ method to return logits for evaluation contexts
+        if pipeline is not None:
             self._patch_pipeline_call(pipeline)
 
-    def _patch_generate_method(self, model: Any) -> None:
-        """
-        Patch the model's generate method to extract parameters from generation_config.
-
-        This handles the case where CausalLMGenerator expects max_new_tokens as a direct
-        parameter but the pipeline passes it inside generation_config.
-
-        Parameters
-        ----------
-        model : Any
-            The model whose generate method to patch.
-        """
-        # Store the original generate method
-        original_generate = model.generate
-
-        def patched_generate(*args, **kwargs):
-            # Extract generation_config if present
-            generation_config = kwargs.get("generation_config")
-
-            if generation_config:
-                # Extract max_new_tokens from generation_config if not directly provided
-                if "max_new_tokens" not in kwargs and hasattr(generation_config, "max_new_tokens"):
-                    kwargs["max_new_tokens"] = generation_config.max_new_tokens
-
-                # Extract other parameters that CausalLMGenerator might need
-                if "temperature" not in kwargs and hasattr(generation_config, "temperature"):
-                    kwargs["temperature"] = generation_config.temperature
-
-                if "top_k" not in kwargs and hasattr(generation_config, "top_k"):
-                    kwargs["top_k"] = generation_config.top_k
-
-            # Call the original generate method with the extracted parameters
-            return original_generate(*args, **kwargs)
-
-        # Replace the generate method with our patched version
-        model.generate = patched_generate
-
     def _patch_pipeline_call(self, pipeline: Any) -> None:
         """
         Patch the pipeline's __call__ method to return logits for evaluation contexts.
 
-        When the input is tokenized tensors (evaluation context), we bypass the pipeline's
+        When the input is tokenized tensors (evaluation context), bypass the pipeline's
         text generation and return raw logits needed for perplexity calculation.
 
         Parameters
         ----------
         pipeline : Any
             The pipeline whose __call__ method to patch.
         """
-        # Store the original __call__ method
         original_call = pipeline.__call__
 
         def patched_call(*args, **kwargs):
-            # Check if we're being called with string inputs (normal generation) or tensor inputs (evaluation context)
             inputs = args[0] if len(args) > 0 else kwargs.get("inputs", kwargs.get("text_inputs"))
 
-            # If input is tensor or we detect evaluation context, return logits
+            # If input is tensor, return logits from the underlying model forward
             if hasattr(inputs, "shape") and hasattr(inputs, "dtype"):
-                # This is a tensor input - likely from evaluation pipeline
-                # Use the model's forward pass to get logits instead of text generation
                 try:
-                    # Prepare inputs for the model
                     if hasattr(pipeline, "model") and hasattr(inputs, "to"):
                         device = next(pipeline.model.parameters()).device
                         inputs = inputs.to(device)
-
-                        # Call the model's forward method directly to get logits
                         with torch.no_grad():
                             outputs = pipeline.model(input_ids=inputs)
-
-                        # Return logits in the format expected by perplexity metric
-                        if hasattr(outputs, "logits"):
-                            return outputs.logits
-                        else:
-                            return outputs
-
+                        return outputs.logits if hasattr(outputs, "logits") else outputs
                 except Exception as e:
                     pruna_logger.warning(f"Failed to get logits from model forward pass: {e}")
                     # Fallback to original pipeline behavior
                     pass
 
-            # For string inputs or fallback, use the original pipeline behavior
             return original_call(*args, **kwargs)
 
-        # Replace the pipeline's __call__ method with our patched version
         pipeline.__call__ = patched_call
 
     def prepare_inputs(
@@ -144,9 +88,9 @@ def prepare_inputs(
         """
         Prepare the inputs for the pipeline.
 
-        For text generation pipelines, this normally converts tokenized tensors back to strings.
-        However, for evaluation contexts (like perplexity), we keep tensors as-is since
-        the patched pipeline will handle them directly.
+        For text generation pipelines, string inputs are passed through. For evaluation
+        contexts (tensor token ids), tensors are passed through so the patched pipeline
+        can return logits for metrics like perplexity.
 
         Parameters
         ----------
@@ -160,27 +104,13 @@ def prepare_inputs(
         """
         x, _ = batch
 
-        # If x is already a string or list of strings, return as is
-        if isinstance(x, (str, list)) and all(isinstance(item, str) for item in (x if isinstance(x, list) else [x])):
-            return x
-
-        # If x is a tensor, we need to decide whether to convert to strings or keep as tensor
-        if isinstance(x, torch.Tensor):
-            # For evaluation contexts, we keep tensors as-is so the patched __call__
-            # can return logits instead of generated text
-            # The patched __call__ method will detect tensor inputs and handle accordingly
-            return x
-
-        # If no tokenizer available or x is not a tensor, return as-is
         return x
 
     def process_output(self, output: Any) -> Any:
         """
         Handle the output of the pipeline.
 
-        With our patched pipeline, the output is either:
-        - Logits tensor (for evaluation with tensor inputs)
-        - Generated text (for normal generation with string inputs)
+        Normalize common pipeline outputs. Returns generated text when available.
 
         Parameters
         ----------
@@ -192,14 +122,20 @@ def process_output(self, output: Any) -> Any:
         Any
             The processed output - pass through since patched __call__ handles the logic.
         """
-        # The patched pipeline __call__ method already handles the tensor vs string logic
-        # So we just pass through the output as-is
+        # HuggingFace text-generation pipeline returns list[dict]
+        if (
+            isinstance(output, list)
+            and len(output) > 0
+            and isinstance(output[0], dict)
+            and "generated_text" in output[0]
+        ):
+            return [o["generated_text"] for o in output]
         return output
 
     def log_model_info(self) -> None:
         """Log information about the inference handler."""
         pruna_logger.info(
             "Detected transformers pipeline. Using PipelineHandler.\n"
-            "- Tensor inputs will be converted to strings for pipeline processing.\n"
-            "- Pipeline outputs will be processed to extract generated text."
+            "- Token ids will be decoded to strings for pipeline processing.\n"
+            "- Pipeline outputs will be normalized to generated text when applicable."
         )

src/pruna/engine/handler/handler_utils.py

@@ -50,15 +50,17 @@ def register_inference_handler(model: Any) -> InferenceHandler:
 
     model_module = model._orig_mod.__module__ if hasattr(model, "_orig_mod") else model.__module__
 
-    # Check if it's a transformers pipeline
-    if hasattr(model, "__class__") and "Pipeline" in model.__class__.__name__:
+    # Prefer diffusers handler first to avoid routing diffusers pipelines to generic pipeline handler
+    if "diffusers" in model_module:
+        return DiffuserHandler(call_signature=inspect.signature(model.__call__))
+
+    # Transformers models and pipelines
+    if "transformers" in model_module and "Pipeline" in type(model).__name__:
         # Specific check for text generation pipelines
-        if "TextGeneration" in model.__class__.__name__:
+        if "TextGeneration" in type(model).__name__:
             return PipelineHandler(pipeline=model)
-        # For other pipelines, fallback to standard pipeline handler
+        # For other transformers pipelines, use PipelineHandler
         return PipelineHandler(pipeline=model)
-    elif "diffusers" in model_module:
-        return DiffuserHandler(call_signature=inspect.signature(model.__call__))
     elif "transformers" in model_module:
         return TransformerHandler()
     else:

tests/algorithms/test_combinations.py

@@ -48,7 +48,7 @@ def prepare_smash_config(self, smash_config: SmashConfig, device: str) -> None:
         ("sd_tiny_random", dict(quantizer="hqq_diffusers", compiler="torch_compile"), False, 'cmmd'),
         ("flux_tiny_random", dict(quantizer="hqq_diffusers", compiler="torch_compile"), False, 'cmmd'),
         ("sd_tiny_random", dict(quantizer="diffusers_int8", compiler="torch_compile"), False, 'cmmd'),
-        ("tiny_llama", dict(quantizer="gptq", compiler="torch_compile", torch_compile_fullgraph=False), True, 'perplexity'),
+        ("tiny_llama", dict(quantizer="gptq", compiler="torch_compile"), True, 'perplexity'),
         ("llama_3_tiny_random_as_pipeline", dict(quantizer="llm_int8", compiler="torch_compile"), True, 'perplexity'),
         ("flux_tiny_random", dict(cacher="pab", quantizer="hqq_diffusers"), False, 'cmmd'),
         ("flux_tiny_random", dict(cacher="pab", quantizer="diffusers_int8"), False, 'cmmd'),