From 33b0dd91256df224553f415c38027e36d1d3de0e Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Thu, 29 Oct 2026 17:43:22 -0700
Subject: [PATCH 1/7] Calibrate and score AutoQuantize formats in parallel
 shards

AutoQuantize calibrates and scores every searched format in one process,
so its runtime is the sum over formats. The gradient scorer replays each
candidate against the same unquantized layer output and gradient, so a
(group, format) score does not depend on the other formats: the formats
can be split across processes without changing any score.

auto_quantize(format_shard=(index, count)) calibrates and scores only
every count-th searched format (the fixed recipe on every shard, as it
shapes the scoring reference) and saves the shard's scores without
solving. merge_auto_quantize_checkpoints() checks that the shards ran
the same search, unions their calibration states and scores, and writes
an ordinary checkpoint that restores, solves and re-solves like the
unsharded search.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 CHANGELOG.rst                                 |   1 +
 modelopt/torch/quantization/algorithms.py     | 154 ++++++++++++++++--
 modelopt/torch/quantization/model_quant.py    |  33 +++-
 .../unit/torch/quantization/test_autoquant.py |  43 ++++-
 4 files changed, 216 insertions(+), 15 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 70d228bd9ec..062c3e90f78 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -19,6 +19,7 @@ Changelog
 - 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}]``), calibrating and scoring each from the original weights and keeping the selected candidates' codes in the returned model. A search with such candidates that fails before calibrating any of them leaves the model unconverted, so it can be retried on the same model.
 - ``local_hessian`` now calibrates the GGML IQ formats, which it previously left unchanged: like imatrix it encodes each weight once, without layerwise calibration, and pins the codes, but picks each group's local scale by the layer-output error under its inputs' 256-column Hessian, damped by the new ``perc_damp`` (default 0.3). Use ``{method: local_hessian}`` in place of ``{method: imatrix}``, alone or after GPTQ; it helps most at IQ2_XXS and IQ2_XS, and like imatrix it skips the max-calibration forward pass when no activation quantizer needs one.
+- Add ``format_shard=(index, count)`` to ``mtq.auto_quantize`` to calibrate and score the searched formats in parallel processes, and ``mtq.merge_auto_quantize_checkpoints`` to combine the shard checkpoints into one that solves (and re-solves for other budgets) exactly like the unsharded search. Gradient and KL-divergence methods only.
 - GPTQ now also quantizes HuggingFace fused MoE experts, each with its own Hessian from the tokens routed to it; they were previously skipped. The built-in IQ1 and IQ2 recipes therefore run GPTQ on the experts too, and large MoEs need layerwise calibration to hold one layer's expert Hessians at a time.
 - GPTQ no longer runs a max-calibration forward when no activation quantizer needs one, and ``hf_ptq.py`` now sizes the default layerwise calibration batch to what one decoder layer's forward fits in GPU memory, up to 32,768 tokens (64 samples at ``--calib_seq 512``), instead of 1. Layerwise GPTQ on HF fused-MoE models runs over 10x faster; pass ``--batch_size 1`` for the old batching.
 - Add a checkpoint-specific NVIDIA Nemotron 3.5 Super VL AutoQuantize recipe with a 4.9-effective-bit NVFP4-MSE/FP8/BF16 search and fixed NVFP4 MTP experts. Select it under ``models/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/auto_quantize/``.
diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py
index 1fb6c354a55..7cb29141c87 100644
--- a/modelopt/torch/quantization/algorithms.py
+++ b/modelopt/torch/quantization/algorithms.py
@@ -593,6 +593,8 @@ def __init__(
         self.cost_weight = cost_weight
         self.allow_no_quant = allow_no_quant
         self.is_fixed = fixed_recipe is not None
+        # A format shard scores only its own formats; None scores every choice.
+        self._scored_recipes: set | None = None
 
         # Module hashes depend on object identity, so sets can produce different orders per rank.
         self.quant_modules = list(dict.fromkeys(quant_modules or []))
@@ -698,6 +700,13 @@ def restore_active_quantizers(self) -> None:
         assert isinstance(active, QuantRecipe)
         self._apply_quantizer_choice(active)
 
+    @property
+    def scored_choices(self) -> list[QuantRecipe]:
+        """Choices whose sensitivity this process scores."""
+        if self._scored_recipes is None:
+            return self.choices
+        return [recipe for recipe in self.choices if recipe in self._scored_recipes]
+
     @property
     def solver_choices(self) -> list[QuantRecipe]:
         """Return choices exposed to the LP after removing an internal no-quant baseline."""
@@ -894,6 +903,7 @@ def default_search_config(self):
             "cost_model": COST_MODEL_WEIGHT,
             "cost": {},
             "active_moe_expert_ratio": None,
+            "format_shard": None,
         }
 
     @property
@@ -913,6 +923,7 @@ def default_state_dict(self) -> SearchStateDict:
             "disabled_layers": None,
             "candidate_stats": defaultdict(dict),
             "quantizer_states": {},
+            "format_shard": None,
             "best": {"recipe": {}, "constraints": {}, "score": float("inf"), "is_satisfied": False},
         }
 
@@ -926,6 +937,15 @@ def sanitize_search_config(self, config: SearchConfig | None) -> SearchConfig:
         assert config["forward_step"] is not None, (
             "`forward_step` must be provided for `auto_quantize`."
         )
+        if config["format_shard"] is not None:
+            index, count = config["format_shard"]
+            if not (isinstance(index, int) and isinstance(count, int) and 0 <= index < count):
+                raise ValueError("format_shard must be (index, count) with 0 <= index < count.")
+            if self.method_name not in ("gradient", "kl_div"):
+                raise ValueError("format_shard supports method='gradient' or 'kl_div'.")
+            if config["checkpoint"] is None:
+                raise ValueError("format_shard needs a checkpoint to save the shard's scores.")
+            config["format_shard"] = (index, count)
         return config
 
     def validate_search_input(self, constraints, config) -> None:
@@ -1279,18 +1299,14 @@ def initialize_candidate_stats(self):
                 continue
 
             formats, raw_scores, scores, costs = [], [], [], []
-            prev_score = float("inf")
+            scored = set(hparam.scored_choices)
             for recipe in hparam.solver_choices:
                 formats.append(recipe)
-
-                score = hparam.get_score(recipe)
-                cost = hparam.get_cost(recipe)
-
-                raw_scores.append(score)
-                score = min(score, prev_score)  # TODO: Should we get rid of this?
-                scores.append(score)
-                costs.append(cost)
-                prev_score = score
+                # A format shard leaves other shards' formats unscored; their merge fills them.
+                unscored = recipe not in scored and not recipe.is_no_quant
+                raw_scores.append(None if unscored else hparam.get_score(recipe))
+                costs.append(hparam.get_cost(recipe))
+            scores = _monotone_scores(raw_scores)
 
             self.candidate_stats[name]["formats"] = formats
             self.candidate_stats[name]["scores"] = scores
