[None][fix] Honour producer ignored layers for fused MoE and any quant format - #17551
[None][fix] Honour producer ignored layers for fused MoE and any quant format#17551joerowell wants to merge 2 commits into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (3)
WalkthroughQuantization configuration loading now merges producer exclusions with stable deduplication. Laguna applies layer-specific ChangesQuantization exclusion handling
Estimated code review effort: 3 (Moderate) | ~20 minutes Mergeability Score: ⚪ Minimal · up to The change is localized to quantization behavior and test-list updates; no concrete correctness, runtime, or deployment issue is identified at the current head, so no actionable merge-blocking risk remains beyond normal checks. Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant ModelConfig
participant LagunaAttention
participant g_proj_quant_config
participant Linear
ModelConfig->>LagunaAttention: provide quant_config and layer_idx
LagunaAttention->>g_proj_quant_config: resolve g_proj quantization
g_proj_quant_config->>g_proj_quant_config: match excluded g_proj module
g_proj_quant_config-->>LagunaAttention: return resolved QuantConfig
LagunaAttention->>Linear: construct g_proj with resolved QuantConfig
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 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.
Actionable comments posted: 4
🤖 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 `@tensorrt_llm/_torch/models/modeling_laguna.py`:
- Around line 217-237: Annotate all newly added functions: in
tensorrt_llm/_torch/models/modeling_laguna.py:217-237, add a precise type for
model_config in g_proj_quant_config and retain its Optional[QuantConfig] return
type; in tests/unittest/_torch/test_hf_quant_config.py:67-162, annotate
_fp8_block_scales_config and every new test function, using None for procedures;
and in tests/unittest/_torch/models/test_laguna_gproj_quant.py:36-89, annotate
_model_config and every test function, using None for procedures.
In `@tensorrt_llm/_torch/models/modeling_utils.py`:
- Around line 618-638: The MoE exclusion candidates in the module-handling logic
only cover expert index 0, so exact exclusions for other experts are missed.
Update the logic around the MoE/VanillaMoE branch and its base-name handling to
match exclusions for every expert index and gate/up/down projection, including
.backend-wrapped modules; add regression coverage for a nonzero exact expert
path and a wrapped module.
In `@tests/unittest/_torch/models/test_laguna_gproj_quant.py`:
- Around line 45-89: Add unittest/_torch/models/test_laguna_gproj_quant.py to
the l0_cpu.yml test list so its four tests are collected by CPU CI; do not add a
QA-list entry or modify the existing test implementations.
In `@tests/unittest/_torch/test_hf_quant_config.py`:
- Around line 77-162: Add the five new test functions in
tests/unittest/_torch/test_hf_quant_config.py to the CPU-only test list in
l0_cpu.yml, using the existing test-list format and preserving their unittest
path. Do not add a QA-list entry; verify the updated selection with pytest
tests/unittest/.
🪄 Autofix
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: 866dd3ca-98cc-4f17-aa66-68b21511b2c3
📒 Files selected for processing (5)
tensorrt_llm/_torch/model_config.pytensorrt_llm/_torch/models/modeling_laguna.pytensorrt_llm/_torch/models/modeling_utils.pytests/unittest/_torch/models/test_laguna_gproj_quant.pytests/unittest/_torch/test_hf_quant_config.py
| elif isinstance(module, (MoE, VanillaMoE)): | ||
| # Fused MoE: a checkpoint's ignore list names the | ||
| # per-expert weights (re:...experts.[0-9]+.gate_proj), | ||
| # but the fused module has no per-expert child, so such | ||
| # a rule never matches and those experts get quantized | ||
| # despite being bf16 in the checkpoint. Offer | ||
| # representative per-expert names, mirroring the Linear | ||
| # expansion above. | ||
| # | ||
| # Strip a trailing ".backend": ConfigurableMoE wraps the | ||
| # module that actually owns the weights, and excluding | ||
| # only the wrapper does nothing because create_weights() | ||
| # delegates to the backend. Without this the backend | ||
| # stays quantized and the output is still degenerate. | ||
| base = (name[:-len('.backend')] | ||
| if name.endswith('.backend') else name) | ||
| candidates += [ | ||
| f'{base}.0.gate_proj', | ||
| f'{base}.0.up_proj', | ||
| f'{base}.0.down_proj', | ||
| ] |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Match literal exclusions for every expert index.
Lines 634-637 only test expert 0. An exact producer exclusion such as ...experts.1.up_proj does not match. The fused MoE then remains quantized although that expert has no quantization scales.
Match all expert indices, or detect exclusions that target any {base}.<expert>.(gate_proj|up_proj|down_proj) path. Add regression coverage for an exact nonzero expert index and a .backend-wrapped module.
🤖 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 `@tensorrt_llm/_torch/models/modeling_utils.py` around lines 618 - 638, The MoE
exclusion candidates in the module-handling logic only cover expert index 0, so
exact exclusions for other experts are missed. Update the logic around the
MoE/VanillaMoE branch and its base-name handling to match exclusions for every
expert index and gate/up/down projection, including .backend-wrapped modules;
add regression coverage for a nonzero exact expert path and a wrapped module.
…t format Signed-off-by: Joe Rowell <joerowell4@gmail.com>
ab2b98d to
37e78a7
Compare
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
The load_hf_quant_config merge is the right place for this and the switch from set() to dict.fromkeys fixes a real nondeterminism, but the fused-MoE half — the change the title leads with — has no test and rests on an assumption worth making explicit (see inline).
Two things at the whole-PR level:
Scope. Three independent fixes are bundled here: the general producer-ignore merge, the fused-MoE exclusion matching, and two Laguna-specific fixes (g_proj quant config, weight_scale_inv double-rename). The Laguna weight_scale_inv fix in particular is unrelated to exclusions and isn't mentioned in the PR title. Splitting the Laguna changes out would make each easier to bisect if one regresses.
Test coverage. apply_quant_config_exclude_modules is pure Python over a module tree and is directly testable — build a small nn.Module containing a ConfigurableMoE-shaped wrapper, set exclude_modules to a per-expert regex, and assert the backend's quant_config.quant_algo ends up None and _weights_created is False. Without it, the .backend strip and the per-expert candidate names are only exercised on real checkpoints. Please add coverage for that path.
Also worth confirming: tests/unittest/_torch/models/test_laguna_gproj_quant.py is a new directory-level file — check whether tests/unittest/_torch/models/ is already collected by an existing CI stage or whether the new file needs registering in tests/integration/test_lists/test-db/.
I couldn't run the suite in my checkout (missing deps), so the assertions above are from reading, not from a green run.
Notes on lines outside the diff:
- tensorrt_llm/_torch/model_config.py:580: This MXFP8
ignored_layershandling is now fully subsumed by the general merge at line 613 — and it differs subtly (it dropsdefault_excludesemantics and only merges whenhf_exclude_modulesis None vs. not). Leaving both means two places to keep in sync. Reduce this branch to just settingexclude_modulesfrommodules_to_not_convertand let the general merge addignored_layers.
| if name.endswith('.backend') else name) | ||
| candidates += [ | ||
| f'{base}.0.gate_proj', | ||
| f'{base}.0.up_proj', |
There was a problem hiding this comment.
Using expert 0 as representative makes this an all-or-nothing decision for the whole fused module, and it fails silently in both directions:
- Ignore list covers experts 0–3 of 128 →
...experts.0.gate_projmatches → the entire fused module is de-quantized, but experts 4–127 are FP8-with-scales in the checkpoint. Loading then either trips on shapes/scales or produces garbage. - Ignore list covers experts 1–127 but not 0 → nothing matches → the whole layer stays quantized even though most experts are BF16.
Uniform experts.[0-9]+.… rules are the common case, but partial lists exist and there's nothing here that says "we generalized from expert 0." Probe at least a second index (num_experts - 1) and logger.warning when the two disagree, so a partial list is visible rather than silently rounded to all-or-nothing.
| # only the wrapper does nothing because create_weights() | ||
| # delegates to the backend. Without this the backend | ||
| # stays quantized and the output is still degenerate. | ||
| base = (name[:-len('.backend')] |
There was a problem hiding this comment.
Worth noting in the comment that the .backend strip is only needed for the synthesized per-expert candidates. For plain rules, is_module_excluded_from_quantization already walks ancestors (...experts.backend → ...experts), so a rule naming the wrapper reaches the backend without this. As written the comment claims excluding the wrapper "does nothing," which reads as broader than what actually holds.
| else: | ||
| quant_config.exclude_modules = default_exclude | ||
|
|
||
| # Honour the producer's "leave these layers unquantized" list, whatever |
There was a problem hiding this comment.
The comment attributes ignored_layers to modelopt, but modelopt configs never reach here — line 520 returns early via _build_modelopt_quant_config, and read_modelopt_quant_config renames modelopt's ignore to exclude_modules itself. Drop modelopt from the list so a future reader doesn't assume this block covers that path.
| alignment, so a quantised g_proj trips Linear.__init__'s assert under | ||
| tp_size > 1. __post_init__'s generic exclusion pass runs too late. | ||
| """ | ||
| quant_config = model_config.get_quant_config() |
There was a problem hiding this comment.
get_quant_config() is called with no name, so this bypasses per_layer_quant_configs. For a MIXED_PRECISION checkpoint, self.quant_config is the global config and the per-layer entry for this layer is never consulted — the exclusion check then runs against the wrong config. The pre-existing call site had the same behaviour so this isn't a regression, but since the function's whole job is resolving one specific layer's config, passing name (and falling back on ValueError) would make it correct for that case too.
Signed-off-by: Joe Rowell <joerowell4@gmail.com>
zhaoyangwang-nvidia
left a comment
zhaoyangwang-nvidia
left a comment
There was a problem hiding this comment.
Approve with nits.
| # g_proj whose out dim is not a multiple of the block alignment, and | ||
| # MoE experts the producer kept in bf16 for quality). Merged on top of | ||
| # whatever per-format defaults were set above. | ||
| producer_ignored = list(hf_quant_config.get("ignored_layers", []) or []) |
There was a problem hiding this comment.
The new cross-format merge already folds in ignored_layers, so the mxfp8-specific read at L581-586 is now dead — it can no longer change the result after dedup. Please drop it so the producer-ignore list is read in exactly one place; otherwise a future edit to one site silently diverges from the other.
| tp_size > 1. __post_init__'s generic exclusion pass runs too late. | ||
| """ | ||
| quant_config = model_config.get_quant_config() | ||
| if quant_config is None or layer_idx is None: |
There was a problem hiding this comment.
layer_idx cannot be None here: LagunaAttention.__init__ already indexes per_layer_heads[layer_idx] (L257) and layer_types[layer_idx] (L263) before g_proj is built, and the only caller LagunaDecoderLayer types it as int. Falling back to the original quantized config on None also silently reinstates the exact TP>1 assert this PR fixes, and test_missing_layer_idx_does_not_match_a_literal_none pins that as intended behaviour. Suggest tightening the parameter to layer_idx: int and dropping both the guard and that test — or raising explicitly if you do want to defend against it.
| if quant_config is None or layer_idx is None: | ||
| return quant_config | ||
|
|
||
| name = f"model.layers.{layer_idx}.self_attn.g_proj" |
There was a problem hiding this comment.
This hardcodes the module path that named_modules() produces for g_proj, so the two must be kept in sync by hand. If the attribute name or the model.layers. prefix ever changes, this silently stops matching and the only symptom is the Linear.__init__ assert under TP>1 — far from the cause. Consider passing the module's own prefix in from the call site, or at minimum noting the coupling in the docstring.
|
One interaction to preserve when this conflicts with #17725: recursively adding producer |
5 participants
Dev Engineer Review
ignored_layersandignoreentries intoquant_config.exclude_modulesfor all quantization formats.g_projhandling that disables incompatible weight quantization while preserving KV-cache quantization.weight_scale_invrenaming.l0_cputest configuration with valid test paths.QA Engineer Review
test_excluded_g_proj_is_not_quantised.test_unexcluded_g_proj_keeps_the_model_quant_config.test_unquantised_model_is_passed_through.test_missing_layer_idx_does_not_match_a_literal_none.tests/unittest/_torch/test_hf_quant_config.py.tests/integration/test_lists/test-db/l0_cpu.yml.