Enhance save_pretrained

[rahul-tuli]

This PR improves the save_pretrained method wrapper to use a more maintainable and robust approach for extending the function signature with compression parameters. The current version uses an approach with weakref that can be difficult to maintain and debug.

This is an initial implementation addressing ticket feedback welcome!

Changes

• Replaces the weakref-based approach with a direct method wrapping using inspect
• Adds proper signature preservation for enhanced save_pretrained method
• Creates a separate helper function to generate the enhanced signature
• Fixes a typo (changes _overriden to _overridden)
• Adds unit tests for the modified function
• Adds documentation explaining the enhanced parameters (docs/save_pretrained.md)

Benefits of New Approach

1. More Robust: The new implementation properly preserves and displays the original function signature

   Before:

   # Function signature is not properly handled or displayed
   # Checking help for the method would show limited information:
   >>> help(model.save_pretrained)
   Help on method save_pretrained:
   
   save_pretrained(save_directory, **kwargs) method of transformers.models.llama.modeling_llama.LlamaForCausalLM instance
       Save a model and its configuration file to a directory...
       # No information about compression parameters in help output

   After:

   # Proper signature with parameters visible in help and IDE tools:
   >>> help(model.save_pretrained)
   Help on method save_pretrained:
   
   save_pretrained(self, save_directory, *, sparsity_config=None, 
                  quantization_format=None, save_compressed=True,
                  skip_sparsity_compression_stats=True,
                  disable_sparse_compression=False, **kwargs) method of transformers.models.llama.modeling_llama.LlamaForCausalLM instance
       Wrapper around PreTrainedModel.save_pretrained() that adds compression
       functionality. The compression format is saved to the model's config file
       
       :param sparsity_config: Optional sparsity compression configuration...
       # Full parameter documentation available

2. Better Maintainability: Code is more modular with a separate signature creation function

   Before:

   # Complex nested functions making it hard to track state
   def save_pretrained_compressed(save_pretrained_method):
       # ...
       def save_pretrained_wrapper(...):
           # ...
       return save_pretrained_wrapper

   After:

   # Clean separation of concerns
   def _create_compression_signature(orig_sig):
       # Create and return enhanced signature
       
   def modify_save_pretrained(model):
       # Create wrapper function
       # Apply signature
       # Attach to model

3. Improved Error Handling: Somewhat better handling of common pitfalls

   Before:

   # Implicit save_directory handling
   original_save_pretrained.__get__(model, model_class)(
       save_directory,
       state_dict=compressed_state_dict,
       safe_serialization=safe_serialization,
       **kwargs,
   )

   After:

   # Explicit save_directory extraction with clear error message
   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"
       )

4. Improved Documentation: documentation of parameters and usage

   """
   NOTE: If adding parameters here, also update _create_compression_signature()
   to maintain signature consistency.
   
    For more information on the compression parameterrs and model saving in
    llmcompressor, refer to docs/save_pretrained.md
   """

   Plus a dedicated docs/save_pretrained.md file with examples.

Testing

Added tests that verify:

• Function signature contains all expected parameters
• Functionality works correctly for both compressed and uncompressed saving
• Parameters are correctly passed through to the underlying implementation

This is an initial implementation to improve the developer experience, and I'm open to suggestions for further refinements.

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[Copilot]

Pull Request Overview

This PR enhances the save_pretrained method by replacing the weakref-based approach with an inspect-based implementation, ensuring the signature is maintained and error handling improved. Key changes include:

• Replacing weakref wrapping with a direct method wrapping using inspect.
• Introducing a helper to generate an enhanced function signature.
• Adding unit tests and updated documentation to validate and explain the new behavior.

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
tests/llmcompressor/transformers/sparsification/test_compress_tensor_utils.py	Adds unit tests to verify the modified save_pretrained signature and functionality.
src/llmcompressor/transformers/sparsification/compressed_tensors_utils.py	Updates modify_save_pretrained to use inspect and adds a helper for signature updates.
docs/save_pretrained.md	Provides documentation on the enhanced save_pretrained parameters and usage examples.

[Copilot]

Pull Request Overview

This PR enhances the save_pretrained functionality by replacing the weakref‐based wrapping with a direct method wrapping that preserves and augments the original function signature with compression parameters. Key changes include:

• Replacing weakref-based method wrapping with an inspect-based approach.
• Adding a helper function (_create_compression_signature) that creates an enhanced signature.
• Updating tests and documentation to verify and explain the new behavior.

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
tests/llmcompressor/transformers/sparsification/test_compress_tensor_utils.py	Updated tests to validate the enhanced save_pretrained signature and its functionality.
src/llmcompressor/transformers/sparsification/compressed_tensors_utils.py	Refactored modify_save_pretrained to use a direct wrapping with a preserved signature and removed weakref.
docs/save_pretrained.md	Added documentation to explain the new compression parameters and usage examples.

[brian-dellabetta]

awesome! nice work

[brian-dellabetta]

95% sure ._overridden is not saved with the model, just want to confirm that this change won't break any previously saved models that had ._overriden

[brian-dellabetta]

very nice!