Speed up compressed-tensors load-time matching (for Kimi models)#1999
Speed up compressed-tensors load-time matching (for Kimi models)#1999rohansjoshi wants to merge 1 commit into
Conversation
compress_model/decompress_model pass ~69k exact module names as targets; the stock match_named_modules nested loop is O(modules x targets), which made the load-time "Compressing model..." phase take ~100 min on Kimi-K2.7-scale MoE checkpoints. Bucket exact-name targets into a set for O(1) membership (regex/class targets unchanged) and rebind the name in every module that imported it. Verified identical match output; 69,120 modules now match in <1 s. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Rohan Joshi <rohjoshi@nvidia.com>
📝 WalkthroughWalkthroughChangesHF PTQ matcher patch
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant PTQScript
participant example_utils
participant compressed_tensors
participant ImportedModules
PTQScript->>example_utils: patch_match_named_modules()
example_utils->>compressed_tensors: wrap match_named_modules
example_utils->>ImportedModules: rebind imported matcher references
compressed_tensors-->>PTQScript: optimized module matching
Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 3
🧹 Nitpick comments (1)
examples/hf_ptq/example_utils.py (1)
70-71: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMove standard-library imports to module scope.
logging,contextlib, andsysare neither optional nor heavy. Keep only the optional compressed-tensors import local, with a brief justification.As per coding guidelines, “Keep imports at the top of Python source and test files; use local imports only for justified circular dependencies, optional dependencies, or unusually heavy imports, with a brief explanatory comment.”
Also applies to: 122-138
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/hf_ptq/example_utils.py` around lines 70 - 71, Move the standard-library imports logging, contextlib, and sys to module scope, and retain only the optional compressed_tensors.utils.match import inside the relevant function. Add a brief comment explaining that this local import is kept because compressed-tensors is optional, covering the code around the referenced import and the additional occurrence near lines 122-138.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/hf_ptq/example_utils.py`:
- Line 150: The patched-module status message in match_named_modules must be
rank-aware because multinode_ptq.py executes it on every rank. Replace the
direct print with the repository’s print_rank_0 helper, preserving the existing
message content and flush behavior.
- Around line 70-73: Update the import handling in the compressed-tensors
optimization setup to catch only the exception indicating the optional package
is absent; allow other import failures from compressed_tensors.utils.match to
propagate instead of silently returning. Preserve the no-op behavior when
compressed-tensors is genuinely unavailable.
- Line 114: Update the matcher call in the regex-target handling to use the
public _m.match_name(name, t, None) helper instead of the private _m._match_name
symbol, preserving the existing matching arguments and behavior.
---
Nitpick comments:
In `@examples/hf_ptq/example_utils.py`:
- Around line 70-71: Move the standard-library imports logging, contextlib, and
sys to module scope, and retain only the optional compressed_tensors.utils.match
import inside the relevant function. Add a brief comment explaining that this
local import is kept because compressed-tensors is optional, covering the code
around the referenced import and the additional occurrence near lines 122-138.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b0a7a50b-6e74-45b5-8f8d-246ace55e55c
📒 Files selected for processing (3)
examples/hf_ptq/example_utils.pyexamples/hf_ptq/hf_ptq.pyexamples/hf_ptq/multinode_ptq.py
| try: | ||
| import compressed_tensors.utils.match as _m | ||
| except Exception: | ||
| return |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟡 Minor | ⚡ Quick win
Do not silently disable the optimization on arbitrary import failures.
Catching every exception makes this a no-op on broken/incompatible compressed-tensors imports, sending users back to the slow load path without an actionable error. Catch only an absent optional package; re-raise other failures.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/hf_ptq/example_utils.py` around lines 70 - 73, Update the import
handling in the compressed-tensors optimization setup to catch only the
exception indicating the optional package is absent; allow other import failures
from compressed_tensors.utils.match to propagate instead of silently returning.
Preserve the no-op behavior when compressed-tensors is genuinely unavailable.
| break | ||
| if mt is None: | ||
| for t in re_targets: | ||
| if _m._match_name(name, t, None): |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant file and locate matcher-related symbols/usages.
git ls-files examples/hf_ptq/example_utils.py
ast-grep outline examples/hf_ptq/example_utils.py --view expanded || true
rg -n "_match_name|match_name|compressed_tensors|patch_match_named_modules" examples/hf_ptq/example_utils.py -n -C 3 || true
# Search the repository for any references to the matcher helper API.
rg -n "match_name\(|_match_name\(" -S . || trueRepository: NVIDIA/Model-Optimizer
Length of output: 5174
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant file section directly if needed.
sed -n '1,220p' examples/hf_ptq/example_utils.py | cat -nRepository: NVIDIA/Model-Optimizer
Length of output: 10829
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect upstream API documentation/source via web if local repo lacks the dependency source.
python3 - <<'PY'
print("No-op placeholder: awaiting shell output from repository inspection.")
PYRepository: NVIDIA/Model-Optimizer
Length of output: 229
🌐 Web query:
compressed-tensors utils.match match_name _match_name current API
💡 Result:
In the current API of the compressed-tensors library, the module matching logic has been refactored significantly to improve performance and maintainability [1][2]. The older functions and internal naming conventions, such as _match_name, have been largely replaced or consolidated [3][2]. Key updates to the matching API include: - Consolidation into match_named_modules: Much of the logic previously scattered across various helper functions has been centralized into match_named_modules [1][2]. This function is designed to efficiently determine if a module or parameter name matches the targets or ignore lists defined in a quantization configuration [1]. - Performance Improvements: Recent updates (e.g., as of April 2026) have rewritten match_named_modules to optimize matching, moving away from nested loops (O(N*M) complexity) to more efficient lookup structures like sets and precompiled regexes [1]. This is particularly important for large models, such as Mixture-of-Experts (MoE) architectures, which may contain thousands of modules [1][4]. - Deprecation of Older Helpers: Older utilities like expand_target_names, is_target, find_name_or_class_matches, and _find_matches have been replaced by the consolidated matching logic [2]. If you are working with the library, you should favor using match_named_modules from src/compressed_tensors/utils/match.py rather than relying on legacy internal underscore-prefixed functions like _match_name, which are generally considered internal and subject to change or removal [3][1][2].
Citations:
- 1: [match] optimized match_named_modules vllm-project/compressed-tensors#697
- 2: Refactor module / parameter matching logic vllm-project/compressed-tensors#406
- 3: [Help Wanted] Refactor module / parameter matching logic vllm-project/compressed-tensors#405
- 4: match_named_modules O(N×M) causes multi-minute silent load tail on fine-grained MoE checkpoints (Qwen3-30B-A3B-AWQ, ~18K Linear modules) vllm-project/compressed-tensors#695
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Look for dependency pins and local imports that can reveal the expected API.
rg -n "compressed-tensors|compressed_tensors" -S pyproject.toml requirements*.txt .github examples || trueRepository: NVIDIA/Model-Optimizer
Length of output: 1348
🌐 Web query:
site:github.com/vllm-project/compressed-tensors _match_name match_name compressed_tensors.utils.match
💡 Result:
In the vllm-project/compressed-tensors library, match_name is a utility function used for matching module or tensor names against specified targets [1]. It is part of the library's utility infrastructure, specifically residing within the compressed_tensors.utils.match module [2][3]. The library utilizes several matching functions to handle module and parameter targeting, including match_name, match_named_modules, match_targets, and is_narrow_match [3][4][1]. These utilities are essential for operations like applying quantization configurations, dequantization, and managing model compression [3][1]. Recent developments in the library have focused on optimizing these matching processes. For instance, the match_named_modules function was recently refactored to replace inefficient quadratic O(N x M) logic with optimized lookup structures, significantly improving performance for large models such as MoE (Mixture of Experts) architectures [2][5]. The match_name function is often used within these broader workflows, such as in converter logic, to verify specific target matches when identifying quantizable tensors or dependencies [1].
Citations:
- 1: [convert] compressed-tensors dequantizer vllm-project/compressed-tensors#689
- 2: [match] optimized match_named_modules vllm-project/compressed-tensors#697
- 3: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/quantization/lifecycle/apply.py
- 4: [Help Wanted] Refactor module / parameter matching logic vllm-project/compressed-tensors#405
- 5: match_named_modules O(N×M) causes multi-minute silent load tail on fine-grained MoE checkpoints (Qwen3-30B-A3B-AWQ, ~18K Linear modules) vllm-project/compressed-tensors#695
🌐 Web query:
current compressed-tensors matcher API match_name _match_name site:github.com/vllm-project/compressed-tensors
💡 Result:
In the compressed-tensors library, the primary API for module matching is match_named_modules, which is located in src/compressed_tensors/utils/match.py [1][2]. As of the current library structure, match_named_modules is the established function used to traverse model modules and match them against specified target or ignore criteria [1][3]. Historical development discussions have referenced internal helper methods or naming variations such as match_name, match_class, and _find_matches during refactoring efforts [4], but match_named_modules serves as the public-facing utility for this operation [1][3]. Recent updates, such as those in Pull Request #697, have optimized match_named_modules from an O(N×M) complexity to near-linear performance by implementing one-time lookup structures for target and ignore lists [1][3]. The function signature typically follows this pattern [2]: def match_named_modules( model: torch.nn.Module, targets: Iterable[str] | None, ignore: Iterable[str] | None = None,) -> Generator[tuple[str, torch.nn.Module]]: The library uses these utilities to facilitate the efficient application of quantization and compression configurations across model components [3][5].
Citations:
- 1: [match] optimized match_named_modules vllm-project/compressed-tensors#697
- 2: [Python 3.10+] Modernize codebase with Python 3.10+ idioms and features vllm-project/compressed-tensors#492
- 3: match_named_modules O(N×M) causes multi-minute silent load tail on fine-grained MoE checkpoints (Qwen3-30B-A3B-AWQ, ~18K Linear modules) vllm-project/compressed-tensors#695
- 4: [Help Wanted] Refactor module / parameter matching logic vllm-project/compressed-tensors#405
- 5: https://github.com/vllm-project/compressed-tensors
Use the public matcher helper for regex targets.
_m._match_name(...) is not part of the current matcher API, so re: targets can hit AttributeError here. Call _m.match_name(name, t, None) instead.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/hf_ptq/example_utils.py` at line 114, Update the matcher call in the
regex-target handling to use the public _m.match_name(name, t, None) helper
instead of the private _m._match_name symbol, preserving the existing matching
arguments and behavior.
| if getattr(_mod, "match_named_modules", None) is _orig: | ||
| _mod.match_named_modules = fast_match_named_modules | ||
| rebound.append(_modname) | ||
| print(f"[match-speedup] patched match_named_modules; rebound importers: {rebound}", flush=True) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Make this multinode log rank-aware.
examples/hf_ptq/multinode_ptq.py invokes the patch during import on every rank, so each rank emits this message. Route it through the repository’s rank-zero logging helper.
As per coding guidelines, “Develop distributed code with rank-aware logging such as print_rank_0 or warn_rank_0.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/hf_ptq/example_utils.py` at line 150, The patched-module status
message in match_named_modules must be rank-aware because multinode_ptq.py
executes it on every rank. Replace the direct print with the repository’s
print_rank_0 helper, preserving the existing message content and flush behavior.
Source: Coding guidelines
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1999 +/- ##
==========================================
+ Coverage 75.97% 76.69% +0.72%
==========================================
Files 518 518
Lines 58269 58269
==========================================
+ Hits 44267 44692 +425
+ Misses 14002 13577 -425
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
1 participant
What does this PR do?
Loading Kimi models (which use Compressed-Tensors weights) is very slow with
hf_ptq/multinode_ptqdue to a CPU-bound phase involving matching module names with an inefficient algorithm. For Kimi K2.6, compress_model/decompress_model pass ~69k exact module names as targets; the stock match_named_modules nested loop is O(modules x targets), which made the load-time "Compressing model..." phase take ~100 min on Kimi-K2.6/2.7-scale MoE checkpoints.Idea of this PR: Bucket exact-name targets into a set for O(1) membership (regex/class targets unchanged) and rebind the name in every module that imported it. Verified identical match output; 69,120 modules now match in <1 s.
Type of change: ?
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: ✅ / ❌ / N/AAdditional Information
Summary by CodeRabbit