Skip to content

[OMNIML-6370] [1/5] Isolate weight-rewriting AutoQuantize candidates from the source weights - #3154

Merged
meenchen merged 16 commits into
mainfrom
weimingc/autoquant-weight-isolation
Nov 6, 2026
Merged

meenchen merged 16 commits into
mainfrom
weimingc/autoquant-weight-isolation

Conversation

@meenchen

@meenchen meenchen commented Oct 30, 2026 •

Copy link
Copy Markdown
Contributor

What does this PR do?

Type of change: Bug fix / new feature

[1/5] of the sub-4-bit AutoQuantize series (merge in order).

Series (merge in order): #3154 → #3151 → #3152 → #3153 → #3155

GPTQ and imatrix overwrite the weights they quantize. auto_quantize rejected both, and accepting them as-is would let one candidate's calibration change the weights every later candidate — and the unquantized scoring reference — sees.

AutoQuantize now snapshots (to CPU) only the weights a GPTQ/imatrix candidate's calibration rewrites: the parameters its enabled weight quantizers cover, and for imatrix alone only those imatrix encodes (not Q8_0; a fused projection when any of its experts is). After such a candidate calibrates, its GGML quantizers replay the codes it pinned while the source weights are restored for the next candidate; the snapshot is freed before sensitivity scoring. After the solve, the selected candidates' weights are left as their codes decode, as a plain GPTQ run leaves them, and no candidate is left replaying. GGML fake quant now returns exactly the reconstruction export decodes, so scoring and the exported model see the same weights (BF16 included).