@@ -1313,6 +1329,30 @@ def _run_func(self, func, num_iters=1, desc=""):
         ):
             func(self.model, data)
 
+    def _owned_recipes(self, quant_recipe_hparams, search_recipes):
+        """The searched formats this format shard calibrates and scores; None without sharding.
+
+        Formats are dealt round-robin in canonical (compression, then config) order. The fixed
+        recipe is calibrated by every shard because it shapes the scoring reference.
+        """
+        if self.format_shard is None:
+            return None
+        index, count = self.format_shard
+        configurable = [hparam for hparam in quant_recipe_hparams if hparam.is_configurable]
+        searched = [
+            recipe
+            for recipe in search_recipes
+            if not recipe.is_no_quant and any(recipe in hparam.choices for hparam in configurable)
+        ]
+        if count > len(searched):
+            raise ValueError(
+                f"format_shard count {count} exceeds {len(searched)} searched formats."
+            )
+        owned = set(searched[index::count])
+        for hparam in configurable:
+            hparam._scored_recipes = owned
+        return owned
+
     def _snapshot_source_weights(self, quant_recipe_hparams, search_recipes):
         """CPU copies of the weights candidates' calibrations will rewrite, or None if none will.
 
@@ -1542,7 +1582,26 @@ def before_search(self):
         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)
+        format_shard = self.config["format_shard"]
+        restored_format_shard = self.format_shard
+        if restored_format_shard is not None:
+            restored_format_shard = tuple(restored_format_shard)
+        if has_restored_calibration_or_scores and restored_format_shard != format_shard:
+            raise ValueError(
+                f"Checkpoint holds format shard {restored_format_shard}, not {format_shard}. "
+                "Combine shard checkpoints with merge_auto_quantize_checkpoints() first."
+            )
+        self.format_shard = format_shard
+        owned_recipes = self._owned_recipes(quant_recipe_hparams, search_recipes)
+
+        source_weights = self._snapshot_source_weights(
+            quant_recipe_hparams,
+            [
+                r
+                for r in search_recipes
+                if owned_recipes is None or r in {*owned_recipes, fixed_recipe}
+            ],
+        )
         # Free the weights conversion replaced before calibration and scoring.
         self.restore_unconverted = None
         QuantRecipe.disable_folding_pqs_to_weights()
@@ -1553,6 +1612,8 @@ def before_search(self):
             for recipe in search_recipes:
                 if recipe == QuantRecipe(quant_cfg=None):  # No-quant format
                     continue
+                if owned_recipes is not None and recipe not in {*owned_recipes, fixed_recipe}:
+                    continue
 
                 for hparam in quant_recipe_hparams:
                     hparam.set_calibration_recipe(recipe)
@@ -1721,6 +1782,13 @@ def run_search_with_stats(self, max_weight_size, verbose=False):
     def run_search(self):
         """Search for the best per-layer quantization configuration and return the best model and configuration."""
         verbose = self.config["verbose"]
+        if self.format_shard is not None:
+            if verbose:
+                print_rank_0(
+                    f"AutoQuantize: format shard {self.format_shard} scored; merge the shard "
+                    "checkpoints with merge_auto_quantize_checkpoints() to solve."
+                )
+            return
         assert "effective_bits" in self.constraints and (
             set(self.constraints) <= AUTO_QUANTIZE_CONSTRAINT_KEYS
         ), (
@@ -1786,6 +1854,66 @@ def run_search(self):
         QuantRecipe.fold_pqs_to_weights(self.model)
 
 
+def _monotone_scores(raw_scores):
+    """Scores never increasing with compression; formats left unscored (None) stay None."""
+    scores, prev_score = [], float("inf")
+    for score in raw_scores:
+        if score is not None:
+            score = min(score, prev_score)  # TODO: Should we get rid of this?
+            prev_score = score
+        scores.append(score)
+    return scores
+
+
+_SHARD_INVARIANT_STATE = (
+    "method",
+    "cost_model",
+    "cost",
+    "active_moe_expert_ratio",
+    "quantization_formats_signature",
+    "fixed_quantization_config_signature",
+    "module_search_space_signature",
+    "resolved_search_setup_signature",
+)
+
+
+def merge_auto_quantize_shards(shard_states: Sequence[dict]) -> dict:
+    """Combine the states of every format shard of one search into an unsharded search state."""
+    ids = [tuple(state.get("format_shard") or ()) for state in shard_states]
+    if sorted(ids) != [(index, len(ids)) for index in range(len(ids))]:
+        raise ValueError(
+            f"Merge needs one checkpoint per format shard (0..count-1, count), got {ids}."
+        )
+    shards = [state for _, state in sorted(zip(ids, shard_states), key=lambda pair: pair[0])]
+    merged = copy.deepcopy(shards[0])
+    for state in shards[1:]:
+        for key in _SHARD_INVARIANT_STATE:
+            if state.get(key) != merged.get(key):
+                raise ValueError(f"Format shards were run with different search setups ({key}).")
+        if state["candidate_stats"].keys() != merged["candidate_stats"].keys():
+            raise ValueError("Format shards were run with different quantization groups.")
+        merged["quantizer_states"].update(state["quantizer_states"])
+        for name, stats in state["candidate_stats"].items():
+            target = merged["candidate_stats"][name]
+            if [str(f) for f in stats["formats"]] != [str(f) for f in target["formats"]]:
+                raise ValueError(f"Format shards list different formats for {name}.")
+            for i, score in enumerate(stats["raw_scores"]):
+                current = target["raw_scores"][i]
+                if score is None or current == score:
+                    continue
+                # Every shard scores BF16 and fixed groups alike (0); a sharded format only once.
+                if current is not None:
+                    raise ValueError(f"Two format shards scored {target['formats'][i]} for {name}.")
+                target["raw_scores"][i] = score
+    for name, stats in merged["candidate_stats"].items():
+        if any(score is None for score in stats["raw_scores"]):
+            raise ValueError(f"No format shard scored every format of {name}.")
+        stats["scores"] = _monotone_scores(stats["raw_scores"])
+    merged["format_shard"] = None
+    merged["best"] = {"recipe": {}, "constraints": {}, "score": float("inf"), "is_satisfied": False}
+    return merged
+
+
 def _get_auto_quantize_score(grad_output, output_diff):
     x = grad_output.float() * output_diff.float()
     return x.clamp(-1e10, 1e10).square().sum()
@@ -1976,7 +2104,7 @@ def forward(self, module: nn.Module, *args, **kwargs):
             return output
 
         output_diffs = self._replay_candidates(
-            module, base, lambda hparam: hparam.choices, *args, **kwargs
+            module, base, lambda hparam: hparam.scored_choices, *args, **kwargs
         )
         self._register_candidate_score_hook(module, base, output_diffs)
         return output
@@ -2267,7 +2395,7 @@ def set_to_unquantized():
             ):
                 if not isinstance(hparam, QuantRecipeHparam):
                     continue
-                for recipe in hparam.choices:
+                for recipe in hparam.scored_choices:
                     if recipe == QuantRecipe(quant_cfg=None):
                         continue
                     hparam.active = recipe
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index 1b117d97461..4024e9b4172 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -42,7 +42,7 @@
     set_quantizer_attributes_partial,
     set_quantizer_by_cfg,
 )
-from modelopt.torch.utils import atomic_print
+from modelopt.torch.utils import atomic_print, safe_load, safe_save
 
 from . import model_calib
 from ._auto_quantize_cost import COST_MODEL_KV_CACHE
@@ -52,6 +52,7 @@
     _algorithm_methods,
     _refuse_layerwise_dirs,
     _rewrites_weights,
+    merge_auto_quantize_shards,
 )
 from .algorithms import get_auto_quantize_config as _get_auto_quantize_config
 from .config import QuantizeAlgoCfgType
@@ -75,6 +76,7 @@
     "enable_quantizer",
     "fold_weight",
     "get_auto_quantize_config",
+    "merge_auto_quantize_checkpoints",
     "postprocess_amax",
     "print_quant_summary",
     "quantize",
@@ -844,6 +846,7 @@ def auto_quantize(
     checkpoint: str | None = None,
     module_search_spaces: list[dict[str, Any]] | None = None,
     fixed_quantization_config: dict[str, Any] | str | None = None,
+    format_shard: tuple[int, int] | None = None,
 ):
     r"""Perform optimal per-layer quantization by searching for the best quantization formats per-layer.
 
