diff --git a/alto/models/patcher.py b/alto/models/patcher.py index 461b762e..df8e329d 100644 --- a/alto/models/patcher.py +++ b/alto/models/patcher.py @@ -62,12 +62,6 @@ def forward(ctx, x, scale, zero_point, args, g_idx, global_scale): # Dispatches to the REAL packed Triton kernel (quantize -> # bit-packed bytes -> dequantize), not the torch fake-quant # emulation -- this is the intentional, validated design - # (commit 95b1114: LLaMA-3.2-1B wikitext loss 2.1141 == the - # fake-quant path's 2.1140). Requires a CUDA tensor (Triton - # kernel launch). num_bits / group_size are fixed by the - # format rather than honoured, so they are validated here - # instead of forwarded -- a recipe that disagrees with the - # format would otherwise be silently ignored. from alto.kernels.mx import MX_QUANT_BIT, convert_to_mx, convert_from_mx assert args.group_size in (None, 16), \ @@ -75,8 +69,9 @@ def forward(ctx, x, scale, zero_point, args, g_idx, global_scale): assert args.num_bits == MX_QUANT_BIT[target_dtype], \ (f"{target_dtype} packed kernel is fixed at num_bits " f"{MX_QUANT_BIT[target_dtype]}, got {args.num_bits}") - packed = convert_to_mx(x, target_dtype=target_dtype) - return convert_from_mx(packed, target_dtype, x.dtype, x.shape) + axis = getattr(args, "block_axis", -1) + packed = convert_to_mx(x, target_dtype=target_dtype, axis=axis) + return convert_from_mx(packed, target_dtype, x.dtype, x.shape, axis=axis) return original_fake_quantize(x, scale, zero_point, args, g_idx, global_scale) @staticmethod diff --git a/alto/models/petr/__init__.py b/alto/models/petr/__init__.py new file mode 100644 index 00000000..85eac522 --- /dev/null +++ b/alto/models/petr/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT diff --git a/alto/models/petr/quantize.py b/alto/models/petr/quantize.py new file mode 100644 index 00000000..9ab93e19 --- /dev/null +++ b/alto/models/petr/quantize.py @@ -0,0 +1,150 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Fully-dynamic MX6 / MX9 recipes. + +Model-agnostic: the part of a model to leave alone is named by pattern, not by +an attribute name a particular architecture happens to use, so the same entry +point works on a detector, an LLM, or anything else built out of ``nn.Conv2d`` +and ``nn.Linear``. + +Mixed precision comes from ``ignore``. A model already runs in some dtype, and +the modules a recipe skips keep it, so "MX6 here, BF16 there" is one recipe with +the BF16 half excluded. + +Scales come from the tensor being quantized on each forward pass, so no +calibration data, observers or persisted scales are involved. + +The public surface is :func:`apply_mx_quantization`; building a recipe and +driving its modifiers through ``convert -> initialize -> pre_step -> post_step +-> finalize`` is internal. +""" + +import re +import torch.nn as nn +from compressed_tensors.utils import match_named_modules +from alto.config import Recipe + +_MX_FORMAT_BITS = {"mx6": 5, "mx9": 8} + +# Axis the MX blocks run along, per module type. A Conv2d is blocked along its +# input channels; +_MX_BLOCK_AXIS = {nn.Conv2d: 1, nn.Linear: -1} + + +def _mx_quant_args(mx_format: str, block_axis: int) -> dict: + """Fully-dynamic MX quantization args. + + ``num_bits`` is required by ``compressed_tensors`` for bookkeeping and + follows the format's fixed physical layout; the Triton codec still owns the + actual packed representation. + """ + try: + num_bits = _MX_FORMAT_BITS[mx_format] + except KeyError as exc: + supported = ", ".join(sorted(_MX_FORMAT_BITS)) + raise ValueError( + f"unsupported MX format {mx_format!r}; choose one of: {supported}" + ) from exc + + return { + "num_bits": num_bits, + "type": "int", + "symmetric": True, + "strategy": "tensor", + "dynamic": True, + "format": mx_format, + "block_axis": block_axis, + } + + +def _create_mx_recipe(mx_format: str, ignore_patterns=()) -> Recipe: + """Build an MX W+A recipe for every Conv2d and Linear. + + Args: + mx_format: ``"mx6"`` or ``"mx9"``. + ignore_patterns: ``re:`` patterns naming the modules left out of the + recipe, so they keep the model's dtype. + """ + config_groups = { + # Split per module type because MX blocks run along a different axis for + # each. + f"group_{cls.__name__.lower()}": { + "targets": [cls.__name__], + "weights": dict(_mx_quant_args(mx_format, axis)), + "input_activations": dict(_mx_quant_args(mx_format, axis)), + } + for cls, axis in _MX_BLOCK_AXIS.items() + } + # ignore sits on the modifier, not on a scheme: compressed_tensors applies it + # across every config group. + modifier = {"sequential": False, "config_groups": config_groups} + if ignore_patterns: + modifier["ignore"] = list(ignore_patterns) + return Recipe.from_dict( + { + "quantization_stage": { + "quantization_modifiers": {"QuantizationModifier": modifier} + } + } + ) + + +def _apply_recipe(model, recipe): + """Drive a recipe's modifiers through their lifecycle against ``model``.""" + modifiers = recipe.modifiers + if not modifiers: + raise ValueError(f"recipe {recipe} produced no modifiers") + + model_parts = [model] + for modifier in modifiers: + modifier.convert(model) + for modifier in modifiers: + modifier.initialize(model_parts) + for modifier in modifiers: + modifier.pre_step(model_parts) + for modifier in modifiers: + modifier.post_step(model_parts) + for modifier in modifiers: + modifier.finalize(model_parts) + return model + + +def apply_mx_quantization(model, mx_format: str, ignore=()): + """Apply an MX6 or MX9 W+A recipe in place. + + Every Conv2d and Linear is quantized unless ``ignore`` excludes it; excluded + modules keep the dtype the model already runs in. + + Raises: + ValueError: if the recipe would reach nothing, or if ``ignore`` was given + but excludes nothing. Neither fails on its own -- a pattern that + matches no module quantizes the whole model quietly -- and both mean + the run will not measure what was asked for. Both are checked before + the model is touched, so a rejected call leaves it unmodified. + """ + # A bare name is matched against the module name exactly, which would select + # only the container -- not a Conv2d or Linear, and so exclude nothing. Widen + # it to the subtree; re: patterns pass through. + patterns = [ + p if p.startswith("re:") else f"re:{re.escape(p)}($|\\.)" for p in ignore + ] + recipe = _create_mx_recipe(mx_format, patterns) + + # Match_named_modules is the matcher apply_quantization_config itself uses, + # so the scope is checked exactly, before the model is touched. + targets = [cls.__name__ for cls in _MX_BLOCK_AXIS] + total = [name for name, _ in match_named_modules(model, targets)] + kept = [name for name, _ in match_named_modules(model, targets, patterns)] + if not kept: + raise ValueError( + f"the {mx_format} recipe matched no {'/'.join(targets)} in " + f"{type(model).__name__}") + if ignore and len(kept) == len(total): + raise ValueError( + f"ignore={list(ignore)} excluded nothing: all {len(total)} module(s) " + f"would be quantized. A pattern is a module name, covering that " + f"module and everything under it, or 're:' plus a regex; names look " + f"like {total[0]!r}") + + return _apply_recipe(model, recipe) diff --git a/alto/modifiers/quantization/format_registry.py b/alto/modifiers/quantization/format_registry.py index 05000285..509e53a3 100644 --- a/alto/modifiers/quantization/format_registry.py +++ b/alto/modifiers/quantization/format_registry.py @@ -3,17 +3,23 @@ # SPDX-License-Identifier: MIT """Runtime patch that wires emulated formats into the standard quant path. -Importing this module injects a real ``format`` field into +Importing this module injects real ``format`` and ``block_axis`` fields into ``compressed_tensors.QuantizationArgs`` so recipe values like ``format: mx9`` -survive pydantic parsing and become readable via ``getattr(args, "format", None)`` -(by default unknown fields are silently dropped). +survive pydantic parsing and become readable via +``getattr(args, "format", None)`` (by default unknown fields are silently +dropped). + +``block_axis`` is the tensor axis the packed formats group elements along, and +defaults to the last one. A recipe only needs to set it where that is the wrong +axis, such as MX on Conv2d weights, whose blocks run along the input channels +rather than along the kernel width. The actual ``fake_quantize`` dispatch (``args.format == "mx9"`` -> mx9) lives in ``alto.models.patcher.ModelPatcher.patch_fake_quantize`` where the single wrap of ``compressed_tensors...forward.fake_quantize`` already happens. ``inject_format_field()`` is called at the top of this package's ``__init__`` (before -``QuantizationModifier`` is imported) so the field exists before the modifier +``QuantizationModifier`` is imported) so the fields exist before the modifier compiles its nested ``QuantizationScheme`` schema. """ @@ -21,9 +27,11 @@ _FORMAT_FIELD_INJECTED = False +DEFAULT_BLOCK_AXIS = -1 + def inject_format_field() -> None: - """Add ``format: Optional[str] = None`` to ``QuantizationArgs`` (idempotent).""" + """Add the ALTO fields to ``QuantizationArgs`` (idempotent).""" global _FORMAT_FIELD_INJECTED if _FORMAT_FIELD_INJECTED: return @@ -31,10 +39,16 @@ def inject_format_field() -> None: from pydantic.fields import FieldInfo from compressed_tensors.quantization import QuantizationArgs, QuantizationConfig, QuantizationScheme - if "format" not in QuantizationArgs.model_fields: - QuantizationArgs.model_fields["format"] = FieldInfo( - annotation=Optional[str], default=None - ) + added = False + for name, field in ( + ("format", FieldInfo(annotation=Optional[str], default=None)), + ("block_axis", FieldInfo(annotation=int, default=DEFAULT_BLOCK_AXIS)), + ): + if name not in QuantizationArgs.model_fields: + QuantizationArgs.model_fields[name] = field + added = True + + if added: QuantizationArgs.model_rebuild(force=True) # QuantizationArgs is nested inside these models. Rebuild them as well so # recipe dictionaries with weights/input_activations.format are accepted diff --git a/tests/unittest/mx9_mx6/test_mx_recipe_scope.py b/tests/unittest/mx9_mx6/test_mx_recipe_scope.py new file mode 100644 index 00000000..07fef3c3 --- /dev/null +++ b/tests/unittest/mx9_mx6/test_mx_recipe_scope.py @@ -0,0 +1,227 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Scoping an MX recipe with ``ignore``, and the guards around it. + +The model is a plain three-stage stack with normalization layers interleaved, not +any real architecture, because the API must not know about one. + +Run with: + pytest tests/unittest/mx9_mx6/test_mx_recipe_scope.py +""" + +import pytest +import torch +import torch.nn as nn + +from alto.models.petr.quantize import apply_mx_quantization + + +class Stack(nn.Module): + """Three named parts, each mixing quantizable leaves with normalization.""" + + def __init__(self): + super().__init__() + self.stem = nn.Sequential( + nn.Conv2d(16, 32, 3, padding=1, bias=False), + nn.BatchNorm2d(32), + nn.ReLU(), + ) + self.trunk = nn.Sequential( + nn.Conv2d(32, 32, 3, padding=1, bias=False), + nn.BatchNorm2d(32), + nn.ReLU(), + ) + self.head = nn.Sequential( + nn.Linear(32, 64, bias=False), + nn.LayerNorm(64), + nn.Linear(64, 16, bias=False), + ) + + def forward(self, x): + x = self.trunk(self.stem(x)) + return self.head(x.mean(dim=(2, 3))) + + +requires_cuda = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="the MX codec is a Triton kernel and needs a GPU", +) + +pytestmark = requires_cuda + + +@pytest.fixture +def model(): + torch.manual_seed(0) + return Stack().cuda().to(torch.bfloat16).eval() + + +@pytest.fixture +def inputs(): + generator = torch.Generator(device="cpu").manual_seed(1) + return torch.randn(2, 16, 8, 8, generator=generator).cuda().to(torch.bfloat16) + + +def quantized_names(model): + return { + name + for name, module in model.named_modules() + if getattr(module, "quantization_scheme", None) is not None + } + + +def test_uniform_recipe_covers_every_leaf(model, inputs): + apply_mx_quantization(model, "mx6") + + assert quantized_names(model) == {"stem.0", "trunk.0", "head.0", "head.2"} + with torch.no_grad(): + model(inputs) + + +def test_normalization_is_never_quantized(model): + """Normalization owns a ``weight`` and would be quantized if ever targeted.""" + apply_mx_quantization(model, "mx6") + + for name in ("stem.1", "trunk.1", "head.1"): + assert getattr(model.get_submodule(name), "quantization_scheme", None) is None + + +def test_ignored_modules_keep_their_dtype(model, inputs): + """The BF16 half: no scheme attached, weights bit-for-bit untouched.""" + head_weight = model.head[0].weight.detach().clone() + + apply_mx_quantization(model, "mx6", ignore=["re:^head\\..*"]) + + assert quantized_names(model) == {"stem.0", "trunk.0"} + assert getattr(model.head[0], "quantization_scheme", None) is None + assert torch.equal(model.head[0].weight, head_weight) + assert model.head[0].weight.dtype == torch.bfloat16 + with torch.no_grad(): + model(inputs) # the mixed graph still runs + + +def test_bare_name_covers_the_subtree(model): + """compressed_tensors alone would match the container, which holds no leaf.""" + apply_mx_quantization(model, "mx9", ignore=["head"]) + + assert quantized_names(model) == {"stem.0", "trunk.0"} + + +def test_bare_name_can_also_be_a_single_leaf(model): + apply_mx_quantization(model, "mx9", ignore=["head.0", "head.2"]) + + assert quantized_names(model) == {"stem.0", "trunk.0"} + + +def test_bare_name_does_not_match_a_prefix_of_a_sibling(model): + """``stem`` must not also swallow a hypothetical ``stem_extra``.""" + torch.manual_seed(0) + wider = Stack().cuda().to(torch.bfloat16).eval() + wider.add_module("stem_extra", nn.Conv2d(16, 16, 1, bias=False).cuda().to(torch.bfloat16)) + + apply_mx_quantization(wider, "mx6", ignore=["stem"]) + + assert "stem_extra" in quantized_names(wider) + assert "stem.0" not in quantized_names(wider) + + +def test_ignore_that_excludes_nothing_is_rejected(model): + """A typo'd pattern would otherwise quantize the whole model silently.""" + with pytest.raises(ValueError, match="excluded nothing"): + apply_mx_quantization(model, "mx6", ignore=["no_such_part"]) + + +def test_recipe_that_matches_nothing_is_rejected(): + model = nn.Sequential(nn.BatchNorm2d(8), nn.ReLU()).cuda() + with pytest.raises(ValueError, match="matched no"): + apply_mx_quantization(model, "mx6") + + +def test_unsupported_format_is_rejected(model): + """The format is validated while the recipe is built, before the scope is.""" + with pytest.raises(ValueError, match="unsupported MX format"): + apply_mx_quantization(model, "mx7") + + assert quantized_names(model) == set() + + +def test_scheme_carries_the_mx_format(model): + """``format`` is what the patched fake_quantize dispatches the codec on.""" + apply_mx_quantization(model, "mx9", ignore=["re:^head\\..*"]) + + scheme = model.stem[0].quantization_scheme + assert scheme.weights.format == "mx9" + assert scheme.input_activations.format == "mx9" + assert getattr(model.head[0], "quantization_scheme", None) is None + + +def test_a_rejected_ignore_leaves_the_model_untouched(model): + """Scope is validated before anything is wrapped, so the failure is atomic.""" + with pytest.raises(ValueError, match="excluded nothing"): + apply_mx_quantization(model, "mx6", ignore=["no_such_part"]) + + assert quantized_names(model) == set() + + +def test_block_axis_reaches_the_kernel_per_module_type(model, inputs, monkeypatch): + """The axis in the scheme has to arrive at the codec, or blocks run the wrong way. + + Conv2d tensors are 4D here and must block along axis 1, their input channels; + Linear tensors are 2D and block along the last axis. Blocking a 3x3 weight + along the last axis instead would put three elements in a 16-element block. + """ + import alto.kernels.mx as mx + + seen = [] + convert_to_mx = mx.convert_to_mx + + def spy(data_hp, target_dtype, axis=-1): + seen.append((tuple(data_hp.shape), axis)) + return convert_to_mx(data_hp, target_dtype, axis=axis) + + monkeypatch.setattr(mx, "convert_to_mx", spy) + + apply_mx_quantization(model, "mx6") + with torch.no_grad(): + model(inputs) + + assert {(len(shape), axis) for shape, axis in seen} == {(4, 1), (2, -1)} + assert (tuple(model.stem[0].weight.shape), 1) in seen + assert (tuple(model.head[0].weight.shape), -1) in seen + + +@pytest.mark.xfail( + strict=True, + reason="known gap: freeze_module_quantization only promotes COMPRESSED to " + "FROZEN, and a fully-dynamic recipe never becomes COMPRESSED, so finalize " + "leaves the status at CALIBRATION. Drop this marker once that is fixed.", +) +def test_finalized_modules_are_frozen(model): + """A finished lifecycle should not still advertise itself as calibrating.""" + from compressed_tensors.quantization import QuantizationStatus + + apply_mx_quantization(model, "mx6") + + assert model.stem[0].quantization_status == QuantizationStatus.FROZEN + + +def test_scoped_run_stays_closer_to_the_unquantized_model(model, inputs): + """Quantizing fewer layers has to move the output less, or ignore is inert.""" + with torch.no_grad(): + reference = model(inputs).clone() + + torch.manual_seed(0) + uniform = Stack().cuda().to(torch.bfloat16).eval() + apply_mx_quantization(uniform, "mx6") + + torch.manual_seed(0) + scoped = Stack().cuda().to(torch.bfloat16).eval() + apply_mx_quantization(scoped, "mx6", ignore=["re:^head\\..*"]) + + with torch.no_grad(): + uniform_out = uniform(inputs) + scoped_out = scoped(inputs) + + assert not torch.equal(uniform_out, scoped_out) + assert (scoped_out - reference).norm() < (uniform_out - reference).norm()