diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 25c9d772fe3..404efda9520 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 ``_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..38b5e53a18d 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,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,7 +1597,35 @@ 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() @@ -1550,7 +1633,7 @@ 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 @@ -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 diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 108cf599161..4152c421212 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, restore_training_modes +from modelopt.torch.utils import atomic_print, restore_training_modes, 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", @@ -828,6 +830,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. @@ -1023,6 +1026,14 @@ 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; 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. @@ -1085,6 +1096,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 @@ -1230,6 +1243,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, @@ -1250,11 +1264,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. @@ -1282,6 +1297,37 @@ 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. 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. 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. + 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, map_location="cpu", mmap=True) for path in shard_checkpoints] + ) + 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. + 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): """Build a flat quant config from auto_quantize search_state. 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) diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py index 9c25a1117e2..af187ae281b 100644 --- a/tests/unit/torch/quantization/test_autoquant.py +++ b/tests/unit/torch/quantization/test_autoquant.py @@ -868,6 +868,99 @@ 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, 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] + + 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}, + 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, + num_calib_steps=calib_steps, + 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), 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")) + # 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 + with pytest.raises(FileExistsError): + 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]) + _, 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"]) + assert state["best"]["recipe"] == expected["best"]["recipe"] + + def test_auto_quantize_disable_layers(): model = TransformerBlock() @@ -1083,7 +1176,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):