[Bugfix] Fix Whisper compatibility with CPU offloading

[kylesayrs]

Purpose

• Fixes Problem faced with quantizing whisper #1573

Background

The whisper model includes these lines where the embedding weight is used as a bare addend. The model can do this because the inputs_embeds are outputs of conv layers, which always output the same shape. self.embed_positions acts essentially as an unlearned bias buffer.

embed_pos = self.embed_positions.weight
hidden_states = inputs_embeds + embed_pos

However, this naked usage of torch.nn.Embedding means that the forward function is never called, and the weight is never onloaded, breaking CPU offloading compatibility with the whisper model. This is likely not a concern for most users, since the whisper models are quite small and don't usually require CPU offloading.

Changes

• Support CPU offloading (which is required by sequential onloading) by modifying the whisper model definition
  • An upstream PR will be requested for future transformers versions

Testing

• Ran whisper example to completion

[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 delivers a crucial bugfix for Whisper model compatibility when CPU offloading is active, especially during quantization. It achieves this by introducing a new, modular framework for preparing and patching various transformers models, ensuring their correct behavior across diverse hardware configurations. The changes lay the groundwork for more robust model support within the library.

Highlights

• Core Bugfix for Whisper CPU Offloading: This PR directly addresses a compatibility issue with the Whisper model when CPU offloading is enabled, particularly during quantization. It introduces a patched WhisperEncoder that explicitly calls embed_positions during its forward pass, ensuring that embedding weights are correctly loaded onto the device before use, which was the root cause of the bug.
• New Model Preparation Framework: A new llmcompressor.modeling package has been introduced, featuring a prepare_for_quantization function. This function provides a generic and extensible mechanism to traverse transformers.PreTrainedModel instances and replace specific submodules (e.g., WhisperEncoder, DeepseekV3MoE) with custom, patched versions, thereby improving model compatibility for various operations like quantization.
• DeepseekV3 MoE Layer Patching: As part of the new model preparation framework, a DeepseekV3MoECalibrate module has been added. This module is designed to correctly handle DeepseekV3's Mixture of Experts (MoE) layers, suggesting a broader effort to ensure robust model compatibility across different architectures within the llmcompressor library.
• General Module Traversal Utility: A new utility function, module_bfs, has been added to llmcompressor.utils. This function enables breadth-first traversal and modification of torch.nn.Module hierarchies, serving as a foundational component for the new model preparation and patching logic.
• Codebase Structure and Formatting: The pyproject.toml configuration has been updated to reflect the new llmcompressor.modeling package structure by adjusting isort and ruff exclusion paths. Additionally, the debug.py script received minor reformatting for improved readability.

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 introduces a fix for Whisper compatibility with CPU offloading by patching the WhisperEncoder. It also adds a new modeling sub-package with utilities for preparing models for quantization, including specific patches for Whisper and DeepSeek V3. The changes are generally well-structured. Key recommendations include adding comprehensive docstrings and type hints to the new modules for better maintainability, clarifying a TODO comment, and a specific refinement to the Whisper patch for improved robustness regarding tensor device handling.

[kylesayrs]

Upstream fix: huggingface/transformers#38994