@@ -1039,6 +1042,13 @@ def forward_backward_step(model, batch) -> None:
             active while searched modules are scored, is calibrated only with its own algorithm,
             and remains part of the effective-bits numerator and denominator. This is one
             integrated AutoQuantize operation, not staged PTQ followed by AutoQuantize.
+        format_shard: Optional ``(index, count)`` to calibrate and score only every ``count``-th
+            searched format starting at ``index`` (in canonical compression order), so the
+            formats can run in ``count`` parallel processes on the same model, data and
+            arguments. Each shard saves its scores to ``checkpoint`` without solving; combine the
+            shard checkpoints with :func:`merge_auto_quantize_checkpoints` and pass the result as
+            ``checkpoint`` to solve. Scores are per (group, format), so the merged search is
+            the unsharded one. Gradient and KL-divergence methods only.
 
     Returns: A tuple (model, state_dict) where ``model`` is the searched and quantized model and
         ``state_dict`` contains the history and detailed stats of the search procedure.
@@ -1101,6 +1111,8 @@ def forward_backward_step(model, batch) -> None:
 
     method_name, method_config = _parse_auto_quantize_method(method)
     is_kv_search = constraints is not None and constraints.get("cost_model") == COST_MODEL_KV_CACHE
+    if is_kv_search and format_shard is not None:
+        raise ValueError("format_shard is not supported for KV-cache AutoQuantize.")
     if is_kv_search:
         assert constraints is not None
         assert quantization_formats is not None
@@ -1246,6 +1258,7 @@ def forward_backward_step(model, batch) -> None:
         "disabled_layers": disabled_layers,
         "verbose": verbose,
         "checkpoint": checkpoint,
+        "format_shard": format_shard,
     }
     if method_config:
         # Only the selected method's declared options are accepted; core inputs (loaders,
@@ -1298,6 +1311,24 @@ def forward_backward_step(model, batch) -> None:
     return model, searcher.state_dict()
 
 
