diff --git a/src/mobius/weights/__init__.py b/src/mobius/weights/__init__.py new file mode 100644 index 000000000..a3dd7744d --- /dev/null +++ b/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", +] diff --git a/src/mobius/weights/_codecs.py b/src/mobius/weights/_codecs.py new file mode 100644 index 000000000..9a796b6a9 --- /dev/null +++ b/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)) diff --git a/src/mobius/weights/_codecs_test.py b/src/mobius/weights/_codecs_test.py new file mode 100644 index 000000000..e65ffaf0b --- /dev/null +++ b/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") diff --git a/src/mobius/weights/_records.py b/src/mobius/weights/_records.py new file mode 100644 index 000000000..cc274f576 --- /dev/null +++ b/src/mobius/weights/_records.py @@ -0,0 +1,121 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Typed logical weights used between checkpoint readers and model adapters.""" + +from __future__ import annotations + +__all__ = [ + "FloatWeight", + "PackedWeight", + "WeightBundle", + "WeightRecord", +] + +import dataclasses +from collections.abc import Iterator, Mapping +from types import MappingProxyType + +import torch + +from mobius._component_manifest import ComponentDescriptor + + +@dataclasses.dataclass(frozen=True) +class FloatWeight: + """One ordinary floating-point checkpoint tensor.""" + + value: torch.Tensor + source_key: str + + +@dataclasses.dataclass(frozen=True) +class PackedWeight: + """One logical affine-quantized weight grouped from checkpoint sidecars.""" + + qweight: torch.Tensor + scales: torch.Tensor + zero_points: torch.Tensor | None + qweight_key: str + scales_key: str + zero_points_key: str | None + method: str + + @property + def source_keys(self) -> tuple[str, ...]: + """Checkpoint keys consumed by this logical packed weight.""" + keys = [self.qweight_key, self.scales_key] + if self.zero_points_key is not None: + keys.append(self.zero_points_key) + return tuple(keys) + + def as_state_dict(self) -> dict[str, torch.Tensor]: + """Reconstruct the source sidecars for a compatibility codec.""" + tensors = { + self.qweight_key: self.qweight, + self.scales_key: self.scales, + } + if self.zero_points_key is not None and self.zero_points is not None: + tensors[self.zero_points_key] = self.zero_points + return tensors + + +WeightStorage = FloatWeight | PackedWeight + + +@dataclasses.dataclass(frozen=True) +class WeightRecord: + """A named logical weight owned by exactly one package component.""" + + name: str + component: str + storage: WeightStorage + + @property + def is_quantized(self) -> bool: + """Whether this record stores an existing packed weight.""" + return isinstance(self.storage, PackedWeight) + + @property + def source_keys(self) -> tuple[str, ...]: + """Checkpoint keys represented by this record.""" + if isinstance(self.storage, FloatWeight): + return (self.storage.source_key,) + return self.storage.source_keys + + +@dataclasses.dataclass(frozen=True) +class WeightBundle(Mapping[str, WeightRecord]): + """Immutable records routed to one component descriptor.""" + + component: ComponentDescriptor + records: Mapping[str, WeightRecord] + + def __post_init__(self) -> None: + for name, record in self.records.items(): + if name != record.name: + raise ValueError( + f"weight bundle key {name!r} does not match record name {record.name!r}" + ) + if record.component != self.component.name: + raise ValueError( + f"weight {name!r} belongs to component " + f"{record.component!r}, expected {self.component.name!r}" + ) + object.__setattr__(self, "records", MappingProxyType(dict(self.records))) + + def __getitem__(self, name: str) -> WeightRecord: + return self.records[name] + + def __iter__(self) -> Iterator[str]: + return iter(self.records) + + def __len__(self) -> int: + return len(self.records) + + @property + def source_keys(self) -> frozenset[str]: + """All checkpoint keys represented by this bundle.""" + return frozenset( + source_key for record in self.records.values() for source_key in record.source_keys + ) diff --git a/src/mobius/weights/_records_test.py b/src/mobius/weights/_records_test.py new file mode 100644 index 000000000..f15aae40a --- /dev/null +++ b/src/mobius/weights/_records_test.py @@ -0,0 +1,47 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for typed logical checkpoint records.""" + +from __future__ import annotations + +import pytest +import torch + +from mobius._component_manifest import ComponentDescriptor +from mobius.weights import FloatWeight, WeightBundle, WeightRecord + + +def _component() -> ComponentDescriptor: + return ComponentDescriptor( + name="decoder", + module_attribute_path="decoder", + role="decoder", + source_paths=("model.layers",), + ) + + +def test_bundle_tracks_source_keys(): + record = WeightRecord( + name="model.norm.weight", + component="decoder", + storage=FloatWeight( + value=torch.ones(4), + source_key="model.norm.weight", + ), + ) + bundle = WeightBundle(_component(), {record.name: record}) + + assert bundle.source_keys == frozenset({"model.norm.weight"}) + assert bundle["model.norm.weight"] is record + + +def test_bundle_rejects_wrong_component(): + record = WeightRecord( + name="model.norm.weight", + component="vision_encoder", + storage=FloatWeight(torch.ones(4), "model.norm.weight"), + ) + + with pytest.raises(ValueError, match="belongs to component"): + WeightBundle(_component(), {record.name: record})