Add typed checkpoint weight records

src/mobius/weights/__init__.py

@@ -0,0 +1,28 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Typed checkpoint records and quantization format codecs."""
+
+from __future__ import annotations
+
+from mobius.weights._codecs import (
+    QuantizationCodec,
+    QuantizationCodecRegistry,
+    codec_registry,
+)
+from mobius.weights._records import (
+    FloatWeight,
+    PackedWeight,
+    WeightBundle,
+    WeightRecord,
+)
+
+__all__ = [
+    "FloatWeight",
+    "PackedWeight",
+    "QuantizationCodec",
+    "QuantizationCodecRegistry",
+    "WeightBundle",
+    "WeightRecord",
+    "codec_registry",
+]

src/mobius/weights/_codecs.py

@@ -0,0 +1,199 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Quantization format codecs for existing packed checkpoint weights."""
+
+from __future__ import annotations
+
+__all__ = [
+    "QuantizationCodec",
+    "QuantizationCodecRegistry",
+    "codec_registry",
+]
+
+from collections.abc import Mapping
+from typing import Protocol
+
+import torch
+
+from mobius._component_manifest import ComponentDescriptor
+from mobius._configs import QuantizationConfig
+from mobius._weight_utils import preprocess_quantized_weights
+from mobius.weights._records import (
+    FloatWeight,
+    PackedWeight,
+    WeightBundle,
+    WeightRecord,
+)
+
+
+class QuantizationCodec(Protocol):
+    """Groups and normalizes one producer's existing packed weight layout."""
+
+    method: str
+
+    def group(
+        self,
+        component: ComponentDescriptor,
+        state_dict: Mapping[str, torch.Tensor],
+        config: QuantizationConfig,
+    ) -> WeightBundle:
+        """Group checkpoint sidecars into typed logical records."""
+        ...
+
+    def normalize(
+        self,
+        record: WeightRecord,
+        config: QuantizationConfig,
+    ) -> dict[str, torch.Tensor]:
+        """Convert one packed record to Mobius's canonical parameter layout."""
+        ...
+
+
+class QuantizationCodecRegistry:
+    """Registry keyed by serialized ``quant_method``."""
+
+    def __init__(self) -> None:
+        self._codecs: dict[str, QuantizationCodec] = {}
+
+    def register(self, codec: QuantizationCodec) -> None:
+        if not codec.method:
+            raise ValueError("quantization codec method must not be empty")
+        if codec.method in self._codecs:
+            raise ValueError(f"quantization codec {codec.method!r} is already registered")
+        self._codecs[codec.method] = codec
+
+    def get(self, method: str) -> QuantizationCodec:
+        try:
+            return self._codecs[method]
+        except KeyError:
+            raise KeyError(
+                f"No quantization codec registered for {method!r}. "
+                f"Available methods: {sorted(self._codecs)}"
+            ) from None
+
+    def __contains__(self, method: str) -> bool:
+        return method in self._codecs
+
+
+class _LegacyAffineCodec:
+    """Typed facade over the existing Olive/GPTQ/AWQ normalization helpers."""
+
+    def __init__(self, method: str):
+        self.method = method
+
+    @staticmethod
+    def _logical_name(qweight_key: str) -> tuple[str, str, str, str | None]:
+        if qweight_key.endswith("_qweight"):
+            stem = qweight_key[: -len("_qweight")]
+            logical_name = stem if stem.endswith(".weight") else f"{stem}.weight"
+            return (
+                logical_name,
+                f"{stem}_scales",
+                qweight_key,
+                f"{stem}_qzeros",
+            )
+        if qweight_key.endswith(".qweight"):
+            stem = qweight_key[: -len(".qweight")]
+            return (
+                f"{stem}.weight",
+                f"{stem}.scales",
+                qweight_key,
+                f"{stem}.qzeros",
+            )
+        raise ValueError(f"{qweight_key!r} is not a packed qweight key")
+
+    def group(
+        self,
+        component: ComponentDescriptor,
+        state_dict: Mapping[str, torch.Tensor],
+        config: QuantizationConfig,
+    ) -> WeightBundle:
+        if config.quant_method != self.method:
+            raise ValueError(
+                f"codec {self.method!r} cannot group quant_method {config.quant_method!r}"
+            )
+
+        records: dict[str, WeightRecord] = {}
+        consumed: set[str] = set()
+        qweight_keys = sorted(
+            key for key in state_dict if key.endswith(("_qweight", ".qweight"))
+        )
+        for qweight_key in qweight_keys:
+            logical_name, scales_key, _, zero_points_key = self._logical_name(qweight_key)
+            if scales_key not in state_dict:
+                raise ValueError(
+                    f"Packed weight {qweight_key!r} is missing scales {scales_key!r}"
+                )
+            zero_points = state_dict.get(zero_points_key)
+            if not config.sym and zero_points is None:
+                raise ValueError(
+                    f"Asymmetric packed weight {qweight_key!r} is missing "
+                    f"zero points {zero_points_key!r}"
+                )
+            storage = PackedWeight(
+                qweight=state_dict[qweight_key],
+                scales=state_dict[scales_key],
+                zero_points=zero_points,
+                qweight_key=qweight_key,
+                scales_key=scales_key,
+                zero_points_key=zero_points_key if zero_points is not None else None,
+                method=self.method,
+            )
+            if logical_name in records:
+                raise ValueError(
+                    f"Checkpoint declares logical weight {logical_name!r} more than once"
+                )
+            records[logical_name] = WeightRecord(
+                name=logical_name,
+                component=component.name,
+                storage=storage,
+            )
+            consumed.update(storage.source_keys)
+
+        orphan_sidecars = sorted(
+            key
+            for key in state_dict
+            if key.endswith(("_scales", "_qzeros", ".scales", ".qzeros"))
+            and key not in consumed
+        )
+        if orphan_sidecars:
+            raise ValueError(
+                f"Packed checkpoint sidecars have no matching qweight: {orphan_sidecars}"
+            )
+
+        for key, value in state_dict.items():
+            if key in consumed:
+                continue
+            if key in records:
+                raise ValueError(f"Float and packed checkpoint values both target {key!r}")
+            records[key] = WeightRecord(
+                name=key,
+                component=component.name,
+                storage=FloatWeight(value=value, source_key=key),
+            )
+        return WeightBundle(component=component, records=records)
+
+    def normalize(
+        self,
+        record: WeightRecord,
+        config: QuantizationConfig,
+    ) -> dict[str, torch.Tensor]:
+        if not isinstance(record.storage, PackedWeight):
+            raise TypeError(f"weight {record.name!r} is not packed")
+        if record.storage.method != self.method:
+            raise ValueError(
+                f"weight {record.name!r} uses {record.storage.method!r}, "
+                f"not codec {self.method!r}"
+            )
+        return preprocess_quantized_weights(
+            record.storage.as_state_dict(),
+            config,
+            tie_embeddings=False,
+            qmoe_target_path=None,
+        )
+
+
+codec_registry = QuantizationCodecRegistry()
+for _method in ("olive", "gptq", "awq"):
+    codec_registry.register(_LegacyAffineCodec(_method))