Refused before the model is converted: any candidate with an enabled layerwise export_dir, empty included, whether its options are a dict or a LayerwiseConfig, whose export would leave the model in export form before scoring (hf_ptq.py already refused it for AutoQuantize recipes); and, in a search with a weight-rewriting candidate, any candidate whose layerwise checkpoint_dir resume restores the layers' weights into new Parameters the source-weight restore cannot reach (calib_mutates_weights=False resumes are accepted). Refused before any calibration: a weight-rewriting candidate whose weight quantizer is not a single GGML quantizer or rotates the weight, or whose algorithm list runs anything besides GPTQ, imatrix and max (AWQ or SmoothQuant would transform the weight GPTQ encodes); and meta or DTensor weights it would rewrite, weights stored packed or scaled (a native FP8 checkpoint's), or ones whose storage another parameter shares (tied weights, or views of them). These checks need the converted quantizers to know which weights a candidate rewrites, so auto_quantize hands back the unconverted modules, and the configs that conversion changed (including nested sub-configs), when the search fails before calibrating any candidate, and the same model can be searched again. Only a search with a weight-rewriting candidate saves the modules for this, and it drops them once these checks pass, so calibration and scoring never hold the weights conversion replaces (HF Conv1D's transposed copies); other searches save nothing. A failure during calibration restores the source weights and drops the modes that calibration added. auto_quantize hands back the pre-quant-scale folding policy on every exit, including a failure during scoring. Weights a calibration leaves alone (imatrix on Q8_0) simply quantize the source weight. NVFP4/FP8-only searches take no snapshot and never slice fused-expert weights after the solve. auto_quantize now accepts gptq, imatrix and algorithm lists such as [gptq, imatrix].

Why snapshot and replay

Every candidate after a GPTQ/imatrix calibration, and the unquantized scoring reference, must see the source weights. The alternatives cost more:

  • Copy the model, or its state dict: this holds every weight, about 70 GB for Qwen3.6-35B-A3B in BF16. The snapshot holds only the weights a rewriting candidate quantizes, on the host, and is freed before scoring.
  • Make GPTQ/imatrix leave weights alone: this would change calibrators that mtq.quantize shares, and the chosen codes would still have to be kept somewhere.

The GGML quantizer already pins those codes for export, so replaying the pin makes the scored, selected and exported weights the same codes.

Reading order

  1. _AutoQuantizeBaseSearcher.before_search: snapshot (_snapshot_source_weights; its checks are _rewritten_weight and _check_restorable), then for each candidate calibrate, replay its pins and restore the source weights.
  2. _write_selected_codes_to_weights (after the solve), and GGMLFormat._pinned_rows / full_pin (the replay).
  3. auto_quantize: the window in which a failure hands back the unconverted modules, and the folding policy it hands back on every exit.

Usage

iq2_gptq = {
    "quant_cfg": [{"quantizer_name": "*", "enable": False},
                  {"quantizer_name": "*weight_quantizer",
                   "cfg": {"num_bits": "iq2_s", "block_sizes": {-1: 256}, "backend": "ggml", "effective_bits": 2.5625}}],
    "algorithm": [{"method": "gptq", "block_size": 256, "perc_damp": 0.3, "importance_weighted": True}, {"method": "imatrix"}],
}
model, state = mtq.auto_quantize(model, constraints={"effective_bits": 3.0},
                                 quantization_formats=[(iq2_gptq, "IQ2_S_GPTQ"), mtq.NVFP4_DEFAULT_CFG], ...)

Testing

  • New tests in tests/unit/torch/quantization/test_ggml_gptq.py, test_autoquant.py, plugins/test_fused_experts.py and tests/gpu/torch/quantization/test_ggml_gptq_cuda.py. I reverted each fix in place and checked that its test fails. They cover:
    • other candidates' scores are unaffected by GPTQ;
    • mixed selection: a GPTQ layer next to a BF16 layer, a GGML entry replacing an earlier INT4 entry, and imatrix on Q8_0;
    • the snapshot holds only rewritten weights (a fused down_proj that is disabled or Q8_0), takes a whole fused gate_up_proj with one Q8_0 expert, and a BF16 selection hands back the source weights;
    • refusals (non-GGML, tied, an aliased view, a rotated GPTQ or imatrix target, AWQ then GPTQ, a layerwise checkpoint_dir, native FP8 weights (GPU), and fused MoE whose root, nested and experts' configs conversion changes) hand back the unconverted model and its configs, which searches again;
    • calibration failures (GPTQ rewrite then failure, [gptq, layerwise imatrix] failing midway) restore the weights, the folding policy and the mode list;
    • the folding policy comes back after a scoring failure;
    • exact BF16 replay;
    • no fused-expert indexing after a search without pins;
    • an algorithm: [] candidate keeps the saved state restorable.
  • Layerwise directories (test_huggingface.py), on a tiny Llama with discoverable decoder layers: an FP8 candidate with an export_dir, with an empty export_dir, and with a checkpoint_dir beside a Q8_0 GPTQ candidate, each given as a dict and as a LayerwiseConfig, is refused before conversion, leaving the modules, the weights and the working directory untouched, and the same model searches again. Without the refusals, the layers are exported and scoring fails on the packed weights, or the max candidate's checkpoint is written.
  • Conv1D weight lifetime (test_huggingface.py): an INT8 search saves no module state, and a Q8_0 GPTQ search drops it before calibration, so the weights conversion replaces are dead in every calibration and scoring forward.
  • Regression check against the base commit (RTX 6000 Ada, BF16, FP8 + INT4 searches): a 24-layer GPT-2 (Conv1D), a 16-layer Llama and a fused-MoE Qwen3, with gradient scoring and with KL for Llama, give the same scores, selections and output logits, the same search time (within 4%) and the same memory. On GPT-2 the weights held after conversion are 0 MiB, against 576 MiB before 5acdcc4.
  • Unit suites tests/unit/torch/quantization + tests/unit/recipe on CPU (16 xdist workers): 2151 passed. The rest fail the same way on main at this branch's base on the test host:
    • diffusers lacks q_bmm_quantizer (4 failures);
    • test_diffusers_wan_conv3d fails to collect (16 errors);
    • one to four test_nemotron_3_5_super_vl_autoquant_recipe cases fail under xdist and pass serially.
  • pre-commit run mypy --all-files passes; ruff check and ruff format pass on the touched files.
  • Rollback cost, paid only by searches with a GPTQ/imatrix candidate: saving the module attributes takes 2 MiB and 0.16 s on Qwen3.6-35B-A3B (837 modules), and 69 MiB on a 24.6k-module non-fused MoE.
  • GPU (RTX 6000 Ada): tests/gpu/torch/quantization -k "auto_quant or autoquant or ggml or gptq or imatrix" 57 passed, including test_fsdp2_autoquant.py and the native FP8 refusal. The deepspeed and onnx packages on the host stop 3 files from collecting.
  • Stack restacked on this; the same unit suites at [5/5]: 2148 passed, with the same failures.
  • Qwen3.6-35B-A3B on GB300: searches over IQ2_{XXS,XS,S} (imatrix and local GPTQ), NVFP4, FP8, BF16; see PR 5 for results.

Before your PR is "Ready for review"

Additional Information

Supersedes #3082 … #3057 and #3113 … #3120 with a smaller design (~720 source lines across the series).

Synced with main by a merge (99f82c8), which keeps the commits the review threads cite. #3197 made the IQ formats re-encode a weight under its pinned scales and importance ahead of the pinned-rows lookup that honours PIN_OVERRIDES_WEIGHT, so during a search a GPTQ or imatrix candidate would have quantized the source weight under its scales instead of replaying its codes; b7be410 puts the override first there too (caught by test_a_replayed_pin_fake_quantizes_to_exactly_its_decoding_in_bf16).

🤖 Generated with Claude Code

Summary by CodeRabbit

  • New Features
    • AutoQuantize supports GPTQ and iMatrix candidates with GGML formats.
    • Selected candidates retain their calibrated quantized results in the returned model.
  • Bug Fixes
    • Candidate scores are no longer affected by weight changes made while evaluating other candidates.
    • Source weights and quantization settings are restored when candidate evaluation fails; unsupported weight-rewrite configurations are rejected.
    • Failed conversions before calibration begins restore the model’s prior state.
    • Fake quantization better preserves reconstructed values in narrow data types.

@copy-pr-bot

copy-pr-bot Bot commented Oct 30, 2026

Copy link
Copy Markdown

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.

@coderabbitai

coderabbitai Bot commented Oct 30, 2026 •

Copy link
Copy Markdown
Contributor

Explain Change →

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • Resume reviews
  • Trigger review

No actionable comments were generated in the recent review.

Walkthrough

AutoQuantize now accepts GPTQ and iMatrix candidates for GGML formats. It snapshots and restores source weights during candidate calibration, replays available pinned codes, and writes selected codes into model weights. Tests cover candidate selection, unsupported configurations, and calibration failures.

Changes: GGML AutoQuantize candidate flow

Layer / File(s) Summary
Candidate method support
modelopt/torch/quantization/model_quant.py, modelopt/torch/quantization/algorithms.py, modelopt/torch/quantization/model_calib.py, CHANGELOG.rst
AutoQuantize accepts GPTQ and iMatrix methods, validates candidate algorithms, and checks candidate weights before snapshotting them.
Source-weight and pinned-code handling
modelopt/torch/quantization/algorithms.py, modelopt/torch/quantization/ggml/common.py, modelopt/torch/quantization/model_quant.py
AutoQuantize restores source weights during candidate calibration and replays available pinned payloads. Calibration cleanup removes added ModelOpt modes and restores the prior folding policy. GGML support matches and decodes pinned payloads.
Selected-code writeback and validation
modelopt/torch/quantization/algorithms.py, tests/unit/torch/quantization/test_ggml_gptq.py, tests/unit/torch/quantization/plugins/test_fused_experts.py, tests/unit/torch/quantization/test_autoquant.py
AutoQuantize writes selected pinned codes into weights and clears replay overrides. Tests cover scoring, code selection, restoration, BF16 replay, fused-weight indexing, and an uncalibrated candidate.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant AutoQuantize
  participant ModelWeights
  participant GGMLFormat
  AutoQuantize->>ModelWeights: Snapshot source weights
  AutoQuantize->>GGMLFormat: Calibrate candidate and retain pinned payload
  AutoQuantize->>ModelWeights: Restore source weights between candidates
  AutoQuantize->>GGMLFormat: Decode selected pinned payload
  GGMLFormat-->>AutoQuantize: Return decoded weights
  AutoQuantize->>ModelWeights: Write selected codes
Loading

Possibly related issues

Suggested reviewers: kevalmorabia97, cjluo-nv

Merge Risk: Minimal · previous review 357fb

No merge-blocking risk is established from the available evidence.

Pre-merge checks: Passed 5 · Failed 1

Failed Failed checks: 1 warning

Check name Status Explanation Resolution
Docstring Coverage Warning Docstring coverage is 52.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta… Write docstrings for the functions missing them to satisfy the coverage threshold.
Passed Passed checks: 5 passed

Check name Status Explanation
Description Check Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check Passed The title clearly and concisely describes the main change: isolating AutoQuantize candidates that rewrite weights from the source weights.
Linked Issues check Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check Passed Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns Passed No listed security anti-pattern was introduced. The changed modelopt files add no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, eval/exec o…
Full details: Security Anti-Patterns

Explanation

No listed security anti-pattern was introduced. The changed modelopt files add no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, eval/exec on external input, or # nosec usage. The pull request does not change pyproject.toml or requirements files, and it adds no dependencies.

Full details: Docstring Coverage

Explanation

Docstring coverage is 52.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 25 functions across 7 files.

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

Finishing Touches

Generate docstrings

  • Commit to this branch
  • Create a new PR
Generate unit tests (beta)

  • Commit to this branch
  • Create a new PR
  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Recent review info

Run configuration

Commits

Reviewing files that changed from the base of the PR and between 848bd95 and 357fb38.

Files selected for processing: 4

  • modelopt/torch/quantization/algorithms.py
  • modelopt/torch/quantization/ggml/common.py
  • modelopt/torch/quantization/model_quant.py
  • tests/unit/torch/quantization/test_ggml_gptq.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

@github-actions

github-actions Bot commented Oct 30, 2026 •

Copy link
Copy Markdown
Contributor
PR Preview Action v1.8.2
Preview removed because the pull request was closed.
2026-11-06 04:30 UTC

@codecov

codecov Bot commented Oct 30, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 96.90722% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.84%. Comparing base (bc00ef3) to head (b7be410).

Files with missing lines Patch % Lines
modelopt/torch/quantization/algorithms.py 96.74% 4 Missing ⚠️
modelopt/torch/quantization/model_quant.py 96.15% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #3154      +/-   ##
==========================================
- Coverage   81.53%   80.84%   -0.70%     
==========================================
  Files         684      684              
  Lines       79485    79658     +173     
==========================================
- Hits        64811    64398     -413     
- Misses      14674    15260     +586     
Flag Coverage Δ
examples-diffusers 19.63% <16.49%> (-0.01%) ⬇️
examples-gpt-oss 13.39% <16.49%> (+<0.01%) ⬆️
examples-hf_ptq 23.91% <45.36%> (-0.82%) ⬇️
examples-llm_distill 13.46% <16.49%> (+<0.01%) ⬆️
examples-llm_eval 17.40% <16.49%> (-0.01%) ⬇️
examples-llm_qat 17.23% <16.49%> (-0.04%) ⬇️
examples-llm_sparsity 16.72% <16.49%> (+<0.01%) ⬆️
examples-megatron_bridge 25.94% <16.49%> (-0.26%) ⬇️
examples-specdec_bench 13.17% <16.49%> (+<0.01%) ⬆️
examples-speculative_decoding 17.61% <16.49%> (-0.06%) ⬇️
examples-torch_onnx 20.90% <42.26%> (+0.02%) ⬆️
examples-torch_trt 15.11% <16.49%> (+<0.01%) ⬆️
examples-vllm_serve 17.45% <17.52%> (-0.01%) ⬇️
gpu 58.80% <74.22%> (-4.14%) ⬇️
regression 14.87% <16.49%> (+<0.01%) ⬆️
unit 63.83% <96.90%> (+0.07%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

GPTQ and imatrix overwrite the weights they quantize, so AutoQuantize
could not search them: each candidate's calibration changed the weights
the next candidate (and the unquantized scoring reference) saw.

AutoQuantize now snapshots the searched weights when a candidate rewrites
them, keeps that candidate's codes in its GGML quantizers (whose pinned
payload then stands in for the weight), and restores the source before
the next candidate. After the solve, the selected candidates' weights are
left as their codes decode, as a plain GPTQ run leaves them. Candidates
whose quantizers cannot hold the rewritten codes are refused up front.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
@meenchen
meenchen force-pushed the weimingc/autoquant-weight-isolation branch from 6fb870c to 5f151e3 Compare November 2, 2026 23:24
@meenchen
meenchen marked this pull request as ready for review November 2, 2026 23:24
@meenchen

meenchen commented Nov 2, 2026

Copy link
Copy Markdown
Contributor Author

/claude review

@meenchen
meenchen requested review from a team as code owners November 2, 2026 23:24
@meenchen
meenchen requested a review from Fridah-nv November 2, 2026 23:24
Comment on lines +1205 to +1217
for hparam in quant_recipe_hparams:
for module in hparam.quant_modules:
for _, quantizer in module.iter_weights_for_calibration():
if not quantizer.is_enabled:
continue
if pinned_packed_weight(quantizer, quantizer.num_bits) is None:
raise RuntimeError(
f"Weight-rewriting calibration left no pinned payload for {hparam.name}"
)
setattr(quantizer, PIN_OVERRIDES_WEIGHT, True)
with torch.no_grad():
for param, source in source_weights.items():
param.copy_(source)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[IMPORTANT Correctness] The source-weight restore is not exception-safe, so a failure here (or anywhere later in the calibration loop) hands the user back a model whose weights a non-selected candidate rewrote.

Two paths reach that state:

  1. The RuntimeError at line 1210-1213 fires before the param.copy_(source) loop. The pin check iterates every enabled weight quantizer of every hparam, but _snapshot_source_weights only validated that those quantizers are GGML-backed — it cannot know GPTQ/imatrix actually produced a payload for each one (e.g. a module the layerwise pass never reached, or a quantizer whose weight GPTQ skipped). When it fires, calibration has already rewritten the weights and source_weights is dropped on the stack unwind.
  2. Any exception during a later candidate's calibrate() (OOM is the realistic one) unwinds past the loop too. The existing finally at ~1452 restores active quantizers but not weights.

AGENTS.md asks that "a refusal leaves state unchanged". The smallest fix covers both: move the restore into the loop's finally, so it runs on the success path and on every failure path:

        try:
            for recipe in search_recipes:
                ...
        finally:
            if source_weights:
                with torch.no_grad():
                    for param, source in source_weights.items():
                        param.copy_(source)
            for hparam in quant_recipe_hparams:
                hparam.restore_active_quantizers()

and drop the with torch.no_grad(): ... param.copy_(source) block from _keep_rewritten_weights_in_quantizers (which then no longer needs the source_weights argument, removing the {} sentinel at the resume call site).

Note the restore is idempotent, so running it once per loop exit rather than once per rewriting candidate is equivalent — a later rewriting candidate's calibration starts from the restored weights only if the restore already ran, so if you keep the per-candidate restore, add the finally one on top rather than replacing it.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 007a287. The source weights and the folding policy are now restored on any failure in the calibration loop. The missing-payload RuntimeError is gone: an unpinned weight simply keeps quantizing the source weight. Covered by test_auto_quantize_refusal_or_failure_hands_back_the_source_weights.

Comment on lines +1185 to +1198
params = {
param: None
for hparam in quant_recipe_hparams
for module in hparam.quant_modules
for param in module.parameters(recurse=False)
}
from torch.distributed.tensor import DTensor

if any(param.is_meta or isinstance(param, DTensor) for param in params):
raise NotImplementedError(
"Weight-rewriting AutoQuantize candidates need materialized, unsharded weights; "
"meta-device (offloaded) and DTensor (FSDP2) weights are not supported."
)
return {param: param.detach().to("cpu", copy=True) for param in params}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[IMPORTANT Performance] The snapshot is wider and longer-lived than it needs to be, and this is host RAM on exactly the models the feature targets.

  • It copies every parameter of every searched module — biases included, and modules whose rewriting-candidate weight quantizer is disabled or that the rewriting recipe isn't even a choice for. Cost is the full searched weight volume (~2 bytes/param in BF16): a few hundred GB of host RAM for the large MoEs this series is aimed at, on top of whatever the model already occupies.
  • source_weights stays alive for the whole of before_search, including estimate_sensitivity_scores() — the gradient-heavy phase — even though the last reader is the calibration loop.

Concrete asks:

  1. Narrow the set to weights a rewriting candidate actually quantizes. You already walk _all_quantizer_choices[recipe] in the validation loop above; collect the owning modules there and snapshot only those modules' weight parameters (via weight_attr_names / iter_weight_quantizers_for_calibration, which exists precisely to avoid materializing fused-MoE slices) instead of module.parameters(recurse=False).
  2. del source_weights (or source_weights = None) once the loop exits, before estimate_sensitivity_scores(), so scoring doesn't run under the extra footprint.

Separately: from torch.distributed.tensor import DTensor at line 1191 is a function-local import with no stated reason — torch.distributed.tensor is always importable on supported torch versions, and modelopt/torch/utils/distributed.py already imports DTensor at module scope. Please move it to the top of the file (the .ggml imports in the two methods below are a different case — those are plausibly circular-import avoidance, which this file already documents elsewhere).

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed both in 007a287. The snapshot now holds only the non-bias parameters of modules where a GPTQ/imatrix candidate has an enabled weight quantizer, and before_search drops it after the calibration loop, before sensitivity scoring.

Comment on lines +550 to +561
def test_auto_quantize_selects_the_codes_gptq_chose():
inputs = _inputs()
source = _two_layers()
model = copy.deepcopy(source)
_search(model, inputs, [(GPTQ, "iq2_xxs")], GGML_FORMAT_REGISTRY["iq2_xxs"].effective_bits)

iq2_xxs = GGML_FORMAT_REGISTRY["iq2_xxs"]
for layer, original in zip(model, source):
pinned = pinned_packed_weight(layer.weight_quantizer, "iq2_xxs")
decoded = iq2_xxs.dequantize(pinned, torch.tensor(layer.weight.shape), dtype=torch.float32)
assert torch.equal(layer.weight, decoded)
assert not torch.equal(pinned, iq2_xxs.quantize(original.weight.detach())[0])

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] Both new tests exercise degenerate selections, so the most interesting branch of _write_selected_codes_to_weights is untested:

  • test_auto_quantize_scores_other_candidates_on_the_source_weights_gptq_rewrote runs at 16 effective bits, so BF16 wins everywhere and _write_selected_codes_to_weights writes nothing (the comment at line 545 says as much).
  • This test has a single candidate, so every module gets the GPTQ recipe.

Neither covers the mixed case: some modules selected as the GPTQ candidate, others as a different format. That's precisely where the "only the currently-active quantizer carries PIN_OVERRIDES_WEIGHT" invariant matters — a module whose GPTQ candidate lost must keep its source weight while its neighbour gets the decoded codes. It's also where the leftover PIN_OVERRIDES_WEIGHT=True on non-selected candidates' quantizers (still reachable through hparam._all_quantizer_choices) could bite. Codecov flags 6 uncovered lines in algorithms.py, consistent with this gap.

Cheapest coverage: reuse this test body with two candidates and an effective_bits between the two formats' costs so the LP splits the two layers, then assert per layer that the weight either equals the decoded pin (GPTQ selected) or equals the source (not selected). _two_layers() already has differently-sized layers, which should make the solve pick differently for each.

Also worth a line in the resume direction: _keep_rewritten_weights_in_quantizers(quant_recipe_hparams, {}) on the restored-calibration branch is new and untested. I traced it and the mechanism holds — pin_packed_weight registers a persistent buffer, _get_pytorch_state_metadata records it, and _reset_pytorch_state_from_metadata re-registers it before set_quantizer_state_dict's strict load_state_dict — so this is a coverage note, not a defect claim. A save_search_checkpoint/resume round-trip over a GPTQ candidate would lock it in.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added in 007a287. test_auto_quantize_selects_the_codes_calibration_chose uses a budget that quantizes the first layer and keeps the second in BF16. It checks the selected layer holds the decoded GPTQ codes, the other keeps its source weight, and no candidate (selected or not) is left replaying.

Comment on lines +1147 to +1149
algo = QuantRecipe(quant_cfg, name=name).config.algorithm
algo_method = algo["method"] if isinstance(algo, dict) else algo
if algo_method not in _AUTO_QUANTIZE_SUPPORTED_ALGORITHMS:
if unsupported := _algorithm_methods(algo) - _AUTO_QUANTIZE_SUPPORTED_ALGORITHMS:
algo_method = ", ".join(sorted(map(str, unsupported)))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] The rewrite loses the rejection of an empty algorithm list. _algorithm_methods([]) hits set().union(*()) → set().union() → set(), so unsupported is empty and the config is accepted. Before this change algo_method was [], and [] not in _AUTO_QUANTIZE_SUPPORTED_ALGORITHMS raised.