+def merge_auto_quantize_checkpoints(shard_checkpoints: Sequence[str], checkpoint: str) -> None:
+    """Combine the checkpoints of every ``format_shard`` of one search into one checkpoint.
+
+    Pass ``checkpoint`` to :func:`auto_quantize` with the shards' other arguments (and no
+    ``format_shard``) to restore every format's calibration and scores and solve. Under
+    distributed scoring, merge each rank's files separately.
+
+    Args:
+        shard_checkpoints: The checkpoint file of each shard ``0..count-1``, in any order.
+        checkpoint: Output file; must not exist.
+    """
+    if os.path.exists(checkpoint):
+        raise FileExistsError(f"{checkpoint} exists; merge into a new checkpoint file.")
+    merged = merge_auto_quantize_shards([safe_load(path) for path in shard_checkpoints])
+    os.makedirs(os.path.dirname(checkpoint) or ".", exist_ok=True)
+    safe_save(merged, checkpoint)
+
+
 def get_auto_quantize_config(search_state, constraints=None, verbose=False):
     """Build a flat quant config from auto_quantize search_state.
 
diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py
index 9c25a1117e2..dcb296e06c1 100644
--- a/tests/unit/torch/quantization/test_autoquant.py
+++ b/tests/unit/torch/quantization/test_autoquant.py
@@ -868,6 +868,47 @@ def loss_func(output):
     assert torch.allclose(output_ref, output_test)
 
 
+@pytest.mark.parametrize("method", ["gradient", "kl_div"])
+def test_format_shards_merge_into_the_unsharded_search(tmp_path, method):
+    source = TransformerBlock()
+    data = [source.get_input() for _ in range(2)]
+    formats = [mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_DEFAULT_CFG, mtq.FP8_DEFAULT_CFG]
+
+    def search(checkpoint, format_shard=None):
+        return mtq.auto_quantize(
+            copy.deepcopy(source),
+            constraints={"effective_bits": 6.0},
+            quantization_formats=formats,
+            data_loader=data,
+            forward_step=lambda model, batch: model(batch),
+            loss_func=(lambda output, batch: output.sum()) if method == "gradient" else None,
+            num_calib_steps=2,
+            num_score_steps=2,
+            checkpoint=str(checkpoint),
+            format_shard=format_shard,
+            method=method,
+        )
+
+    _, expected = search(tmp_path / "full.pth")
+    shards = [tmp_path / f"shard{i}.pth" for i in range(2)]
+    for i, shard in enumerate(shards):
+        _, state = search(shard, (i, 2))
+        assert not state["best"]["recipe"]
+    # Shards hold disjoint formats: 2 of the 3 searched formats on shard 0, 1 on shard 1.
+    calibrated = [set(safe_load(shard)["quantizer_states"]) for shard in shards]
+    assert not calibrated[0] & calibrated[1] and len(calibrated[0] | calibrated[1]) == 3
+
+    merged = tmp_path / "merged.pth"
+    model_quant.merge_auto_quantize_checkpoints(list(reversed(shards)), str(merged))
+    with pytest.raises(ValueError, match="format shard"):
+        search(shards[0])
+    _, state = search(merged)
+
+    for name, stats in expected["candidate_stats"].items():
+        assert state["candidate_stats"][name]["raw_scores"] == pytest.approx(stats["raw_scores"])
+    assert state["best"]["recipe"] == expected["best"]["recipe"]
+
+
 def test_auto_quantize_disable_layers():
     model = TransformerBlock()
 
@@ -1083,7 +1124,7 @@ def test_gradient_scoring_tracks_reused_module_invocations():
 
     class TestHparam:
         is_configurable = True
-        choices = [no_quant_recipe, quant_recipe]
+        choices = scored_choices = [no_quant_recipe, quant_recipe]
         active = no_quant_recipe
 
     class ScoreModule(torch.nn.Module):

From 46eefbc0ddc8a14bffd0e673a7f1ed1ef4cb9f04 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Tue, 3 Nov 2026 22:27:18 -0800
Subject: [PATCH 2/7] Type scored_choices like solver_choices

mypy rejected the property's return types (return-value, misc). Behavior is unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 modelopt/torch/quantization/algorithms.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py
index 7cb29141c87..9320fedf0e8 100644
--- a/modelopt/torch/quantization/algorithms.py
+++ b/modelopt/torch/quantization/algorithms.py
@@ -703,9 +703,12 @@ def restore_active_quantizers(self) -> None:
     @property
     def scored_choices(self) -> list[QuantRecipe]:
         """Choices whose sensitivity this process scores."""
-        if self._scored_recipes is None:
-            return self.choices
-        return [recipe for recipe in self.choices if recipe in self._scored_recipes]
+        recipes: list[QuantRecipe] = []
+        for recipe in self.choices:
+            assert isinstance(recipe, QuantRecipe)
+            if self._scored_recipes is None or recipe in self._scored_recipes:
+                recipes.append(recipe)
+        return recipes
 
     @property
     def solver_choices(self) -> list[QuantRecipe]:

From de2946894c3728bbe42ee8d878b909bea5dd5e47 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Thu, 5 Nov 2026 21:35:00 -0800
Subject: [PATCH 3/7] Calibrate the formats groups hold fixed in every format
 shard

A one-format module_search_spaces rule without no-quant keeps its format active while the
other groups are scored, so that format's calibration shapes their scores. A format shard
calibrated only its own formats and the fixed recipe, so a held format was calibrated by the
shard that owned it, or by none, and the other shards scored against the uncalibrated group:
the merged scores no longer matched the unsharded search. Every shard now calibrates the
formats non-configurable groups hold.

The merge also copies only the containers it fills instead of deep-copying the first shard's
calibrated tensors.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 modelopt/torch/quantization/algorithms.py     | 44 ++++++++++---------
 modelopt/torch/quantization/model_quant.py    |  9 ++--
 .../unit/torch/quantization/test_autoquant.py | 12 ++++-
 3 files changed, 39 insertions(+), 26 deletions(-)

diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py
index 9320fedf0e8..65d24bd2917 100644
--- a/modelopt/torch/quantization/algorithms.py
+++ b/modelopt/torch/quantization/algorithms.py
@@ -1332,14 +1332,15 @@ def _run_func(self, func, num_iters=1, desc=""):
         ):
             func(self.model, data)
 
-    def _owned_recipes(self, quant_recipe_hparams, search_recipes):
-        """The searched formats this format shard calibrates and scores; None without sharding.
+    def _other_shards_recipes(self, quant_recipe_hparams, search_recipes) -> set:
+        """The searched formats this format shard leaves to the others; none without sharding.
 
-        Formats are dealt round-robin in canonical (compression, then config) order. The fixed
-        recipe is calibrated by every shard because it shapes the scoring reference.
+        Formats are dealt round-robin in canonical (compression, then config) order, each
+        calibrated and scored by its own shard. Every shard calibrates the formats groups hold
+        fixed, since those shape the reference the other groups are scored against.
         """
         if self.format_shard is None:
