[Transform] Spinquant with R1 and R2

[brian-dellabetta]

Prerequisites

• [Transform] [WIP] Transforms HEAD neuralmagic/compressed-tensors#391
• [Core] [Offloading] Fix saving offloaded submodules huggingface/transformers#39280

[gemini-code-assist]

Summary of Changes

Hello @brian-dellabetta, 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 significantly enhances the TransformModifier by introducing support for predefined transformation configurations, known as presets. This allows users to easily apply complex transformation schemes like QUIP and SpinQuant, streamlining the process of applying advanced model compression techniques. The changes also include an updated example demonstrating the new functionality and improved validation for the modifier.

Highlights

• Enhanced TransformModifier Flexibility: The TransformModifier now accepts either a preset_config string to load predefined transformation schemes (like QUIP or SpinQuant) or a direct config object for custom transformation setups, making it more versatile and user-friendly.
• Introduction of Predefined Transformation Presets: New modules have been added under src/llmcompressor/modifiers/transform/presets to define and expose QUIP, QUIP_ONLINE, LLAMA_SPINQUANT, and LLAMA_SPINQUANT_R1R2 configurations. These presets simplify the application of complex transformation strategies based on research papers.
• Updated Llama-3 Example: The llama3_example.py script has been revised to showcase the usage of the TransformModifier with a preset_config (specifically LLAMA_SPINQUANT_R1R2) and to use QuantizationModifier instead of GPTQModifier. The example also now uses a smaller Llama model for faster execution and includes a dispatch_for_generation call.

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 configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

The pull request significantly enhances the TransformModifier by introducing a robust preset configuration system and improving module targeting. The refactoring to use Pydantic for configuration validation greatly improves maintainability and prevents invalid states. The changes to use regex for module targeting in the presets (spinquant.py and quip.py) are a notable improvement for flexibility and robustness.

[brian-dellabetta]

@kylesayrs why are we excluding these? wouldn't we want them to persist in json?