Add: Failing test

rahul-tuli · rahul-tuli · commit cfa2abc6cd4f · 2025-04-23T15:59:21.000-04:00
Refactor: modify_save_pretrained

Signed-off-by: Rahul Tuli &lt;rtuli@redhat.com&gt;
diff --git a/src/llmcompressor/transformers/sparsification/compressed_tensors_utils.py b/src/llmcompressor/transformers/sparsification/compressed_tensors_utils.py
@@ -1,6 +1,6 @@
+import inspect
 import os
 import re
-import weakref
 from functools import wraps
 from typing import Dict, Optional
 
@@ -33,115 +33,101 @@
 __all__ = ["modify_save_pretrained"]
 
 
-def modify_save_pretrained(model: PreTrainedModel):
+def modify_save_pretrained(model: PreTrainedModel) -> None:
     """
     Overrides a PreTrainedModel's save_pretrained() method with a wrapped version that
-    supports compression. The new save_pretrained function performs the following saving
-    operations:
+    also supports compression params. The modified save_pretrained function performs the
+    following operations:
 
     1. Saves the model state, potentially in a compressed format
     2. Saves the recipe, appending any current recipes to existing recipe files
     3. Copies any necessary python files from the model cache
-    """
 
-    def save_pretrained_compressed(save_pretrained_method):
-        if getattr(save_pretrained_method, "_overridden", False):
-            # `model.save_pretrained` has already been replaced, return.
-            return save_pretrained_method
-
-        # Keep a weak reference to the model class and unbound save_pretrained
-        # method so we can call the original
-        model_ref = weakref.ref(save_pretrained_method.__self__)
-        original_save_pretrained = save_pretrained_method.__func__
-        model_class = model_ref().__class__
-        del save_pretrained_method
-
-        @wraps(original_save_pretrained)
-        def save_pretrained_wrapper(
-            save_directory: str,
-            sparsity_config: Optional[SparsityCompressionConfig] = None,
-            quantization_format: Optional[str] = None,
-            save_compressed: bool = True,
-            safe_serialization: bool = True,
-            skip_sparsity_compression_stats: bool = True,
-            disable_sparse_compression: bool = False,
-            **kwargs,
-        ):
-            """
-            Wrapper around PreTrainedModel.save_pretrained(), adds functionality for
-            saving models in a compressed format on disk. The compression format is
-            saved to the model's config file
-
-            :param save_directory: output directory to save model to
-            :param sparsity_config: optional sparsity config to compress model with,
-                if no config is provided it will be inferred from the model
-            :param quantization_format: optional compression format for quantized
-                models. If none is provided it will be inferred from the model
-            :param save_compressed: whether or not to compress the model on disk
-            :param skip_sparsity_compression_stats: whether to skip the calculation of
-                sparsity statistics (such as global sparsity and sparsity structure)
-                when saving a model in dense format
-            :param disable_sparse_compression: whether to skip sparse compression
-                during save, default is False
-            :param kwargs: additional kwargs to pass on to model.save_pretrained
-            """
-
-            # HACK: Override the dtype_byte_size function in transformers to
-            # support float8 types. Fix is posted upstream
-            # https://github.com/huggingface/transformers/pull/30488
-            transformers.modeling_utils.dtype_byte_size = new_dtype_byte_size
-
-            # state_dict gets passed in as a kwarg for FSDP models
-            state_dict = kwargs.pop("state_dict", None)
-            if state_dict is None:
-                logger.info("Fetching state_dict - this may take some time")
-                state_dict = get_state_dict_offloaded_model(model)
-
-            logger.info("Fetching compressor")
-            compressor = get_model_compressor(
-                model=model,
-                sparsity_config=sparsity_config,
-                quantization_format=quantization_format,
-                save_compressed=save_compressed,
-                skip_sparsity_compression_stats=skip_sparsity_compression_stats,
-                state_dict=state_dict,
-                disable_sparse_compression=disable_sparse_compression,
+    :param model: The model whose save_pretrained method will be modified
+    """
+    original = model.save_pretrained
+    # Avoid double-wrapping if already modified
+    if getattr(original, "_overridden", False):
+        return
+
+    # Create enhanced signature with compression parameters
+    orig_sig = inspect.signature(original)
+    sig_with_compression_params = _create_compression_signature(orig_sig)
+
+    @wraps(original)
+    def save_pretrained_wrapper(
+        *args,
+        sparsity_config: Optional[SparsityCompressionConfig] = None,
+        quantization_format: Optional[str] = None,
+        save_compressed: bool = True,
+        skip_sparsity_compression_stats: bool = True,
+        disable_sparse_compression: bool = False,
+        **kwargs,
+    ):
+        """
+        Wrapper around PreTrainedModel.save_pretrained() that adds compression
+        functionality. The compression format is saved to the model's config file
+
+        NOTE: If adding parameters here, also update _create_compression_signature()
+        to maintain signature consistency.
+
+        :param sparsity_config: Optional sparsity compression configuration.
+            If None and `skip_sparsity_compression_stats` is False, a sparsity
+            config will be inferred from the model.
+        :param quantization_format: Optional format string for quantization
+        :param save_compressed: Whether to save the model in compressed format
+        :param skip_sparsity_compression_stats: Whether to skip calculating
+            sparsity stats.
+        :param disable_sparse_compression: Whether to disable sparse compression
+            entirely
+        """
+        # HACK: Override the dtype_byte_size function in transformers to
+        # support float8 types. Fix is posted upstream
+        # https://github.com/huggingface/transformers/pull/30488
+        transformers.modeling_utils.dtype_byte_size = new_dtype_byte_size
+
+        # Extract save_directory from args or kwargs
+        save_directory = args[0] if args else kwargs.get("save_directory")
+        if save_directory is None:
+            raise ValueError(
+                "`save_directory` must be provided as first positional arg or kwarg"
             )
 
-            if compressor is None:
-                # model is not compressed or quantized, save as normal
-                original_save_pretrained_func = original_save_pretrained.__get__(
-                    model, model_class
-                )
-                original_save_pretrained_func(
-                    save_directory, state_dict=state_dict, **kwargs
-                )
-                return
-
-            # make sure we're on the main process when saving
-            if state_dict is not None and len(state_dict) > 0:
-                compressed_state_dict = compressor.compress(model, state_dict)
-                logger.info("Saving compressed model to disk")
-                original_save_pretrained.__get__(model, model_class)(
-                    save_directory,
-                    state_dict=compressed_state_dict,
-                    safe_serialization=safe_serialization,
-                    **kwargs,
-                )
-                compressor.update_config(save_directory)
-
-            # update existing recipe
-            update_and_save_recipe(model.name_or_path, save_directory)
-
-            # copy python files from cache dir to save_path if any
-            copy_python_files_from_model_cache(model, save_directory)
-
-        save_pretrained_wrapper._overriden = True
-        return save_pretrained_wrapper
-
-    # wrap save_pretrained if not already
-    if not getattr(model.save_pretrained, "_overriden", False):
-        model.save_pretrained = save_pretrained_compressed(model.save_pretrained)
+        # Get state_dict or fetch it if not provided
+        state_dict = kwargs.pop("state_dict", None)
+        if state_dict is None:
+            logger.info("Fetching state_dict – this may take some time")
+            state_dict = get_state_dict_offloaded_model(model)
+
+        logger.info("Fetching compressor")
+        compressor = get_model_compressor(
+            model=model,
+            sparsity_config=sparsity_config,
+            quantization_format=quantization_format,
+            save_compressed=save_compressed,
+            skip_sparsity_compression_stats=skip_sparsity_compression_stats,
+            state_dict=state_dict,
+            disable_sparse_compression=disable_sparse_compression,
+        )
+
+        if compressor is None:
+            # No compression needed
+            original(*args, state_dict=state_dict, **kwargs)
+        else:
+            # Compress and save
+            compressed_state_dict = compressor.compress(model, state_dict)
+            logger.info("Saving compressed model to disk")
+            original(*args, state_dict=compressed_state_dict, **kwargs)
+            compressor.update_config(save_directory)
+
+        # These operations happen regardless of compression
+        update_and_save_recipe(model.name_or_path, save_directory)
+        copy_python_files_from_model_cache(model, save_directory)
+
+    # Apply compression signature
+    save_pretrained_wrapper.__signature__ = sig_with_compression_params
+    save_pretrained_wrapper._overridden = True
+    model.save_pretrained = save_pretrained_wrapper
 
 
 # HACK: Override the dtype_byte_size function in transformers to support float8 types
@@ -306,3 +292,59 @@ def update_and_save_recipe(model_stub: str, save_directory: str):
     # save recipe
     recipe_path = os.path.join(save_directory, RECIPE_FILE_NAME)
     recipe.yaml(recipe_path)
+
+
+def _create_compression_signature(orig_sig: inspect.Signature) -> inspect.Signature:
+    """
+    Creates an enhanced signature with compression parameters.
+
+    :param orig_sig: Original function signature
+    :return: Enhanced signature with compression parameters
+    """
+    # Define compression parameters
+    compression_params = [
+        inspect.Parameter(
+            name="sparsity_config",
+            kind=inspect.Parameter.KEYWORD_ONLY,
+            default=None,
+            annotation=Optional[SparsityCompressionConfig],
+        ),
+        inspect.Parameter(
+            name="quantization_format",
+            kind=inspect.Parameter.KEYWORD_ONLY,
+            default=None,
+            annotation=Optional[str],
+        ),
+        inspect.Parameter(
+            name="save_compressed",
+            kind=inspect.Parameter.KEYWORD_ONLY,
+            default=True,
+            annotation=bool,
+        ),
+        inspect.Parameter(
+            name="skip_sparsity_compression_stats",
+            kind=inspect.Parameter.KEYWORD_ONLY,
+            default=True,
+            annotation=bool,
+        ),
+        inspect.Parameter(
+            name="disable_sparse_compression",
+            kind=inspect.Parameter.KEYWORD_ONLY,
+            default=False,
+            annotation=bool,
+        ),
+    ]
+
+    # Only add parameters that don't exist in the original signature
+    existing_params = orig_sig.parameters.keys()
+    new_params = []
+
+    for param in orig_sig.parameters.values():
+        if param.kind == inspect.Parameter.VAR_KEYWORD:
+            # Add compression params before **kwargs
+            new_params.extend(
+                [p for p in compression_params if p.name not in existing_params]
+            )
+        new_params.append(param)
+
+    return orig_sig.replace(parameters=new_params)
diff --git a/tests/llmcompressor/transformers/sparsification/test_compress_tensor_utils.py b/tests/llmcompressor/transformers/sparsification/test_compress_tensor_utils.py
@@ -1,3 +1,4 @@
+import inspect
 import math
 import os
 import shutil
@@ -8,7 +9,11 @@
 from accelerate.accelerator import get_state_dict_offloaded_model
 from compressed_tensors import QUANTIZATION_CONFIG_NAME, CompressionFormat
 from compressed_tensors.compressors import ModelCompressor
-from compressed_tensors.config import BitmaskConfig, DenseSparsityConfig
+from compressed_tensors.config import (
+    BitmaskConfig,
+    DenseSparsityConfig,
+    SparsityCompressionConfig,
+)
 from compressed_tensors.quantization import (
     QuantizationConfig,
     QuantizationStatus,
@@ -708,3 +713,64 @@ def test_correct_compressor_inferred(
         )
     else:
         assert compressor.sparsity_config.format == expected_sparsity_compressor
+
+
+@pytest.mark.parametrize(
+    "sparse_uncompressed_model",
+    ["nm-testing/llama2.c-stories15M-pruned_50.2of4-uncompressed"],
+)
+@pytest.mark.parametrize("save_compressed", [True, False])
+def test_modify_save_pretrained(sparse_uncompressed_model, save_compressed, tmp_path):
+    """
+    Test if the `modify_save_pretrained` function correctly modifies the model's
+    `save_pretrained` method.
+    """
+    model = AutoModelForCausalLM.from_pretrained(sparse_uncompressed_model)
+
+    modify_save_pretrained(model)
+
+    # Get the actual function object (handle both bound and unbound methods)
+    modified_func = getattr(
+        model.save_pretrained,
+        "__func__",
+        model.save_pretrained,
+    )
+
+    # Check that the method was properly modified
+    assert hasattr(model, "save_pretrained")
+    assert callable(model.save_pretrained)
+    assert getattr(modified_func, "_overridden", True)
+
+    # Verify the signature contains expected compression parameters
+    expected_params = {
+        "sparsity_config",
+        "quantization_format",
+        "save_compressed",
+        "skip_sparsity_compression_stats",
+        "disable_sparse_compression",
+    }
+    sig = inspect.signature(model.save_pretrained)
+    actual_params = set(sig.parameters.keys())
+
+    # Check that all expected parameters are present
+    assert expected_params.issubset(
+        actual_params
+    ), f"Missing parameters: {expected_params - actual_params}"
+
+    # Test the actual functionality
+    save_dir = tmp_path / "compressed_model"
+    model.save_pretrained(
+        save_dir,
+        save_compressed=save_compressed,
+        skip_sparsity_compression_stats=not save_compressed,
+    )
+
+    # Verify the model was saved correctly
+    assert (save_dir / "recipe.yaml").exists()
+
+    # Additional checks when saving in compressed format
+    if save_compressed:
+        # Verify we can load a compressor from the saved model config
+        compressor = ModelCompressor.from_pretrained(save_dir)
+        assert compressor is not None
+        assert isinstance(compressor.sparsity_config, SparsityCompressionConfig)