-            return None
+            return set()
         index, count = self.format_shard
         configurable = [hparam for hparam in quant_recipe_hparams if hparam.is_configurable]
         searched = [
@@ -1354,7 +1355,13 @@ def _owned_recipes(self, quant_recipe_hparams, search_recipes):
         owned = set(searched[index::count])
         for hparam in configurable:
             hparam._scored_recipes = owned
-        return owned
+        held = {
+            recipe
+            for hparam in quant_recipe_hparams
+            if not hparam.is_configurable
+            for recipe in hparam.choices
+        }
+        return set(searched) - owned - held
 
     def _snapshot_source_weights(self, quant_recipe_hparams, search_recipes):
         """CPU copies of the weights candidates' calibrations will rewrite, or None if none will.
@@ -1595,16 +1602,10 @@ def before_search(self):
                 "Combine shard checkpoints with merge_auto_quantize_checkpoints() first."
             )
         self.format_shard = format_shard
-        owned_recipes = self._owned_recipes(quant_recipe_hparams, search_recipes)
-
-        source_weights = self._snapshot_source_weights(
-            quant_recipe_hparams,
-            [
-                r
-                for r in search_recipes
-                if owned_recipes is None or r in {*owned_recipes, fixed_recipe}
-            ],
-        )
+        other_shards_recipes = self._other_shards_recipes(quant_recipe_hparams, search_recipes)
+        calibration_recipes = [r for r in search_recipes if r not in other_shards_recipes]
+
+        source_weights = self._snapshot_source_weights(quant_recipe_hparams, calibration_recipes)
         # Free the weights conversion replaced before calibration and scoring.
         self.restore_unconverted = None
         QuantRecipe.disable_folding_pqs_to_weights()
@@ -1612,11 +1613,9 @@ def before_search(self):
         # Iterate over the search recipes and calibrate the quantizers for each recipe
         calibrated_new = False
         try:
-            for recipe in search_recipes:
+            for recipe in calibration_recipes:
                 if recipe == QuantRecipe(quant_cfg=None):  # No-quant format
                     continue
-                if owned_recipes is not None and recipe not in {*owned_recipes, fixed_recipe}:
-                    continue
 
                 for hparam in quant_recipe_hparams:
                     hparam.set_calibration_recipe(recipe)
@@ -1888,7 +1887,12 @@ def merge_auto_quantize_shards(shard_states: Sequence[dict]) -> dict:
             f"Merge needs one checkpoint per format shard (0..count-1, count), got {ids}."
         )
     shards = [state for _, state in sorted(zip(ids, shard_states), key=lambda pair: pair[0])]
-    merged = copy.deepcopy(shards[0])
+    # Copy the containers the merge fills, not the calibrated tensors (GGML codes) they hold.
+    merged = {
+        **shards[0],
+        "quantizer_states": dict(shards[0]["quantizer_states"]),
+        "candidate_stats": copy.deepcopy(shards[0]["candidate_stats"]),
+    }
     for state in shards[1:]:
         for key in _SHARD_INVARIANT_STATE:
             if state.get(key) != merged.get(key):
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index 4024e9b4172..617e3513b12 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -1045,10 +1045,11 @@ def forward_backward_step(model, batch) -> None:
         format_shard: Optional ``(index, count)`` to calibrate and score only every ``count``-th
             searched format starting at ``index`` (in canonical compression order), so the
             formats can run in ``count`` parallel processes on the same model, data and
-            arguments. Each shard saves its scores to ``checkpoint`` without solving; combine the
-            shard checkpoints with :func:`merge_auto_quantize_checkpoints` and pass the result as
-            ``checkpoint`` to solve. Scores are per (group, format), so the merged search is
-            the unsharded one. Gradient and KL-divergence methods only.
+            arguments; formats a group holds fixed are calibrated by every shard. Each shard
+            saves its scores to ``checkpoint`` without solving; combine the shard checkpoints
+            with :func:`merge_auto_quantize_checkpoints` and pass the result as ``checkpoint``
+            to solve. Scores are per (group, format), so the merged search is the unsharded
+            one. Gradient and KL-divergence methods only.
 
     Returns: A tuple (model, state_dict) where ``model`` is the searched and quantized model and
         ``state_dict`` contains the history and detailed stats of the search procedure.
diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py
index dcb296e06c1..4e21083d80b 100644
--- a/tests/unit/torch/quantization/test_autoquant.py
+++ b/tests/unit/torch/quantization/test_autoquant.py
@@ -879,6 +879,14 @@ def search(checkpoint, format_shard=None):
             copy.deepcopy(source),
             constraints={"effective_bits": 6.0},
             quantization_formats=formats,
+            # o_proj holds FP8, so its calibrated input scale shapes the mlp's scores.
+            module_search_spaces=[
+                {
+                    "module_name_patterns": ["*o_proj*"],
+                    "quantization_formats": [mtq.FP8_DEFAULT_CFG],
+                    "allow_no_quant": False,
+                }
+            ],
             data_loader=data,
             forward_step=lambda model, batch: model(batch),
             loss_func=(lambda output, batch: output.sum()) if method == "gradient" else None,
@@ -894,9 +902,9 @@ def search(checkpoint, format_shard=None):
     for i, shard in enumerate(shards):
         _, state = search(shard, (i, 2))
         assert not state["best"]["recipe"]
-    # Shards hold disjoint formats: 2 of the 3 searched formats on shard 0, 1 on shard 1.
+    # Each searched format is calibrated by one shard, o_proj's held FP8 by both.
     calibrated = [set(safe_load(shard)["quantizer_states"]) for shard in shards]
-    assert not calibrated[0] & calibrated[1] and len(calibrated[0] | calibrated[1]) == 3
+    assert len(calibrated[0] & calibrated[1]) == 1 and len(calibrated[0] | calibrated[1]) == 3
 
     merged = tmp_path / "merged.pth"
     model_quant.merge_auto_quantize_checkpoints(list(reversed(shards)), str(merged))

From 81bfb30f5bc22235152a531c7d43ca30b0db3d98 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Fri, 6 Nov 2026 10:29:37 -0800
Subject: [PATCH 4/7] Merge AutoQuantize format shards without the GPUs that
 saved them

A shard's checkpoint holds its calibration on the GPUs that produced it, so
merge_auto_quantize_checkpoints failed on any host with fewer GPUs: a CPU node,
or a one-GPU process merging shards that each spanned a node's GPUs. Memory-map
the shard files on the CPU instead. The restore moves each calibration onto the
model's devices, and the merge reads only the tensors it keeps.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 modelopt/torch/quantization/model_quant.py    | 10 ++++--
 .../torch/quantization/test_ggml_gptq_cuda.py | 35 +++++++++++++++++++
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index 617e3513b12..f37278dacae 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -1316,8 +1316,10 @@ def merge_auto_quantize_checkpoints(shard_checkpoints: Sequence[str], checkpoint
     """Combine the checkpoints of every ``format_shard`` of one search into one checkpoint.
 
     Pass ``checkpoint`` to :func:`auto_quantize` with the shards' other arguments (and no
-    ``format_shard``) to restore every format's calibration and scores and solve. Under
-    distributed scoring, merge each rank's files separately.
+    ``format_shard``) to restore every format's calibration and scores and solve. The shard files
+    are memory-mapped on the CPU, so any host can merge them, and the restore moves each
+    calibration onto the model's devices. Under distributed scoring, merge each rank's files
+    separately.
 
     Args:
         shard_checkpoints: The checkpoint file of each shard ``0..count-1``, in any order.
