[Transform] Online Rotations

[kylesayrs]

No description provided.

[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 integrates the foundational components for applying 'online rotations' (specifically R1 and R2 from the SpinQuant paper) into the llmcompressor framework. It primarily introduces a new SpinQuantModifier that leverages novel model transformation utilities, such as embedding normalization and norm-linear fusion, to prepare models for more effective quantization. Additionally, it refines the handling of tied word embeddings, ensuring compatibility and robustness across various model configurations.

Highlights

• New Feature: SpinQuantModifier: Introduced a new SpinQuantModifier to apply 'offline' rotations (R1 and R2) from the SpinQuant paper. These rotations transform model weights and activations to improve quantization accuracy without introducing runtime overhead.
• Model Transformation Utilities: Added new utilities for normalizing embedding layers and fusing norm layers into subsequent linear layers. These are crucial preprocessing steps for applying SpinQuant rotations, ensuring transform invariance.
• Improved Tied Word Embedding Handling: Refactored and enhanced the utility for untying word embeddings. The updated implementation is more robust, correctly handling cases where embeddings are tied, especially with offloaded parameters, and centralizes the untying logic.
• Example Usage and Integration: Provided new example scripts (compress_model.py, spinquant_example.py) demonstrating how to use the SpinQuantModifier for model compression. The modifier is also integrated into the data-free pipeline for seamless application.

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

This pull request introduces support for SpinQuant online rotations, a technique for improving quantization performance. It adds a new SpinQuantModifier, along with utilities for model transformation like layer fusion and embedding normalization. The changes also include updates to the data-free pipeline, improvements to handling tied word embeddings, and new example scripts and tests.

My review identified a critical bug in a Pydantic validator within the new SpinQuantModifier that prevents it from being used. I've also pointed out a few medium-severity issues, including a required argument missing in a script, brittle directory name construction, a documentation typo, and a maintainability concern with a hardcoded pipeline selection. Addressing these points will improve the correctness and robustness of the new features.

[gemini-code-assist]

The validate_not_implemented validator for randomize and learnable will always raise a NotImplementedError, even when the default value False is used. This makes the SpinQuantModifier unusable as it cannot be instantiated. The validation should only fail if these flags are explicitly set to True.

Suggested change

	@field_validator("randomize", "learnable", mode="before")
	def validate_not_implemented(cls, value, info: ValidationInfo):
	raise NotImplementedError(f"{info.field_name} is not supported right now")
	@field_validator("randomize", "learnable", mode="before")
	def validate_not_implemented(cls, value, info: ValidationInfo):
	if value:
	raise NotImplementedError(f"{info.field_name} is not supported right now")
	return value

[gemini-code-assist]

The --model_id argument is essential for this script to run. Without it, args.model_id will be None, causing AutoModelForCausalLM.from_pretrained to fail. Please make this argument required to prevent runtime errors and add a more descriptive help message.

Suggested change

	parser.add_argument("--model_id", type=str, help="Model stub to compress")
	parser.add_argument("--model_id", type=str, required=True, help="Hugging Face model ID to compress")

[gemini-code-assist]

The current method of constructing SAVE_DIR will include "None" in the directory name if --transform_type or --scheme are not provided. This can lead to confusing directory names. It's better to build the save directory name conditionally, only including the parts that are actually provided.

Suggested change

	SAVE_DIR = MODEL_ID.split("/")[1] + f"-{args.transform_type}-{args.scheme}"
	SAVE_DIR = "-".join([p for p in (MODEL_ID.split("/")[1], args.transform_type, args.scheme) if p])

[gemini-code-assist]

There appears to be a copy-paste error in the docstring for transform_type. The value "random-matrix" is listed twice. Based on the Literal type hint for transform_type on line 64, one of these should likely be "random-hadamard". Please correct the docstring to accurately reflect the available options.

Suggested change

	`"random-matrix"` has more performance cost, but supports a much larger set of
	sizes.
	`"random-matrix"` has the greatest performance cost, but supports any size
	`"random-hadamard"` has more performance cost, but supports a much larger set of
	sizes.
	`"random-matrix"` has the greatest performance cost, but supports any size

[gemini-code-assist]

This logic for selecting the datafree pipeline for SpinQuantModifier is hardcoded. While this works, it's not very extensible, as acknowledged by the TODO comment. For better maintainability, consider adding a method to the Modifier base class, such as requires_calibration_data(). SpinQuantModifier could then override this to return False. This would make the pipeline selection logic more robust and easier to manage as new modifiers are introduced.