HAQ Poc - #2138
Draft
juhi10071998 wants to merge 7 commits into
Draft
Conversation
…xed-kernel lookup First HAQ implementation slice: a self-contained latency look-up table for hardware-aware AutoQuantize, decoupled from the solver so it is unit-testable and reusable by the offline benchmark-to-LUT converter. - FixedKernelPolicy/KernelSelector: versioned, explicit (op_kind, recipe_id) -> backend selection; never falls back to another backend. enable_w4a4_proxy (W4A16 only) authorizes the documented bsz=1 low-M W4A4->W4A16 proxy and forces proxy provenance. - canonicalize_benchmark_csv: parses the sectioned combined_results.csv and emits haq_latency_v1 rows. Strict: exactly one successful row per (group, M, declared recipe); zero/failed/ambiguous are aggregated coverage errors, never minimized across backends. - LatencyLUT: validates the canonical CSV, exposes a SHA-256 digest, and does exact (profile, m, group_pattern, recipe_id) lookup + source-pattern group matching. 23 synthetic tests, ruff-clean, mypy-clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com>
Add the LatencyCostModel and constraint schema for hardware-aware AutoQuantize,
preserving effective-bits defaults:
- constraints['cost_model'] = 'latency' with a top-level constraints['latency']
{relative_to_min >= 1.0} budget and constraints['cost'] {lut_path,
deployment_profile, positive-int m}.
- latency and effective_bits are mutually exclusive; the 'latency' block is
rejected for every non-latency cost model; cost-excluded patterns still apply.
- weight/active_moe paths and defaults unchanged.
13 constraint-validation tests; existing cost-model tests unaffected.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Juhi Mittal <juhim@nvidia.com>
Wire the haq_latency_v1 LUT into the AutoQuantize gradient searcher so the LP minimizes sensitivity subject to a summed-kernel-latency budget, leaving the effective-bits and active-MoE paths unchanged. - QuantRecipe stores its stable format name (self.name), the LUT recipe_id join key (NONE / FP8_DEFAULT_CFG / W4A16_NVFP4_CFG). - before_search loads/validates the LUT (cost_model: latency), records profile/M/relative_to_min/digest, rejects latency + non-gradient scoring. - _candidate_cost returns per-(group, recipe) LUT latency (0 for cost-excluded); _verify_latency_coverage aggregates every missing/failed/ambiguous row into one pre-calibration error (never substitutes a backend). - run_search budget = relative_to_min * summed per-group minimum latency; drops the near-budget lower-bound workaround for latency; reports selected/min/budget latency, profile, M, LUT digest, kernel policy in best['constraints']. Two CPU e2e tests; full test_autoquant.py: 111 passed. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com>
Extend ModelOptAutoQuantizeRecipe so a recipe can drive the latency-aware AutoQuantize search: - cost_model gains 'latency'; AutoQuantizeCost gains lut_path/deployment_profile/m; a top-level 'latency' block carries relative_to_min (>= 1.0). - A model validator enforces latency<->effective_bits exclusivity and the required latency fields; to_mtq_constraints() serializes to the mtq dict, dropping effective_bits for latency and emitting the latency budget + cost keys. - hf_ptq translation uses to_mtq_constraints() so the recipe reaches mtq.auto_quantize unchanged for both effective-bits and latency. 5 schema tests added; recipe loader + hf_ptq arg tests: 231 passed. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com>
Latency counterpart to w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml: same disabled and cost-excluded layers and same FP8 + NVFP4-weight-only + BF16 candidates, but cost_model: latency against the batch-size-1 M=1 TRT-LLM LUT so the two objectives are compared on one setup. W4A16 is priced by the documented W4A4 proxy. (Recipe validated via modelopt.recipe.loader.load_recipe + schema unit tests; the check-modelopt-recipes hook was skipped as it rebuilds nvidia-modelopt and fails on a root-squashed egg-info left by an earlier container editable install.) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com>
conv1d, shared_expert_gate, and lm_head are disabled (BF16) but the batch-size-1 GEMM/MoE benchmark never measured them, so the latency LUT has no rows for them. Effective-bits prices them from weight numel; latency cannot. Cost-exclude them (the design's sanctioned explicit-exclusion path) so the strict coverage check passes. Validated: the Qwen3.6-35B latency smoke now runs end-to-end. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com>
M=32 variant of the M=1 latency recipe (same disabled/excluded/candidates and the same TRT-LLM LUT, only cost.m=32). M=32 is more compute-bound than the M=1 decode point; comparing the two shows how the latency-optimal allocation shifts as the workload leaves the weight-memory-bound regime (at M=32 the MoE experts strongly favor W4A16 over FP8, unlike M=1). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.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. |
Contributor
📝 WalkthroughWalkthroughAutoQuantize now supports hardware-aware latency optimization. The change adds latency constraints, validated latency LUTs, latency-based candidate costing, relative latency budgets, search metadata, Qwen3.6 MoE recipes, and comprehensive tests. ChangesLatency-Based AutoQuantize
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant AutoQuantizeSearcher
participant LatencyLUT
participant GradientSolver
participant SearchHistory
AutoQuantizeSearcher->>LatencyLUT: Load and validate latency LUT
AutoQuantizeSearcher->>LatencyLUT: Request candidate latency by group and recipe
LatencyLUT-->>AutoQuantizeSearcher: Return latency costs or coverage errors
AutoQuantizeSearcher->>GradientSolver: Apply latency budget and candidate costs
GradientSolver-->>AutoQuantizeSearcher: Return selected quantization resolution
AutoQuantizeSearcher->>SearchHistory: Record latency and LUT metadata
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
Contributor
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
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: 8
🧹 Nitpick comments (4)
tests/unit/torch/quantization/test_auto_quantize_latency.py (2)
426-453: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove the
ALL_COLUMNSimport to module scope.Line 429 imports
ALL_COLUMNSinside_minimal_canonical_csv. The same module is already imported at Line 23. No circular import or optional dependency applies.Based on path instructions: "Imports inside functions or test methods without explicit justification. Imports belong at the top of the file so import errors surface at collection time, not mid-test."
♻️ Proposed change
from modelopt.torch.quantization._auto_quantize_latency import ( + ALL_COLUMNS, RECIPE_FP8,def _minimal_canonical_csv( *, schema_version=SCHEMA_VERSION, latency_us="5.0", cost_is_proxy="False" ) -> str: - from modelopt.torch.quantization._auto_quantize_latency import ALL_COLUMNS - values = {🤖 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 `@tests/unit/torch/quantization/test_auto_quantize_latency.py` around lines 426 - 453, Move the ALL_COLUMNS import from inside _minimal_canonical_csv to the module-level imports, reusing the existing _auto_quantize_latency module import and leaving the helper’s CSV construction unchanged.Source: Path instructions
232-234: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe test name does not match the behavior it validates, and the named behavior is untested.
test_selector_rejects_proxy_on_wrong_recipeasserts thatKernelSelector(enable_w4a4_proxy=True)withoutproxy_reasonraises. That is the__post_init__check atmodelopt/torch/quantization/_auto_quantize_latency.pyLines 184-185, not a wrong-recipe rejection.The wrong-recipe branch at Lines 493-497 of that module, which appends
"enable_w4a4_proxy is only valid for W4A16_NVFP4_CFG", has no test.Rename this test and add coverage for the canonicalizer branch.
As per coding guidelines: "Tests must exercise the behavior they claim to validate."
♻️ Proposed change
-def test_selector_rejects_proxy_on_wrong_recipe(): +def test_selector_requires_proxy_reason(): with pytest.raises(ValueError, match="proxy_reason"): KernelSelector(backend="nvfp4_trtllm", enable_w4a4_proxy=True) + + +def test_canonicalize_rejects_proxy_on_wrong_recipe(raw_csv): + policy = FixedKernelPolicy( + kernel_policy_id="p", + selectors={ + "gemm": { + RECIPE_FP8: KernelSelector( + backend="fp8_trtllm", enable_w4a4_proxy=True, proxy_reason="poc" + ) + } + }, + ) + rows, problems = _canonicalize(raw_csv, policy) + assert rows == [] + assert any("enable_w4a4_proxy is only valid" in p for p in problems)🤖 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 `@tests/unit/torch/quantization/test_auto_quantize_latency.py` around lines 232 - 234, Rename test_selector_rejects_proxy_on_wrong_recipe to describe the missing-proxy-reason validation, then add a separate test covering the canonicalizer’s wrong-recipe branch: enable the proxy with a non-W4A16_NVFP4_CFG recipe and assert the resulting validation error includes “enable_w4a4_proxy is only valid for W4A16_NVFP4_CFG”.Source: Coding guidelines
modelopt/torch/quantization/_auto_quantize_latency.py (2)
722-761: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winGroup-pattern resolution rescans the whole LUT on every candidate pricing call.
match_group_patternbuilds itsseenindex by iteratingself._rowson each invocation, and the searcher calls it once per(hparam, recipe)pair. For a large model the total work is O(groups x recipes x rows), and_verify_latency_coveragerepeats the same resolution a second time before calibration.
modelopt/torch/quantization/_auto_quantize_latency.py#L722-L761: precompute the(deployment_profile, m) -> {group_pattern: source_module_patterns}index once inLatencyLUT.__init__and read it inmatch_group_patterninstead of scanningself._rows.modelopt/torch/quantization/algorithms.py#L1413-L1430: resolvegroup_patternonce perhparamand cache it, so_candidate_costperforms only thelookupcall for each recipe.🤖 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 `@modelopt/torch/quantization/_auto_quantize_latency.py` around lines 722 - 761, In modelopt/torch/quantization/_auto_quantize_latency.py lines 722-761, precompute a (deployment_profile, m) to group_pattern/source_module_patterns index during LatencyLUT.__init__, then update match_group_pattern to read that index without rescanning self._rows. In modelopt/torch/quantization/algorithms.py lines 1413-1430, resolve and cache group_pattern once per hparam so _candidate_cost only performs the recipe-specific lookup; apply the change at both listed sites.
219-220: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd docstrings to the public APIs.
Ruff does not enforce pydocstyle in this private module because
*/_[a-zA-Z]*ignoresDrules. The coding guidelines still require docstrings for the listed methods and properties. Document the"mixed"sentinel returned byselection_policyandkernel_policy_id.🤖 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 `@modelopt/torch/quantization/_auto_quantize_latency.py` around lines 219 - 220, Add docstrings to the public APIs in the latency quantization module, including selector and the properties selection_policy and kernel_policy_id. Document their purpose, parameters or return values as appropriate, and explicitly describe the "mixed" sentinel returned by selection_policy and kernel_policy_id.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
`@modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_latency-trtllm_m1.yaml`:
- Around line 43-46: Replace the developer-specific absolute lut_path in
modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_latency-trtllm_m1.yaml
lines 43-46 with a repository-relative path to the committed LUT or a documented
override placeholder, and correct the Line 23
outputs/qwen36_trtllm_kernel_policy.yaml reference. Apply the identical lut_path
change in
modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_latency-trtllm_m32.yaml
lines 40-43 so both recipes resolve the same artifact.
In `@modelopt/recipe/config.py`:
- Around line 240-285: Update _validate_latency_cost_model to reject explicitly
provided effective_bits and unsupported latency cost keys using
model_fields_set, including cost.active_moe_expert_ratio, instead of silently
dropping them in to_mtq_constraints. Preserve the existing required latency cost
validation and ensure invalid fields fail at recipe validation before
serialization.
In `@modelopt/torch/quantization/_auto_quantize_latency.py`:
- Around line 52-61: Update the module’s __all__ declaration to include the
consumer-imported public symbols RECIPE_FP8, RECIPE_NONE, RECIPE_W4A16_NVFP4,
SCHEMA_VERSION, normalize_layer_indices, ALL_COLUMNS, and parse_benchmark_csv,
preserving the existing exports.
- Around line 607-682: Update LatencyRow construction in from_csv to validate
every required integer conversion for m, tp, and ep within the existing per-row
validation flow, catching invalid or missing values and appending row-numbered
messages to problems before continuing. Ensure malformed integer fields never
escape as bare ValueError or TypeError, while preserving successful row parsing
and aggregated LatencyCoverageError behavior.
- Around line 236-258: Update load_fixed_kernel_policy to reject unknown
top-level keys and require selectors to be a non-empty mapping, raising
ValueError for missing or malformed coverage configuration instead of
constructing an empty policy. Update _parse_selector to validate that each
selector entry contains backend and raise ValueError when it is absent, while
preserving existing valid-policy parsing.
In `@modelopt/torch/quantization/algorithms.py`:
- Line 1154: Persist the latency identity triple in default_state_dict and
validate it in before_search alongside cost_model, rejecting checkpoints whose
LUT digest, deployment profile, or m differs from the active configuration.
Ensure restored latency costs cannot be used after an identity mismatch. Add
class-level defaults for _latency_profile, _latency_m, _latency_relative_to_min,
and _lut_digest, matching the existing _latency_lut default.
- Around line 1552-1565: Update _resolve_best_recipe to detect latency-search
state before treating constraints["effective_bits"] or
search_state["cost_denominator"] as weight budgets. Either reject re-solving for
latency states with a clear error, or route it through the existing latency
budget and LUT metadata so recipe selection uses latency values correctly;
preserve the current weight-budget behavior for non-latency searches.
In `@tests/unit/recipe/test_loader.py`:
- Around line 1891-1955: Move the AutoQuantizeConstraints and pytest imports
used by the five new tests to module scope alongside the existing imports in the
test module. Remove the corresponding imports from
test_autoquantize_constraints_latency_to_mtq,
test_autoquantize_constraints_effective_bits_to_mtq_unchanged,
test_autoquantize_constraints_latency_requires_fields,
test_autoquantize_constraints_latency_block_rejected_for_other_cost_models, and
test_autoquantize_constraints_relative_to_min_lower_bound.
---
Nitpick comments:
In `@modelopt/torch/quantization/_auto_quantize_latency.py`:
- Around line 722-761: In modelopt/torch/quantization/_auto_quantize_latency.py
lines 722-761, precompute a (deployment_profile, m) to
group_pattern/source_module_patterns index during LatencyLUT.__init__, then
update match_group_pattern to read that index without rescanning self._rows. In
modelopt/torch/quantization/algorithms.py lines 1413-1430, resolve and cache
group_pattern once per hparam so _candidate_cost only performs the
recipe-specific lookup; apply the change at both listed sites.
- Around line 219-220: Add docstrings to the public APIs in the latency
quantization module, including selector and the properties selection_policy and
kernel_policy_id. Document their purpose, parameters or return values as
appropriate, and explicitly describe the "mixed" sentinel returned by
selection_policy and kernel_policy_id.
In `@tests/unit/torch/quantization/test_auto_quantize_latency.py`:
- Around line 426-453: Move the ALL_COLUMNS import from inside
_minimal_canonical_csv to the module-level imports, reusing the existing
_auto_quantize_latency module import and leaving the helper’s CSV construction
unchanged.
- Around line 232-234: Rename test_selector_rejects_proxy_on_wrong_recipe to
describe the missing-proxy-reason validation, then add a separate test covering
the canonicalizer’s wrong-recipe branch: enable the proxy with a
non-W4A16_NVFP4_CFG recipe and assert the resulting validation error includes
“enable_w4a4_proxy is only valid for W4A16_NVFP4_CFG”.
🪄 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: 22091def-c436-429e-ba2d-b4f0fd5fb3e2
📒 Files selected for processing (10)
examples/hf_ptq/hf_ptq.pymodelopt/recipe/config.pymodelopt/torch/quantization/_auto_quantize_cost.pymodelopt/torch/quantization/_auto_quantize_latency.pymodelopt/torch/quantization/algorithms.pymodelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_latency-trtllm_m1.yamlmodelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_latency-trtllm_m32.yamltests/unit/recipe/test_loader.pytests/unit/torch/quantization/test_auto_quantize_latency.pytests/unit/torch/quantization/test_autoquant.py
Comment on lines
+43
to
+46
| cost: | ||
| lut_path: /home/scratch.juhim_coreai/code/dev_notes/Autoquant_research/outputs/qwen36_sm100_tp1_ep1_haq_latency_v1.csv | ||
| deployment_profile: qwen36_sm100_tp1_ep1_decode | ||
| m: 1 |
Contributor
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Both shipped recipes reference a LUT in a developer scratch home directory. The root cause is one unshipped artifact, qwen36_sm100_tp1_ep1_haq_latency_v1.csv, referenced by absolute path. LatencyCostModel.normalize_cost_constraints accepts any non-empty string, so the failure surfaces only when LatencyLUT.from_csv opens the file, and the path also embeds a personal account name.
modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_latency-trtllm_m1.yaml#L43-L46: replace the absolutelut_pathwith a repository-relative path to a committed LUT, or a documented placeholder the user must override; also fix the Line 23 reference tooutputs/qwen36_trtllm_kernel_policy.yaml.modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_latency-trtllm_m32.yaml#L40-L43: apply the identicallut_pathchange so both recipes resolve the same committed artifact.
📍 Affects 2 files
modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_latency-trtllm_m1.yaml#L43-L46(this comment)modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_latency-trtllm_m32.yaml#L40-L43
🤖 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
`@modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_latency-trtllm_m1.yaml`
around lines 43 - 46, Replace the developer-specific absolute lut_path in
modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_latency-trtllm_m1.yaml
lines 43-46 with a repository-relative path to the committed LUT or a documented
override placeholder, and correct the Line 23
outputs/qwen36_trtllm_kernel_policy.yaml reference. Apply the identical lut_path
change in
modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_latency-trtllm_m32.yaml
lines 40-43 so both recipes resolve the same artifact.
Source: Path instructions
Comment on lines
+240
to
+285
| @model_validator(mode="after") | ||
| def _validate_latency_cost_model(self): | ||
| if self.cost_model == "latency": | ||
| if self.latency is None: | ||
| raise ValueError("cost_model: latency requires a 'latency' budget block.") | ||
| missing = [ | ||
| key | ||
| for key in ("lut_path", "deployment_profile", "m") | ||
| if self.cost is None or getattr(self.cost, key) is None | ||
| ] | ||
| if missing: | ||
| raise ValueError( | ||
| f"cost_model: latency requires cost.{{{', '.join(missing)}}} to be set." | ||
| ) | ||
| elif self.latency is not None: | ||
| raise ValueError( | ||
| f"The 'latency' budget block is only valid with cost_model: latency, " | ||
| f"got cost_model={self.cost_model!r}." | ||
| ) | ||
| return self | ||
|
|
||
| def to_mtq_constraints(self) -> dict: | ||
| """Serialize to the dict passed to ``mtq.auto_quantize(constraints=...)``. | ||
|
|
||
| The 'latency' cost model drops ``effective_bits`` (mutually exclusive) and | ||
| emits the ``latency`` budget plus the latency ``cost`` keys; other cost | ||
| models keep the historical effective-bits dict. | ||
| """ | ||
| if self.cost_model == "latency": | ||
| assert self.latency is not None # guaranteed by _validate_latency_cost_model | ||
| cost = { | ||
| key: getattr(self.cost, key) | ||
| for key in ("lut_path", "deployment_profile", "m") | ||
| if self.cost is not None and getattr(self.cost, key) is not None | ||
| } | ||
| return { | ||
| "cost_model": "latency", | ||
| "latency": {"relative_to_min": self.latency.relative_to_min}, | ||
| "cost": cost, | ||
| } | ||
| out: dict = {"effective_bits": self.effective_bits, "cost_model": self.cost_model} | ||
| if self.cost is not None: | ||
| cost = self.cost.model_dump(exclude_none=True) | ||
| if cost: | ||
| out["cost"] = cost | ||
| return out |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Reject explicitly set effective_bits (and non-latency cost keys) under cost_model: latency instead of dropping them.
_validate_latency_cost_model accepts a recipe that sets both cost_model: latency and effective_bits. to_mtq_constraints then drops effective_bits silently. The MTQ layer disagrees: normalize_auto_quantize_constraints in modelopt/torch/quantization/_auto_quantize_cost.py (Lines 313-318) raises "'effective_bits' and cost_model: latency are mutually exclusive." for the same input. A recipe author gets no signal that the value was ignored.
The same applies to cost.active_moe_expert_ratio: LatencyCostModel.supported_cost_keys rejects it, but to_mtq_constraints never forwards it, so the rejection never fires.
Use model_fields_set to detect explicit values and fail at the recipe boundary.
🐛 Proposed fix
`@model_validator`(mode="after")
def _validate_latency_cost_model(self):
if self.cost_model == "latency":
+ if "effective_bits" in self.model_fields_set:
+ raise ValueError(
+ "'effective_bits' and cost_model: latency are mutually exclusive. Provide "
+ "the budget via latency.relative_to_min instead."
+ )
+ if self.cost is not None and self.cost.active_moe_expert_ratio is not None:
+ raise ValueError(
+ "cost.active_moe_expert_ratio is not valid with cost_model: latency."
+ )
if self.latency is None:
raise ValueError("cost_model: latency requires a 'latency' budget block.")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| @model_validator(mode="after") | |
| def _validate_latency_cost_model(self): | |
| if self.cost_model == "latency": | |
| if self.latency is None: | |
| raise ValueError("cost_model: latency requires a 'latency' budget block.") | |
| missing = [ | |
| key | |
| for key in ("lut_path", "deployment_profile", "m") | |
| if self.cost is None or getattr(self.cost, key) is None | |
| ] | |
| if missing: | |
| raise ValueError( | |
| f"cost_model: latency requires cost.{{{', '.join(missing)}}} to be set." | |
| ) | |
| elif self.latency is not None: | |
| raise ValueError( | |
| f"The 'latency' budget block is only valid with cost_model: latency, " | |
| f"got cost_model={self.cost_model!r}." | |
| ) | |
| return self | |
| def to_mtq_constraints(self) -> dict: | |
| """Serialize to the dict passed to ``mtq.auto_quantize(constraints=...)``. | |
| The 'latency' cost model drops ``effective_bits`` (mutually exclusive) and | |
| emits the ``latency`` budget plus the latency ``cost`` keys; other cost | |
| models keep the historical effective-bits dict. | |
| """ | |
| if self.cost_model == "latency": | |
| assert self.latency is not None # guaranteed by _validate_latency_cost_model | |
| cost = { | |
| key: getattr(self.cost, key) | |
| for key in ("lut_path", "deployment_profile", "m") | |
| if self.cost is not None and getattr(self.cost, key) is not None | |
| } | |
| return { | |
| "cost_model": "latency", | |
| "latency": {"relative_to_min": self.latency.relative_to_min}, | |
| "cost": cost, | |
| } | |
| out: dict = {"effective_bits": self.effective_bits, "cost_model": self.cost_model} | |
| if self.cost is not None: | |
| cost = self.cost.model_dump(exclude_none=True) | |
| if cost: | |
| out["cost"] = cost | |
| return out | |
| `@model_validator`(mode="after") | |
| def _validate_latency_cost_model(self): | |
| if self.cost_model == "latency": | |
| if "effective_bits" in self.model_fields_set: | |
| raise ValueError( | |
| "'effective_bits' and cost_model: latency are mutually exclusive. Provide " | |
| "the budget via latency.relative_to_min instead." | |
| ) | |
| if self.cost is not None and self.cost.active_moe_expert_ratio is not None: | |
| raise ValueError( | |
| "cost.active_moe_expert_ratio is not valid with cost_model: latency." | |
| ) | |
| if self.latency is None: | |
| raise ValueError("cost_model: latency requires a 'latency' budget block.") | |
| missing = [ | |
| key | |
| for key in ("lut_path", "deployment_profile", "m") | |
| if self.cost is None or getattr(self.cost, key) is None | |
| ] | |
| if missing: | |
| raise ValueError( | |
| f"cost_model: latency requires cost.{{{', '.join(missing)}}} to be set." | |
| ) | |
| elif self.latency is not None: | |
| raise ValueError( | |
| f"The 'latency' budget block is only valid with cost_model: latency, " | |
| f"got cost_model={self.cost_model!r}." | |
| ) | |
| return self | |
| def to_mtq_constraints(self) -> dict: | |
| """Serialize to the dict passed to ``mtq.auto_quantize(constraints=...)``. | |
| The 'latency' cost model drops ``effective_bits`` (mutually exclusive) and | |
| emits the ``latency`` budget plus the latency ``cost`` keys; other cost | |
| models keep the historical effective-bits dict. | |
| """ | |
| if self.cost_model == "latency": | |
| assert self.latency is not None # guaranteed by _validate_latency_cost_model | |
| cost = { | |
| key: getattr(self.cost, key) | |
| for key in ("lut_path", "deployment_profile", "m") | |
| if self.cost is not None and getattr(self.cost, key) is not None | |
| } | |
| return { | |
| "cost_model": "latency", | |
| "latency": {"relative_to_min": self.latency.relative_to_min}, | |
| "cost": cost, | |
| } | |
| out: dict = {"effective_bits": self.effective_bits, "cost_model": self.cost_model} | |
| if self.cost is not None: | |
| cost = self.cost.model_dump(exclude_none=True) | |
| if cost: | |
| out["cost"] = cost | |
| return out |
🤖 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 `@modelopt/recipe/config.py` around lines 240 - 285, Update
_validate_latency_cost_model to reject explicitly provided effective_bits and
unsupported latency cost keys using model_fields_set, including
cost.active_moe_expert_ratio, instead of silently dropping them in
to_mtq_constraints. Preserve the existing required latency cost validation and
ensure invalid fields fail at recipe validation before serialization.
Comment on lines
+52
to
+61
| __all__ = [ | ||
| "FixedKernelPolicy", | ||
| "KernelSelector", | ||
| "LatencyCoverageError", | ||
| "LatencyLUT", | ||
| "LatencyRow", | ||
| "canonicalize_benchmark_csv", | ||
| "load_fixed_kernel_policy", | ||
| "write_canonical_csv", | ||
| ] |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Add the symbols that consumers already import to __all__.
tests/unit/torch/quantization/test_auto_quantize_latency.py imports RECIPE_FP8, RECIPE_NONE, RECIPE_W4A16_NVFP4, SCHEMA_VERSION, normalize_layer_indices, and ALL_COLUMNS from this module. None of them appear in __all__. parse_benchmark_csv is also public by name and absent.
Declare the full public surface, or rename the internal-only names with a leading underscore.
As per coding guidelines: "Define each module's public API with __all__ = [...]."
♻️ Proposed change
__all__ = [
+ "ALL_COLUMNS",
"FixedKernelPolicy",
"KernelSelector",
"LatencyCoverageError",
"LatencyLUT",
"LatencyRow",
+ "RECIPE_FP8",
+ "RECIPE_NONE",
+ "RECIPE_NVFP4",
+ "RECIPE_W4A16_NVFP4",
+ "SCHEMA_VERSION",
"canonicalize_benchmark_csv",
"load_fixed_kernel_policy",
+ "normalize_layer_indices",
+ "parse_benchmark_csv",
"write_canonical_csv",
]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| __all__ = [ | |
| "FixedKernelPolicy", | |
| "KernelSelector", | |
| "LatencyCoverageError", | |
| "LatencyLUT", | |
| "LatencyRow", | |
| "canonicalize_benchmark_csv", | |
| "load_fixed_kernel_policy", | |
| "write_canonical_csv", | |
| ] | |
| __all__ = [ | |
| "ALL_COLUMNS", | |
| "FixedKernelPolicy", | |
| "KernelSelector", | |
| "LatencyCoverageError", | |
| "LatencyLUT", | |
| "LatencyRow", | |
| "RECIPE_FP8", | |
| "RECIPE_NONE", | |
| "RECIPE_NVFP4", | |
| "RECIPE_W4A16_NVFP4", | |
| "SCHEMA_VERSION", | |
| "canonicalize_benchmark_csv", | |
| "load_fixed_kernel_policy", | |
| "normalize_layer_indices", | |
| "parse_benchmark_csv", | |
| "write_canonical_csv", | |
| ] |
🤖 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 `@modelopt/torch/quantization/_auto_quantize_latency.py` around lines 52 - 61,
Update the module’s __all__ declaration to include the consumer-imported public
symbols RECIPE_FP8, RECIPE_NONE, RECIPE_W4A16_NVFP4, SCHEMA_VERSION,
normalize_layer_indices, ALL_COLUMNS, and parse_benchmark_csv, preserving the
existing exports.
Source: Coding guidelines
Comment on lines
+236
to
+258
| def load_fixed_kernel_policy(source: str | Path | dict[str, Any]) -> FixedKernelPolicy: | ||
| """Load a fixed-kernel policy from a YAML/JSON file path or an in-memory dict.""" | ||
| if isinstance(source, dict): | ||
| data = source | ||
| else: | ||
| text = Path(source).read_text() | ||
| # yaml.safe_load parses JSON as well, so a single loader covers both. | ||
| data = yaml.safe_load(text) | ||
| if not isinstance(data, dict): | ||
| raise ValueError("Fixed-kernel policy must be a mapping.") | ||
|
|
||
| selectors: dict[str, dict[str, KernelSelector]] = {} | ||
| for op_kind, recipe_selectors in (data.get("selectors") or {}).items(): | ||
| if not isinstance(recipe_selectors, dict): | ||
| raise ValueError(f"selectors[{op_kind!r}] must be a mapping.") | ||
| selectors[op_kind] = { | ||
| recipe_id: _parse_selector(raw) for recipe_id, raw in recipe_selectors.items() | ||
| } | ||
| return FixedKernelPolicy( | ||
| kernel_policy_id=data.get("kernel_policy_id", ""), | ||
| mode=data.get("mode", SELECTION_POLICY_FIXED_KERNEL), | ||
| selectors=selectors, | ||
| ) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
A policy file with a typo or a missing selectors key produces an empty LUT with zero reported problems.
load_fixed_kernel_policy reads data.get("selectors") or {} and ignores unknown top-level keys. A file that spells the key selector, or omits it, yields a FixedKernelPolicy with no selectors. canonicalize_benchmark_csv then hits if not recipe_ids: continue at Line 452 for both op kinds and returns ([], []) — no rows and no coverage problems.
That contradicts the module contract: "missing coverage is always a fatal, aggregated error."
_parse_selector also indexes raw["backend"] directly, so a selector entry without backend raises a bare KeyError instead of a ValueError.
Validate the top-level keys, require a non-empty selectors mapping, and report a missing backend as a ValueError.
🐛 Proposed fix
def _parse_selector(raw: dict[str, Any]) -> KernelSelector:
known = {"backend", "kernel_source", "enable_w4a4_proxy", "proxy_reason"}
+ if not isinstance(raw, dict):
+ raise ValueError(f"Kernel selector must be a mapping, got {type(raw).__name__}.")
unknown = set(raw) - known
if unknown:
raise ValueError(f"Unknown kernel selector fields: {sorted(unknown)}.")
+ if "backend" not in raw:
+ raise ValueError("Kernel selector requires a 'backend' field.")
return KernelSelector( if not isinstance(data, dict):
raise ValueError("Fixed-kernel policy must be a mapping.")
+ unknown_keys = set(data) - {"kernel_policy_id", "mode", "selectors"}
+ if unknown_keys:
+ raise ValueError(f"Unknown fixed-kernel policy fields: {sorted(unknown_keys)}.")
+ if not data.get("selectors"):
+ raise ValueError("Fixed-kernel policy requires a non-empty 'selectors' mapping.")🤖 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 `@modelopt/torch/quantization/_auto_quantize_latency.py` around lines 236 -
258, Update load_fixed_kernel_policy to reject unknown top-level keys and
require selectors to be a non-empty mapping, raising ValueError for missing or
malformed coverage configuration instead of constructing an empty policy. Update
_parse_selector to validate that each selector entry contains backend and raise
ValueError when it is absent, while preserving existing valid-policy parsing.
Comment on lines
+607
to
+682
| for i, record in enumerate(reader, start=2): | ||
| if record.get("schema_version") != SCHEMA_VERSION: | ||
| problems.append( | ||
| f"Row {i}: schema_version={record.get('schema_version')!r}, " | ||
| f"expected {SCHEMA_VERSION!r}." | ||
| ) | ||
| continue | ||
| try: | ||
| latency_us = float(record["latency_us"]) | ||
| except (TypeError, ValueError): | ||
| problems.append(f"Row {i}: non-numeric latency_us={record.get('latency_us')!r}.") | ||
| continue | ||
| if not (latency_us > 0.0 and latency_us < float("inf")): | ||
| problems.append( | ||
| f"Row {i}: latency_us must be finite and positive, got {latency_us}." | ||
| ) | ||
| continue | ||
| try: | ||
| source_patterns = json.loads(record["source_module_patterns"]) | ||
| except (TypeError, ValueError, json.JSONDecodeError): | ||
| problems.append( | ||
| f"Row {i}: source_module_patterns is not valid JSON: " | ||
| f"{record.get('source_module_patterns')!r}." | ||
| ) | ||
| continue | ||
| if not isinstance(source_patterns, list) or not all( | ||
| isinstance(p, str) for p in source_patterns | ||
| ): | ||
| problems.append(f"Row {i}: source_module_patterns must be a JSON list of strings.") | ||
| continue | ||
|
|
||
| cost_is_proxy = _parse_bool_cell(record.get("cost_is_proxy", "")) | ||
| proxy_reason = record.get("proxy_reason") or None | ||
| measured_runtime_format = record.get("measured_runtime_format") or None | ||
| if cost_is_proxy and not (proxy_reason and measured_runtime_format): | ||
| problems.append( | ||
| f"Row {i}: cost_is_proxy=True requires both measured_runtime_format " | ||
| "and proxy_reason." | ||
| ) | ||
| continue | ||
|
|
||
| rows.append( | ||
| LatencyRow( | ||
| schema_version=SCHEMA_VERSION, | ||
| deployment_profile=record["deployment_profile"], | ||
| group_pattern=record["group_pattern"], | ||
| source_module_patterns=source_patterns, | ||
| recipe_id=record["recipe_id"], | ||
| runtime_format=record["runtime_format"], | ||
| m=int(record["m"]), | ||
| latency_us=latency_us, | ||
| backend=record["backend"], | ||
| with_quant=_parse_bool_cell(record["with_quant"]), | ||
| op_kind=record["op_kind"], | ||
| timing_scope=record["timing_scope"], | ||
| selection_policy=record["selection_policy"], | ||
| kernel_policy_id=record["kernel_policy_id"], | ||
| tp=int(record["tp"]), | ||
| ep=int(record["ep"]), | ||
| hardware=record["hardware"], | ||
| n=_parse_int_cell(record.get("n", "")), | ||
| k=_parse_int_cell(record.get("k", "")), | ||
| h=_parse_int_cell(record.get("h", "")), | ||
| f=_parse_int_cell(record.get("f", "")), | ||
| local_experts=_parse_int_cell(record.get("local_experts", "")), | ||
| top_k=_parse_int_cell(record.get("top_k", "")), | ||
| benchmark_provenance=record.get("benchmark_provenance") or None, | ||
| measured_runtime_format=measured_runtime_format, | ||
| cost_is_proxy=cost_is_proxy, | ||
| proxy_reason=proxy_reason, | ||
| ) | ||
| ) | ||
|
|
||
| if problems: | ||
| raise LatencyCoverageError(problems) | ||
| return cls(rows, digest) |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Integer fields bypass the aggregated coverage error.
from_csv collects latency_us, source_module_patterns, schema_version, and proxy-provenance failures into problems and raises one LatencyCoverageError. The integer conversions at Lines 656, 664, and 665 are not guarded. A row with m=abc, an empty tp, or a truncated record raises a bare ValueError or TypeError from int(...), so the caller sees an unaggregated error with no row number.
Convert the integers inside the same validated block.
🐛 Proposed fix
+ try:
+ m_value = int(record["m"])
+ tp_value = int(record["tp"])
+ ep_value = int(record["ep"])
+ except (KeyError, TypeError, ValueError):
+ problems.append(
+ f"Row {i}: m, tp, and ep must be integers; got "
+ f"m={record.get('m')!r} tp={record.get('tp')!r} ep={record.get('ep')!r}."
+ )
+ continue
+
cost_is_proxy = _parse_bool_cell(record.get("cost_is_proxy", ""))- m=int(record["m"]),
+ m=m_value,- tp=int(record["tp"]),
- ep=int(record["ep"]),
+ tp=tp_value,
+ ep=ep_value,📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| for i, record in enumerate(reader, start=2): | |
| if record.get("schema_version") != SCHEMA_VERSION: | |
| problems.append( | |
| f"Row {i}: schema_version={record.get('schema_version')!r}, " | |
| f"expected {SCHEMA_VERSION!r}." | |
| ) | |
| continue | |
| try: | |
| latency_us = float(record["latency_us"]) | |
| except (TypeError, ValueError): | |
| problems.append(f"Row {i}: non-numeric latency_us={record.get('latency_us')!r}.") | |
| continue | |
| if not (latency_us > 0.0 and latency_us < float("inf")): | |
| problems.append( | |
| f"Row {i}: latency_us must be finite and positive, got {latency_us}." | |
| ) | |
| continue | |
| try: | |
| source_patterns = json.loads(record["source_module_patterns"]) | |
| except (TypeError, ValueError, json.JSONDecodeError): | |
| problems.append( | |
| f"Row {i}: source_module_patterns is not valid JSON: " | |
| f"{record.get('source_module_patterns')!r}." | |
| ) | |
| continue | |
| if not isinstance(source_patterns, list) or not all( | |
| isinstance(p, str) for p in source_patterns | |
| ): | |
| problems.append(f"Row {i}: source_module_patterns must be a JSON list of strings.") | |
| continue | |
| cost_is_proxy = _parse_bool_cell(record.get("cost_is_proxy", "")) | |
| proxy_reason = record.get("proxy_reason") or None | |
| measured_runtime_format = record.get("measured_runtime_format") or None | |
| if cost_is_proxy and not (proxy_reason and measured_runtime_format): | |
| problems.append( | |
| f"Row {i}: cost_is_proxy=True requires both measured_runtime_format " | |
| "and proxy_reason." | |
| ) | |
| continue | |
| rows.append( | |
| LatencyRow( | |
| schema_version=SCHEMA_VERSION, | |
| deployment_profile=record["deployment_profile"], | |
| group_pattern=record["group_pattern"], | |
| source_module_patterns=source_patterns, | |
| recipe_id=record["recipe_id"], | |
| runtime_format=record["runtime_format"], | |
| m=int(record["m"]), | |
| latency_us=latency_us, | |
| backend=record["backend"], | |
| with_quant=_parse_bool_cell(record["with_quant"]), | |
| op_kind=record["op_kind"], | |
| timing_scope=record["timing_scope"], | |
| selection_policy=record["selection_policy"], | |
| kernel_policy_id=record["kernel_policy_id"], | |
| tp=int(record["tp"]), | |
| ep=int(record["ep"]), | |
| hardware=record["hardware"], | |
| n=_parse_int_cell(record.get("n", "")), | |
| k=_parse_int_cell(record.get("k", "")), | |
| h=_parse_int_cell(record.get("h", "")), | |
| f=_parse_int_cell(record.get("f", "")), | |
| local_experts=_parse_int_cell(record.get("local_experts", "")), | |
| top_k=_parse_int_cell(record.get("top_k", "")), | |
| benchmark_provenance=record.get("benchmark_provenance") or None, | |
| measured_runtime_format=measured_runtime_format, | |
| cost_is_proxy=cost_is_proxy, | |
| proxy_reason=proxy_reason, | |
| ) | |
| ) | |
| if problems: | |
| raise LatencyCoverageError(problems) | |
| return cls(rows, digest) | |
| for i, record in enumerate(reader, start=2): | |
| if record.get("schema_version") != SCHEMA_VERSION: | |
| problems.append( | |
| f"Row {i}: schema_version={record.get('schema_version')!r}, " | |
| f"expected {SCHEMA_VERSION!r}." | |
| ) | |
| continue | |
| try: | |
| latency_us = float(record["latency_us"]) | |
| except (TypeError, ValueError): | |
| problems.append(f"Row {i}: non-numeric latency_us={record.get('latency_us')!r}.") | |
| continue | |
| if not (latency_us > 0.0 and latency_us < float("inf")): | |
| problems.append( | |
| f"Row {i}: latency_us must be finite and positive, got {latency_us}." | |
| ) | |
| continue | |
| try: | |
| source_patterns = json.loads(record["source_module_patterns"]) | |
| except (TypeError, ValueError, json.JSONDecodeError): | |
| problems.append( | |
| f"Row {i}: source_module_patterns is not valid JSON: " | |
| f"{record.get('source_module_patterns')!r}." | |
| ) | |
| continue | |
| if not isinstance(source_patterns, list) or not all( | |
| isinstance(p, str) for p in source_patterns | |
| ): | |
| problems.append(f"Row {i}: source_module_patterns must be a JSON list of strings.") | |
| continue | |
| try: | |
| m_value = int(record["m"]) | |
| tp_value = int(record["tp"]) | |
| ep_value = int(record["ep"]) | |
| except (KeyError, TypeError, ValueError): | |
| problems.append( | |
| f"Row {i}: m, tp, and ep must be integers; got " | |
| f"m={record.get('m')!r} tp={record.get('tp')!r} ep={record.get('ep')!r}." | |
| ) | |
| continue | |
| cost_is_proxy = _parse_bool_cell(record.get("cost_is_proxy", "")) | |
| proxy_reason = record.get("proxy_reason") or None | |
| measured_runtime_format = record.get("measured_runtime_format") or None | |
| if cost_is_proxy and not (proxy_reason and measured_runtime_format): | |
| problems.append( | |
| f"Row {i}: cost_is_proxy=True requires both measured_runtime_format " | |
| "and proxy_reason." | |
| ) | |
| continue | |
| rows.append( | |
| LatencyRow( | |
| schema_version=SCHEMA_VERSION, | |
| deployment_profile=record["deployment_profile"], | |
| group_pattern=record["group_pattern"], | |
| source_module_patterns=source_patterns, | |
| recipe_id=record["recipe_id"], | |
| runtime_format=record["runtime_format"], | |
| m=m_value, | |
| latency_us=latency_us, | |
| backend=record["backend"], | |
| with_quant=_parse_bool_cell(record["with_quant"]), | |
| op_kind=record["op_kind"], | |
| timing_scope=record["timing_scope"], | |
| selection_policy=record["selection_policy"], | |
| kernel_policy_id=record["kernel_policy_id"], | |
| tp=tp_value, | |
| ep=ep_value, | |
| hardware=record["hardware"], | |
| n=_parse_int_cell(record.get("n", "")), | |
| k=_parse_int_cell(record.get("k", "")), | |
| h=_parse_int_cell(record.get("h", "")), | |
| f=_parse_int_cell(record.get("f", "")), | |
| local_experts=_parse_int_cell(record.get("local_experts", "")), | |
| top_k=_parse_int_cell(record.get("top_k", "")), | |
| benchmark_provenance=record.get("benchmark_provenance") or None, | |
| measured_runtime_format=measured_runtime_format, | |
| cost_is_proxy=cost_is_proxy, | |
| proxy_reason=proxy_reason, | |
| ) | |
| ) | |
| if problems: | |
| raise LatencyCoverageError(problems) | |
| return cls(rows, digest) |
🤖 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 `@modelopt/torch/quantization/_auto_quantize_latency.py` around lines 607 -
682, Update LatencyRow construction in from_csv to validate every required
integer conversion for m, tp, and ep within the existing per-row validation
flow, catching invalid or missing values and appending row-numbered messages to
problems before continuing. Ensure malformed integer fields never escape as bare
ValueError or TypeError, while preserving successful row parsing and aggregated
LatencyCoverageError behavior.
| self.active_moe_expert_ratio = self.config["active_moe_expert_ratio"] | ||
| self.disabled_layers = self.config["disabled_layers"] | ||
| self.cost_denominator = getattr(self, "cost_denominator", None) | ||
| self._load_latency_cost_model() |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Record the latency LUT identity in the checkpoint and reject a mismatch.
candidate_stats["costs"] now hold measured latencies that are valid only for one (lut_path digest, deployment_profile, m) triple. before_search compares a restored checkpoint against cost_model and active_moe_expert_ratio only (Lines 1136-1147). It does not compare the latency identity.
A user who reruns with the same checkpoint path but a different LUT file, a different deployment_profile, or a different m gets stale costs restored at Line 1331, and the search returns silently. _latency_resolution then reports lut_digest, deployment_profile, and m from the NEW configuration while the costs came from the OLD one. The recorded provenance contradicts the actual result.
Persist the triple in the state dict and fail on a mismatch, matching the existing cost-model guard.
Also note that _latency_profile, _latency_m, _latency_relative_to_min, and _lut_digest receive no class-level default, unlike _latency_lut at Line 650. Every current read sits behind an is_latency branch, so add the defaults to keep that invariant explicit.
🐛 Proposed fix
method_name: str | None = None
# Loaded only for cost_model: latency (see _load_latency_cost_model).
_latency_lut: LatencyLUT | None = None
+ _latency_profile: str | None = None
+ _latency_m: int | None = None
+ _latency_relative_to_min: float | None = None
+ _lut_digest: str | None = None def _load_latency_cost_model(self) -> None:
"""Load and validate the latency LUT when the latency cost model is active."""
if self.config["cost_model"] != COST_MODEL_LATENCY:
self._latency_lut = None
return self._latency_relative_to_min = float(latency_budget[LATENCY_RELATIVE_TO_MIN_KEY])
self._lut_digest = lut.digest
+ latency_identity = (lut.digest, self._latency_profile, self._latency_m)
+ restored_identity = getattr(self, "latency_identity", None)
+ if self.candidate_stats and restored_identity is not None:
+ if tuple(restored_identity) != latency_identity:
+ raise ValueError(
+ "Checkpoint latency LUT identity does not match the current search config: "
+ f"checkpoint={tuple(restored_identity)}, current={latency_identity}. "
+ "Use a different checkpoint path."
+ )
+ self.latency_identity = latency_identityAdd the key to default_state_dict so it is serialized:
"cost_model": "weight",
"cost": {},
+ "latency_identity": None,
"active_moe_expert_ratio": None,Also applies to: 1393-1411
🤖 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 `@modelopt/torch/quantization/algorithms.py` at line 1154, Persist the latency
identity triple in default_state_dict and validate it in before_search alongside
cost_model, rejecting checkpoints whose LUT digest, deployment profile, or m
differs from the active configuration. Ensure restored latency costs cannot be
used after an identity mismatch. Add class-level defaults for _latency_profile,
_latency_m, _latency_relative_to_min, and _lut_digest, matching the existing
_latency_lut default.
Comment on lines
+1552
to
+1565
| if is_latency: | ||
| self.best["constraints"] = self._latency_resolution( | ||
| best_constraints, minimum_latency, budget | ||
| ) | ||
| if verbose: | ||
| self._print_latency_summary(best_recipe, best_constraints, minimum_latency, budget) | ||
| else: | ||
| effective_bits_from_search = (best_constraints / total_weight_size) * 16 | ||
| if verbose: | ||
| effective_bits_from_search = self._print_recipe_summary( | ||
| best_recipe, best_constraints, total_weight_size | ||
| ) | ||
| else: | ||
| effective_bits_from_search = (best_constraints / total_weight_size) * 16 | ||
| self.best["constraints"] = {"effective_bits": effective_bits_from_search} |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Find consumers of best["constraints"], effective_bits reporting, and cost_denominator.
rg -nP -C4 --type=py '\bbest\b\s*\[\s*["'\'']constraints["'\'']\s*\]'
rg -nP -C4 --type=py 'search_history\s*\[\s*["'\'']best["'\'']\s*\]'
rg -nP -C4 --type=py '\bcost_denominator\b'Repository: NVIDIA/Model-Optimizer
Length of output: 160
🏁 Script executed:
#!/bin/bash
set -u
printf '%s\n' '--- candidate files ---'
git ls-files | rg '(^|/)(algorithms|.*quant.*|.*search.*|.*export.*|.*report.*)' | head -200
printf '%s\n' '--- all exact identifier references ---'
rg -n -F -C3 '["constraints"]' .
rg -n -F -C3 "['constraints']" .
rg -n -F -C3 'effective_bits' .
rg -n -F -C3 'cost_denominator' .
rg -n -F -C3 'search_history' .
rg -n -F -C3 '_latency_resolution' modelopt/torch/quantization/algorithms.pyRepository: NVIDIA/Model-Optimizer
Length of output: 50379
🏁 Script executed:
#!/bin/bash
set -u
printf '%s\n' '--- algorithms structure and relevant implementation ---'
ast-grep outline modelopt/torch/quantization/algorithms.py
sed -n '1360,1615p' modelopt/torch/quantization/algorithms.py
printf '%s\n' '--- auto-quantize public and restore paths ---'
ast-grep outline modelopt/torch/quantization/model_quant.py
sed -n '560,760p' modelopt/torch/quantization/model_quant.py
rg -n -C5 'get_auto_quantize_config|search_state|state_dict|load_state|search_history|best.*constraints|constraints.*best' modelopt/torch/quantization/model_quant.py modelopt/torch/quantization/algorithms.py
printf '%s\n' '--- latency-specific tests and result consumers ---'
sed -n '1650,1740p' tests/unit/torch/quantization/test_autoquant.py
rg -n -C4 'cost_model.*latency|minimum_latency_us|deployment_profile|search_history.*best|best.*constraints' tests modelopt examples docs --glob '*.py' --glob '*.rst' --glob '*.md' | head -300Repository: NVIDIA/Model-Optimizer
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -u
printf '%s\n' '--- exact cost_denominator references ---'
rg -n -C6 '\bcost_denominator\b' modelopt tests examples docs --glob '*.py' --glob '*.rst' --glob '*.md'
printf '%s\n' '--- exact best-constraints references ---'
rg -n -C4 'search_state\["best"\]\["constraints"\]|search_history\["best"\]\["constraints"\]|self\.best\["constraints"\]|best\["constraints"\]' modelopt tests examples docs --glob '*.py' --glob '*.rst' --glob '*.md'
printf '%s\n' '--- resolver and config generation ---'
sed -n '2150,2315p' modelopt/torch/quantization/algorithms.py
printf '%s\n' '--- searcher checkpoint/state implementation ---'
rg -n -C8 'def (state_dict|load_search_checkpoint|search|load_state|default_state_dict)|checkpoint|cost_denominator' modelopt/torch/opt/searcher.py modelopt/torch/quantization/algorithms.py | head -350
printf '%s\n' '--- complete latency test region ---'
sed -n '1680,1740p' tests/unit/torch/quantization/test_autoquant.pyRepository: NVIDIA/Model-Optimizer
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -u
printf '%s\n' '--- solver path used by _resolve_best_recipe ---'
sed -n '880,1085p' modelopt/torch/quantization/algorithms.py
sed -n '1600,1745p' modelopt/torch/quantization/algorithms.py
sed -n '1935,2025p' modelopt/torch/quantization/algorithms.py
printf '%s\n' '--- latency config normalization and constants ---'
rg -n -C8 'COST_MODEL_LATENCY|LATENCY_|normalize_auto_quantize_constraints|constraints.*latency' modelopt/torch/quantization --glob '*.py'
sed -n '1035,1175p' modelopt/torch/quantization/algorithms.py
printf '%s\n' '--- export/reporting consumers outside the searcher ---'
rg -n -C5 'best.*constraints|constraints.*best|effective_bits.*search|cost_denominator|search_history' modelopt/torch/export modelopt/torch/quantization examples docs tests --glob '*.py' --glob '*.rst' --glob '*.md' | rg -v 'algorithms.py|test_autoquant.py' | head -300
printf '%s\n' '--- static check of resolver assumptions ---'
python3 - <<'PY'
from pathlib import Path
p = Path("modelopt/torch/quantization/algorithms.py")
text = p.read_text()
start = text.index("def _resolve_best_recipe")
end = text.index("\ndef _match_quantizer_cfg", start)
chunk = text[start:end]
for i, line in enumerate(chunk.splitlines(), 1):
if any(token in line for token in ("cost_model", "cost_denominator", "effective_bits", "max_weight_size", "run_search_with_stats")):
print(f"{start + text[:start].count(chr(10)) + i}: {line}")
PYRepository: NVIDIA/Model-Optimizer
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -u
printf '%s\n' '--- exact run_search_with_stats implementations ---'
rg -n -A75 '^ def run_search_with_stats' modelopt/torch/quantization/algorithms.py
printf '%s\n' '--- concise resolver assumption extract ---'
python3 - <<'PY'
from pathlib import Path
text = Path("modelopt/torch/quantization/algorithms.py").read_text()
start = text.index("def _resolve_best_recipe")
end = text.index("\ndef _match_quantizer_cfg", start)
for n, line in enumerate(text[start:end].splitlines(), text[:start].count("\n") + 1):
if any(x in line for x in (
'effective_bits =', 'cost_denominator', 'max_weight_size',
'searcher.cost_model', 'searcher.cost =', 'searcher.config =',
'run_search_with_stats',
)):
print(f"{n}: {line}")
print("--- public latency/re-solve coverage ---")
for path in Path("tests").rglob("*.py"):
body = path.read_text(errors="ignore")
if "get_auto_quantize_config" in body and ("latency" in body or "cost_model" in body):
print(path)
for n, line in enumerate(body.splitlines(), 1):
if "get_auto_quantize_config" in line or "cost_model" in line and "latency" in line:
print(f" {n}: {line.strip()}")
PYRepository: NVIDIA/Model-Optimizer
Length of output: 13300
Reject or implement latency-aware re-solving
_resolve_best_recipe() always treats constraints["effective_bits"] and search_state["cost_denominator"] as weight-budget values. For latency searches, these values represent latency, so re-solving silently uses an invalid budget and can select the wrong recipe. Reject re-solving latency states or apply the latency budget and LUT metadata.
🤖 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 `@modelopt/torch/quantization/algorithms.py` around lines 1552 - 1565, Update
_resolve_best_recipe to detect latency-search state before treating
constraints["effective_bits"] or search_state["cost_denominator"] as weight
budgets. Either reject re-solving for latency states with a clear error, or
route it through the existing latency budget and LUT metadata so recipe
selection uses latency values correctly; preserve the current weight-budget
behavior for non-latency searches.
Comment on lines
+1891
to
+1955
| def test_autoquantize_constraints_latency_to_mtq(): | ||
| from modelopt.recipe.config import AutoQuantizeConstraints | ||
|
|
||
| c = AutoQuantizeConstraints( | ||
| cost_model="latency", | ||
| latency={"relative_to_min": 1.2}, | ||
| cost={"lut_path": "/tmp/lut.csv", "deployment_profile": "b100_tp1_ep1_decode", "m": 1}, | ||
| ) | ||
| out = c.to_mtq_constraints() | ||
| assert out["cost_model"] == "latency" | ||
| assert "effective_bits" not in out | ||
| assert out["latency"] == {"relative_to_min": 1.2} | ||
| assert out["cost"] == { | ||
| "lut_path": "/tmp/lut.csv", | ||
| "deployment_profile": "b100_tp1_ep1_decode", | ||
| "m": 1, | ||
| } | ||
|
|
||
|
|
||
| def test_autoquantize_constraints_effective_bits_to_mtq_unchanged(): | ||
| from modelopt.recipe.config import AutoQuantizeConstraints | ||
|
|
||
| c = AutoQuantizeConstraints( | ||
| effective_bits=6.0, cost_model="active_moe", cost={"active_moe_expert_ratio": 0.03125} | ||
| ) | ||
| out = c.to_mtq_constraints() | ||
| assert out == { | ||
| "effective_bits": 6.0, | ||
| "cost_model": "active_moe", | ||
| "cost": {"active_moe_expert_ratio": 0.03125}, | ||
| } | ||
| assert "latency" not in out | ||
|
|
||
|
|
||
| def test_autoquantize_constraints_latency_requires_fields(): | ||
| import pytest | ||
|
|
||
| from modelopt.recipe.config import AutoQuantizeConstraints | ||
|
|
||
| with pytest.raises(ValueError, match="latency"): | ||
| AutoQuantizeConstraints(cost_model="latency") # no latency block, no cost | ||
| with pytest.raises(ValueError, match=r"lut_path|deployment_profile|m"): | ||
| AutoQuantizeConstraints(cost_model="latency", latency={"relative_to_min": 1.2}) | ||
|
|
||
|
|
||
| def test_autoquantize_constraints_latency_block_rejected_for_other_cost_models(): | ||
| import pytest | ||
|
|
||
| from modelopt.recipe.config import AutoQuantizeConstraints | ||
|
|
||
| with pytest.raises(ValueError, match="only valid with cost_model: latency"): | ||
| AutoQuantizeConstraints(cost_model="weight", latency={"relative_to_min": 1.2}) | ||
|
|
||
|
|
||
| def test_autoquantize_constraints_relative_to_min_lower_bound(): | ||
| import pytest | ||
|
|
||
| from modelopt.recipe.config import AutoQuantizeConstraints | ||
|
|
||
| with pytest.raises(ValueError, match="relative_to_min"): | ||
| AutoQuantizeConstraints( | ||
| cost_model="latency", | ||
| latency={"relative_to_min": 0.9}, | ||
| cost={"lut_path": "/tmp/lut.csv", "deployment_profile": "p", "m": 1}, | ||
| ) |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Move the imports to module scope.
All five new tests import AutoQuantizeConstraints inside the function body, and four also import pytest locally. No circular import or optional dependency applies here. pytest is already used across this module.
Import errors must surface at collection time, not mid-test.
Based on path instructions: "Imports inside functions or test methods without explicit justification. Imports belong at the top of the file so import errors surface at collection time, not mid-test."
♻️ Proposed change
Add to the module-level imports:
from modelopt.recipe.config import AutoQuantizeConstraintsThen remove the local imports:
def test_autoquantize_constraints_latency_to_mtq():
- from modelopt.recipe.config import AutoQuantizeConstraints
-
c = AutoQuantizeConstraints( def test_autoquantize_constraints_latency_requires_fields():
- import pytest
-
- from modelopt.recipe.config import AutoQuantizeConstraints
-
with pytest.raises(ValueError, match="latency"):🧰 Tools
🪛 ast-grep (0.45.0)
[info] 1896-1896: Do not hardcode temporary file or directory names
Context: "/tmp/lut.csv"
Note: [CWE-377] Insecure Temporary File.
(hardcoded-tmp-file)
[info] 1903-1903: Do not hardcode temporary file or directory names
Context: "/tmp/lut.csv"
Note: [CWE-377] Insecure Temporary File.
(hardcoded-tmp-file)
[info] 1953-1953: Do not hardcode temporary file or directory names
Context: "/tmp/lut.csv"
Note: [CWE-377] Insecure Temporary File.
(hardcoded-tmp-file)
🤖 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 `@tests/unit/recipe/test_loader.py` around lines 1891 - 1955, Move the
AutoQuantizeConstraints and pytest imports used by the five new tests to module
scope alongside the existing imports in the test module. Remove the
corresponding imports from test_autoquantize_constraints_latency_to_mtq,
test_autoquantize_constraints_effective_bits_to_mtq_unchanged,
test_autoquantize_constraints_latency_requires_fields,
test_autoquantize_constraints_latency_block_rejected_for_other_cost_models, and
test_autoquantize_constraints_relative_to_min_lower_bound.
Source: Path instructions
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2138 +/- ##
==========================================
- Coverage 78.73% 77.59% -1.14%
==========================================
Files 522 523 +1
Lines 60342 61401 +1059
==========================================
+ Hits 47508 47647 +139
- Misses 12834 13754 +920
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:
|
| latency: | ||
| relative_to_min: 1.2 | ||
| cost: | ||
| lut_path: /home/scratch.juhim_coreai/code/dev_notes/Autoquant_research/outputs/qwen36_sm100_tp1_ep1_haq_latency_v1.csv |
Collaborator
There was a problem hiding this comment.
internal paths need to be fixed here
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: ?
Usage
# Add a code snippet demonstrating how to use thisTesting
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
New Features
Bug Fixes
Tests