@@ -1325,7 +1327,9 @@ def merge_auto_quantize_checkpoints(shard_checkpoints: Sequence[str], checkpoint
     """
     if os.path.exists(checkpoint):
         raise FileExistsError(f"{checkpoint} exists; merge into a new checkpoint file.")
-    merged = merge_auto_quantize_shards([safe_load(path) for path in shard_checkpoints])
+    merged = merge_auto_quantize_shards(
+        [safe_load(path, map_location="cpu", mmap=True) for path in shard_checkpoints]
+    )
     os.makedirs(os.path.dirname(checkpoint) or ".", exist_ok=True)
     safe_save(merged, checkpoint)
 
diff --git a/tests/gpu/torch/quantization/test_ggml_gptq_cuda.py b/tests/gpu/torch/quantization/test_ggml_gptq_cuda.py
index e522b37152d..5407c991a39 100644
--- a/tests/gpu/torch/quantization/test_ggml_gptq_cuda.py
+++ b/tests/gpu/torch/quantization/test_ggml_gptq_cuda.py
@@ -132,3 +132,38 @@ def search(algorithm):
         assert torch.equal(param, original)
     # The same model searches again with a calibration that leaves its weights alone.
     search("max")
+
+
+def test_format_shards_saved_on_gpus_merge_on_a_host_without_them(tmp_path, monkeypatch):
+    torch.manual_seed(0)
+    source = torch.nn.Sequential(torch.nn.Linear(512, 256), torch.nn.Linear(256, 256)).cuda()
+    inputs = _correlated_inputs()
+    iq1_s = _iq1_s_config({"method": "gptq", "block_size": 256})
+    iq1_s["quant_cfg"][1]["cfg"]["effective_bits"] = GGML_FORMAT_REGISTRY["iq1_s"].effective_bits
+
+    def search(checkpoint, format_shard=None):
+        return mtq.auto_quantize(
+            copy.deepcopy(source),
+            constraints={"effective_bits": 6.0},
+            quantization_formats=[(iq1_s, "IQ1_S"), mtq.INT8_DEFAULT_CFG],
+            data_loader=[inputs],
+            forward_step=lambda m, x: m(x),
+            num_calib_steps=1,
+            num_score_steps=1,
+            method="kl_div",
+            checkpoint=str(checkpoint),
+            format_shard=format_shard,
+        )
+
+    reference, expected = search(tmp_path / "full.pth")
+    shards = [tmp_path / f"shard{i}.pth" for i in range(2)]
+    for i, shard in enumerate(shards):
+        search(shard, (i, 2))
+    # The shards hold GPU tensors (IQ1_S's pinned codes among them); merge them on a CPU node.
+    with monkeypatch.context() as patch:
+        patch.setattr(torch.cuda, "is_available", lambda: False)
+        mtq.merge_auto_quantize_checkpoints(shards, str(tmp_path / "merged.pth"))
+
+    model, state = search(tmp_path / "merged.pth")
+    assert state["best"]["recipe"] == expected["best"]["recipe"]
+    torch.testing.assert_close(model(inputs), reference(inputs), rtol=0, atol=0)

From 1f44a436d100fb3a1617723780dd8790dc02dae8 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Fri, 6 Nov 2026 14:32:35 -0800
Subject: [PATCH 5/7] Refuse merging differently scored shards and hand back a
 refused shard's model

- Shards record num_calib_steps and num_score_steps, and the merge refuses
  shards whose counts differ: scores sum over the batches, so shards scored
  over different counts are not comparable. Keeping the weights and batches
  the same stays with the caller, as the merge's docstring now says.
- A format shard refused before any calibration, e.g. for a count the search
  cannot split, hands back the unconverted model, as weight-rewriting searches
  already did, so the caller can retry on the same model.
- The merged checkpoint is created exclusively, so a file or symlink that
  appears at the output path after the existence check is refused rather
  than written through.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 modelopt/torch/quantization/algorithms.py     |  7 ++++++
 modelopt/torch/quantization/model_quant.py    | 19 +++++++++------
 .../unit/torch/quantization/test_autoquant.py | 23 +++++++++++++++----
 3 files changed, 38 insertions(+), 11 deletions(-)

diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py
index 65d24bd2917..e6276ae2fb1 100644
--- a/modelopt/torch/quantization/algorithms.py
+++ b/modelopt/torch/quantization/algorithms.py
@@ -927,6 +927,8 @@ def default_state_dict(self) -> SearchStateDict:
             "candidate_stats": defaultdict(dict),
             "quantizer_states": {},
             "format_shard": None,
+            "num_calib_steps": None,
+            "num_score_steps": None,
             "best": {"recipe": {}, "constraints": {}, "score": float("inf"), "is_satisfied": False},
         }
 
@@ -1602,6 +1604,9 @@ def before_search(self):
                 "Combine shard checkpoints with merge_auto_quantize_checkpoints() first."
             )
         self.format_shard = format_shard
+        # Scores sum over the batches, so shards merge only with the same workload.
+        self.num_calib_steps = self.config["num_calib_steps"]
+        self.num_score_steps = self.config["num_score_steps"]
         other_shards_recipes = self._other_shards_recipes(quant_recipe_hparams, search_recipes)
         calibration_recipes = [r for r in search_recipes if r not in other_shards_recipes]
 
@@ -1876,6 +1881,8 @@ def _monotone_scores(raw_scores):
     "fixed_quantization_config_signature",
     "module_search_space_signature",
     "resolved_search_setup_signature",
+    "num_calib_steps",
+    "num_score_steps",
 )
 
 
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index f37278dacae..ee642d576b0 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -1280,11 +1280,12 @@ def forward_backward_step(model, batch) -> None:
     searcher.validate_search_input(search_constraints, search_config)
 
     # A weight-rewriting candidate's refusals need the converted quantizers to tell which weights it
-    # rewrites. A failure before any candidate calibrates then hands back the unconverted modules,
-    # so the caller can retry on the same model, which a converted model would refuse. The saved
-    # modules hold every weight conversion replaces (HF Conv1D's), so other searches save none and
-    # the searcher drops them once those checks pass.
-    if rewrites_weights:
+    # rewrites, as does a format shard's check of its count against the searched formats. A failure
+    # before any candidate calibrates then hands back the unconverted modules, so the caller can
+    # retry on the same model, which a converted model would refuse. The saved modules hold every
+    # weight conversion replaces (HF Conv1D's), so other searches save none and the searcher drops
+    # them once those checks pass.
+    if rewrites_weights or format_shard is not None:
         searcher.restore_unconverted = _save_module_attributes(model)
     # Calibration turns pre-quant-scale folding off for the whole process. Hand back the policy
     # found here on every exit, also after a failure while scoring, so later calibrations fold.
@@ -1319,7 +1320,9 @@ def merge_auto_quantize_checkpoints(shard_checkpoints: Sequence[str], checkpoint
     ``format_shard``) to restore every format's calibration and scores and solve. The shard files
     are memory-mapped on the CPU, so any host can merge them, and the restore moves each
     calibration onto the model's devices. Under distributed scoring, merge each rank's files
-    separately.
+    separately. Shards that recorded different search setups (formats, groups, method, cost, or
+    calibration and scoring step counts) are refused; running every shard on the same weights
+    and batches is up to the caller.
 
     Args:
         shard_checkpoints: The checkpoint file of each shard ``0..count-1``, in any order.
@@ -1331,7 +1334,9 @@ def merge_auto_quantize_checkpoints(shard_checkpoints: Sequence[str], checkpoint
         [safe_load(path, map_location="cpu", mmap=True) for path in shard_checkpoints]
     )
     os.makedirs(os.path.dirname(checkpoint) or ".", exist_ok=True)
-    safe_save(merged, checkpoint)
+    # Created exclusively: a file or symlink that appeared since the check is refused, not written.
+    with open(checkpoint, "xb") as f:
+        safe_save(merged, f)
 
 
 def get_auto_quantize_config(search_state, constraints=None, verbose=False):
diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py
index 4e21083d80b..cbbf7392a52 100644
--- a/tests/unit/torch/quantization/test_autoquant.py
+++ b/tests/unit/torch/quantization/test_autoquant.py
@@ -874,9 +874,9 @@ def test_format_shards_merge_into_the_unsharded_search(tmp_path, method):
     data = [source.get_input() for _ in range(2)]
     formats = [mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_DEFAULT_CFG, mtq.FP8_DEFAULT_CFG]
 
-    def search(checkpoint, format_shard=None):
+    def search(checkpoint, format_shard=None, model=None, score_steps=2):
         return mtq.auto_quantize(
-            copy.deepcopy(source),
+            copy.deepcopy(source) if model is None else model,
             constraints={"effective_bits": 6.0},
             quantization_formats=formats,
             # o_proj holds FP8, so its calibrated input scale shapes the mlp's scores.
@@ -891,22 +891,37 @@ def search(checkpoint, format_shard=None):
             forward_step=lambda model, batch: model(batch),
             loss_func=(lambda output, batch: output.sum()) if method == "gradient" else None,
             num_calib_steps=2,
-            num_score_steps=2,
+            num_score_steps=score_steps,
             checkpoint=str(checkpoint),
             format_shard=format_shard,
             method=method,
         )
 
     _, expected = search(tmp_path / "full.pth")
+    # A refused shard count hands the model back, so shard 0 then runs on that same model.
+    model = copy.deepcopy(source)
+    with pytest.raises(ValueError, match="format_shard count"):
+        search(tmp_path / "refused.pth", (0, 99), model=model)
     shards = [tmp_path / f"shard{i}.pth" for i in range(2)]
     for i, shard in enumerate(shards):
-        _, state = search(shard, (i, 2))
+        _, state = search(shard, (i, 2), model=model if i == 0 else None)
         assert not state["best"]["recipe"]
     # Each searched format is calibrated by one shard, o_proj's held FP8 by both.
     calibrated = [set(safe_load(shard)["quantizer_states"]) for shard in shards]
     assert len(calibrated[0] & calibrated[1]) == 1 and len(calibrated[0] | calibrated[1]) == 3
 
+    # Scores sum over batches, so a shard scored over fewer is refused.
+    fewer = tmp_path / "fewer_steps.pth"
+    search(fewer, (1, 2), score_steps=1)
+    with pytest.raises(ValueError, match="num_score_steps"):
+        model_quant.merge_auto_quantize_checkpoints([shards[0], fewer], str(tmp_path / "x.pth"))
+
     merged = tmp_path / "merged.pth"
+    merged.symlink_to(tmp_path / "elsewhere.pth")  # planted after the existence check
+    with pytest.raises(FileExistsError):
+        model_quant.merge_auto_quantize_checkpoints(shards, str(merged))
+    assert not (tmp_path / "elsewhere.pth").exists()
+    merged.unlink()
     model_quant.merge_auto_quantize_checkpoints(list(reversed(shards)), str(merged))
     with pytest.raises(ValueError, match="format shard"):
         search(shards[0])

From cb6dea0a7f39b61b2c91a687424f673864f2bd8d Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Fri, 6 Nov 2026 16:34:19 -0800
Subject: [PATCH 6/7] Refuse resuming or merging shards whose step counts are
 not their own

- Restored calibration and scores keep the step counts they were computed
  with. A shard resumed under other counts is refused, so its merge compares
  the counts its scores came from; an unsharded search still reuses them, e.g.
  to re-solve another budget.
- The merge refuses shards that recorded no counts, as earlier versions saved
  them, since missing counts compare equal whatever the workloads.
- A search interrupted between calibration and scoring now resumes: safe_save
  stores candidate_stats as a dict, which scoring after the resume indexed by
  name as a defaultdict.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 modelopt/torch/quantization/algorithms.py     | 29 ++++++++++++++++---
 .../unit/torch/quantization/test_autoquant.py | 25 ++++++++++++++--
 2 files changed, 47 insertions(+), 7 deletions(-)

diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py
index e6276ae2fb1..38b5e53a18d 100644
--- a/modelopt/torch/quantization/algorithms.py
+++ b/modelopt/torch/quantization/algorithms.py
@@ -957,7 +957,10 @@ def validate_search_input(self, constraints, config) -> None:
         """Validate method-specific inputs before quantizing the model."""
 
     def load_search_checkpoint(self) -> bool:
-        return super().load_search_checkpoint(strict=False)
+        loaded = super().load_search_checkpoint(strict=False)
+        # safe_save stores the defaultdict as a dict, which scoring after a resume fills by name.
+        self.candidate_stats = defaultdict(dict, self.candidate_stats)
+        return loaded
 
     @staticmethod
     def _is_auto_quantize_module(module):
@@ -1604,9 +1607,21 @@ def before_search(self):
                 "Combine shard checkpoints with merge_auto_quantize_checkpoints() first."
             )
         self.format_shard = format_shard
-        # Scores sum over the batches, so shards merge only with the same workload.
-        self.num_calib_steps = self.config["num_calib_steps"]
-        self.num_score_steps = self.config["num_score_steps"]
+        # Scores sum over the batches, so restored calibration and scores keep the step counts they
+        # were computed with, which a shard's merge compares; a shard refuses to reuse them under
+        # other counts. An unsharded search reuses them as before, e.g. to re-solve a budget.
+        reused = {
+            "num_calib_steps": bool(self.quantizer_states),
+            "num_score_steps": bool(self.candidate_stats),
+        }
+        for key, is_reused in reused.items():
+            if not is_reused:
+                setattr(self, key, self.config[key])
+            elif format_shard is not None and getattr(self, key) != self.config[key]:
+                raise ValueError(
+                    f"This shard's checkpoint was computed with {key}={getattr(self, key)}, not "
+                    f"{self.config[key]}; resume it with the same counts or a new checkpoint."
+                )
         other_shards_recipes = self._other_shards_recipes(quant_recipe_hparams, search_recipes)
         calibration_recipes = [r for r in search_recipes if r not in other_shards_recipes]
 
@@ -1894,6 +1909,12 @@ def merge_auto_quantize_shards(shard_states: Sequence[dict]) -> dict:
             f"Merge needs one checkpoint per format shard (0..count-1, count), got {ids}."
         )
     shards = [state for _, state in sorted(zip(ids, shard_states), key=lambda pair: pair[0])]
+    for key in ("num_calib_steps", "num_score_steps"):
+        # Unrecorded counts would compare equal across shards whose scores scale differently.
+        if any(state.get(key) is None for state in shards):
+            raise ValueError(
+                f"Every format shard must record {key}; rerun shards saved without it."
+            )
     # Copy the containers the merge fills, not the calibrated tensors (GGML codes) they hold.
     merged = {
         **shards[0],
diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py
index cbbf7392a52..54e7f2fa23d 100644
--- a/tests/unit/torch/quantization/test_autoquant.py
+++ b/tests/unit/torch/quantization/test_autoquant.py
@@ -874,7 +874,7 @@ def test_format_shards_merge_into_the_unsharded_search(tmp_path, method):
     data = [source.get_input() for _ in range(2)]
     formats = [mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_DEFAULT_CFG, mtq.FP8_DEFAULT_CFG]
 
-    def search(checkpoint, format_shard=None, model=None, score_steps=2):
+    def search(checkpoint, format_shard=None, model=None, calib_steps=2, score_steps=2):
         return mtq.auto_quantize(
             copy.deepcopy(source) if model is None else model,
             constraints={"effective_bits": 6.0},
@@ -890,7 +890,7 @@ def search(checkpoint, format_shard=None, model=None, score_steps=2):
             data_loader=data,
             forward_step=lambda model, batch: model(batch),
             loss_func=(lambda output, batch: output.sum()) if method == "gradient" else None,
-            num_calib_steps=2,
+            num_calib_steps=calib_steps,
             num_score_steps=score_steps,
             checkpoint=str(checkpoint),
             format_shard=format_shard,
@@ -915,6 +915,25 @@ def search(checkpoint, format_shard=None, model=None, score_steps=2):
     search(fewer, (1, 2), score_steps=1)
     with pytest.raises(ValueError, match="num_score_steps"):
         model_quant.merge_auto_quantize_checkpoints([shards[0], fewer], str(tmp_path / "x.pth"))
+    # So are shards that recorded no counts, as earlier versions saved them.
+    unlabeled = [tmp_path / f"unlabeled{i}.pth" for i in range(2)]
+    for shard, path in zip(shards, unlabeled):
+        state = safe_load(shard)
+        del state["num_calib_steps"], state["num_score_steps"]
+        safe_save(state, path)
+    with pytest.raises(ValueError, match="num_calib_steps"):
+        model_quant.merge_auto_quantize_checkpoints(unlabeled, str(tmp_path / "y.pth"))
+    # A shard resumes its calibration and scores only under the counts that produced them.
+    interrupted = safe_load(shards[0])
+    interrupted["candidate_stats"] = {}  # as saved between calibration and scoring
+    safe_save(interrupted, tmp_path / "interrupted.pth")
+    with pytest.raises(ValueError, match="num_calib_steps"):
+        search(tmp_path / "interrupted.pth", (0, 2), calib_steps=1)
+    with pytest.raises(ValueError, match="num_score_steps"):
+        search(shards[0], (0, 2), score_steps=1)
+    resumed = search(tmp_path / "interrupted.pth", (0, 2))[1]["candidate_stats"]
+    for name, stats in safe_load(shards[0])["candidate_stats"].items():
+        assert resumed[name]["raw_scores"] == stats["raw_scores"]
 
     merged = tmp_path / "merged.pth"
     merged.symlink_to(tmp_path / "elsewhere.pth")  # planted after the existence check
@@ -925,7 +944,7 @@ def search(checkpoint, format_shard=None, model=None, score_steps=2):
     model_quant.merge_auto_quantize_checkpoints(list(reversed(shards)), str(merged))
     with pytest.raises(ValueError, match="format shard"):
         search(shards[0])
-    _, state = search(merged)
+    _, state = search(merged, score_steps=1)  # unsharded, the counts may change to re-solve
 
     for name, stats in expected["candidate_stats"].items():
         assert state["candidate_stats"][name]["raw_scores"] == pytest.approx(stats["raw_scores"])

From 92a18baa25dd79ffd6d93544cfafd3b5c4ae6483 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Fri, 6 Nov 2026 19:54:24 -0800
Subject: [PATCH 7/7] Remove a failed merge's partial checkpoint so the merge
 can be retried

The merged checkpoint is created exclusively before it is written, so a write
that failed (a full disk, an interrupt) left a partial file behind: the retry
then refused the existing path, and auto_quantize would try to restore from
it. Remove the file when the write fails; a path that already existed still
fails the exclusive open and is left alone.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 modelopt/torch/quantization/model_quant.py      |  9 +++++++--
 tests/unit/torch/quantization/test_autoquant.py | 12 +++++++++++-
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index 972ce5d6fec..4152c421212 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -1319,8 +1319,13 @@ def merge_auto_quantize_checkpoints(shard_checkpoints: Sequence[str], checkpoint
     )
     os.makedirs(os.path.dirname(checkpoint) or ".", exist_ok=True)
     # Created exclusively: a file or symlink that appeared since the check is refused, not written.
-    with open(checkpoint, "xb") as f:
-        safe_save(merged, f)
+    f = open(checkpoint, "xb")
+    try:
+        with f:
+            safe_save(merged, f)
+    except BaseException:
+        os.unlink(checkpoint)  # only this call's partial file, so the merge can be retried
+        raise
 
 
 def get_auto_quantize_config(search_state, constraints=None, verbose=False):
diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py
index 54e7f2fa23d..af187ae281b 100644
--- a/tests/unit/torch/quantization/test_autoquant.py
+++ b/tests/unit/torch/quantization/test_autoquant.py
@@ -869,7 +869,7 @@ def loss_func(output):
 
 
 @pytest.mark.parametrize("method", ["gradient", "kl_div"])
-def test_format_shards_merge_into_the_unsharded_search(tmp_path, method):
+def test_format_shards_merge_into_the_unsharded_search(tmp_path, monkeypatch, method):
     source = TransformerBlock()
     data = [source.get_input() for _ in range(2)]
     formats = [mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_DEFAULT_CFG, mtq.FP8_DEFAULT_CFG]
@@ -941,6 +941,16 @@ def search(checkpoint, format_shard=None, model=None, calib_steps=2, score_steps
         model_quant.merge_auto_quantize_checkpoints(shards, str(merged))
     assert not (tmp_path / "elsewhere.pth").exists()
     merged.unlink()
+
+    def disk_full(obj, f):  # writes part of the checkpoint, then fails
+        f.write(b"partial")
+        raise OSError("No space left on device")
+
+    with monkeypatch.context() as patch:
+        patch.setattr(model_quant, "safe_save", disk_full)
+        with pytest.raises(OSError, match="No space"):
+            model_quant.merge_auto_quantize_checkpoints(shards, str(merged))
+    assert not merged.exists()  # so the merge can be retried
     model_quant.merge_auto_quantize_checkpoints(list(reversed(shards)), str(merged))
     with pytest.raises(ValueError, match="format shard"):
         search(shards[0])