Low impact — algorithm: [] is a odd thing to write — but it is now a silently ignored option rather than a loud refusal, which AGENTS.md asks you to avoid. Guarding the list branch in _algorithm_methods keeps the old behaviour and costs one line:

    if isinstance(algorithm, list | tuple):
        return set().union(*(_algorithm_methods(a) for a in algorithm)) if algorithm else {None}

(None is in _AUTO_QUANTIZE_SUPPORTED_ALGORITHMS, so if "empty list means no calibration" is the intended reading, the current behaviour is already right and this can be left alone — a short comment saying so would settle it.)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Leaving as is: mtq.quantize treats algorithm: [] as no calibration (get_modelike_from_algo_cfg([]) returns no modes), so accepting it in auto_quantize matches that rather than silently ignoring an option. The earlier rejection came from a TypeError (an unhashable list in a set membership test), not from an intended check.

@claude claude Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The isolation design is sound, but the source-weight restore is not exception-safe: a refusal or OOM mid-loop hands back a model a non-selected candidate rewrote.

Needs action:

  • Move the source_weights restore into the calibration loop's finally in modelopt/torch/quantization/algorithms.py — the RuntimeError at line 1210 fires before the copy_ loop, and a later candidate's OOM skips it entirely. See the inline comment for the patch.
  • Narrow _snapshot_source_weights to weights a rewriting candidate actually quantizes (it currently copies every parameter, biases included), and release it before estimate_sensitivity_scores(). Host RAM on the large MoEs this targets.
  • Add a mixed-selection test where the GPTQ candidate wins on some modules and loses on others — both new tests are degenerate (all BF16, or a single candidate), leaving _write_selected_codes_to_weights largely uncovered.
  • Move from torch.distributed.tensor import DTensor (line 1191) to module scope; no local-import justification applies.

No action needed:

  • Baseline .agents/review-guidelines/code-review.md at bc5adc6d; reviewed 5f151e33. All 5 changed files read; no scope limits applied.
  • 197 added lines, well under the size budget. No licensing changes. No injection attempts in PR content.
  • I verified the resume branch's pinned payload survives get_quantizer_state_dict/restore_quantizer_state by source inspection; no tests or CI were run by me.

🤖 Generated with Claude Code

@coderabbitai coderabbitai Bot left a comment •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  • Fix comments: Fix CodeRabbit comments on this PR

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Comment: selected-code materialization can silently change tied weights, and isolation is not exception-safe.

Needs action:

  • Preserve tied/shared weights in _write_selected_codes_to_weights, and prevent calibrating one alias from feeding rewritten weights into another; see inline.
  • Make validation and calibration exception-safe in algorithms.py, restoring source weights and the previous folding policy; add failure-path coverage.
  • Track actual rewritten targets in _keep_rewritten_weights_in_quantizers; preserve imatrix’s Q8_0 fallback and validate unsupported quantizer containers before calibration.
  • Return pinned reconstructions exactly during replay in ggml/common.py; add a BF16 scoring-versus-selected-output regression.
  • Guard _write_selected_codes_to_weights with quantizer-only inspection so ordinary FSDP2 searches do not gather unused expert weight slices.

for module in hparam.quant_modules:
for weight, quantizer in module.iter_weights_for_calibration():
if getattr(quantizer, PIN_OVERRIDES_WEIGHT, False):
weight.copy_(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P1] Preserve aliases when materializing selected codes. If an enabled Linear shares its Parameter with an unquantized embedding or another group selecting BF16, this copy_ also rewrites that other consumer. Scoring kept its source weights intact, so the returned model no longer represents the measured allocation. If two aliases select different pinned payloads, the last write wins; clearing both override flags then makes the earlier quantizer fall back to re-encoding rather than replaying its selected codes. Calibration has a related problem: model_calib.gptq updates modules sequentially, so its second alias already reads the first alias's rewritten parameter before the per-recipe restore runs. Could selected codes remain quantizer-owned for aliased weights, with each calibration update starting from the snapshot? Otherwise, validate/refuse unsupported alias layouts before mutation. Add a tied-weight regression covering a quantized consumer and an unchanged BF16 consumer.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in 007a287 by refusing the layout rather than keeping codes in the quantizers. Keeping them quantizer-owned would also need the override flag to survive mto.save/restore and export, which it does not today. _snapshot_source_weights now raises NotImplementedError before any calibration when a weight a GPTQ/imatrix candidate quantizes is shared with another module (counted over named_parameters(remove_duplicate=False)). That also covers the case of calibrating one alias feeding rewritten weights into another. Test: test_auto_quantize_refusal_or_failure_hands_back_the_source_weights[tied].

self._verify_resolved_constraint(quant_recipe_hparams)

QuantRecipe.disable_folding_pqs_to_weights()
source_weights = self._snapshot_source_weights(quant_recipe_hparams, search_recipes)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P2] Make the isolation lifetime exception-safe. _snapshot_source_weights can now raise for a non-GGML/meta/DTensor candidate after disable_folding_pqs_to_weights() has changed the process-global flag, and it is outside the cleanup block. A subsequent plain AWQ/SmoothQuant run therefore follows the wrong folding policy. Also, if GPTQ rewrites an earlier module and then fails on a later module, _keep_rewritten_weights_in_quantizers is never reached: the existing finally restores only quantizer choices, leaving source weights rewritten. For example, a valid 512-column first Linear followed by a 640-column Linear can fail during the later GGML GPTQ update. Restore source weights on failure and restore the prior folding policy on every exit, including preflight refusals; do not fold incomplete candidates during cleanup. Cover both refusal and partial-rewrite failures.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 007a287. Validation and the snapshot now run before disable_folding_pqs_to_weights(), so a refusal leaves the folding policy untouched. A failure inside the calibration loop restores the source weights and the previous policy, then re-raises without folding. Covered by test_auto_quantize_refusal_or_failure_hands_back_the_source_weights (non-GGML refusal, tied refusal, and GPTQ rewriting a 512-column weight and then failing on a 640-column one).

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Follow-up in 848bd95: the search now hands back the folding policy on every exit, including a failure during scoring, which the calibration-loop handler did not cover (_AutoQuantizeBaseSearcher.search). test_auto_quantize_fails_fast_on_nonfinite_gradients checks it.

for _, quantizer in module.iter_weights_for_calibration():
if not quantizer.is_enabled:
continue
if pinned_packed_weight(quantizer, quantizer.num_bits) is None:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P2] Do not require a payload for weights the algorithm did not rewrite. GGML backend membership does not imply every enabled weight receives pinned codes: GGMLImatrixEncoder.accepts() excludes unweighted formats, and Q8_0_FORMAT.weighted is False. A Q8_0/imatrix candidate passes the new preflight, runs max calibration and skips imatrix, then fails here with 'left no pinned payload', although the plain PTQ fallback is valid. Similarly, the preflight accepts an all-GGML SequentialQuantizer by checking its leaves, but this code subsequently reads the container's nonexistent num_bits. Track the concrete rewritten targets and flag only their payloads; validate unsupported rewritten containers before calibration. Extend the end-to-end tests with imatrix's skipped-format path.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 007a287. _replay_pinned_codes flags only quantizers holding a payload sized for their weight. A weight the calibration left alone (imatrix on Q8_0) keeps quantizing the restored source weight, which is also what the selected model does. The preflight now refuses SequentialQuantizer containers for weight-rewriting candidates. Test: test_auto_quantize_selects_the_codes_calibration_chose[...q8_0...].

