[AWQ] use match_modules_set and fix logic #2070

HDCharles · 2025-11-25T21:06:53Z

Summary:

modified _set_resolved_mappings
- get smoothing and balance layers at same time using match_modules_set
- (bugfix) correct logic so that if any balance layers are incompatible, that matching is skipped
- added warnings
- get rid of tqdm and skip counting @kylesayrs
- check for modules that map to multiple names and warn if detected
- remove hardcoded handling for single balance layer by updating get_lowest_common_module to handle that
modified get_lowest_common parent
- renamed to get_lowest_common_module to reflect new behavior and bugfix
  - single module returns single module [abc]->abc (matches how _set_resolved_mappings handled single modules)
  - parent and child returns parent [x.abc, x.abc.a]->x.abc (rather than "" like before and aligns with single module behavior, could see argument for this returning x rather than x.abc which is why i renamed it)
  - substring handled correctly [abc, ab] -> "" (rather than ->ab like before)
updated test_base
- added test case for _set_resolved_mappings to check that partially skipped matches are handled correctly
- added test cases for test_get_lowest_common_module to check that the behavior matches the cases I outlined above
- imported Linear and used that instead of torch.nn.Linear

TEST PLAN: # (hits both relevent tests)
pytest /home/HDCharles/repos/llm-compressor/tests/llmcompressor/modifiers/awq/test_base.py

Summary: modified _set_resolved_mappings to get smoothing and balance layers at same time. Signed-off-by: HDCharles <[email protected]>

github-actions · 2025-11-25T21:07:03Z

👋 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 · 2025-11-25T21:07:07Z

Summary of Changes

Hello @HDCharles, 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 introduces a small but significant refactor to the _set_resolved_mappings method within the AWQ modifier. The core change involves replacing a sequential matching process with a more integrated approach using match_modules_set. This modification aims to streamline the identification of related smoothing and balance layers, making the mapping resolution more efficient and coherent by finding these layers together rather than in separate steps.

Highlights

Refactoring _set_resolved_mappings: The _set_resolved_mappings method in the AWQ modifier has been refactored to improve how smoothing and balance layers are identified and associated.
Introduction of match_modules_set: The refactoring leverages the match_modules_set utility function, allowing for the simultaneous matching of coherent sets of smoothing and balance layers.
Improved Module Lookup Efficiency: A module_to_name mapping is now pre-built at the start of the method to optimize module name lookups, enhancing performance.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

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.

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

Summary Signed-off-by: HDCharles <[email protected]>

gemini-code-assist

Code Review

This pull request refactors _set_resolved_mappings to use a new match_modules_set utility, simplifying the process of finding related layers for AWQ smoothing. The change improves readability by abstracting the complex matching logic. However, the new implementation introduces a potential bug when dealing with models that have shared modules, as it may resolve to an incorrect module name. I've added a comment detailing this issue.

src/llmcompressor/modifiers/awq/base.py

kylesayrs · 2025-11-25T23:03:59Z

src/llmcompressor/modifiers/awq/base.py

+                    if (
+                        isinstance(smooth_layer, torch.nn.Linear)
+                        and isinstance(balance_layer, torch.nn.Linear)
+                        and balance_name.endswith(".o_proj")
+                        and (
+                            (
+                                smooth_name.endswith(".v_proj")
+                                and smooth_layer.out_features
+                                != balance_layer.in_features
+                            )
+                            or (
+                                smooth_name.endswith(".qkv_proj")
+                                and smooth_layer.out_features
+                                != 3 * balance_layer.in_features
                            )
-                        ):
-                            num_skipped_mappings += 1
-                            continue
+                        )
+                    ):
+                        num_skipped_mappings += 1
+                        continue


Can we replace and generalize this to

if any(smooth_layer.weight.size(0) % layer.weight.size(-1) != 0 for layer in balance_layers): continue

doesn't that not work for the qkv_proj? I think mlp layers can do similar with gate_up_proj too.

src/llmcompressor/modifiers/awq/base.py

Summary Signed-off-by: HDCharles <[email protected]>

kylesayrs

🔥

[AWQ] small refactor to use match_modules_set

69cf95b

Summary: modified _set_resolved_mappings to get smoothing and balance layers at same time. Signed-off-by: HDCharles <[email protected]>

HDCharles requested review from brian-dellabetta, fynnsu and kylesayrs November 25, 2025 21:06

HDCharles requested a review from dsikka November 25, 2025 21:07

formatting

dc8e3ae

Summary Signed-off-by: HDCharles <[email protected]>

HDCharles added ready When a PR is ready for review awq For any issue / PR related to AWQ support labels Nov 25, 2025

gemini-code-assist bot reviewed Nov 25, 2025

View reviewed changes

src/llmcompressor/modifiers/awq/base.py Outdated Show resolved Hide resolved

kylesayrs reviewed Nov 25, 2025

View reviewed changes

fixing logic and test update

351568d

Summary Signed-off-by: HDCharles <[email protected]>

HDCharles requested a review from kylesayrs November 26, 2025 18:27

HDCharles marked this pull request as draft November 26, 2025 20:27

updates to get_lowest_common_x

dea5eab

Summary Signed-off-by: HDCharles <[email protected]>

HDCharles changed the title ~~[AWQ] small refactor to use match_modules_set~~ [AWQ] use match_modules_set and fix logic Nov 27, 2025

format

728b8c0

Summary Signed-off-by: HDCharles <[email protected]>

HDCharles marked this pull request as ready for review November 27, 2025 03:26

kylesayrs approved these changes Nov 27, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[AWQ] use match_modules_set and fix logic #2070

[AWQ] use match_modules_set and fix logic #2070

HDCharles commented Nov 25, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Nov 25, 2025

Uh oh!

gemini-code-assist bot commented Nov 25, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kylesayrs Nov 25, 2025

Uh oh!

HDCharles Nov 26, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kylesayrs left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

[AWQ] use match_modules_set and fix logic #2070

Are you sure you want to change the base?

[AWQ] use match_modules_set and fix logic #2070

Conversation

HDCharles commented Nov 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Nov 25, 2025

Uh oh!

gemini-code-assist bot commented Nov 25, 2025

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kylesayrs Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

HDCharles Nov 26, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kylesayrs left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

HDCharles commented Nov 25, 2025 •

edited

Loading