[1/2] Reload versioned quantizer recipes in vLLM fake-quant serving

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

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 = {

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+)*$")

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)),
     )
 
 

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

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)

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)

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