diff --git a/examples/vllm_serve/vllm_ptq_utils.py b/examples/vllm_serve/vllm_ptq_utils.py index d4cfab33a4b..01411e3da1b 100644 --- a/examples/vllm_serve/vllm_ptq_utils.py +++ b/examples/vllm_serve/vllm_ptq_utils.py @@ -28,6 +28,10 @@ import modelopt.torch.quantization as mtq from modelopt.recipe import ModelOptPTQRecipe, load_recipe +from modelopt.torch.export.plugins.vllm_fakequant_hf import ( + QUANTIZER_RECIPE_FORMAT, + QUANTIZER_RECIPE_VERSION, +) from modelopt.torch.quantization.plugins.vllm import VllmMLAAttention @@ -331,6 +335,22 @@ def get_quant_config(quant_config: dict[str, Any], model: Any) -> dict[str, Any] if quant_config["recipe_path"]: recipe_path = Path(quant_config["recipe_path"]) raw_recipe = yaml.safe_load(recipe_path.read_text()) if recipe_path.is_file() else None + if isinstance(raw_recipe, dict) and raw_recipe.get("format") == QUANTIZER_RECIPE_FORMAT: + # Versioned per-quantizer recipe, e.g. written by Megatron-Bridge's HF export. + if raw_recipe.get("version") != QUANTIZER_RECIPE_VERSION: + raise ValueError( + f"Unsupported {QUANTIZER_RECIPE_FORMAT} version {raw_recipe.get('version')!r}; " + f"expected {QUANTIZER_RECIPE_VERSION}" + ) + raw_recipe = raw_recipe.get("quantizers") + if not ( + isinstance(raw_recipe, dict) + and raw_recipe + and all(isinstance(name, str) and "_quantizer" in name for name in raw_recipe) + ): + raise ValueError( + f"'quantizers' of {recipe_path} must map quantizer names to states" + ) if ( isinstance(raw_recipe, dict) and raw_recipe diff --git a/examples/vllm_serve/vllm_reload_utils.py b/examples/vllm_serve/vllm_reload_utils.py index 8bb7b0175a1..07fbdca0660 100644 --- a/examples/vllm_serve/vllm_reload_utils.py +++ b/examples/vllm_serve/vllm_reload_utils.py @@ -373,12 +373,8 @@ def quantizer_recipe_to_quant_cfg(recipe: dict[str, Any], model: Any) -> dict[st # Uncaptured quantizers stay disabled; exported weights are already folded. entries = [{"quantizer_name": "*", "enable": False}] for name, state in mapped_recipe.items(): - entry = {"quantizer_name": name, "enable": not state["_disabled"]} - cfg = { - key[1:]: value - for key, value in state.items() - if key in ("_num_bits", "_axis", "_block_sizes", "_backend", "_backend_extra_args") - } + cfg = TensorQuantizer.attribute_config_from_modelopt_state(state) + entry = {"quantizer_name": name, "enable": cfg.pop("enable")} if cfg: entry["cfg"] = cfg entries.append(entry) @@ -678,6 +674,28 @@ def _validate_quantizer_ranges(model, saved_quant_dict, current_state_dict) -> N ) +def _alias_state_dict_keys(model: Any, state_dict: dict[str, Any]) -> dict[str, list[str]]: + """Map each state-dict key of a module registered under several names to its other names' keys. + + vLLM can register one module under two names (e.g. the MoE runner also holds the shared expert + as ``experts._shared_experts._layer``), so the state dict lists its tensors once per name. + """ + names_by_module: dict[Any, list[str]] = {} + for name, module in model.named_modules(remove_duplicate=False): + names_by_module.setdefault(module, []).append(name) + aliases = { + name: names for names in names_by_module.values() if len(names) > 1 for name in names + } + alias_keys = {} + for key in state_dict: + module_name, _, attr = key.rpartition(".") + if module_name in aliases: + alias_keys[key] = [ + f"{other}.{attr}" for other in aliases[module_name] if other != module_name + ] + return alias_keys + + def load_state_dict_from_path(quantizer_file_path: str, model: Any) -> dict[str, Any]: """Overlay mapped, TP-sharded quantizer tensors on the model's state dict.""" # Load on CPU to avoid failures when the checkpoint was saved from a different GPU mapping. @@ -692,6 +710,12 @@ def load_state_dict_from_path(quantizer_file_path: str, model: Any) -> dict[str, saved_quant_dict = convert_dict_to_vllm(saved_quant_dict, map_fun=map_fun) current_state_dict = model.state_dict() + # The checkpoint names one of a module's names; load it under all of them, or a stale alias + # entry would overwrite the loaded value in ``model.load_state_dict``. + for key, aliases in _alias_state_dict_keys(model, current_state_dict).items(): + if key in saved_quant_dict: + for alias in aliases: + saved_quant_dict.setdefault(alias, saved_quant_dict[key]) _validate_quantizer_ranges(model, saved_quant_dict, current_state_dict) checkpoint_quant_keys = [key for key in saved_quant_dict if "quantizer" in key] local_checkpoint_quantizers = { diff --git a/modelopt/torch/export/plugins/vllm_fakequant_hf.py b/modelopt/torch/export/plugins/vllm_fakequant_hf.py index 270c03894eb..e8a33927acf 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_hf.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_hf.py @@ -61,12 +61,19 @@ ) __all__ = [ + "QUANTIZER_RECIPE_FORMAT", + "QUANTIZER_RECIPE_VERSION", "export_hf_vllm_fq_checkpoint", "infer_quantizer_prefix_remap", "is_weight_quantizer_state_key", "merge_amax_tensors_for_group", ] +# Header of a versioned exported quantizer recipe (``quant_recipe.yaml``); its ``quantizers`` map HF +# quantizer names to properties-only ModelOpt quantizer states. +QUANTIZER_RECIPE_FORMAT = "modelopt_quantizer_recipe" +QUANTIZER_RECIPE_VERSION = 1 + # Matches ``…weight_quantizer``, ``…weight_quantizer.0``, ``…w13_weight_quantizer.0``, # and the plural fused-experts form ``…weight_quantizers.0`` (per-expert ModuleList). _WEIGHT_QUANTIZER_STATE_KEY = re.compile(r"(?:^|\.)(?:\w+_)?weight_quantizers?(?:\.\d+)*$") diff --git a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py index de873f6d673..30723624739 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py @@ -25,6 +25,10 @@ import torch import yaml +from modelopt.torch.export.plugins.vllm_fakequant_hf import ( + QUANTIZER_RECIPE_FORMAT, + QUANTIZER_RECIPE_VERSION, +) from modelopt.torch.export.quant_format import QUANTIZATION_NONE from modelopt.torch.export.unified_export_megatron import GPTModelExporter from modelopt.torch.quantization.nn import GroupedQuantizer, SequentialQuantizer, TensorQuantizer @@ -141,7 +145,9 @@ def gather_mcore_vllm_fq_quantizer_recipe( quantizer_state_by_name: dict[str, dict], save_directory: str | os.PathLike, ) -> None: - """Sync each rank's captured quantizer configs and save as ``quant_recipe.yaml``. + """Sync each rank's captured quantizer configs and save as a versioned ``quant_recipe.yaml``. + + Removes both sidecars instead when no rank captured a quantizer. Args: quantizer_state_by_name: HF-prefixed quantizer name -> resolved config, collected in @@ -153,10 +159,23 @@ def gather_mcore_vllm_fq_quantizer_recipe( DistributedProcessGroup(None), _merge_quantizer_states, ) + if not merged: + # Nothing to reload (the loader rejects an empty recipe): write no sidecars, like + # Megatron-Bridge, and drop those of an earlier export. + for name in ("quantizer_state.pth", "quant_recipe.yaml"): + _save_quantizer_state( + Path(save_directory) / name, lambda path: path.unlink(missing_ok=True) + ) + return + recipe = { + "format": QUANTIZER_RECIPE_FORMAT, + "version": QUANTIZER_RECIPE_VERSION, + "quantizers": merged, + } _save_quantizer_state( Path(save_directory) / "quant_recipe.yaml", - lambda path: path.write_text(yaml.safe_dump(merged, sort_keys=False)), + lambda path: path.write_text(yaml.safe_dump(recipe, sort_keys=False)), ) diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index 454d18bdaf6..482f71f8119 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -289,6 +289,24 @@ def _constant_amax_setter(val): if self.is_mx_format: self._pass_through_bwd = True + @staticmethod + def attribute_config_from_modelopt_state(state: dict[str, Any]) -> dict[str, Any]: + """Return the :class:`QuantizerAttributeConfig` fields of a properties-only quantizer state. + + Inverts :meth:`set_from_attribute_config` for the state saved by + ``get_modelopt_state(properties_only=True)``; keys that are not config fields are ignored. + """ + inverse_getters: dict[str, tuple[str, Callable]] = { + "_disabled": ("enable", lambda val: not val), + "_dynamic": ("type", lambda val: "dynamic" if val else "static"), + } + attribute_cfg = {} + for key, val in state.items(): + attribute, getter = inverse_getters.get(key, (key.removeprefix("_"), lambda v: v)) + if attribute in QuantizerAttributeConfig.model_fields: + attribute_cfg[attribute] = getter(val) + return attribute_cfg + def dequantize(self, inputs: BaseQuantizedTensor | QTensorWrapper): """De-quantize a real quantized tensor to a given dtype.""" qtensor = inputs.get_qtensor() if isinstance(inputs, QTensorWrapper) else inputs diff --git a/tests/examples/vllm_serve/test_fakequant_example.py b/tests/examples/vllm_serve/test_fakequant_example.py index 3e9e22a2daf..ea02423f198 100644 --- a/tests/examples/vllm_serve/test_fakequant_example.py +++ b/tests/examples/vllm_serve/test_fakequant_example.py @@ -20,6 +20,7 @@ import importlib.util import os import sys +import warnings from pathlib import Path from types import SimpleNamespace from unittest.mock import Mock, patch @@ -34,6 +35,7 @@ import modelopt.torch.quantization as mtq from modelopt.torch.opt.config_loader import load_config +from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.conversion import set_quantizer_by_cfg from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.plugins.vllm import VllmMLAAttention @@ -245,23 +247,53 @@ def test_recipe_path_rejects_manual_quant_config(conflicting_field): @pytest.mark.parametrize( - ("contents", "error"), + ("contents", "error", "match"), [ - ("", ValueError), - ("quantize: [", yaml.YAMLError), - ("- not\n- a mapping\n", ValueError), + ("", ValueError, None), + ("quantize: [", yaml.YAMLError, None), + ("- not\n- a mapping\n", ValueError, None), + ("format: modelopt_quantizer_recipe\nversion: 2\nquantizers: {}\n", ValueError, "version"), + # Rejected as a quantizer recipe, not handed to the PTQ recipe parser. + ( + "format: modelopt_quantizer_recipe\nversion: 1\nquantizers: {}\n", + ValueError, + "quantizer names", + ), ], ) -def test_invalid_recipe_yaml_is_rejected(tmp_path, contents, error): +def test_invalid_recipe_yaml_is_rejected(tmp_path, contents, error, match): ptq_utils = _load_example_module("vllm_ptq_utils") recipe_path = tmp_path / "recipe.yaml" recipe_path.write_text(contents) config = {"recipe_path": str(recipe_path), "quant_cfg": None, "kv_quant_cfg": None} - with pytest.raises(error): + with pytest.raises(error, match=match): ptq_utils.get_quant_config(config, model=None) +def test_versioned_recipe_restores_full_quantizer_state(monkeypatch, tmp_path): + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + ptq_utils = _load_example_module("vllm_ptq_utils") + source = TensorQuantizer( + QuantizerAttributeConfig(num_bits=(4, 3), axis=None, use_constant_amax=True) + ) + recipe = { + "format": ptq_utils.QUANTIZER_RECIPE_FORMAT, + "version": ptq_utils.QUANTIZER_RECIPE_VERSION, + "quantizers": {"k_bmm_quantizer": source.get_modelopt_state(properties_only=True)}, + } + recipe_path = tmp_path / "quant_recipe.yaml" + recipe_path.write_text(yaml.safe_dump(recipe)) + model = torch.nn.Module() + model.k_bmm_quantizer = TensorQuantizer() + model.v_bmm_quantizer = TensorQuantizer() + config = {"recipe_path": str(recipe_path), "quant_cfg": None, "kv_quant_cfg": None} + mtq.quantize(model, {**ptq_utils.get_quant_config(config, model), "algorithm": None}) + assert model.k_bmm_quantizer.is_enabled and model.k_bmm_quantizer._use_constant_amax + assert model.k_bmm_quantizer.num_bits == (4, 3) + assert not model.v_bmm_quantizer.is_enabled + + def _calibration_worker( num_blocks: int, *, @@ -956,6 +988,26 @@ def test_recipe_metadata_and_tensor_reload_share_module_mapping( assert converted["modelopt_state_dict"][0][1]["metadata"]["quantizer_state"].keys() == expected +@pytest.mark.parametrize("checkpoint_name_first", [True, False]) +def test_aliased_quantizer_keeps_the_loaded_range(reload_utils, tmp_path, checkpoint_name_first): + # vLLM's MoE runner also registers the shared expert as ``experts._shared_experts._layer``; + # either name may come first in ``named_modules``. + model = torch.nn.Module() + shared_expert = torch.nn.Module() + shared_expert.input_quantizer = TensorQuantizer(amax=9.0) + if checkpoint_name_first: + model.shared_expert = shared_expert + model.experts = torch.nn.Module() + model.experts.shared = shared_expert + model.shared_expert = shared_expert + path = tmp_path / "quantizer_state.pth" + torch.save({"shared_expert.input_quantizer._amax": torch.tensor(1.001)}, path) + with warnings.catch_warnings(): + warnings.simplefilter("error") + model.load_state_dict(reload_utils.load_state_dict_from_path(str(path), model)) + assert torch.equal(model.shared_expert.input_quantizer.amax, torch.tensor(1.001)) + + def test_constant_amax_kv_recipe_reloads_without_calibration_buffers(reload_utils, tmp_path): model = torch.nn.Module() model.input_quantizer = TensorQuantizer(amax=9.0) diff --git a/tests/gpu_megatron/torch/export/test_vllm_fakequant_megatron_export.py b/tests/gpu_megatron/torch/export/test_vllm_fakequant_megatron_export.py index 8e95163a8ff..5f0b947890d 100644 --- a/tests/gpu_megatron/torch/export/test_vllm_fakequant_megatron_export.py +++ b/tests/gpu_megatron/torch/export/test_vllm_fakequant_megatron_export.py @@ -35,6 +35,10 @@ import modelopt.torch.quantization as mtq from modelopt.torch.export import export_mcore_gpt_to_hf_vllm_fq +from modelopt.torch.export.plugins.vllm_fakequant_hf import ( + QUANTIZER_RECIPE_FORMAT, + QUANTIZER_RECIPE_VERSION, +) from modelopt.torch.export.plugins.vllm_fakequant_megatron import ( VllmFqGPTModelExporter, gather_mcore_vllm_fq_quantized_state_dict, @@ -45,7 +49,7 @@ def _assert_exported_quantizers(export_dir, expected_names, amax=1.001, disabled_names=()): state = torch.load(export_dir / "quantizer_state.pth", weights_only=True, map_location="cpu") - recipe = yaml.safe_load((export_dir / "quant_recipe.yaml").read_text()) + recipe = yaml.safe_load((export_dir / "quant_recipe.yaml").read_text())["quantizers"] assert expected_names <= recipe.keys() assert {name + "._amax" for name in expected_names} <= state.keys() for name in expected_names: @@ -259,7 +263,7 @@ def forward_loop(model): (source / "config.json").write_text(json.dumps(config)) tied_dir = tmp_path / "tied_export" export_mcore_gpt_to_hf_vllm_fq(model, source, export_dir=str(tied_dir)) - tied_recipe = yaml.safe_load((tied_dir / "quant_recipe.yaml").read_text()) + tied_recipe = yaml.safe_load((tied_dir / "quant_recipe.yaml").read_text())["quantizers"] assert not any(name.startswith("lm_head.") for name in tied_recipe) assert json.loads((tied_dir / "config.json").read_text())["tie_word_embeddings"] @@ -534,10 +538,19 @@ def _test_cross_rank_quantizer_merge(tmp_path, rank, size): None, {1: {name + "._amax": tensor}}, tmp_path ) if error is None: - assert yaml.safe_load((tmp_path / "quant_recipe.yaml").read_text()) == {name: recipe} + assert yaml.safe_load((tmp_path / "quant_recipe.yaml").read_text()) == { + "format": QUANTIZER_RECIPE_FORMAT, + "version": QUANTIZER_RECIPE_VERSION, + "quantizers": {name: recipe}, + } state = torch.load(tmp_path / "quantizer_state.pth", weights_only=True) torch.testing.assert_close(state[name + "._amax"], tensor, rtol=0, atol=0) + # No quantizer on any rank: no sidecars, and those of the export above are removed. + gather_mcore_vllm_fq_quantizer_recipe({}, tmp_path) + assert not (tmp_path / "quant_recipe.yaml").exists() + assert not (tmp_path / "quantizer_state.pth").exists() + failure_dir = tmp_path / "write_failure" if rank == 0: (failure_dir / "quant_recipe.yaml").mkdir(parents=True) @@ -700,7 +713,7 @@ def count_qdq(module, args, output): _assert_grouped_weights(export_dir, expected_weights) quantizer_state = torch.load(export_dir / "quantizer_state.pth", weights_only=True) - recipe = yaml.safe_load((export_dir / "quant_recipe.yaml").read_text()) + recipe = yaml.safe_load((export_dir / "quant_recipe.yaml").read_text())["quantizers"] assert not any("weight_quantizer" in key for key in quantizer_state) assert {key.rsplit(".", 1)[0] for key in quantizer_state} <= recipe.keys() assert not any("{}" in key for key in recipe) diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py index 26f048d0d1d..31352e42bb5 100644 --- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py +++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py @@ -18,6 +18,7 @@ import numpy as np import pytest import torch +import yaml from _test_utils.torch.quantization.models import SimpleLinear from _test_utils.torch.quantization.tensor_quant_common import FakeTensorQuantTester @@ -623,3 +624,27 @@ def test_preserve_quantizer_attributes_context_restores_type_and_properties(): quantizer.get_modelopt_state(properties_only=True) for quantizer in model.weight_quantizer ] assert restored_states == original_states + + +@pytest.mark.parametrize( + "attribute_cfg", + [ + { + "num_bits": (2, 1), + "axis": None, + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + }, + {"num_bits": (4, 3), "axis": None, "use_constant_amax": True}, + {"num_bits": 8, "axis": 0, "unsigned": True, "narrow_range": True, "type": "dynamic"}, + {"num_bits": "custom", "backend": "custom_backend", "backend_extra_args": {"offset": 0.25}}, + {"enable": False}, + ], +) +def test_attribute_config_from_modelopt_state_round_trips_yaml(attribute_cfg): + quantizer = TensorQuantizer(QuantizerAttributeConfig(**attribute_cfg)) + state = quantizer.get_modelopt_state(properties_only=True) + restored_cfg = TensorQuantizer.attribute_config_from_modelopt_state( + yaml.safe_load(yaml.safe_dump(state)) + ) + restored = TensorQuantizer(QuantizerAttributeConfig(**restored_cfg)) + assert restored.get_modelopt_state(properties_only=True) == state