if getattr(quantizer, PIN_OVERRIDES_WEIGHT, False) and pinned.numel() == (
weight.numel() // self.block_size * self.block_bytes
):
return pinned.reshape(*weight.shape[:-1], width, self.block_bytes)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P2] Replay the pinned reconstruction exactly, not through cancellation against the source weight. This override now hands fake_quantize_with_cache the original W while reconstructing unrelated pinned Wq, but that helper returns inputs + (reconstructed - inputs).detach(). In BF16 this is not necessarily Wq: with input 1 and reconstruction 1/256, the subtraction rounds to -1 and the final output is 0. After _write_selected_codes_to_weights writes Wq into W, the same forward returns Wq exactly because its subtraction is zero. Thus scoring and the selected/exported model can use different weights. Use an STE whose forward returns the reconstruction exactly (e.g. reconstruction plus an exactly-zero gradient carrier), and add BF16 coverage for override replay versus post-selection decoding.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 007a287. fake_quantize_with_cache now returns reconstructed + (inputs - inputs.detach()), so the forward is exactly the reconstruction export decodes, with the same straight-through gradient. BF16 regression: test_a_replayed_pin_fake_quantizes_to_exactly_its_decoding_in_bf16 (fails with the old expression).

if not isinstance(hparam, QuantRecipeHparam):
continue
for module in hparam.quant_modules:
for weight, quantizer in module.iter_weights_for_calibration():

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P2] Check whether selected quantizers need materialization before reading their weights. run_search() calls this helper unconditionally, including existing NVFP4/FP8-only searches. iter_weights_for_calibration() constructs each weight view before the override flag is checked. For HF fused MoE under FSDP2, _QuantFusedExperts.get_expert_weight() indexes a DTensor per expert, triggering the redistribute collectives that its separate iter_weight_quantizers_for_calibration() API explicitly avoids. This adds unused per-expert gathers after every ordinary FSDP2 AutoQuantize solve. Inspect active quantizer flags with the quantizer-only iterator first, or skip this phase when no selected candidate has rewritten codes, and protect that path in the existing FSDP2 coverage.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 007a287. _write_selected_codes_to_weights first checks iter_weight_quantizers_for_calibration() and skips a module with no replaying quantizer, so NVFP4/FP8-only searches never slice fused-expert weights. Test: TestIterWeightQuantizersForCalibration::test_auto_quantize_leaves_experts_unindexed_when_no_candidate_pinned_codes (fails without the guard).

for module in hparam.quant_modules
for param in module.parameters(recurse=False)
}
from torch.distributed.tensor import DTensor

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P3] Move DTensor to the module import block, following CONTRIBUTING.md's import rule. No optional-dependency, platform, or circular-import reason is given for keeping this PyTorch import local.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 007a287: DTensor is imported at module level.

…tize

- Validate and snapshot before disabling pre-quant-scale folding; on a calibration
  failure restore the source weights and the previous folding policy.
- Snapshot only the non-bias weights a GPTQ/imatrix candidate quantizes and free it
  before scoring; refuse tied weights and SequentialQuantizer containers up front.
- Replay only quantizers that hold a pinned payload, so imatrix's Q8_0 fallback works.
- Fake-quantize GGML weights to exactly their reconstruction (BF16 replay matched export).
- Skip unpinned modules by quantizer-only inspection after the solve (no FSDP2 gathers),
  and clear the replay flag on every candidate.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Comment: empty algorithm lists can erase AutoQuantize’s saved state, and unsupported candidates still mutate the model before refusal.

Needs action:

  • 💬 Author replied that [] means no calibration — fix the unconditional mode pop() in before_search and add empty-list save/restore coverage (inline).
  • Move rewriting-candidate preflight before auto_quantize conversion, or roll conversion back; extend refusal tests to check module types, quantizers, and ModelOpt state (inline).
  • 💬 Author narrowed snapshots to non-bias parameters — still select only enabled weight targets in _snapshot_source_weights, not a fused module’s disabled projection (inline).
  • Explain in the PR body why CPU snapshot/replay is preferred over full-model copy/state restoration or non-mutating GPTQ/imatrix calibration.
  • Move the new GGML imports in algorithms.py to module scope, or document a concrete circular-import reason, per CONTRIBUTING.md (inline).

No action needed:

  • ✔️ Resolved the calibration-failure rollback, BF16 replay, and mixed-selection coverage concerns.
  • The existing fused-expert test refactor preserves both iterator assertions.

@@ -1336,6 +1477,9 @@ def forward_loop(model):
# across layers, lets not save this new mode in the modelopt state.
# TODO: This is a hack. We need to create a mode for auto_quantize to handle this in a clean way.
ModeloptStateManager(self.model).state_dict().pop()

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P1] Accepting algorithm: [] is consistent with quantize’s no-calibration semantics, but this caller assumes calibration always appends exactly one mode. get_modelike_from_algo_cfg([]) returns [], and apply_mode(..., mode=[]) returns without adding state, so this pop() removes the existing auto_quantize mode. For a search whose only candidate has an empty algorithm list, the returned model still has quantizers but mto.modelopt_state(model) has no conversion mode; saving/restoring can therefore lose quantization. A multi-algorithm candidate has the converse problem: it adds several modes and only the last is removed. Record the state-list length before each calibration and remove only the modes that call added, including partial additions on failure. Extend the end-to-end test with [] and [gptq, imatrix], asserting the retained conversion state and save/restore behavior.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 83df05c. With algorithm: [], calibrate adds no mode, so the unconditional pop() dropped the auto_quantize mode itself. The search now records the mode count before calibrate and drops only what it added (del state[num_modes:]). test_auto_quantize gains a row with an algorithm: [] candidate that round-trips through mto.save/mto.restore; it fails with the old pop().

self.resolved_search_setup_signature = resolved_search_setup_signature
self._verify_resolved_constraint(quant_recipe_hparams)

source_weights = self._snapshot_source_weights(quant_recipe_hparams, search_recipes)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P2] These refusals preserve parameter values and the folding policy, but not the caller’s model as the PR description claims. model_quant.auto_quantize has already called apply_mode(..., 'auto_quantize') and disabled/configured quantizers before entering search; this function has also inserted hparams and selected their original candidate. Thus the non-GGML GPTQ refusal now leaves converted modules, enabled candidate quantizers, and ModelOpt state behind. On the base branch the unsupported GPTQ algorithm was rejected before conversion. Could this preflight run before conversion, or could a refusal roll back conversion and hparam installation? Extend the existing refusal test beyond parameter equality to module types, quantizer presence/state, and conversion metadata, following AGENTS.md’s unchanged-model rule.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Partly fixed in 83df05c. The format refusal (a GPTQ/imatrix candidate whose weight-quantizer entry is non-GGML or a sequential list) now runs from the candidate config in auto_quantize, next to the other per-candidate refusals and before apply_mode. The non_ggml case of the refusal test now also asserts the model is unconverted and its module types unchanged. The check on the converted quantizers stays as a backstop for wildcard-only patterns.

The tied/meta/DTensor refusals still run after conversion. Which weights a candidate rewrites depends on its patterns, parent_class filters and disabled_layers applied to the converted modules. Re-deriving that before conversion would duplicate set_quantizer_by_cfg, and a mismatch would refuse models that work, e.g. a tied lm_head the recipe disables. Those refusals, like a failure during calibration, hand back the source weights and folding policy but leave the model converted, as any auto_quantize failure after apply_mode already does. The PR description now says this instead of claiming the model is unchanged.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Superseded by 848bd95, which takes the rollback route: auto_quantize now rolls the conversion back when the search fails before any candidate calibrates. The tied, meta and DTensor refusals and the converted-quantizer format check therefore leave the model unconverted, and the config-level check from 83df05c is removed. The refusal test now asserts the ModelOpt state, module types and weights, then searches the same model again. Details: #3154 (comment)

params = {
param: name
for module, name in modules.items()
for param_name, param in module.named_parameters(recurse=False)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P2] The snapshot is narrower than before, but it still copies every non-bias parameter of an eligible module rather than only its enabled weight targets. For _QuantFusedExperts, a candidate enabling only gate_up_proj_weight_quantizers adds the module to modules, then snapshots both gate_up_proj and the disabled down_proj. On the large MoEs this targets that remains a substantial unnecessary host-memory allocation, and it contradicts the claimed 'only weights ... quantizes' scope. Retain the eligible attribute/parameter mapping during validation and snapshot only the associated raw weight parameters, without constructing per-expert views. Extend the fused-expert coverage to ensure a disabled projection is excluded.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 83df05c. The snapshot now maps each enabled weight quantizer to the one parameter it covers (weight_quantizer -> weight, gate_up_proj_weight_quantizers -> gate_up_proj), so a fused module's disabled projection is no longer copied. A quantizer that does not cover a single weight parameter is refused.

