[OMNIML-5570] 2/2 Compose GEMM and KV-cache AutoQuant workflows - #2273
[OMNIML-5570] 2/2 Compose GEMM and KV-cache AutoQuant workflows#2273meenchen wants to merge 2 commits into
Conversation
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. 🗂️ Base branches to auto review (3)
Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Comment |
meenchen
changed the title
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## agent/kv-cache-autoquant-core #2273 +/- ##
=================================================================
- Coverage 70.49% 69.84% -0.65%
=================================================================
Files 526 526
Lines 61588 61551 -37
=================================================================
- Hits 43417 42992 -425
- Misses 18171 18559 +388
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:
|
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (claude-opus-5) — DM the bot to share feedback.
Stacked composition PR (targets agent/kv-cache-autoquant-core, not main). The orchestration itself is small and readable, and the happy paths for both new shipped recipes are covered by tests. Three things I'd want resolved before this lands:
1. README now contradicts the shipped recipes and the new test. The diff rewrites the KV-AutoQuant paragraph to say "The shipped canary recipe searches calibrated FP8 K/V ... Each candidate uses max calibration so its persistent K/V scales are present in the unified HF checkpoint." All three shipped KV recipes (kv_fp8_nvfp4_cast_…, fp8_ptq_then_kv_…, nvfp4_fp8_gradient_then_kv_…) still use algorithm: (None) + constant_amax: 448.0, i.e. calibration-free cast candidates — which is exactly what the new loader test test_builtin_kv_autoquantize_recipes_use_calibration_free_cast_candidates asserts, and what kv_cache_auto_quant._validate_deployable_candidate treats as the constant-amax branch. Either revert this doc edit or change the recipes.
2. get_quant_config fail-closed guard removed with no replacement signal. The NotImplementedError for "uniform quantized weights + mixed-precision KV map" is dropped so the fixed-FP8-PTQ→mixed-KV flow can export. But the PR body itself says this combination "remains gated until the runtime's uniform-weight ModelOpt configuration consumes kv_cache_quantized_layers" — so the export now silently produces a checkpoint no released runtime can load. At minimum emit a warn() on that branch. Note this also affects non-AutoQuant flows: any uniform-weight PTQ run where only some KV-eligible layers are quantized (e.g. MTP layers excluded) now falls into the elif and exports kv_cache_quant_algo: MIXED_PRECISION + a per-layer map instead of the plain uniform algo; that path has no test.
3. Design gate (5 directories) is unaddressed in the PR body. The recipe schema already expresses stages as quantize (fixed PTQ) + auto_quantize; this adds a third special-cased field kv_auto_quantize plus three new cross-field validators, a second --kv_auto_quantize_checkpoint CLI flag, and a primary_is_kv / primary_uses_kv_checkpoint branch in the runner. The obvious in-repo alternative — a single ordered stage list (auto_quantize: list[AutoQuantizeConfig], or a generic stages:) with one checkpoint path per stage index — would collapse the validators and the checkpoint-attr branching and generalizes past two stages. The PR body explains what was built but not why the existing two-field shape couldn't be generalized. Please state the "why not a stage list / why not extend mtq.auto_quantize to both domains" rationale in the body.
Also: no negative tests for the three new ModelOptAutoQuantizeRecipe validators, and the checkpoint_attr string-indirection is worth simplifying (details inline).
[{"file": "examples/hf_ptq/README.md", "line": 457, "body": "This contradicts every shipped KV recipe and the new test. kv_fp8_nvfp4_cast_kl_div_at_5p4bits, fp8_ptq_then_kv_… and nvfp4_fp8_gradient_then_kv_… all specify algorithm: (None) with constant_amax: 448.0 — no calibration pass at all — and test_builtin_kv_autoquantize_recipes_use_calibration_free_cast_candidates asserts exactly that (candidate.algorithm is None). The previous wording ("Each candidate uses an explicit constant scale, avoiding an additional calibration pass") was correct. Same for the line above: "calibrated FP8 K/V" should stay "FP8-cast K/V", matching the recipe filenames."}, {"file": "modelopt/torch/export/quant_utils.py", "line": 1777, "body": "Two concerns with dropping the fail-closed guard here:\n\n1. Per the PR body, uniform-weight + mixed-KV checkpoints are still undeployable ("gated until the runtime's uniform-weight ModelOpt configuration consumes kv_cache_quantized_layers"). Previously that raised; now it exports silently. Please emit a warn() on this branch so a user doesn't discover it at deploy time.\n2. This branch is also reached by plain PTQ runs where only some KV-eligible layers are quantized (e.g. MTP/attention layers excluded by the recipe): len(kv_cache_formats) == 1 but all_kv_layers_quantized is False. Those exports now flip from a uniform kv_cache_quant_algo to MIXED_PRECISION + a per-layer map. That's a deployment-visible metadata change for existing recipes and there's no test for it — worth one covering "uniform weight format, uniform KV format, partial KV coverage"."}, {"file": "examples/hf_ptq/hf_ptq.py", "line": 481, "body": "checkpoint_attr: str + getattr(args, checkpoint_attr, None) is stringly-typed indirection that also silently yields None when the attribute is missing (see test_composed_kv_autoquantize_rejects_enabled_actual_kv_quantizers, whose SimpleNamespace has no kv_auto_quantize_checkpoint). Passing the resolved value — checkpoint: str | None = None, computed once in _run_auto_quantize_recipe — is simpler, keeps auto_quantize independent of argparse attribute names, and makes a missing flag a real error rather than a silent no-checkpoint run."}, {"file": "examples/hf_ptq/hf_ptq.py", "line": 935, "body": "The checkpoint flag a user must pass for the same KV search depends on whether a fixed quantize block precedes it: standalone KV recipe → --auto_quantize_checkpoint, fixed-PTQ-then-KV → --kv_auto_quantize_checkpoint. That's surprising and only discoverable from the README. Consider making any kv_effective_bits stage always use --kv_auto_quantize_checkpoint (with a one-release fallback to the old flag), or at least add a comment here explaining the rule.\n\nAlso: _assert_kv_autoquantize_input_is_clean only fires after the fixed PTQ stage has run a full calibration pass. Since the offending config is knowable from the recipe (a quantize block that enables *[kv]_bmm_quantizer under a KV-primary recipe), a cheap pre-check at recipe load/stage start would fail in seconds instead of after calibration."}, {"file": "modelopt/recipe/config.py", "line": 358, "body": "None of the three new error paths (kv_auto_quantize after a KV-domain primary, kv_auto_quantize without kv_effective_bits, auto_quantize.kv_cache + kv_auto_quantize) has a negative test — tests/unit/recipe/test_loader.py only adds happy-path cases. Please add pytest.raises(ValidationError, match=...) coverage for each, consistent with the existing test_load_recipe_autoquantize_fixed_baseline_requires_explicit_search style.\n\nSeparately: primary_is_kv now short-circuits the three pre-existing fixed-baseline checks. That means a KV-primary recipe with both a quantize baseline and module_search_spaces is silently accepted even though _run_auto_quantize_recipe drops fixed_quantize_config after mono_quantize and AutoQuantizeConfig._has_search_space already rejects module_search_spaces for KV. Worth an explicit error rather than relying on the other validator."}]
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (claude-opus-5) — DM the bot to share feedback.
Re-review of the composition PR (still stacked on #2272, targets agent/kv-cache-autoquant-core). Most of the previous round is genuinely fixed:
Resolved
- README contradicted the shipped recipes (critical) — fixed. The paragraph is back to "The shipped canary recipe searches FP8-cast K/V …" / "Each candidate uses an explicit constant scale, avoiding an additional calibration pass", which matches
algorithm: None+constant_amax: 448.0in all three recipes andtest_builtin_kv_autoquantize_recipes_use_calibration_free_cast_candidates. - Fail-closed guard dropped with no signal (critical) — fixed.
get_quant_confignowwarn()s on the uniform-weight + mixed-KV branch, and both new export tests (…exports_mixed_kv_cache_map,…partial_kv_map_with_warning) cover the branch, including the partial-coverage case I asked for. checkpoint_attrstring indirection (minor) — fixed:auto_quantize(..., checkpoint: str | None)is resolved once by the caller.- Flag inconsistency + late clean-input check (minor/critical mix) — fixed:
_resolve_kv_auto_quantize_checkpointgives one rule for every KV-domain search with a deprecated fallback, and_quantize_config_explicitly_enables_kvnow fails beforemono_quantizeruns (test assertsmono_quantizemust not start). - No negative tests for the new validators (critical) — fixed: four
pytest.raises(ValidationError, …)cases, including the new explicit "KV-primary + fixed baseline must not definemodule_search_spaces" error.
Still open
- Design gate (blocker #3) is unchanged. The PR body still explains what was built (an extra
kv_auto_quantizefield, three cross-field validators, a second CLI checkpoint flag, aprimary_is_kvbranch in the runner) but never says why the obvious generalization — one ordered stage list (stages:/auto_quantize: list[AutoQuantizeConfig]) with a checkpoint path per stage — was rejected, nor whymtq.auto_quantize/auto_quantize_kv_cachecouldn't be composed at the API level instead of inhf_ptq.py. Please add that rationale to the body; per the design protocol I can't approve while it's unaddressed. - Two small residuals inline (doc/behavior mismatch on the deprecated-flag scope; dead fallback in the KV pre-check), plus one thing worth a human eye: with
all_kv_layers_quantized(from #2272) plus this PR'swarn()branch, an existing plain-PTQ recipe that excludes some KV-eligible layers (e.g.*mtp*) now exportskv_cache_quant_algo: MIXED_PRECISION+ a per-layer map instead of the uniform algo. It's now tested and warned, but it's a deployment-visible metadata change for shipped recipes and the CHANGELOG is marked N/A here — make sure #2272's entry actually calls it out.
| ) | ||
|
|
||
| kv_auto_quantize: AutoQuantizeConfig | None = ModeloptField( | ||
| default=None, |
There was a problem hiding this comment.
Bot comment.
Design gate, unchanged from the last round: kv_auto_quantize is a third special-cased stage field alongside quantize and auto_quantize, and it brings three new cross-field validators plus a primary_is_kv short-circuit through the pre-existing fixed-baseline checks. The recipe schema already expresses "fixed PTQ then search"; the natural generalization is a single ordered stage list (stages: / auto_quantize: list[AutoQuantizeConfig]) with one checkpoint path per stage index, which collapses all three validators and the checkpoint-attr branching in the runner and generalizes past two stages. Please state in the PR body why that (and/or composing the two searches inside mtq.auto_quantize* rather than in hf_ptq.py) was rejected — the body currently documents the shape but not the choice.
| help=( | ||
| "Path for saving/restoring any KV-cache AutoQuantize search checkpoint. Use a new " | ||
| "path whenever the preceding weight/activation quantization stage changes. Standalone " | ||
| "KV recipes temporarily accept --auto_quantize_checkpoint as a deprecated fallback." |
There was a problem hiding this comment.
Bot comment.
The help text (and the matching README sentence) says "Standalone KV recipes temporarily accept --auto_quantize_checkpoint as a deprecated fallback", but _run_auto_quantize_recipe calls _resolve_kv_auto_quantize_checkpoint for any KV-primary recipe — including the composed fp8_ptq_then_kv_…, which is not standalone. The behavior is the right one (one rule: every KV-domain search prefers the dedicated flag); it's the wording that's narrower than the code. Suggest "KV-primary recipes" in both places.
| continue | ||
| suffix = pattern.rsplit(".", 1)[-1] | ||
| for name in enabled: | ||
| if fnmatch(name, suffix): |
There was a problem hiding this comment.
Bot comment.
entry.get("enable", entry.get("cfg") is not None): QuantizerCfgEntry.enable is bool with default=True, and this function is fed model_dump() output, so "enable" is always present and the fallback is dead. Not a bug (the dumped default True gives the same fail-closed answer), but the expression reads as if an absent enable were possible — entry["enable"] is clearer, or add a comment if you want to keep it tolerant of hand-built dicts.
2 participants
What does this PR do?
Type of change: new feature.
Stacked on #2272. Adds composition of existing GEMM quantization with the standalone KV-cache AutoQuant API:
kv_auto_quantizerecipe stage with independent method, constraints, candidates, and checkpoint path;hf_ptq.pyorchestration that keeps the selected weight/activation QDQ behavior active while freezing its calibration state during KV scoring;The two stages remain independent searches and reuse the public APIs from #2272. This PR does not change either solver, scoring protocol, or checkpoint schema.
Usage
Fixed FP8 GEMM PTQ followed by KV AutoQuantize:
Weight AutoQuantize followed by KV AutoQuantize:
Testing
f94d74a98(tree2564939297095db02c0ffea5848a0c1f99b18d6d).Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
agent/kv-cache-autoquant-coreso reviewers see only the nine composition-specific files.--auto_quantize_checkpointand--kv_auto_quantize_checkpointare intentionally separate. Changing the preceding GEMM configuration requires a fresh KV checkpoint path.kv_cache_quantized_layers.