[OMNIML-6370] [2/5] Calibrate and score AutoQuantize formats in parallel shards

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/``.

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,16 @@ 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."""
+        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]:
         """Return choices exposed to the LP after removing an internal no-quant baseline."""
@@ -894,6 +906,7 @@ def default_search_config(self):
             "cost_model": COST_MODEL_WEIGHT,
             "cost": {},
             "active_moe_expert_ratio": None,
+            "format_shard": None,
         }
 
     @property
@@ -913,6 +926,9 @@ def default_state_dict(self) -> SearchStateDict:
             "disabled_layers": None,
             "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},
         }
 
@@ -926,13 +942,25 @@ 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:
         """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):
@@ -1279,18 +1307,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 +1337,37 @@ def _run_func(self, func, num_iters=1, desc=""):
         ):
             func(self.model, data)
 
+    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, 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 set()
+        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
+        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.
 
@@ -1542,15 +1597,43 @@ 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
+        # 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]
+
+        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()
 
         # 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
 
@@ -1721,6 +1804,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 +1876,79 @@ 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",
+    "num_calib_steps",
+    "num_score_steps",
+)
+
+
+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])]
+    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],
+        "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):
+                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 +2139,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 +2430,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