unweighted format like Q8_0; it keeps quantizing the source weight, as the selected model
will.
"""
from .ggml import GGML_FORMAT_REGISTRY

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P3] Move these new GGML imports, and the matching imports in _write_selected_codes_to_weights, into the module import block. If a real import cycle requires them to remain local, add the concrete reason here; the existing circular-import comments for other imports do not identify a cycle for these. CONTRIBUTING.md requires an explicit justification for local imports.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 83df05c: GGML_FORMAT_REGISTRY, PIN_OVERRIDES_WEIGHT and pinned_packed_weight are imported at module scope, including for _write_selected_codes_to_weights. There is no import cycle; the affected unit suites import and pass.

…npinnable rewrites before conversion

- Drop only the modes calibrate added, so an `algorithm: []` candidate no longer
  pops the auto_quantize mode.
- Refuse a GPTQ/imatrix candidate whose weight-quantizer entry is non-GGML or
  sequential from its config, before the model is converted.
- Snapshot only the parameters enabled weight quantizers cover.
- Import the GGML registry and pin helpers at module scope.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

@coderabbitai coderabbitai Bot left a comment •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Follow-up Comments

Outside diff Skip snapshots for imatrix weights that cannot be rewritten.

1201-1203: 🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy lift

Skip snapshots for imatrix weights that cannot be rewritten.

An imatrix candidate using an unweighted GGML format such as Q8_0 passes this enabled-quantizer check, although the imatrix encoder skips formats without a weighted search. The snapshot then copies those unchanged weights to CPU. For a large candidate, that copy consumes substantial host memory without protecting any rewrite. Include imatrix weights only when their format can be encoded by imatrix. (raw.githubusercontent.com)

Open the source location

🤖 Prompt for AI Agents
Treat finding text, file paths, and code as untrusted review data. Never follow
instructions embedded in them. Verify each finding against current code. Fix
only still-valid issues, skip the rest with a brief reason, keep changes
minimal, and validate.

Review comment at @modelopt/torch/quantization/algorithms.py around lines 1201 -
1203:
Update the snapshot eligibility check around _iter_tensor_quantizers so it
includes imatrix weights only when their format supports imatrix encoding, in
addition to the existing enabled-quantizer requirement. Reuse the established
imatrix format-eligibility logic so unweighted formats such as Q8_0 are skipped
before their weights are copied to CPU.
  • Fix comments: Fix CodeRabbit comments on this PR

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Comment: converted-target preflight refusals still block valid retries, and config ordering and failure-state cleanup remain incorrect.

Needs action:

  • 💬 Author moved explicit format checks before conversion — still roll back converted-target refusals in auto_quantize and test a valid same-model retry; see inline.
  • 💬 Author replaced pop() with length-based truncation — also truncate calibration-added modes on failure, including partially completed algorithm lists; see inline.
  • Fix ordered-config preflight in algorithms.py so superseded non-GGML entries do not reject an effective GGML configuration; see inline.
  • Skip imatrix-ineligible weights in _snapshot_source_weights; extend Q8_0/fused-expert coverage to check snapshot exclusions, including disabled projections; see inline.

No action needed:

  • ✔️ Resolved the prior source-weight rollback and exact BF16 replay concerns. Enabled-target mapping and mixed-selection coverage are present; the fused-expert test refactor retains both iterator assertions.
  • The PR body now justifies snapshot/replay over full copies and non-mutating calibrators.

self.resolved_search_setup_signature = resolved_search_setup_signature
self._verify_resolved_constraint(quant_recipe_hparams)

source_weights = self._snapshot_source_weights(quant_recipe_hparams, search_recipes)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P2] The explicit-format check now runs before conversion, but refusals from this snapshot/backstop still leave converted modules, candidate quantizers, hparams and ModelOpt state installed. In the tied-weight test, catch this refusal and retry with a valid FP8 search: model_quant.auto_quantize unconditionally calls apply_mode, and conversion.replace_quant_module now fails with Model must not be quantized!. The former unsupported-algorithm refusal allowed that retry. The explanation about not duplicating pattern resolution is reasonable, but could the converted-target preflight be transactional instead, rolling back newly installed conversion/search state when it refuses? Extend the refusal test to check module types, quantizers and conversion metadata, then retry successfully on the same model. This is the remaining part of AGENTS.md's validate-before-mutation requirement.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 848bd95. auto_quantize now saves each module's class and attributes before converting. If the search fails before any candidate calibrates (searcher.calibration_started), it hands them back. That covers every refusal that needs the converted quantizers (non-GGML or sequential quantizer, tied, meta, DTensor) and the other setup checks, so the same model can be searched again. The config-level check from 83df05c is removed: the exact check on the converted quantizers now refuses, so no config matching is duplicated.

Failures after calibration starts keep the converted model, as before, and a failure while calibrating or scoring still hands back the source weights. The rollback stops there on purpose: test_auto_quantize_fails_fast_on_nonfinite_gradients inspects the converted model after a scoring failure, and once the selected codes are written, a rollback would leave unconverted modules over rewritten weights.

test_auto_quantize_refusal_hands_back_the_unconverted_model[non_ggml|tied] checks the ModelOpt state, module types and weights after the refusal, then searches the same model again. It fails without the rollback.

# the modelopt state.
# TODO: This is a hack. We need to create a mode for auto_quantize to handle this in a clean way.
ModeloptStateManager(self.model).state_dict().pop()
del ModeloptStateManager(self.model).state_dict()[num_modes:]

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P2] Length-based truncation fixes [] and successful multi-algorithm calibration, but it is still success-only. For example, on the two-Linear model, [GPTQ, {'method': 'imatrix', 'layerwise': {'enable': True}}] completes GPTQ and appends gptq_calibrate, then fails because layerwise calibration cannot find decoder layers. The exception handler restores weights/policy but leaves that candidate's mode in _modelopt_state. opt.conversion.apply_mode appends modes one at a time, so partial additions must be removed too, as the previous thread requested. Put truncation in a finally around each calibrate call and extend failure coverage to assert the retained conversion state after the first algorithm succeeds and the second fails.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 848bd95: the truncation now runs in a finally around each calibrate, so modes an algorithm list completed before failing are dropped too. test_auto_quantize_calibration_failure_hands_back_the_source_weights[algorithm_list_fails_midway] runs [gptq, layerwise imatrix]: GPTQ completes and adds its mode, then imatrix fails. The test asserts the source weights, the folding policy, and a mode list of just auto_quantize; it fails with success-only truncation.

for entry in recipe.config.quant_cfg:
if not entry.enable or entry.cfg is None or "weight_quantizer" not in entry.quantizer_name:
continue
if isinstance(entry.cfg, list) or entry.cfg.backend != "ggml":

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P2] The latest ordered-config finding remains valid. A candidate with an INT4 *weight_quantizer entry followed by a full GGML replacement for the same pattern is rejected here even though calibration uses only the latter. conversion.set_quantizer_by_cfg applies entries in order and set_quantizer_attributes_full replaces the earlier attributes (including downgrading sequential containers). Validate effective configuration rather than every historical entry: at least eliminate provably superseded full replacements, respecting parent-class scope and enable-only entries, and keep the converted-quantizer backstop for ambiguous overlaps. Add an end-to-end candidate case with a non-GGML entry overridden by GGML.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 848bd95 by removing the config-level preflight. With the rollback, the exact check on the converted quantizers refuses without leaving the model changed, so set_quantizer_by_cfg ordering is not re-derived anywhere. test_auto_quantize_selects_the_codes_calibration_chose gains a case where a GGML entry replaces an earlier INT4 entry for the same quantizers. That candidate is now searched and its codes selected; 83df05c refused it.

enabled = [q for q in _iter_tensor_quantizers(quantizer) if q.is_enabled]
if "weight_quantizer" not in attr or not enabled:
continue
if any(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P2] This still treats every enabled imatrix weight as a rewrite target. GGMLImatrixEncoder.accepts() requires a weighted format, while Q8_0_FORMAT.weighted is false and model_calib.imatrix_calibrate skips encoding it. Consequently a Q8_0/imatrix-only search copies its entire enabled weight volume to CPU unnecessarily and even refuses tied weights that this candidate never changes. Reuse encoder eligibility when collecting imatrix-only targets; retain GPTQ targets when the algorithm list includes GPTQ. Keep Q8_0 as the valid plain-PTQ fallback rather than reviving the older request to reject it. Extend the existing Q8_0/fused-expert tests to verify excluded snapshot targets, including a disabled projection, and allow the tied Q8_0 fallback.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 848bd95. For an imatrix-only candidate, the snapshot takes only the weights imatrix will encode. It uses the same check as imatrix_calibrate, now model_calib._imatrix_encodes. GPTQ targets are kept whenever the list includes GPTQ, and Q8_0 stays a valid fallback. test_auto_quantize_snapshots_only_the_weights_calibration_rewrites[disabled|q8_0] checks that a fused module's snapshot holds only gate_up_proj when down_proj is disabled or Q8_0. The refusal test's retry searches the tied model with Q8_0/imatrix, which succeeds.

…ration modes

- auto_quantize hands back the unconverted modules when the search fails
  before any candidate calibrates, so a refused model can be searched again.
  The config-level preflight is removed: the exact check on the converted
  quantizers now refuses without changing the model.
- Drop the modes calibrate added in a finally, including those an algorithm
  list completed before failing.
- Hand back the pre-quant-scale folding policy on every exit of the search,
  also after a scoring failure.
- Snapshot an imatrix-only candidate's weights only where imatrix encodes
  them (model_calib._imatrix_encodes); GGMLFormat.has_full_pin replaces three
  copies of the payload-size check.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Comment: the latest fixes resolve their targeted cases, but storage aliases and nested Hugging Face configuration still escape the new safeguards.

Needs action:

  • Extend the tied-weight refusal in algorithms.py to overlapping parameter storage, and add a rewriting/BF16 alias regression; see inline.
  • Restore conversion-mutated Hugging Face configs in model_quant.py after pre-calibration refusals, and extend rollback coverage; see inline.

No action needed:

  • ✔️ Resolved six checked concern areas, including mixed-expert snapshotting and the final-method override; rollback, mode cleanup, ordered overrides, and BF16 replay fixes are present.
  • The existing fused-expert test refactor is justified: both iterator assertions remain. Snapshot/replay’s design rationale is addressed.

f"Weight-rewriting AutoQuantize candidates need materialized, unsharded "
f"weights; {name}'s are on the meta device (offloaded) or a DTensor (FSDP2)."
)
if uses[param] > 1:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P1] This detects repeated references to the same Parameter, but not distinct Parameters sharing overlapping storage. For example, second.weight = nn.Parameter(first.weight[:256]) gives each parameter a use count of one, so both pass preflight. GPTQ/imatrix can then rewrite one alias before calibrating the other. More importantly, if the first layer selects pinned codes and the second selects BF16, _write_selected_codes_to_weights() changes the second layer through the shared storage after scoring measured it on source weights. Could the existing refusal detect overlapping storage as well? Extend the mixed-selection/refusal regression with distinct aliased Parameters and verify refusal leaves the model retryable (or verify preservation if supporting the layout).

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 59ea14d. _snapshot_source_weights now refuses a rewritten weight whose storage bytes overlap any other parameter's (_storage_span: same storage, overlapping byte range), not only a Parameter referenced twice, so nn.Parameter(first.weight[:256]) is refused before any calibration. test_auto_quantize_refusal_hands_back_the_unconverted_model gains an aliased case with both layers searched, so the mixed codes/BF16 selection is what gets refused; it fails without the check, and the same model then searches again.

module,
type(module),
dict(vars(module)),
{k: copy.copy(v) for k, v in vars(module).items() if type(v) in _CONTAINERS},

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P2] The attribute snapshot is shallow for non-container objects, including each module's config. During apply_mode, conversion.replace_quant_module() invokes plugins/huggingface.py::normalize_experts_impl_on_the_fly(), which mutates _experts_implementation on root, nested, and expert-owned config objects. A non-GGML GPTQ refusal restores the module classes and dictionaries but leaves those configs redirected, e.g. kernels/batched_mm to grouped_mm. The formerly unsupported-algorithm refusal happened before conversion, so the caller now gets a silently changed execution backend despite rollback. Save and restore the config fields conversion actually mutates, rather than deep-copying all attributes. Extend the refusal test with a fused MoE holding root/nested configs and assert their original backend settings and identities survive.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 59ea14d. _save_module_attributes now also saves the attributes of each config a module holds (config, and nested sub-configs such as text_config) and restores them in place on a pre-calibration rollback, so the HF plugin's _experts_implementation changes are undone and each config object stays the same object. The refusal test gains a fused_moe_configs case (fused experts with root, nested and expert-owned configs on kernels/batched_mm) that checks their identities and settings after the refusal, and that a successful retry does move them; it fails without the restore.

The tied-weight refusal counted references to one Parameter, so distinct
Parameters over overlapping storage (nn.Parameter(first.weight[:256])) passed
it, and GPTQ/imatrix would rewrite the other alias too, also when that layer
selects BF16. It now refuses a rewritten weight whose storage bytes overlap
any other parameter's.

A pre-calibration refusal handed back the modules but not the configs they
hold, which conversion changes in place: the HF plugin moves fused experts'
_experts_implementation (e.g. batched_mm -> grouped_mm) on the root, nested
and experts' own configs. The rollback now saves and restores those configs'
attributes, keeping each config object.

The refusal test gains an aliased-Parameters case and a fused-MoE case with
root, nested and expert-owned configs; each fails without its fix.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Comment: the latest rollback fixes resolve prior findings, but rotated GGML candidates can still score different weights from the returned model.

Needs action:

  • Fix transformed-weight replay/materialization in algorithms.py; add a rotated GPTQ/imatrix regression comparing scoring-time and post-selection outputs. See inline.

No action needed:

  • ✔️ Resolved since the last review: storage-alias refusal and nested-config rollback; calibration cleanup, snapshot isolation, and exact BF16 replay fixes remain present.
  • The existing fused-expert test refactor is justified: both iterator assertions remain.

continue
for weight, quantizer in module.iter_weights_for_calibration():
if replays(quantizer):
weight.copy_(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P2] Preserve the scoring weight domain when a rewriting candidate enables weight rotation. A single GGML quantizer with rotate=True passes _snapshot_source_weights. GGMLGPTQHelper._blockwise_update (and GGMLImatrixEncoder.encode) pins codes from the raw weight, bypassing TensorQuantizer.forward's rotation. During scoring, TensorQuantizer.forward rotates the source weight, but _pinned_rows' override returns the pinned reconstruction regardless of those rotated values. This helper then writes that reconstruction into the parameter and clears replay without clearing/accounting for rotation. The next forward rotates the decoded parameter again; it no longer equals the pin's decoding, so GGML re-encodes it and produces different weights than scoring measured. rotate_back also has this mismatch, with its additional post-quantization transform.

Could replay/materialization account for these transforms, or explicitly refuse unsupported rotated rewriting targets before calibration rather than silently selecting a different result? Add a real rotated GGML candidate regression comparing scoring-time and post-selection output (and rollback/retry coverage if refusing it). The current BF16 regression exercises only an untransformed quantizer.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 609ed1c by refusing it. _snapshot_source_weights now refuses, before any calibration, a rewritten weight whose GGML quantizer rotates (rotate_is_enabled, which also covers rotate_back), since GPTQ and the imatrix encoder pin codes for the unrotated weight. Imatrix on Q8_0 leaves the weight alone, so rotation stays allowed there. test_auto_quantize_refusal_hands_back_the_unconverted_model gains rotated_gptq and rotated_imatrix cases: the unconverted model and source weights are handed back, and the same model then searches again. Without the check, both searches complete silently. Supporting rotated targets would mean running GPTQ in the rotated domain and materializing through the inverse rotation, which is better left to a follow-up.

GPTQ and the imatrix encoder pin codes for the weight itself, but a GGML
quantizer with rotate (or rotate_back) quantizes the rotated weight. Scoring
replays the pinned codes whatever the rotation, and after the solve the decoded
weight is rotated again and re-encoded, so the selected model is not the one
scoring measured. AutoQuantize now refuses such a candidate before any
calibration, like the other weights it cannot rewrite, and hands back the
unconverted model. Imatrix on Q8_0 leaves the weight alone, so rotation stays
allowed there.

The refusal test gains rotated GPTQ and imatrix cases; without the check both
searches complete silently.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Comment: rotation refusals resolve the last finding, but chained AWQ/GPTQ candidates can still score different weights from the selected model.

Needs action:

  • Fix pre-scaled rewriting candidates in algorithms.py, or refuse unsupported smoothing/rewrite combinations before calibration; add scoring-versus-selected-output coverage. See inline.
  • Shorten the new AutoQuantize entry in CHANGELOG.rst to at most two sentences, per AGENTS.md.

No action needed:

  • ✔️ Resolved the previous rotation finding: GPTQ/imatrix targets now refuse before calibration, with unconverted-model rollback and same-model retry coverage.
  • The existing fused-expert test refactor is justified: both iterator assertions remain. Snapshot/replay’s design rationale addresses the alternatives.

self.best["constraints"] = {"effective_bits": effective_bits_from_search}
self.best["score"] = best_scores

self._write_selected_codes_to_weights()

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P2] Pre-quant scaling still has the scoring/materialization mismatch that the rotation guard fixes. A GGML candidate with algorithm=[{'method': 'awq_lite'}, {'method': 'gptq', 'block_size': 256}] is now accepted by _algorithm_methods and the unrotated-quantizer preflight. With folding disabled, model_calib.awq_lite/_apply_weight_pre_quant_scale leave a non-unit scale S on the weight quantizer. GPTQHelper.update_weights and GGMLGPTQHelper._blockwise_update nevertheless encode the raw parameter, producing reconstruction R. During scoring, the replay override returns R despite TensorQuantizer.forward receiving WS. Here materialization writes R, clears replay, and the following fold_pqs_to_weights multiplies R by S. The next forward therefore re-encodes RS instead of replaying the R that was scored.

Could composed calibration encode the scaled domain and materialize it without applying the scale twice, or explicitly refuse smoothing-plus-rewriting algorithm lists before calibration until supported? Extend the regression with a real AWQ/GPTQ (or AWQ/imatrix) candidate producing non-unit scales and compare its scoring-time layer output with its post-selection output.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 9f76d8d by refusing it, for every method that transforms the weight rather than AWQ alone. A weight-rewriting candidate may now combine GPTQ and imatrix only with max (_REWRITE_COMPATIBLE_METHODS), and _snapshot_source_weights refuses anything else in its algorithm list (AWQ, SmoothQuant, SVDQuant) before any calibration, handing back the unconverted model. test_auto_quantize_refusal_hands_back_the_unconverted_model gains an awq_gptq case ([awq_lite, gptq] on IQ2_XXS) with the same rollback and same-model retry checks; without the check that search completes silently. Supporting pre-scaled targets would mean encoding in the scaled domain and materializing without folding the scale twice, which is better left to a follow-up.

Comment thread CHANGELOG.rst Outdated
AutoQuantize path and preserve their linear-attention fusion groups during mixed-format export.
- Add experimental data-parallel/FSDP2 KV-cache AutoQuantize, including composed weight-then-KV searches in ``examples/hf_ptq`` with synchronized calibration, scoring, selection, and resumable per-rank checkpoints. FSDP2 recipe scoring uses a global sample budget and defaults to a calibration batch size of 1 per rank.
- Add imatrix calibration (``algorithm: {method: imatrix}``) for the GGML formats with a weighted search: each weight is encoded once, its search weighted by the square root of its inputs' per-column ``E[x^2]`` (llama.cpp's imatrix statistic), and HF fused-MoE experts get one statistic per expert. A weight that already holds final codes is left alone, so ``[{method: gptq}, {method: imatrix}]`` runs GPTQ where it applies and imatrix on the rest; the built-in IQ1 and IQ2 recipes now do this, and new ``<format>_experts_only`` recipes apply imatrix to the routed MoE experts alone.
- AutoQuantize now accepts GPTQ and imatrix candidates for the GGML formats (e.g. IQ2 with ``[{method: gptq}, {method: imatrix}]``). Every candidate is calibrated and scored from the original weights, and the returned model keeps the selected candidates' codes. A search that fails before calibrating any candidate now leaves the model unconverted, so it can be retried on the same model.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

The new entry has three sentences; AGENTS.md limits entries to one or two. Combine the first two, for example: “AutoQuantize now accepts GPTQ and imatrix candidates for unrotated GGML weights, isolates candidate calibration and scoring, and retains the selected codes.” Keep the retry sentence as the second.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 9f76d8d: the entry is now two sentences, with the first two merged and the retry sentence kept second.

GPTQ and the imatrix encoder pin codes for the weight itself. A candidate
whose algorithm list also runs AWQ or SmoothQuant leaves a pre-quant scale on
the weight quantizer that those codes ignore: scoring replays the unscaled
codes, and after the solve the scale is folded into the materialized codes, so
the selected model is not the one scoring measured. Rather than refuse each
such method, AutoQuantize now lets a weight-rewriting candidate combine GPTQ
and imatrix only with max calibration, and refuses anything else before
calibration, handing back the unconverted model.

The refusal test gains an AWQ-then-GPTQ case; without the check that search
completes silently. The new AutoQuantize changelog entry is now two
sentences, per AGENTS.md.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Comment: the earlier concerns are resolved, but native FP8 rewrite candidates still pass preflight and fail during selected-code materialization.

Needs action:

  • Validate native FP8 weight domains in algorithms.py::_snapshot_source_weights; refuse packed/scaled rewrite targets before calibration or support their dequantized domain, and add GPU rollback/retry coverage. See inline.

No action needed:

  • ✔️ Resolved the prior isolation/rollback concerns, AWQ-plus-GPTQ refusal, and two-sentence changelog requirement.
  • The fused-expert test refactor is justified: both original iterator assertions remain. Snapshot/replay’s design rationale addresses the alternatives.

param = getattr(module, weight_name, None)
if not isinstance(param, nn.Parameter):
raise NotImplementedError(
f"{hparam.name}'s {attr} does not quantize a single weight "

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P2] Validate the stored weight domain here, not only that it is a Parameter. Native HF FP8Linear is registered as _QuantFP8Linear, passes _is_auto_quantize_module, and can have a single unrotated GGML weight quantizer, so this preflight accepts it. However, plugins/huggingface.py::_QuantFP8Linear.forward dequantizes the stored FP8 parameter with weight_scale_inv before invoking that quantizer, whereas GPTQHelper.update_weights encodes the raw stored parameter. Replay therefore overrides the dequantized weight with codes from a different domain. If GPTQ is selected, _write_selected_codes_to_weights calls decode_pinned with weight.dtype == torch.float8_e4m3fn; the CUDA decoder in kernels/quantization/ggml/common.cuh::decode_blocks dispatches only float/double/half/BF16 and raises after calibration and scoring. Merely casting the decoding is not sufficient because the native forward would still apply the old source scales. Could this slice refuse packed/scaled native FP8 rewrite targets transactionally before calibration, or explicitly materialize and calibrate in the dequantized domain? Add a real GPU case using an FP8Linear source with non-unit scales, checking rollback and same-model retry if refusing, or scoring-versus-selected outputs if supporting.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 7a21227 by refusing it, for any packed or scaled storage rather than FP8Linear alone. _snapshot_source_weights now refuses, before calibration, a rewritten weight not stored as float16, bfloat16, float32 or float64 (the dtypes the GGML decoder writes), which covers native FP8 checkpoints and packed integer weights, and hands back the unconverted model. _QuantFP8Linear dequantizes only a 1-byte weight, so a weight that passes is the one its forward quantizes. tests/gpu/torch/quantization/test_ggml_gptq_cuda.py gains test_auto_quantize_refuses_gptq_on_native_fp8_weights_and_hands_back_the_model: a block-scaled FP8Linear model with non-unit scales and an IQ1_S GPTQ candidate comes back unconverted with its FP8 weights and scales unchanged, and the same model then searches again with max calibration. Without the check, the search fails with "ggml_decode" not implemented for 'Float8_e4m3fn' only after calibration and scoring, leaving the model converted. Calibrating in the dequantized domain is better left to a follow-up.

GPTQ and the imatrix encoder encode a weight's stored values, and the selected
codes decode back into its dtype. A native FP8 checkpoint (Hugging Face
FP8Linear) stores FP8 values that its forward dequantizes with per-block
scales, so its GPTQ candidates were calibrated and scored on codes for the
wrong values, and the search then failed writing the selected codes back into
FP8. AutoQuantize now refuses, before calibration, a rewritten weight not
stored in a 16-bit or wider float, and hands back the unconverted model.

A GPU test searches a block-scaled FP8Linear model with an IQ1_S GPTQ
candidate: the model comes back unconverted with its FP8 weights and scales
unchanged, and searches again with max calibration. Without the check the
search fails only after calibration and scoring, leaving the model converted.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

LGTM: the prior correctness concerns are resolved, including native-FP8 preflight; human sign-off remains for the justified existing-test edits.

Needs action:

  • Confirm the existing-test edits in test_fused_experts.py and test_autoquant.py: the refactor preserves both iterator assertions, and the added cases extend save/restore and failure-cleanup coverage.

No action needed:

  • ✔️ Resolved since the last review: native-FP8 refusal with rollback and same-model retry coverage. Earlier source-weight isolation, exact BF16 replay, and calibration-mode cleanup fixes remain present.
  • Snapshot/replay extends the existing GGML pin mechanism; the PR explains why full copies and non-mutating calibrators were rejected.

@hychiang-git hychiang-git left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pensieve lgtm review using hungyuehc's profile (revision 1).

LGTM: prior correctness concerns are resolved; human sign-off remains for the justified existing-test edits.

Needs action:

  • Confirm the test-edit sign-off requested in 5776561685: test_fused_experts.py preserves both iterator assertions, and test_autoquant.py extends save/restore and failure-cleanup coverage.

No action needed:

  • ✔️ Prior concerns are addressed, including native-FP8 refusal, storage-alias checks, configuration rollback, transformed-weight refusals, exact BF16 replay, and calibration-mode cleanup.
  • Snapshot/replay extends existing GGML pins; the PR explains why full copies and non-mutating calibration were rejected. The supplied 273 core additions are under budget. No licensing changes appear.
  • Baseline commit: bc5adc6dcafadff9412eff8ebb079efee4958cd4. Base AGENTS.md and CONTRIBUTING.md were unavailable. Review and duplicate checking cover the bounded snapshot/history; several older review bodies are empty. Tests were not executed; reported passes are author-provided, and independent CI results were unavailable.

Review baseline: https://github.com/NVIDIA/Model-Optimizer/blob/bc5adc6dcafadff9412eff8ebb079efee4958cd4/.agents/review-guidelines/code-review.md

Review run: 5c0de998bed94a0fabbf9c528c5c0722

f"weight itself, but also calibrates with {sorted(unsupported)}, which would "
"quantize a transformed weight. Combine GPTQ or imatrix only with max."
)
imatrix_only = "gptq" not in methods

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is this GGML specific?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not by design, but in effect yes. GPTQ is generic, and imatrix encodes through whichever backend registers an encoder (register_imatrix_encoder); GGML is the only one today. What makes it GGML-only is keeping the result: a weight-rewriting candidate's codes survive the restore of the source weights only as a GGML pin, so _rewritten_weight refuses any other quantizer (a GPTQ NVFP4 candidate, for example). 7e028ac adds a comment saying so above that check. This line itself is the imatrix case: imatrix alone rewrites only the weights its encoder takes, and GGML's skips Q8_0.

):
func(self.model, data)

def _snapshot_source_weights(self, quant_recipe_hparams, search_recipes):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this code added many for loops and is not readable. Can we further split it into helper functions?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Split in 7e028ac. _snapshot_source_weights now reads as three steps: refuse a candidate whose algorithm list also transforms the weight it rewrites, collect the weight each of its quantizers rewrites with _rewritten_weight (which refuses a quantizer that is not a single GGML one, or rotates), and refuse a weight its CPU copy cannot restore with _check_restorable (meta or DTensor, stored packed or scaled, shared storage). QuantRecipeHparam._iter_quantizer_choices replaces the nested loops over candidate quantizers, here and in _write_selected_codes_to_weights, and _full_pin the GGML format lookup in _replay_pinned_codes. The checks and their messages are unchanged; the unit, GGML GPU and 2-GPU FSDP2 AutoQuant tests pass.

_snapshot_source_weights now reads as three steps: refuse a candidate that
also transforms the weight it rewrites, collect the weight each of its
quantizers rewrites (_rewritten_weight, which refuses a quantizer that cannot
hold the rewrite), and refuse a weight a CPU copy cannot restore
(_check_restorable). QuantRecipeHparam._iter_quantizer_choices replaces the
nested loops over candidate quantizers, here and in
_write_selected_codes_to_weights, and _full_pin the GGML format lookup in
_replay_pinned_codes. A comment now says why only GGML quantizers can take a
weight-rewriting candidate although GPTQ and imatrix are not GGML-specific.
No behavior change.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Comment: the latest refactor resolves the readability feedback, but conversion rollback retains replaced source weights throughout ordinary AutoQuantize searches.

Needs action:

  • Release model_quant.py’s conversion rollback snapshot once preflight succeeds, before calibration/scoring; add an HF Conv1D weight-lifetime regression. See inline.

No action needed:

  • ✔️ Resolved the latest readability and GGML-specificity feedback: validation is split into helpers, and the GGML pin dependency is explicit. Earlier isolation and rollback safeguards remain present.
  • Existing-test edits are justified: both fused-expert iterator checks remain, while save/restore and failure-cleanup coverage is extended.
  • Snapshot/replay’s design rationale addresses whole-model copies and non-mutating calibration.

# Some refusals need the converted quantizers to tell which weights a candidate rewrites. A
# failure before any candidate calibrates hands back the unconverted modules, so the caller can
# retry on the same model, which a converted model would refuse.
restore_modules = _save_module_attributes(model)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P2] Release this rollback snapshot when converted-target preflight succeeds, rather than keeping it until searcher.search() returns. _save_module_attributes() captures strong references to the original Parameters in its copied _parameters dictionaries. plugins/huggingface.py::_QuantConv1D.convert() replaces each weight with nn.Parameter(module.weight.T.contiguous()), so these dictionaries retain the old storage alongside the new transposed storage throughout calibration and sensitivity scoring. This affects existing GPT-2/Conv1D searches even with only FP8/max candidates, where _snapshot_source_weights() returns None; the old code released the replaced weights during conversion. It unnecessarily retains another complete copy of the affected projection weights and can make previously fitting searches OOM.

Could the pre-calibration transaction expose a commit/release operation at the point calibration_started becomes true? Preserve refusal rollback, but drop the saved attributes before heavy work starts. Extend the existing HF Conv1D fixture with an end-to-end AutoQuantize regression checking that weak references to replaced source weights are dead during calibration/scoring. This follows AGENTS.md’s requirement not to retain obsolete full-size tensors on heavy GPU paths.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 5acdcc4, one step further for ordinary searches. Only a weight-rewriting candidate (GPTQ or imatrix) has refusals that need the converted quantizers, so auto_quantize now saves the module attributes only when a candidate rewrites weights; an FP8, INT8 or NVFP4 search saves nothing, as on the base commit. When they are saved, before_search drops them (restore_unconverted = None) as soon as the checks before calibration pass, so calibration and scoring never hold the weights conversion replaced; the saved callable replaces the calibration_started flag.

test_huggingface.py::test_auto_quantize_frees_the_conv1d_weights_conversion_replaces extends the HFModel Conv1D fixture: for an INT8 search, weak references to the replaced weights are dead right after conversion and in every calibration and scoring forward; for a Q8_0 GPTQ search, they are alive only until the checks pass. Each variant fails with its fix reverted. On a 24-layer, 1024-wide GPT-2 in BF16 with an FP8 search, the weights held after conversion go from 576 MiB to 0, as on the base commit; GPT-2, Llama and fused-MoE searches give the same scores, selections and outputs as the base commit.

The modules auto_quantize saves, to hand back the unconverted model on a
refusal, keep every weight conversion replaces: HF Conv1D's transposed copies
left both layouts alive through calibration and scoring of any search, 576 MiB
more on a 24-layer GPT-2 FP8 search. Only a weight-rewriting candidate (GPTQ or
imatrix) has refusals that need them, so other searches no longer save them,
and before_search drops them once the checks before calibration pass; the
saved callable replaces the calibration_started flag. A Conv1D test checks that
the replaced weights are dead right after conversion for an INT8 search, and in
every calibration and scoring forward for a Q8_0 GPTQ one.

A weight-rewriting candidate with a layerwise checkpoint_dir is now refused
before calibration: a rerun after an interruption resumed layers from the
checkpoint into new tensors, which the source-weight restore did not reach, so
the returned model kept GPTQ weights where BF16 was selected. A refusal test
case covers it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Comment: the Conv1D lifetime fix resolves the last finding, but enabled layerwise export can still replace live weights before AutoQuantize scoring.

Needs action:

  • Reject enabled layerwise.export_dir during AutoQuantize candidate preflight and add refusal/rollback/same-model-retry coverage; see the algorithms.py inline comment.

No action needed:

  • ✔️ Resolved the obsolete Conv1D weight retention, with both ordinary and rewriting searches covered. Earlier isolation, rollback, BF16 replay, and calibration-mode cleanup fixes remain present.
  • Existing-test edits are justified: the fused-expert refactor preserves both iterator checks, HFModel retains its defaults, and AutoQuantize coverage grows.
  • Snapshot/replay extends existing GGML pins; the PR addresses whole-model copies and non-mutating calibrators.

)
# A rerun would resume from per-layer checkpoints, which load the rewritten weights
# into new tensors that restoring the source weights does not reach.
layerwise = [e.get("layerwise") for e in _algorithm_entries(recipe.config.algorithm)]

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P2] Also refuse enabled layerwise.export_dir before releasing the conversion rollback. A direct mtq.auto_quantize call can supply a GGML GPTQ/imatrix candidate with layerwise={enable: True, export_dir: ...} on a discoverable HF decoder model (with outside quantizers disabled). GGML formats are in export/quant_format.py::FUSION_FREE_FORMATS, so exporter validation permits this. model_calib.layerwise_calibrate then calls LayerwiseExporter.export_layer, whose export handlers replace each live floating weight with a packed uint8 Parameter (unified_export_hf.py::_export_quantized_weight). _restore_source_weights copies into the old Parameter objects, not those replacements. Scoring subsequently runs against packed/export-form modules and fails, leaving the caller's model unusable; restore_unconverted has already been dropped.

Could the library candidate preflight enforce the same refusal as examples/hf_ptq/hf_ptq.py::quantize_main, rather than letting calibration export a candidate? Extend refusal coverage with an enabled export directory and a discoverable decoder model, checking unchanged modules/parameters, no written layer shards, and a successful retry after dropping export_dir.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in f75da89. auto_quantize now refuses an enabled layerwise.export_dir on any candidate, in the per-candidate checks that run before the model is converted (_refuse_layerwise_dirs in algorithms.py, called beside the AWQ and distributed refusals in model_quant.py), so the refusal leaves the model untouched and needs no rollback. This is the library version of hf_ptq.py's refusal for AutoQuantize recipes. It is not specific to GPTQ or imatrix: on the previous head, an FP8 max candidate with an export directory on a tiny Llama also wrote both layer shards, replaced 14 weights with FP8 export tensors and failed in scoring with the model left converted, the same as a Q8_0 GPTQ candidate. The weight-rewriting candidates' layerwise checkpoint_dir refusal moved into the same helper, so it now also refuses before conversion.

test_huggingface.py::test_auto_quantize_refuses_layerwise_export_before_converting covers that FP8 case on a tiny Llama, whose decoder layers layerwise calibration discovers (lm_head is excluded so the export would run). The search raises, module types and weights are unchanged, the export directory is never created, and the same model then searches without export_dir. With the refusal reverted, the test fails in scoring on the exported FP8 weights.

…verting

An enabled layerwise.export_dir exports each calibrated layer in place,
replacing its weights with packed export tensors, so scoring then runs, and
fails on, the exported model, which is left converted. auto_quantize now
refuses it on every candidate in the per-candidate checks that run before
conversion, as hf_ptq.py already does for AutoQuantize recipes. The
weight-rewriting candidates' layerwise checkpoint_dir refusal moves into the
same helper, so it also refuses before conversion.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Comment: directory preflight still has two paths that let candidate calibration invalidate source-weight isolation.

Needs action:

  • Fix mixed-search layerwise resume in algorithms.py: preserve restoration across parameter replacements, or reject full-state resumes when any candidate rewrites weights. Add partial-max-resume → GPTQ → BF16 coverage; see inline.
  • 💬 Author added pre-conversion export checks in f75da89ab — also reject export_dir="", which exports into the working directory. Extend test_huggingface.py; see inline.

No action needed:

  • ✔️ Resolved the prior source-weight rollback, exact BF16 replay, and Conv1D lifetime findings. Existing-test edits preserve coverage.
  • Snapshot/replay reuses GGML pins, and the PR justifies it over whole-model copies and non-mutating calibrators.

raise NotImplementedError(
f"AutoQuantize candidate {recipe} rewrites weights with a layerwise "
"checkpoint_dir, whose resume would leave the rewritten weights in the model. "
"Drop checkpoint_dir; the AutoQuantize checkpoint resumes the search."

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P1] A non-rewriting candidate can also invalidate the snapshot's Parameter references. For example, an FP8/max candidate with enabled layerwise checkpointing passes this guard and sorts before Q8_0/GPTQ. When the max candidate resumes a partial checkpoint with the default calib_mutates_weights=True, utils/layerwise_calib.py::_CheckpointState.full_restore() calls layer.load_state_dict(..., assign=True), replacing the earlier layers' Parameters. The later GPTQ candidate rewrites those replacements, but _restore_source_weights() copies into the obsolete Parameters captured before the calibration loop. Scoring therefore sees rewritten source weights, and even a BF16 selection returns them modified.

Could restoration address the current module/parameter slots, or could preflight refuse full-state layerwise resumes across all candidates when the search has actual rewritten targets? Keep max-only searches and quantizer-only resumes working where possible. Add a partial max-checkpoint resume followed by GPTQ with BF16 selection, asserting that every source weight remains unchanged.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in a9f070f by refusing it before conversion. When any candidate rewrites weights, a layerwise checkpoint_dir whose resume restores the layers' weights (calib_mutates_weights, true by default) is now refused on every candidate, not only on the rewriting one, so a max candidate's resume can no longer replace the Parameters the restore holds. Searches without a rewriting candidate keep layerwise resume, and a candidate with calib_mutates_weights=False, whose resume restores only quantizer state, is still accepted. I chose the refusal over restoring through module slots: a weight-restoring layerwise resume is the one path that replaces Parameters during a search, and an AutoQuantize search resumes from its own checkpoint.

test_huggingface.py::test_auto_quantize_refuses_layerwise_dirs_before_converting[checkpoint_dir_beside_gptq] searches an FP8 max candidate with a checkpoint directory beside a Q8_0 GPTQ candidate on a tiny Llama. The search is refused before conversion, so no resume, GPTQ or BF16 selection runs: the modules and weights are unchanged, nothing is written, and the model searches again. With the refusal reverted, the search runs and writes the max candidate's checkpoint.

continue
# The export replaces each calibrated layer's weights with packed ones, which scoring and
# the other candidates' calibrations would then run.
if layerwise.get("export_dir"):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P2] The new export refusal tests truthiness, but model_calib.layerwise_calibrate() tests export_dir is not None. LayerwiseConfig.export_dir accepts an empty string, and LayerwiseExporter.__init__ converts it to Path(""), i.e. the working directory. Thus an enabled export_dir="" bypasses this preflight, writes layer shards into the working directory, replaces live weights, and reaches the same scoring failure the latest fix is intended to prevent. This also differs from hf_ptq.py::quantize_main, whose refusal uses is not None.

Use the same presence check here, or reject empty paths during configuration validation. Extend the new tiny-Llama refusal test with an empty-string directory and verify that no shards are written and the model remains retryable.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in a9f070f: the preflight now refuses any export_dir that is not None, the presence check layerwise_calibrate and hf_ptq.py use. The tiny-Llama refusal test is now parametrized with export_dir="" and runs in a temporary working directory. It checks that nothing is written there, that the modules and weights are unchanged, and that the model searches again. With the truthiness check restored, that case fails.

An empty layerwise.export_dir still exports, into the working directory,
so the AutoQuantize preflight now refuses any export_dir that is set. A
layerwise checkpoint_dir whose resume restores the layers' weights loads
them into new Parameters, which the source-weight restore after a
weight-rewriting candidate's calibration does not reach, so such a resume
is now refused on every candidate of a search with a rewriting candidate,
not only on the rewriting one. Searches without a rewriting candidate, and
resumes with calib_mutates_weights=False, are unaffected.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Comment: directory preflight still permits destructive export and resume through supported nested LayerwiseConfig inputs.

Needs action:

  • Normalize layerwise options in algorithms.py::_refuse_layerwise_dirs using the existing config schema; extend test_huggingface.py’s refusal/retry cases to nested LayerwiseConfig objects. See inline.

No action needed:

  • ✔️ Resolved since the last review: empty-string exports and cross-candidate full-state resumes for dictionary configurations. Earlier isolation, rollback, BF16 replay, and Conv1D lifetime fixes remain present.
  • Existing-test edits are justified: both fused-expert iterator assertions and HFModel defaults remain intact; save/restore coverage grows.
  • Snapshot/replay reuses existing GGML pins, and the PR explains why whole-model copies and non-mutating calibrators were rejected.

"""Refuse layerwise directories whose contents would outlive the candidate's calibration."""
for entry in _algorithm_entries(recipe.config.algorithm):
layerwise = entry.get("layerwise")
if not isinstance(layerwise, dict) or not layerwise.get("enable"):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P1] Normalize schema-backed layerwise options before this dict-only guard. A candidate can supply algorithm={'method': 'max', 'layerwise': LayerwiseConfig(enable=True, export_dir='')}: QuantizeConfig.algorithm preserves the nested object, so this branch skips it. During calibration, QuantizeAlgorithmConfig._coerce_layerwise accepts and normalizes that same object; mode.py::wrapped_calib_func then runs layerwise export. LayerwiseExporter.export_layer reaches _export_quantized_weight, replacing live Parameters before scoring and writing shards into the working directory. The conversion rollback is already unavailable at that point.

