[AWQ] use match_modules_set and fix logic #2070

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

Depends on vllm-project/compressed-tensors#524

Summary:

modified AWQ _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
- added helper for module_to_name
- remove hardcoded handling for single balance layer by updating get_lowest_common_module to handle that
modified SmoothQuant _resolve_mappings
- brought into alignment with AWQ
  - this is largely a horizontal move though there is handling for situations that would have been missed before like
    - multiple smooth layer matches in a single set
    - parent contexts further than 1 layer away.
- updated mapping definitions to always be tuple(list[str],str) which is always the case but wasn't required unlike in AWQ
removed get_lowest_common_parent
- now we can use CT's get_lowest_common_ancestor_name so only need to check for module_list (it has a lot of bugfixes compared to the get_lowest_common_parent implementation in LLMC)
updated test_base for AWQ and smoothquant
- added test case for _set_resolved_mappings to check that partially skipped matches are handled correctly
- added tests for MoE matching being handled correctly
- added test cases for get_lowest_non_module_list_ancestor
- imported Linear and used that instead of torch.nn.Linear
reverted test_pytorch.py for logarithmic_equalizations and smoothquant
- The test was updated in [Fix] SmoothQuant MoE Support: Smooth All Experts, Not Just expert.0 #2084 by @rahul-tuli to ignore some modules but in general because of the way the new logic works, you need to ignore the whole set.
- if you only ignore one element the matching logic would need to determine whether there's a full set or not somehow which it doesn't do. In the previous logic, this was possible because it was assumed the whole set had to be siblings of the smooth_layer, but the new util is trying to be more flexible and so relaxes this assumption which prevents the same approach from working. If this is a common need, perhaps we can add a util that checks for a context parent context of size N or something.

TEST PLAN:
pytest /home/HDCharles/repos/llm-compressor/tests/llmcompressor/modifiers/awq/test_base.py
pytest /home/HDCharles/repos/llm-compressor/tests/llmcompressor/modifiers/smoothquant/test_base.py

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

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

kylesayrs

🔥

fynnsu

Looks good!

src/llmcompressor/modifiers/awq/base.py

kylesayrs

Sweet

kylesayrs · 2025-12-03T21:22:37Z

src/llmcompressor/modifiers/smoothquant/base.py

+        module_to_name = get_module_to_name_dict(model)
+        for mapping in self.mappings:
+            for *nested_balance_layers, smooth_layers in match_modules_set(
+                model, tree_flatten(mapping)[0], self.ignore


Suggested change

model, tree_flatten(mapping)[0], self.ignore

model, *mapping, self.ignore

I don't think this works [A,[B,C]] -> A,[B,C] we need to get [A,B,C]

Could switch to pytree.tree_leaves tho

do we have a dependency on Jax? i don't have it in my env atm

edit: nvm that's in the same dir, i thought you meant official pytree

src/llmcompressor/utils/pytorch/module.py

fynnsu

Thanks for the pr. Added some comments below

src/llmcompressor/utils/pytorch/module.py

src/llmcompressor/modifiers/transform/spinquant/base.py

src/llmcompressor/modifiers/smoothquant/base.py

src/llmcompressor/modifiers/awq/base.py

Summary: modified _set_resolved_mappings to get smoothing and balance layers at same time. Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

Summary Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

Summary fix smoothquant logic to align with AWQ Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

Summary Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

fynnsu

Thanks for updating. One small comment to resolve and then should be good to go

fynnsu · 2025-12-05T22:50:13Z

src/llmcompressor/modifiers/awq/base.py

    return True


+def get_lowest_ancestor_with_avoid(name: str, model: Module, avoid=torch.nn.Module):


The default for avoid seems risky to me. Maybe you meant it to be torch.nn.ModuleList? Otherwise this will probably just return the root model on every pytorch model unless avoid is explicitly set.

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

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

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

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

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

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

kylesayrs previously approved these changes Nov 27, 2025

View reviewed changes

HDCharles marked this pull request as draft December 2, 2025 18:52

fynnsu previously approved these changes Dec 2, 2025

View reviewed changes

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

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

HDCharles dismissed stale reviews from fynnsu and kylesayrs via 07c657e December 3, 2025 20:28

HDCharles force-pushed the 96_awq_match_module_set branch from b9d3cea to 07c657e Compare December 3, 2025 20:28

HDCharles requested a review from rahul-tuli December 3, 2025 20:57

HDCharles marked this pull request as ready for review December 3, 2025 21:12

kylesayrs reviewed Dec 3, 2025

View reviewed changes

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

kylesayrs previously approved these changes Dec 3, 2025

View reviewed changes

HDCharles dismissed kylesayrs’s stale review via 7a4e7d3 December 3, 2025 21:47

HDCharles force-pushed the 96_awq_match_module_set branch from 7a4e7d3 to 7d28a11 Compare December 3, 2025 21:54

HDCharles requested review from fynnsu and kylesayrs December 3, 2025 21:57

fynnsu reviewed Dec 4, 2025

View reviewed changes

HDCharles force-pushed the 96_awq_match_module_set branch from cf73082 to 8dcd6f0 Compare December 5, 2025 22:17

HDCharles added 15 commits December 5, 2025 22:23

[AWQ] small refactor to use match_modules_set

a8f0573

Summary: modified _set_resolved_mappings to get smoothing and balance layers at same time. Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

formatting

5dc9a70

Summary Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

fixing logic and test update

77a78b9

Summary Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

updates to get_lowest_common_x

f7067e9

Summary Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

format

d5d7d27

Summary Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

fixes

378245c

Summary Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

tests

acf7e5c

Summary Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

fixes and formatting

4ffbed4

Summary fix smoothquant logic to align with AWQ Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

fixing tests to old versions

2e4bf7c

Summary Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

formatting

cc149d4

Summary Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

comments

f9e4cfa

Summary Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

format

1f10f7d

Summary Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

tree flatten -> tree_leaves

056daee

Summary Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

fix infinite loop and the asserts

af7d76e

Summary Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

formatting

320f4ae

Summary Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

HDCharles force-pushed the 96_awq_match_module_set branch from 8a210bf to 320f4ae Compare December 5, 2025 22:23

HDCharles requested a review from fynnsu December 5, 2025 22:24

fynnsu reviewed Dec 5, 2025

View reviewed changes

	model, tree_flatten(mapping)[0], self.ignore
	model, *mapping, self.ignore

		return True


		def get_lowest_ancestor_with_avoid(name: str, model: Module, avoid=torch.nn.Module):

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

fynnsu left a comment

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!

kylesayrs Dec 3, 2025

Choose a reason for hiding this comment

Uh oh!

HDCharles Dec 3, 2025

Choose a reason for hiding this comment

Uh oh!

kylesayrs Dec 4, 2025

Choose a reason for hiding this comment

Uh oh!

fynnsu Dec 4, 2025

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

fynnsu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

fynnsu left a comment

Choose a reason for hiding this comment

Uh oh!

fynnsu Dec 5, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

HDCharles commented Nov 25, 2025 •

edited

Loading

HDCharles Dec 4, 2025 •

edited

Loading