src/mobius/weights/_codecs_test.py

@@ -0,0 +1,132 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Tests for typed quantization format codecs."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from mobius._component_manifest import ComponentDescriptor
+from mobius._configs import QuantizationConfig
+from mobius.weights import PackedWeight, QuantizationCodecRegistry, codec_registry
+
+
+def _component() -> ComponentDescriptor:
+    return ComponentDescriptor(
+        name="decoder",
+        module_attribute_path="decoder",
+        role="decoder",
+        source_paths=("model.layers",),
+    )
+
+
+def _config(method: str = "olive", *, sym: bool = True) -> QuantizationConfig:
+    return QuantizationConfig(
+        bits=4,
+        group_size=16,
+        quant_method=method,
+        sym=sym,
+    )
+
+
+def test_groups_olive_sidecars_into_one_record():
+    state_dict = {
+        "model.q_proj.weight_qweight": torch.zeros(32, 32, dtype=torch.uint8),
+        "model.q_proj.weight_scales": torch.ones(32, 4),
+        "model.norm.weight": torch.ones(64),
+    }
+
+    bundle = codec_registry.get("olive").group(
+        _component(),
+        state_dict,
+        _config(),
+    )
+
+    record = bundle["model.q_proj.weight"]
+    assert isinstance(record.storage, PackedWeight)
+    assert record.source_keys == (
+        "model.q_proj.weight_qweight",
+        "model.q_proj.weight_scales",
+    )
+    assert bundle["model.norm.weight"].is_quantized is False
+
+
+def test_groups_gptq_dotted_sidecars():
+    state_dict = {
+        "model.q_proj.qweight": torch.zeros(8, 32, dtype=torch.int32),
+        "model.q_proj.scales": torch.ones(4, 32),
+    }
+
+    bundle = codec_registry.get("gptq").group(
+        _component(),
+        state_dict,
+        _config("gptq"),
+    )
+
+    assert bundle["model.q_proj.weight"].is_quantized is True
+
+
+def test_rejects_missing_scales():
+    state_dict = {
+        "model.q_proj.weight_qweight": torch.zeros(32, 32, dtype=torch.uint8),
+    }
+
+    with pytest.raises(ValueError, match="missing scales"):
+        codec_registry.get("olive").group(
+            _component(),
+            state_dict,
+            _config(),
+        )
+
+
+def test_rejects_missing_asymmetric_zero_points():
+    state_dict = {
+        "model.q_proj.weight_qweight": torch.zeros(32, 32, dtype=torch.uint8),
+        "model.q_proj.weight_scales": torch.ones(32, 4),
+    }
+
+    with pytest.raises(ValueError, match="missing zero points"):
+        codec_registry.get("olive").group(
+            _component(),
+            state_dict,
+            _config(sym=False),
+        )
+
+
+def test_rejects_orphan_sidecars():
+    with pytest.raises(ValueError, match="no matching qweight"):
+        codec_registry.get("olive").group(
+            _component(),
+            {"model.q_proj.weight_scales": torch.ones(32, 4)},
+            _config(),
+        )
+
+
+def test_compatibility_normalizer_uses_existing_packer():
+    state_dict = {
+        "model.q_proj.weight_qweight": torch.zeros(32, 32, dtype=torch.uint8),
+        "model.q_proj.weight_scales": torch.ones(32, 4),
+    }
+    codec = codec_registry.get("olive")
+    record = codec.group(_component(), state_dict, _config())["model.q_proj.weight"]
+
+    normalized = codec.normalize(record, _config())
+
+    assert normalized["model.q_proj.weight"].shape == (32, 4, 8)
+    assert normalized["model.q_proj.scales"].shape == (32, 4)
+
+
+def test_registry_rejects_duplicate_method():
+    registry = QuantizationCodecRegistry()
+    codec = codec_registry.get("olive")
+    registry.register(codec)
+
+    with pytest.raises(ValueError, match="already registered"):
+        registry.register(codec)
+
+
+def test_registry_reports_unknown_method():
+    with pytest.raises(KeyError, match="Available methods"):
+        QuantizationCodecRegistry().get("unknown")