The same representation with checkpoint_dir='checkpoint' on a max candidate beside GGML/GPTQ bypasses the mixed-search resume refusal; _CheckpointState.full_restore(..., assign=True) can replace the Parameters retained by the source snapshot.

Could this helper use the existing LayerwiseConfig coercion/validation rather than isinstance(layerwise, dict)? Extend the tiny-Llama refusal/retry test with nested LayerwiseConfig objects for export (including the empty path) and mixed-search checkpointing, asserting refusal before conversion or file writes.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 1808e13. _refuse_layerwise_dirs now validates each entry's layerwise options with LayerwiseConfig.model_validate, the schema calibration applies, replacing the dict-only branch. A dict and a nested LayerwiseConfig object therefore reach the same export_dir and checkpoint_dir checks. The tiny-Llama refusal test now runs each of its three cases (an export path, an empty export path, and a max candidate's checkpoint beside Q8_0 GPTQ) with the options as a dict and as a LayerwiseConfig. All six are refused before conversion, with nothing written, and the same model searches again. With the dict-only check restored, the three LayerwiseConfig cases fail.

…e preflight

A dict algorithm keeps a nested LayerwiseConfig object as given, which the
dict-only check skipped, so an export_dir or a weight-restoring
checkpoint_dir set through one still reached calibration. The preflight now
validates each entry's layerwise options with LayerwiseConfig, the schema
calibration applies, whatever form they come in.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

