[Examples] [Bugfix] skip sparsity stats when saving checkpoints

[kylesayrs]

Purpose

• Fix example failure where sparsity stats were attempted to be computed on a compressed model
  • This happened because the get_model_compressor cannot compute sparsity stats of a compressed model

Prerequisites

• [Hotfix] Implement method on dense compressor neuralmagic/compressed-tensors#345

Changes

• When saving checkpoints, call get_model_compressor on the decompressed model (before it is compressed, not after)
• Do not save checkpoints with sparsity stats by default, due to high runtime

Testing

• Ran example to completion with reduced training duration

[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.

[gemini-code-assist]

Summary of Changes

Hello @kylesayrs, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a bug encountered in examples where saving checkpoints of compressed models failed due to an inability to compute sparsity statistics on the compressed state. It resolves this by adjusting the order of operations during checkpoint saving to ensure the model compressor is obtained on the decompressed model. Additionally, it changes the default setting for skipping sparsity statistics during saves to True to improve performance.

Highlights

• Bugfix for Checkpoint Saving: Fixed an issue where attempting to compute sparsity statistics on a compressed model during checkpoint saving would fail. The logic to obtain the model compressor is now performed before the model is potentially compressed for saving, ensuring it operates on the decompressed state.
• Default Behavior Change: Changed the default value for skip_sparsity_compression_stats to True when saving models/checkpoints. This avoids the potentially high runtime cost associated with computing sparsity statistics by default.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configureGemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request addresses a bug by ensuring get_model_compressor is called on the model before it's potentially compressed by save_pretrained when saving checkpoints. It also changes the default to skip sparsity stats computation during this process to avoid high runtime costs. The feedback provided aims to improve the clarity of comments.

[rahul-tuli]

Lgtm, could you also update the save_pretrained readme in the docs folder

[rahul-tuli]

#1528 (comment)

[kylesayrs]

@rahul-tuli I've read through the docs and don't see the section you want updated? Afaict these docs are for users who want help calling save_pretrained themselves. I'm not sure how these changes are related or relevant to that reader?

[rahul-tuli]

@rahul-tuli I've read through the docs and don't see the section you want updated? Afaict these docs are for users who want help calling save_pretrained themselves. I'm not sure how these changes are related or relevant to that reader?

My bad @kylesayrs meant to request a change on #1530
And, missed the filename, thought you were making changes to save_pretrained_wrapper but this is a different pathway altogether. This is good to go from my side!