[AWQ] use match_modules_set and fix logic

[HDCharles]

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]

👋 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 @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.

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

[kylesayrs]

🔥

[fynnsu]

Looks good!

[kylesayrs]

Sweet

[fynnsu]

Thanks for the pr. Added some comments below

[rahul-tuli]

Great Job on this!