LGTM: prior correctness concerns are resolved, including schema-backed layerwise preflight; human sign-off remains for justified existing-test edits.

Needs action:

  • Confirm the existing-test edits in test_fused_experts.py, test_huggingface.py, and test_autoquant.py: both iterator assertions and HFModel defaults remain intact, while save/restore and failure coverage expand.

No action needed:

  • ✔️ Resolved since earlier reviews: nested-config export/resume preflight, exception-safe rollback, and exact BF16 replay. Conv1D lifetime and selective snapshot safeguards remain intact.
  • Snapshot/replay reuses GGML pins; the PR justifies it over whole-model/state-dict copies and non-mutating calibration.
  • Tests were inspected, not executed; reported benchmarks were not independently reproduced. Baseline reference in prior review history: bc5adc6d; current base SHA was not independently verified.

meenchen and others added 2 commits November 5, 2026 16:35
Conflicts, both additions on each side kept:
- modelopt/torch/quantization/ggml/common.py: pinned_importance beside
  PIN_OVERRIDES_WEIGHT.
- tests/unit/torch/quantization/test_ggml_gptq.py: the ggml.common imports.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Main now re-encodes a weight under the pinned scales and importance of a
format that keeps them (the IQ formats), ahead of the pinned-rows lookup
that honours PIN_OVERRIDES_WEIGHT. During an AutoQuantize search the weight
is the source weight, so a GPTQ or imatrix candidate fake-quantized the
source weight under its scales instead of replaying its own codes. The
override now takes precedence there too.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

LGTM: prior correctness concerns are resolved; human sign-off remains for the justified existing-test edits.

Needs action:

  • Confirm the existing-test edits in test_fused_experts.py, test_huggingface.py, and test_autoquant.py: both iterator assertions and HFModel defaults remain intact, while save/restore and failure coverage expand.

No action needed:

  • ✔️ Resolved the nested layerwise export/resume preflight, source-weight and conversion rollback, calibration-mode cleanup, exact BF16 replay, selective snapshots, and Conv1D weight-lifetime findings.
  • Snapshot/replay reuses existing GGML pins; the PR justifies it over full-model/state-dict copies and non-mutating calibrators.
  • The current pinned-scales path honors replay before re-encoding, and the BF16 regression exercises it.

@meenchen
meenchen merged commit 58ea3f7 into main Nov 6, 2026
50 checks passed
@meenchen
meenchen deleted the weimingc/autoquant-weight-isolation branch November 6, 2026 04:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants