diff --git a/src/mobius/_build_context.py b/src/mobius/_build_context.py index 03a0b4317..6ff724e02 100644 --- a/src/mobius/_build_context.py +++ b/src/mobius/_build_context.py @@ -35,8 +35,10 @@ "build_context", "ep_capabilities", "get_build_dtype", + "get_matmul_nbits_accuracy_level", "is_prefill_prefix_pruning_enabled", "prefill_prefix_pruning", + "validate_matmul_nbits_accuracy_level", ] _DEFAULT_CAPABILITIES = EpCapabilities(name="default") @@ -47,6 +49,9 @@ _current_dtype: contextvars.ContextVar[ir.DataType] = contextvars.ContextVar( "mobius_build_dtype", default=ir.DataType.FLOAT ) +_matmul_nbits_accuracy_level: contextvars.ContextVar[int | None] = contextvars.ContextVar( + "mobius_matmul_nbits_accuracy_level", default=None +) _prune_prefill_prefix: contextvars.ContextVar[bool] = contextvars.ContextVar( "mobius_prune_prefill_prefix", default=False ) @@ -56,6 +61,8 @@ def build_context( capabilities: EpCapabilities, dtype: ir.DataType = ir.DataType.FLOAT, + *, + matmul_nbits_accuracy_level: int | None = None, ) -> Iterator[None]: """Activate EP capabilities for the duration of graph construction. @@ -66,6 +73,9 @@ def build_context( capabilities: EP capability descriptor to activate. Typically obtained via ``ep_registry.require(execution_provider)``. dtype: Active build dtype. Defaults to ``ir.DataType.FLOAT``. + matmul_nbits_accuracy_level: Override every emitted ``MatMulNBits`` + accuracy level (0 through 4), regardless of weight bit width. + ``None`` preserves the EP-specific defaults. Example:: @@ -76,13 +86,27 @@ def build_context( with build_context(capabilities, ir.DataType.FLOAT16): pkg = task.build(module, config) """ + validate_matmul_nbits_accuracy_level(matmul_nbits_accuracy_level) capabilities_token = _current_ep.set(capabilities) dtype_token = _current_dtype.set(dtype) + accuracy_token = _matmul_nbits_accuracy_level.set(matmul_nbits_accuracy_level) try: yield finally: _current_ep.reset(capabilities_token) _current_dtype.reset(dtype_token) + _matmul_nbits_accuracy_level.reset(accuracy_token) + + +def validate_matmul_nbits_accuracy_level(level: int | None) -> None: + """Reject values outside the ORT ``MatMulNBits`` accuracy-level domain.""" + if level is not None and (type(level) is not int or not 0 <= level <= 4): + raise ValueError("matmul_nbits_accuracy_level must be an integer from 0 to 4 or None") + + +def get_matmul_nbits_accuracy_level() -> int | None: + """Return the explicit override, or ``None`` to preserve EP defaults.""" + return _matmul_nbits_accuracy_level.get() def ep_capabilities() -> EpCapabilities: diff --git a/src/mobius/_build_context_test.py b/src/mobius/_build_context_test.py index 8d835088c..4cdbc211b 100644 --- a/src/mobius/_build_context_test.py +++ b/src/mobius/_build_context_test.py @@ -4,6 +4,7 @@ from __future__ import annotations import threading +from concurrent.futures import ThreadPoolExecutor import onnx_ir as ir import pytest @@ -12,6 +13,7 @@ build_context, ep_capabilities, get_build_dtype, + get_matmul_nbits_accuracy_level, is_prefill_prefix_pruning_enabled, prefill_prefix_pruning, ) @@ -27,6 +29,67 @@ ) +@pytest.mark.parametrize("level", [None, 0, 1, 2, 3, 4]) +def test_matmul_nbits_accuracy_level_scoping(level): + def raise_in_inner_context(): + with build_context(_CUDA_CAPABILITIES, matmul_nbits_accuracy_level=0): + assert get_matmul_nbits_accuracy_level() == 0 + raise RuntimeError("deliberate") + + assert get_matmul_nbits_accuracy_level() is None + with build_context(_CPU_CAPABILITIES, matmul_nbits_accuracy_level=level): + assert get_matmul_nbits_accuracy_level() == level + with build_context(_CUDA_CAPABILITIES): + assert get_matmul_nbits_accuracy_level() is None + assert get_matmul_nbits_accuracy_level() == level + with pytest.raises(RuntimeError, match="deliberate"): + raise_in_inner_context() + assert get_matmul_nbits_accuracy_level() == level + assert get_matmul_nbits_accuracy_level() is None + + +@pytest.mark.parametrize("level", [-1, 5, True, False, 4.0, "4"]) +def test_matmul_nbits_accuracy_level_rejects_invalid_values_without_leaking(level): + with build_context(_CPU_CAPABILITIES, matmul_nbits_accuracy_level=4): + with ( + pytest.raises(ValueError, match="matmul_nbits_accuracy_level"), + build_context(_CUDA_CAPABILITIES, matmul_nbits_accuracy_level=level), + ): + pytest.fail("Invalid context should not be entered") + assert get_matmul_nbits_accuracy_level() == 4 + assert ep_capabilities() is _CPU_CAPABILITIES + assert get_matmul_nbits_accuracy_level() is None + + +def test_matmul_nbits_accuracy_level_thread_isolation(): + barrier = threading.Barrier(2) + + def build(level): + with build_context(_CPU_CAPABILITIES, matmul_nbits_accuracy_level=level): + barrier.wait(timeout=5) + return get_matmul_nbits_accuracy_level() + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(build, level) for level in (0, 4)] + assert [future.result(timeout=10) for future in futures] == [0, 4] + assert get_matmul_nbits_accuracy_level() is None + + +def test_matmul_nbits_accuracy_level_async_isolation(): + import asyncio + + async def build(level): + with build_context(_CPU_CAPABILITIES, matmul_nbits_accuracy_level=level): + await asyncio.sleep(0) + return get_matmul_nbits_accuracy_level() + + async def run(): + return await asyncio.gather(build(0), build(4)) + + assert asyncio.run(run()) == [0, 4] + assert get_matmul_nbits_accuracy_level() is None + + class TestBuildContextDefaults: def test_default_capabilities_is_default_ep(self): """No context active → returns the 'default' EP descriptor.""" diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py index bb2d3ea82..c3a0afd99 100644 --- a/src/mobius/_builder.py +++ b/src/mobius/_builder.py @@ -134,6 +134,7 @@ def build_from_module( task: str | ModelTask = "text-generation", *, execution_provider: str = "default", + matmul_nbits_accuracy_level: int | None = None, trace_optimization: bool = False, fp8_kv_cache: bool = False, kv_cache_scales: dict[int, tuple[float, float]] | None = None, @@ -149,6 +150,10 @@ def build_from_module( precision and its ``validate`` method runs before build. task: Task name or :class:`ModelTask` instance. execution_provider: Target for EP-aware optimizations. + matmul_nbits_accuracy_level: Optional accuracy level (0 through 4) for + every emitted ``MatMulNBits``, including 2-, 4-, and 8-bit weights. + ``None`` preserves the EP-specific defaults. This does not change + weight quantization or ``GatherBlockQuantized``. trace_optimization: Log optimization diagnostics when true. fp8_kv_cache: Store supported decoder KV caches as FLOAT8E4M3FN. kv_cache_scales: Optional per-layer FP8 key/value scales. @@ -172,7 +177,9 @@ def build_from_module( ) _cast_module_dtype(module, dtype) capabilities = ep_registry.require(execution_provider) - with build_context(capabilities, dtype): + with build_context( + capabilities, dtype, matmul_nbits_accuracy_level=matmul_nbits_accuracy_level + ): package = resolved_task.build(module, config) for name, model in package.items(): diff --git a/src/mobius/_builder_test.py b/src/mobius/_builder_test.py index 354308ed1..5776efbcf 100644 --- a/src/mobius/_builder_test.py +++ b/src/mobius/_builder_test.py @@ -22,6 +22,41 @@ from mobius._model_package import ModelPackage +@pytest.mark.parametrize("level", [None, 0, 4]) +@pytest.mark.parametrize("bits", [4, 8]) +def test_build_from_module_applies_accuracy_override(level, bits): + from mobius._builder import build_from_module + from mobius._configs import QuantizationConfig + from mobius._testing import make_config + from mobius.models import CausalLMModel + + config = make_config( + quantization=QuantizationConfig( + quant_method="olive", bits=bits, group_size=32, sym=True + ) + ) + package = build_from_module( + CausalLMModel(config), + config, + execution_provider="cpu", + matmul_nbits_accuracy_level=level, + ) + nodes = [ + node + for model in package.values() + for node in ir.traversal.RecursiveGraphIterator(model.graph) + if node.op_type == "MatMulNBits" + ] + assert len(nodes) > 1 + expected = level if level is not None else (4 if bits == 4 else None) + for node in nodes: + if expected is None: + assert "accuracy_level" not in node.attributes + else: + assert node.attributes["accuracy_level"].value == expected + assert node.attributes["bits"].value == bits + + def _make_value(name: str) -> ir.Value: return ir.Value(name=name) diff --git a/src/mobius/_component_quantization.py b/src/mobius/_component_quantization.py index b41928916..a7f13038d 100644 --- a/src/mobius/_component_quantization.py +++ b/src/mobius/_component_quantization.py @@ -22,6 +22,7 @@ from mobius._component_manifest import ComponentDescriptor, ComponentManifest from mobius._configs import BaseModelConfig, QuantizationConfig +from mobius._quantization_utils import resolve_matmul_nbits_block_size from mobius._weight_utils import is_packed_quant_key from mobius.components import ( ClippableLinear, @@ -153,7 +154,8 @@ def _linear_layout_matches( ) return ( module._bits == quantization.bits - and module._block_size == quantization.group_size + and module._block_size + == resolve_matmul_nbits_block_size(quantization.group_size, module._k) and (module.zero_points is None) is quantization.sym and ( module.zero_points is None or module.zero_points.dtype == expected_zero_point_dtype @@ -626,7 +628,17 @@ def normalize_component_quantized_weights( component_names = tuple(component_names) manifest = manifest or _default_manifest(module, config, task) mapping = getattr(config, "component_quantization", None) - if mapping is None and not _uses_global_module_rules(config, manifest): + quantization = getattr(config, "quantization", None) + global_per_channel = ( + quantization is not None + and quantization.quant_method in {"gptq", "awq"} + and quantization.group_size == -1 + ) + if ( + mapping is None + and not _uses_global_module_rules(config, manifest) + and not global_per_channel + ): return state_dict routed = _route_component_weights(state_dict, manifest, component_names) result = dict(state_dict) @@ -726,6 +738,9 @@ def normalize_component_quantized_weights( kind="embedding" if isinstance(local_module, QuantizedEmbedding) else "linear", + in_features=local_module._k + if isinstance(local_module, QuantizedLinear) + else None, ) ) diff --git a/src/mobius/_component_quantization_loading_test.py b/src/mobius/_component_quantization_loading_test.py index 927941b85..89ba1f197 100644 --- a/src/mobius/_component_quantization_loading_test.py +++ b/src/mobius/_component_quantization_loading_test.py @@ -210,6 +210,51 @@ def _t5_config(encoder, decoder) -> ArchitectureConfig: ) +@pytest.mark.parametrize("method", ["gptq", "awq"]) +@pytest.mark.parametrize("bits", [4, 8]) +def test_global_per_channel_loading_uses_each_projection_logical_width(method, bits): + quantization = _quantization(quant_method=method, bits=bits, group_size=-1, sym=False) + config = make_config(num_hidden_layers=1, intermediate_size=33, quantization=quantization) + module = CausalLMModel(config) + package = _build(module) + state_dict = {} + expected = {} + for projection, k, n in (("down_proj", 33, 64), ("gate_proj", 64, 33)): + name = f"model.layers.0.mlp.{projection}" + packed = torch.arange(math.ceil(k * bits / 32) * n, dtype=torch.int32).reshape(-1, n) + scales = _floats(1, n) + state_dict.update( + { + f"{name}.qweight": packed, + f"{name}.scales": scales, + f"{name}.qzeros": torch.full((1, n), 0x55555555, dtype=torch.int32), + } + ) + blob_size = math.ceil(k * bits / 8) + weight = ( + torch.stack([(packed >> shift) & 255 for shift in (0, 8, 16, 24)], dim=-1) + .transpose(0, 1) + .reshape(n, -1)[:, :blob_size] + .to(torch.uint8) + .reshape(n, 1, blob_size) + ) + zero = 0x44 if method == "awq" and bits == 4 else (0x54 if method == "awq" else 0x55) + expected[name] = ( + k, + weight, + scales.T, + torch.full((n, 1), zero, dtype=torch.uint8), + ) + + _load(module, package, state_dict) + + for name, (k, weight, scales, zeros) in expected.items(): + _assert_layout(package, "model", f"{name}.weight", bits, k) + _assert_bound(package, "model", f"{name}.weight", weight) + _assert_bound(package, "model", f"{name}.scales", scales) + _assert_bound(package, "model", f"{name}.zero_points", zeros) + + @pytest.mark.parametrize(("method", "bits"), [("gptq", 8), ("awq", 4)]) def test_t5_raw_dotted_scales_survive_canonical_parameter_filter(method, bits): encoder = _quantization(quant_method=method, bits=bits, group_size=32) diff --git a/src/mobius/_component_quantization_test.py b/src/mobius/_component_quantization_test.py index 77352f28e..16458c849 100644 --- a/src/mobius/_component_quantization_test.py +++ b/src/mobius/_component_quantization_test.py @@ -156,6 +156,108 @@ def forward(self, op, x): configure_component_quantization(module, config, _SingleTask()) +def test_per_channel_layout_keeps_specialized_forward(): + class _SpecialQuantizedLinear(QuantizedLinear): + def __init__(self, in_features, out_features, bias=False): + super().__init__(in_features, out_features, block_size=-1, bias=bias) + + def forward(self, op, x): + return super().forward(op, x) + + module = _Projection(_SpecialQuantizedLinear) + original = module.proj + config = ArchitectureConfig( + component_quantization={ + "model": QuantizationConfig(bits=4, group_size=-1, quant_method="olive", sym=True) + } + ) + + configure_component_quantization(module, config, _SingleTask()) + + assert module.proj is original + assert module.proj._block_size == 64 + + +def test_per_channel_component_plan_preserves_output_head_bit_override(): + from mobius._configs import QuantizationOverride + + module = nn.Module() + module.proj = Linear(96, 32, bias=False) + module.lm_head = Linear(96, 64, bias=False) + quantization = QuantizationConfig( + bits=4, + group_size=-1, + quant_method="olive", + sym=True, + quantize_lm_head=True, + overrides={"lm_head": QuantizationOverride(bits=8)}, + ) + config = ArchitectureConfig(component_quantization={"model": quantization}) + task = _SingleTask() + manifest = configure_component_quantization(module, config, task) + packed = torch.arange(32 * 48).reshape(32, 48).to(torch.uint8) + head = torch.arange(64 * 96).reshape(64, 96).to(torch.uint8) + + result = normalize_component_quantized_weights( + { + "proj.weight_qweight": packed, + "proj.weight_scales": torch.ones(32, 1), + "lm_head.weight_qweight": head, + "lm_head.weight_scales": torch.ones(64, 1), + }, + module, + config, + ("model",), + manifest=manifest, + task=task, + ) + + assert (module.proj._bits, module.proj._block_size) == (4, 96) + assert (module.lm_head._bits, module.lm_head._block_size) == (8, 96) + assert result["proj.weight"].shape == (32, 1, 48) + assert result["lm_head.weight"].shape == (64, 1, 96) + torch.testing.assert_close(result["proj.weight"][:, 0, :48], packed) + torch.testing.assert_close(result["lm_head.weight"][:, 0, :96], head) + assert result["proj.scales"].shape == (32, 1) + assert result["lm_head.scales"].shape == (64, 1) + + +@pytest.mark.parametrize("method", ["olive", "gptq", "awq"]) +@pytest.mark.parametrize("bits", [2, 4, 8]) +@pytest.mark.parametrize("group_size", [-1, 33]) +def test_per_channel_component_uses_logical_width_for_packing(method, bits, group_size): + import math + + k, n = 33, 8 + module = nn.Module() + module.proj = Linear(k, n, bias=False) + quantization = QuantizationConfig( + bits=bits, group_size=group_size, quant_method=method, sym=True + ) + config = ArchitectureConfig(component_quantization={"model": quantization}) + task = _SingleTask() + manifest = configure_component_quantization(module, config, task) + if method == "olive": + packed = torch.arange(n * math.ceil(k * bits / 8)).to(torch.uint8).reshape(n, -1) + state_dict = { + "proj.weight_qweight": packed, + "proj.weight_scales": torch.ones(n, 1), + } + expected = packed + else: + packed = torch.arange(math.ceil(k * bits / 32) * n).to(torch.int32).reshape(-1, n) + state_dict = {"proj.qweight": packed, "proj.scales": torch.ones(1, n)} + expected = packed.T.contiguous().view(torch.uint8)[..., : math.ceil(k * bits / 8)] + + result = normalize_component_quantized_weights( + state_dict, module, config, ("model",), manifest=manifest, task=task + ) + + assert module.proj._block_size == k + assert result["proj.weight"].shape == (n, 1, math.ceil(k * bits / 8)) + torch.testing.assert_close(result["proj.weight"][:, 0], expected) + + def test_normalizes_weights_with_component_module_path_routing(): module = _Composite() config = _config() diff --git a/src/mobius/_quantization_utils.py b/src/mobius/_quantization_utils.py new file mode 100644 index 000000000..f3dbbf129 --- /dev/null +++ b/src/mobius/_quantization_utils.py @@ -0,0 +1,15 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Shared layout helpers for affine-quantized weights.""" + +from __future__ import annotations + + +def resolve_matmul_nbits_block_size(block_size: int, in_features: int) -> int: + """Resolve per-channel ``-1`` to one block of exactly K elements.""" + if block_size == -1: + if in_features <= 0: + raise ValueError(f"Per-channel quantization requires K > 0, got {in_features}") + return in_features + return block_size diff --git a/src/mobius/_weight_utils.py b/src/mobius/_weight_utils.py index 098877b7b..1ca957d22 100644 --- a/src/mobius/_weight_utils.py +++ b/src/mobius/_weight_utils.py @@ -23,6 +23,7 @@ import torch from mobius._configs import QuantizationConfig, QuantizedWeightFormat +from mobius._quantization_utils import resolve_matmul_nbits_block_size logger = logging.getLogger(__name__) @@ -884,7 +885,34 @@ def vlm_vision_weights( return renamed -def _reshape_packed_qweight(value: torch.Tensor, blob_size: int) -> torch.Tensor: +def _reshape_per_channel_qweight( + value: torch.Tensor, + bits: int, + *, + in_features: int | None = None, + storage_bits: int = 8, +) -> torch.Tensor: + """Keep one packed row per channel, trimming only producer word padding.""" + k = value.shape[-1] * 8 // bits if in_features is None else in_features + block_size = resolve_matmul_nbits_block_size(-1, k) + blob_size = math.ceil(block_size * bits / 8) + expected_bytes = math.ceil(block_size * bits / storage_bits) * (storage_bits // 8) + if value.shape[-1] != expected_bytes: + raise ValueError( + f"Per-channel packed width {value.shape[-1]} does not match " + f"K={k}, bits={bits}, storage_bits={storage_bits} " + f"(expected {expected_bytes} bytes)" + ) + return value[..., :blob_size].reshape(*value.shape[:-1], 1, blob_size).contiguous() + + +def _reshape_packed_qweight( + value: torch.Tensor, + bits: int, + group_size: int, + *, + in_features: int | None = None, +) -> torch.Tensor: """Transpose and reshape a packed qweight tensor for MatMulNBits. Converts ``[..., K_packed, N]`` int32 to @@ -893,7 +921,12 @@ def _reshape_packed_qweight(value: torch.Tensor, blob_size: int) -> torch.Tensor transposed = value.transpose(-1, -2).contiguous() prefix = transposed.shape[:-2] n = transposed.shape[-2] - packed = transposed.view(torch.uint8) + packed = transposed.reshape(-1).view(torch.uint8).reshape(*prefix, n, -1) + if group_size == -1: + return _reshape_per_channel_qweight( + packed, bits, in_features=in_features, storage_bits=32 + ) + blob_size = group_size * bits // 8 n_blocks = packed.shape[-1] // blob_size return packed.reshape(*prefix, n, n_blocks, blob_size) @@ -1250,6 +1283,8 @@ def preprocess_gptq_weights( state_dict: dict[str, torch.Tensor], bits: int = 4, group_size: int = 128, + *, + in_features: int | None = None, ) -> dict[str, torch.Tensor]: """Rename, transpose and reshape GPTQ weights for MatMulNBits. @@ -1270,7 +1305,9 @@ def preprocess_gptq_weights( Args: state_dict: Model state dict with GPTQ keys. bits: Quantization bit-width (typically 4). - group_size: Number of elements per quantization group. + group_size: Number of elements per quantization group, or ``-1`` for + per-output-channel quantization. + in_features: Logical K for per-channel weights, excluding int32 packing padding. Returns: State dict with renamed, transposed and reshaped weights. @@ -1278,14 +1315,19 @@ def preprocess_gptq_weights( import logging logger = logging.getLogger(__name__) - blob_size = group_size * bits // 8 + if group_size == in_features: + group_size = -1 result: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): if key.endswith(".g_idx"): # g_idx maps element i to its quantization group. For # non-desc_act models this is simply i // group_size. - trivial = torch.arange(value.numel(), dtype=value.dtype) // group_size + trivial = ( + torch.zeros_like(value) + if group_size == -1 + else torch.arange(value.numel(), dtype=value.dtype) // group_size + ) if not torch.equal(value, trivial): logger.warning( "Dropping %s — desc_act models with non-trivial " @@ -1296,7 +1338,9 @@ def preprocess_gptq_weights( if key.endswith(".qweight"): new_key = key.replace(".qweight", ".weight") - result[new_key] = _reshape_packed_qweight(value, blob_size) + result[new_key] = _reshape_packed_qweight( + value, bits, group_size, in_features=in_features + ) elif key.endswith(".qzeros"): new_key = key.replace(".qzeros", ".zero_points") @@ -1307,10 +1351,14 @@ def preprocess_gptq_weights( f"Missing {qw_key} — qweight must be present alongside qzeros for {key}" ) k = state_dict[qw_key].shape[-2] * 32 // bits - n_blocks = math.ceil(k / group_size) + n_blocks = 1 if group_size == -1 else math.ceil(k / group_size) result[new_key] = _reshape_packed_qzeros(value, bits, n_blocks) elif key.endswith(".scales"): + if group_size == -1 and value.shape[-2] != 1: + raise ValueError( + f"Per-channel scales for {key} must have one group, got {value.shape}" + ) result[key] = value.transpose(-1, -2).contiguous() else: @@ -1328,6 +1376,8 @@ def preprocess_olive_weights( tie_word_embeddings: bool = False, embed_key: str = "model.embed_tokens.weight", head_key: str = "lm_head.weight", + *, + in_features: int | None = None, ) -> dict[str, torch.Tensor]: """Rename and reshape Olive-packed quantized weights. @@ -1363,17 +1413,21 @@ def preprocess_olive_weights( Args: state_dict: Model state dict with Olive quantization keys. bits: Quantization bit-width (typically 4). - group_size: Number of elements per quantization group. + group_size: Number of elements per quantization group, or ``-1`` for + per-output-channel linear quantization. quantize_embeddings: Whether the embedding table is quantized. quantize_lm_head: Whether the LM head is quantized. tie_word_embeddings: Whether embedding and LM head share weights. embed_key: Float embedding key used for the tied-head fallback. head_key: Float LM-head key used for the tied-head fallback. + in_features: Logical K for per-channel linear weights. Returns: State dict with renamed and reshaped weights (and, for a float tied head, the aliased ``lm_head.weight``). """ + if group_size == in_features: + group_size = -1 blob_size = group_size * bits // 8 result: dict[str, torch.Tensor] = {} @@ -1440,12 +1494,17 @@ def _rename(key: str, raw_suffix: str, dotted_name: str) -> str: elif key.endswith("_qweight"): if value.dtype != torch.uint8: raise ValueError(f"Olive qweight must be uint8 for {key}, got {value.dtype}") + new_key = _rename(key, "_qweight", "weight") + if group_size == -1: + result[new_key] = _reshape_per_channel_qweight( + value, bits, in_features=in_features + ) + continue if value.shape[-1] % blob_size != 0: raise ValueError( f"Olive qweight packed dimension for {key} ({value.shape[-1]}) " f"must be divisible by blob_size ({blob_size})" ) - new_key = _rename(key, "_qweight", "weight") # Preserve all leading dims so fused expert-major tensors # ``[E, N, packed_K]`` reshape to ``[E, N, n_blocks, blob_size]`` # for the QMoE repacker; 2-D linears ``[N, packed_K]`` are @@ -1457,6 +1516,10 @@ def _rename(key: str, raw_suffix: str, dotted_name: str) -> str: new_key = _rename(key, "_qzeros", "zero_points") result[new_key] = value.contiguous() elif key.endswith("_scales"): + if group_size == -1 and value.shape[-1] != 1: + raise ValueError( + f"Per-channel scales for {key} must have one group, got {value.shape}" + ) new_key = _rename(key, "_scales", "scales") result[new_key] = value else: @@ -1480,6 +1543,8 @@ def preprocess_awq_weights( state_dict: dict[str, torch.Tensor], bits: int = 4, group_size: int = 128, + *, + in_features: int | None = None, ) -> dict[str, torch.Tensor]: """Rename, transpose and reshape AWQ weights for MatMulNBits. @@ -1491,18 +1556,23 @@ def preprocess_awq_weights( Args: state_dict: Model state dict with AWQ keys. bits: Quantization bit-width (typically 4). - group_size: Number of elements per quantization group. + group_size: Number of elements per quantization group, or ``-1`` for + per-output-channel quantization. + in_features: Logical K for per-channel weights, excluding int32 packing padding. Returns: State dict with renamed, transposed and reshaped weights. """ - blob_size = group_size * bits // 8 + if group_size == in_features: + group_size = -1 result: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): if key.endswith(".qweight"): new_key = key.replace(".qweight", ".weight") - result[new_key] = _reshape_packed_qweight(value, blob_size) + result[new_key] = _reshape_packed_qweight( + value, bits, group_size, in_features=in_features + ) elif key.endswith(".qzeros"): new_key = key.replace(".qzeros", ".zero_points") @@ -1513,7 +1583,7 @@ def preprocess_awq_weights( f"Missing {qw_key} — qweight must be present alongside qzeros for {key}" ) k = state_dict[qw_key].shape[-2] * 32 // bits - n_blocks = math.ceil(k / group_size) + n_blocks = 1 if group_size == -1 else math.ceil(k / group_size) # AWQ zero points have an implicit +1 offset; subtract it # before unpacking so MatMulNBits sees the raw value. # For 4-bit, each byte packs TWO nibbles — subtract per-nibble @@ -1530,6 +1600,10 @@ def preprocess_awq_weights( result[new_key] = zp_int16.clamp(min=0).to(torch.uint8) elif key.endswith(".scales"): + if group_size == -1 and value.shape[-2] != 1: + raise ValueError( + f"Per-channel scales for {key} must have one group, got {value.shape}" + ) result[key] = value.transpose(-1, -2).contiguous() else: @@ -1545,6 +1619,7 @@ def preprocess_quantized_weights( tie_embeddings: bool = False, embed_key: str = "model.embed_tokens.weight", head_key: str = "lm_head.weight", + in_features: int | None = None, qmoe_target_path: str | None = None, qmoe_quant_methods: Collection[str] = ("gptq", "awq", "olive"), qmoe_num_experts: int | None = None, @@ -1554,6 +1629,7 @@ def preprocess_quantized_weights( qmoe_source_moe_paths: tuple[str, ...] = (), reject_quantized_embeddings_lm_head: bool = False, defer_non_expert_sidecars: bool = False, + defer_per_channel_sidecars: bool = True, ) -> dict[str, torch.Tensor]: """Apply shared quantization conversion, tying, and QMoE packing. @@ -1564,6 +1640,7 @@ def preprocess_quantized_weights( tie_embeddings: Whether to alias missing float embedding/head weights. embed_key: Canonical float embedding-table key. head_key: Canonical float LM-head key. + in_features: Logical K when normalizing one per-channel linear weight. qmoe_target_path: MoE path replaced by packed QMoE parameter names. ``None`` means this model does not support or require QMoE handling. qmoe_quant_methods: Quantization methods supported by this caller when @@ -1580,6 +1657,8 @@ def preprocess_quantized_weights( defer_non_expert_sidecars: Keep ordinary packed sidecars raw for the component loader to resolve per-projection layouts. Expert packing and float weight tying still run here. + defer_per_channel_sidecars: Defer GPTQ/AWQ per-channel sidecars when + logical K is unknown. Leaf codecs disable this to normalize immediately. Returns: The preprocessed weight dictionary. @@ -1696,7 +1775,14 @@ def _is_packed_sidecar_of(key: str, float_key: str) -> bool: ) deferred = {} - if defer_non_expert_sidecars: + if defer_non_expert_sidecars or ( + defer_per_channel_sidecars + and in_features is None + and quantization is not None + and quantization.quant_method in {"gptq", "awq"} + and quantization.group_size == -1 + ): + # Int32 word padding can only be trimmed once the leaf supplies logical K. deferred = { key: value for key, value in state_dict.items() @@ -1707,11 +1793,17 @@ def _is_packed_sidecar_of(key: str, float_key: str) -> bool: if quantization is not None and quantization.quant_method == "gptq": state_dict = preprocess_gptq_weights( - state_dict, bits=quantization.bits, group_size=quantization.group_size + state_dict, + bits=quantization.bits, + group_size=quantization.group_size, + in_features=in_features, ) elif quantization is not None and quantization.quant_method == "awq": state_dict = preprocess_awq_weights( - state_dict, bits=quantization.bits, group_size=quantization.group_size + state_dict, + bits=quantization.bits, + group_size=quantization.group_size, + in_features=in_features, ) elif quantization is not None and quantization.quant_method == "olive": projection_layouts = { @@ -1762,6 +1854,7 @@ def _is_packed_sidecar_of(key: str, float_key: str) -> bool: tie_word_embeddings=olive_tie, embed_key=embed_key, head_key=head_key, + in_features=in_features, ) if use_qmoe and qmoe_target_path is not None: return_state_dict = stack_per_expert_moe_weights( diff --git a/src/mobius/_weight_utils_test.py b/src/mobius/_weight_utils_test.py index 1d5c48ae8..9d16a1132 100644 --- a/src/mobius/_weight_utils_test.py +++ b/src/mobius/_weight_utils_test.py @@ -5,6 +5,8 @@ from __future__ import annotations +import math + import pytest import torch @@ -602,6 +604,131 @@ def test_missing_qweight_raises(self): with pytest.raises(ValueError, match=r"Missing q_proj\.qweight"): preprocess_gptq_weights(sd, bits=self.BITS, group_size=self.GROUP_SIZE) + def test_per_channel_zero_group_indices_are_trivial(self, caplog): + result = preprocess_gptq_weights( + {"proj.g_idx": torch.zeros(96, dtype=torch.int32)}, group_size=-1 + ) + assert result == {} + assert "desc_act" not in caplog.text + + +@pytest.mark.parametrize("preprocess", [preprocess_gptq_weights, preprocess_awq_weights]) +@pytest.mark.parametrize("bits", [4, 8]) +@pytest.mark.parametrize("prefix", [(), (2,)]) +@pytest.mark.parametrize( + ("k", "group_size"), + [(7, -1), (9, -1), (33, -1), (96, -1), (1536, -1), (6144, -1), (33, 33)], +) +def test_per_channel_dotted_weights_preserve_codes_and_scales( + preprocess, bits, prefix, k, group_size +): + n = 8 + packed = ( + torch.arange(math.prod(prefix) * math.ceil(k * bits / 32) * n) + .reshape(*prefix, -1, n) + .to(torch.int32) + ) + scales = torch.arange(1, math.prod(prefix) * n + 1).reshape(*prefix, 1, n).float() + raw = { + "proj.qweight": packed, + "proj.scales": scales, + "proj.qzeros": torch.full((*prefix, 1, n), 0x55555555, dtype=torch.int32), + } + result = preprocess(raw, bits=bits, group_size=group_size, in_features=k) + weight = result["proj.weight"] + blob_size = math.ceil(k * bits / 8) + assert weight.shape == (*prefix, n, 1, blob_size) + expected = ( + packed.transpose(-1, -2) + .contiguous() + .reshape(-1) + .view(torch.uint8) + .reshape(*prefix, n, -1) + ) + torch.testing.assert_close(weight[..., 0, :], expected[..., :blob_size]) + torch.testing.assert_close(result["proj.scales"], scales.transpose(-1, -2)) + assert result["proj.zero_points"].shape == (*prefix, n, 1) + expected_zero = ( + 0x44 + if preprocess is preprocess_awq_weights and bits == 4 + else (0x54 if preprocess is preprocess_awq_weights else 0x55) + ) + assert (result["proj.zero_points"] == expected_zero).all() + + +@pytest.mark.parametrize("bits", [2, 4, 8]) +@pytest.mark.parametrize("prefix", [(), (2,)]) +@pytest.mark.parametrize( + ("k", "group_size"), + [(7, -1), (9, -1), (33, -1), (96, -1), (1536, -1), (6144, -1), (33, 33)], +) +def test_per_channel_olive_weights_preserve_codes_and_scales(bits, prefix, k, group_size): + n = 8 + blob_size = math.ceil(k * bits / 8) + packed = ( + torch.arange(math.prod(prefix) * n * blob_size) + .reshape(*prefix, n, blob_size) + .to(torch.uint8) + ) + scales = torch.arange(1, math.prod(prefix) * n + 1).reshape(*prefix, n, 1).float() + zeros = torch.full((*prefix, n, 1), 3, dtype=torch.uint8) + result = preprocess_olive_weights( + { + "proj.weight_qweight": packed, + "proj.weight_scales": scales, + "proj.weight_qzeros": zeros, + }, + bits=bits, + group_size=group_size, + in_features=k, + ) + weight = result["proj.weight"] + assert weight.shape == (*prefix, n, 1, blob_size) + torch.testing.assert_close(weight[..., 0, :], packed) + assert result["proj.scales"] is scales + assert result["proj.zero_points"] is zeros + + +@pytest.mark.parametrize("packed_bytes", [16, 18]) +@pytest.mark.parametrize("group_size", [-1, 33]) +def test_per_channel_olive_rejects_mismatched_logical_width(packed_bytes, group_size): + with pytest.raises(ValueError, match="Per-channel packed width"): + preprocess_olive_weights( + {"proj.weight_qweight": torch.zeros(8, packed_bytes, dtype=torch.uint8)}, + bits=4, + group_size=group_size, + in_features=33, + ) + + +@pytest.mark.parametrize("preprocess", [preprocess_gptq_weights, preprocess_awq_weights]) +@pytest.mark.parametrize("packed_words", [4, 6]) +@pytest.mark.parametrize("group_size", [-1, 33]) +def test_per_channel_dotted_rejects_mismatched_logical_width( + preprocess, packed_words, group_size +): + with pytest.raises(ValueError, match="Per-channel packed width"): + preprocess( + {"proj.qweight": torch.zeros(packed_words, 8, dtype=torch.int32)}, + bits=4, + group_size=group_size, + in_features=33, + ) + + +@pytest.mark.parametrize( + ("preprocess", "scales_key", "shape"), + [ + (preprocess_olive_weights, "proj.weight_scales", (8, 2)), + (preprocess_gptq_weights, "proj.scales", (2, 8)), + (preprocess_awq_weights, "proj.scales", (2, 8)), + ], +) +@pytest.mark.parametrize("group_size", [-1, 33]) +def test_per_channel_rejects_multiple_scale_groups(preprocess, scales_key, shape, group_size): + with pytest.raises(ValueError, match="must have one group"): + preprocess({scales_key: torch.ones(shape)}, group_size=group_size, in_features=33) + class TestPreprocessOliveWeights: """Tests for Olive uint8-packed weight preprocessing.""" diff --git a/src/mobius/components/_quantized_linear.py b/src/mobius/components/_quantized_linear.py index 5560e72a4..fb17e4f2e 100644 --- a/src/mobius/components/_quantized_linear.py +++ b/src/mobius/components/_quantized_linear.py @@ -16,13 +16,14 @@ import onnx_ir as ir from onnxscript import OpBuilder, nn -from mobius._build_context import ep_capabilities +from mobius._build_context import ep_capabilities, get_matmul_nbits_accuracy_level +from mobius._quantization_utils import resolve_matmul_nbits_block_size # MatMulNBits packs weights into uint8 blobs. The packed shape depends # on bits and block_size: # packed_weights: [N, n_blocks, blob_size] (uint8) # where n_blocks = ceil(K / block_size) -# blob_size = block_size * bits / 8 +# blob_size = ceil(block_size * bits / 8) # scales: [N, n_blocks] (float16/float32) # zero_points: [N, ceil(n_blocks*bits/8)] (uint8, optional, bit-packed) @@ -58,20 +59,30 @@ def _accuracy_level_attrs(bits: int) -> dict[str, int]: """Return the ``accuracy_level`` attribute for ``MatMulNBits``, if any. - Only emitted for 4-bit weights: ``accuracy_level`` is sourced from - ``EpCapabilities.default_int4_accuracy_level`` and its int8-accumulation - semantics are defined for INT4 ``MatMulNBits``. For 2-bit / 8-bit weights the - attribute is omitted so those paths keep ORT's default behavior. + Without an override, only emitted for 4-bit weights: ``accuracy_level`` is + sourced from ``EpCapabilities.default_int4_accuracy_level`` and its + int8-accumulation semantics are defined for INT4 ``MatMulNBits``. For + 2-bit / 8-bit weights the attribute is omitted so those paths keep ORT's + default behavior. ORT's MLAS CPU ``MatMulNBits`` kernel selects its compute path from the ``accuracy_level`` attribute: unset/0 keeps the highest-precision fp32 dequant + fp32 GEMM path, while ``4`` dynamically quantizes activations to int8 and uses int8 dot-products (SDOT/NEON on ARM, AVX-VNNI on x86) — the same class of kernel llama.cpp uses, and typically 2-4x faster on CPU with - no observable quality loss for Q4 weights. The value is sourced from the - active EP's :attr:`EpCapabilities.default_int4_accuracy_level` (4 for CPU / - WebGPU). When it is 0 the attribute is omitted so ORT keeps its default. + no observable quality loss for Q4 weights. The default value is sourced + from the active EP's :attr:`EpCapabilities.default_int4_accuracy_level` + (4 for CPU / WebGPU). When this EP default is 0 the attribute is omitted + so ORT keeps its default. + + An explicit build override takes precedence for every weight bit width, + including zero, which overrides an EP's nonzero default. It controls + activation computation, not weight quantization. Omitting the override + preserves the default behavior above. """ + override = get_matmul_nbits_accuracy_level() + if override is not None: + return {"accuracy_level": override} if bits != 4: return {} level = ep_capabilities().default_int4_accuracy_level @@ -92,7 +103,9 @@ class QuantizedLinear(nn.Module): in_features: Input dimension (K). out_features: Output dimension (N). bits: Quantization bit-width (2, 4, or 8). - block_size: Number of elements per quantization group. + block_size: Number of elements per quantization group, or ``-1`` for + one scale per output channel. A full-channel block equals K and + need not be a power of two; native kernel support is backend-specific. has_zero_point: Whether asymmetric zero-point is used. zero_point_dtype: Dtype for the zero_points parameter when ``has_zero_point`` is true. ``UINT8`` (default) uses ORT's @@ -117,8 +130,14 @@ def __init__( super().__init__() if bits not in (2, 4, 8): raise ValueError(f"bits must be 2, 4, or 8, got {bits}") - if block_size < 16 or (block_size & (block_size - 1)): - raise ValueError(f"block_size must be a power of 2 >= 16, got {block_size}") + block_size = resolve_matmul_nbits_block_size(block_size, in_features) + if not (block_size == in_features and in_features > 0) and ( + block_size < 16 or (block_size & (block_size - 1)) + ): + raise ValueError( + "block_size must be a power of 2 >= 16 or equal to K > 0, " + f"got {block_size} for K={in_features}" + ) self._bits = bits self._block_size = block_size @@ -126,7 +145,7 @@ def __init__( self._n = out_features n_blocks = math.ceil(in_features / block_size) - blob_size = block_size * bits // 8 + blob_size = math.ceil(block_size * bits / 8) # Packed quantized weight tensor (uint8) self.weight = nn.Parameter( @@ -599,7 +618,8 @@ def make_quantized_linear_factory( Args: bits: Quantization bit-width (typically 4). - block_size: Number of elements per quantization group. + block_size: Number of elements per quantization group, or ``-1`` for + per-output-channel quantization. has_zero_point: Whether to include zero-point parameters. zero_point_dtype: Dtype for the zero_points parameter (see :class:`QuantizedLinear`). diff --git a/src/mobius/components/_quantized_linear_test.py b/src/mobius/components/_quantized_linear_test.py index 9f35c182c..de5243904 100644 --- a/src/mobius/components/_quantized_linear_test.py +++ b/src/mobius/components/_quantized_linear_test.py @@ -21,7 +21,9 @@ BlockQuantizedLinear, ClippableQuantizedLinear, NVFP4QuantizedLinear, + QuantizedEmbedding, QuantizedLinear, + TiedQuantizedLMHead, ) # Test dimensions @@ -32,6 +34,39 @@ BLOCK_SIZE = 32 +@pytest.mark.parametrize("ep", ["default", "cpu", "qnn"]) +@pytest.mark.parametrize("bits", [2, 4, 8]) +@pytest.mark.parametrize("level", [None, 0, 1, 2, 3, 4]) +@pytest.mark.parametrize("kind", ["linear", "clippable", "tied"]) +def test_matmul_nbits_accuracy_override(ep, bits, level, kind): + embedding = QuantizedEmbedding(OUT_FEATURES, IN_FEATURES, bits=bits, block_size=32) + if kind == "tied": + linear = TiedQuantizedLMHead(embedding, IN_FEATURES, OUT_FEATURES) + else: + cls = QuantizedLinear if kind == "linear" else ClippableQuantizedLinear + linear = cls(IN_FEATURES, OUT_FEATURES, bits=bits, block_size=32) + with build_context(ep_registry.require(ep), matmul_nbits_accuracy_level=level): + builder, op, graph = create_test_builder() + x = create_test_input(builder, "x", [1, 4, IN_FEATURES]) + builder.add_output(linear(op, x), "y") + ids = create_test_input(builder, "ids", [1, 4], dtype=ir.DataType.INT64) + builder.add_output(embedding(op, ids), "embeddings") + + node = next(node for node in graph if node.op_type == "MatMulNBits") + assert node.attributes["bits"].value == bits + assert node.attributes["block_size"].value == 32 + expected = level + if expected is None and bits == 4: + expected = ep_registry.require(ep).default_int4_accuracy_level or None + if expected is None: + assert "accuracy_level" not in node.attributes + else: + assert node.attributes["accuracy_level"].value == expected + gather = next(node for node in graph if node.op_type == "GatherBlockQuantized") + assert "accuracy_level" not in gather.attributes + assert gather.attributes["bits"].value == bits + + class TestQuantizedLinearInit: def test_creates_packed_weight(self): ql = QuantizedLinear(IN_FEATURES, OUT_FEATURES, bits=4, block_size=32) @@ -97,6 +132,99 @@ def test_rejects_small_block_size(self): with pytest.raises(ValueError, match="block_size must be a power of 2 >= 16"): QuantizedLinear(IN_FEATURES, OUT_FEATURES, block_size=8) + @pytest.mark.parametrize("bits", [2, 4, 8]) + @pytest.mark.parametrize("k", [1, 7, 9, 16, 33, 64, 96, 1536, 6144]) + @pytest.mark.parametrize("explicit_block_size", [False, True]) + def test_per_channel_uses_exact_input_width(self, bits, k, explicit_block_size): + linear = QuantizedLinear( + k, OUT_FEATURES, bits=bits, block_size=k if explicit_block_size else -1 + ) + assert linear._k == k + assert linear._block_size == k + assert linear.weight.shape == [OUT_FEATURES, 1, math.ceil(k * bits / 8)] + assert linear.scales.shape == [OUT_FEATURES, 1] + + @pytest.mark.parametrize("k", [0, -1]) + def test_per_channel_rejects_nonpositive_width(self, k): + with pytest.raises(ValueError, match="requires K > 0"): + QuantizedLinear(k, OUT_FEATURES, block_size=-1) + + @pytest.mark.parametrize("block_size", [0, -2, -16]) + def test_other_nonpositive_block_sizes_remain_invalid(self, block_size): + with pytest.raises(ValueError, match="block_size must be a power of 2 >= 16"): + QuantizedLinear(IN_FEATURES, OUT_FEATURES, block_size=block_size) + + +@pytest.mark.parametrize( + ("bits", "inline", "k"), + [ + *((bits, False, k) for bits in (4, 8) for k in (16, 32, 64, 128, 256)), + *((4, True, k) for k in (1, 7, 9, 33, 96, 1536, 6144)), + ], +) +@pytest.mark.parametrize("asymmetric", [False, True]) +def test_per_channel_olive_export_matches_dequantized_weights( + tmp_path, k, bits, inline, asymmetric +): + import numpy as np + import onnxruntime as ort + import torch + from onnx_ir.passes.common import InlinePass + + from mobius._weight_utils import preprocess_olive_weights + from mobius.functions import register_function_bodies + + n = 8 + rng = np.random.default_rng(7) + values_per_byte = 8 // bits + packed_k = math.ceil(k / values_per_byte) + codes = rng.integers(0, 1 << bits, (n, packed_k * values_per_byte), dtype=np.uint8) + packed = np.bitwise_or.reduce( + codes.reshape(n, packed_k, values_per_byte) + << (np.arange(values_per_byte, dtype=np.uint8) * bits), + axis=-1, + ) + scales = rng.uniform(0.01, 0.05, (n, 1)).astype(np.float32) + zeros = ( + rng.integers(0, 1 << bits, (n, 1), dtype=np.uint8) + if asymmetric + else np.full((n, 1), 1 << (bits - 1), dtype=np.uint8) + ) + raw = { + "proj.weight_qweight": torch.from_numpy(packed), + "proj.weight_scales": torch.from_numpy(scales), + } + if asymmetric: + raw["proj.weight_qzeros"] = torch.from_numpy(zeros) + weights = preprocess_olive_weights(raw, bits=bits, group_size=-1, in_features=k) + linear = QuantizedLinear(k, n, bits=bits, block_size=-1, has_zero_point=asymmetric) + builder, op, graph = create_test_builder() + x = create_test_input(builder, "x", [2, 3, k]) + builder.add_output(linear(op, x), "y") + graph.opset_imports["com.microsoft"] = 1 + for name, weight in weights.items(): + value = graph.initializers[name.removeprefix("proj.")] + assert list(value.shape) == list(weight.shape) + value.const_value = ir.tensor(weight.numpy()) + node = next(node for node in graph if node.op_type == "MatMulNBits") + assert node.attributes["K"].value == k + assert node.attributes["block_size"].value == k + model = ir.Model(graph, ir_version=10) + if inline: + register_function_bodies(model) + InlinePass()(model) + assert not any(node.op_type == "MatMulNBits" for node in graph) + path = tmp_path / "per-channel.onnx" + ir.save(model, path) + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + session = ort.InferenceSession(str(path), options, providers=["CPUExecutionProvider"]) + inputs = rng.standard_normal((2, 3, k)).astype(np.float32) + actual = session.run(None, {"x": inputs})[0] + dequantized = (codes[:, :k].astype(np.float32) - zeros.astype(np.float32)) * scales + np.testing.assert_allclose(actual, inputs @ dequantized.T, rtol=1e-5, atol=1e-4) + class TestQuantizedLinearForward: def test_graph_has_matmulnbits_node(self): @@ -172,7 +300,7 @@ def test_cpu_ep_emits_accuracy_level_4(self): pytest.fail("MatMulNBits node not found") def test_cpu_ep_omits_accuracy_level_for_non_int4(self): - """accuracy_level is INT4-specific: 8-bit weights keep ORT's default.""" + """Without an override, 8-bit weights preserve the existing ORT default.""" with build_context(ep_registry.require("cpu")): ql = QuantizedLinear(IN_FEATURES, OUT_FEATURES, bits=8, block_size=32) b, op, graph = create_test_builder() diff --git a/src/mobius/functions/matmul_nbits.py b/src/mobius/functions/matmul_nbits.py index f2c5ee43c..00969e3b3 100644 --- a/src/mobius/functions/matmul_nbits.py +++ b/src/mobius/functions/matmul_nbits.py @@ -11,7 +11,7 @@ ``MatMulNBits`` node and forces the quantized MatMuls onto CPU. **Op semantics (4-bit):** the weight ``B`` is stored packed as -``(N, n_blocks, block_size/2)`` ``uint8`` — two 4-bit values per byte, low +``(N, n_blocks, ceil(block_size/2))`` ``uint8`` — two 4-bit values per byte, low nibble first. ``scales`` is ``(N, n_blocks)`` and the optional ``zero_points`` is ``(N, ceil(n_blocks/2))`` ``uint8`` (also 4-bit packed along the block axis). Element ``(n, k)`` dequantizes as ``(w[n,k] - zp[n, k//block]) * scale[n, k//block]`` @@ -60,7 +60,7 @@ def matmul_nbits(*, has_zero_points: bool = True) -> ir.Function: Inputs: A: (..., K) activation. - B: (N, n_blocks, block_size/2) uint8 packed 4-bit weight. + B: (N, n_blocks, ceil(block_size/2)) uint8 packed 4-bit weight. scales: (N, n_blocks) dequant scales. zero_points: (N, ceil(n_blocks/2)) uint8 packed 4-bit zero points. @@ -95,6 +95,14 @@ def body(op, a_input, b_packed, scales_input, zero_points_input): keep_n = op.Constant(value=ir.tensor(np.array([0, -1], dtype=np.int64))) w_blocks = op.Reshape(w_inter, keep_n_nb) w_2d = op.Reshape(w_blocks, keep_n) + # The last storage block can extend past K (including per-channel + # blocks). Discard padding before dequantization and the matmul. + w_2d = op.Slice( + w_2d, + op.Constant(value=ir.tensor(np.array([0], dtype=np.int64))), + op.Shape(a_input, start=-1), + op.Constant(value=ir.tensor(np.array([1], dtype=np.int64))), + ) w4 = op.Cast(w_2d, to=ir.DataType.UINT4) if has_zero_points: diff --git a/src/mobius/functions/matmul_nbits_test.py b/src/mobius/functions/matmul_nbits_test.py index 454f6795a..79a6fbb4b 100644 --- a/src/mobius/functions/matmul_nbits_test.py +++ b/src/mobius/functions/matmul_nbits_test.py @@ -19,6 +19,7 @@ from __future__ import annotations import json +import math import tempfile from pathlib import Path @@ -65,8 +66,8 @@ def _build_matmulnbits_model( Returns ``(model, feeds, packed_weight, scales, packed_zero_points)``. """ - nb = k_in // block - blob = block // 2 # 4-bit → 2 values per byte + nb = math.ceil(k_in / block) + blob = math.ceil(block / 2) # 4-bit → 2 values per byte rng = np.random.default_rng(seed) numpy_dtype = np.float16 if dtype == ir.DataType.FLOAT16 else np.float32 packed = rng.integers(0, 256, (n_out, nb, blob), dtype=np.uint8) @@ -205,6 +206,34 @@ def test_inline_matches_native_op_odd_blocks(self): got = _run(inlined, feeds2) np.testing.assert_allclose(got, reference, rtol=0, atol=0) + @pytest.mark.parametrize( + ("k", "block"), + [(7, 7), (33, 33), (96, 96), (1536, 1536), (7, 16), (33, 64), (65, 32)], + ) + @pytest.mark.parametrize("has_zero_points", [False, True]) + def test_inline_trims_single_block_padding(self, k, block, has_zero_points): + kwargs = { + "n_out": 8, + "k_in": k, + "block": block, + "has_zero_points": has_zero_points, + } + inlined, feeds, packed, scales, packed_zeros = _build_matmulnbits_model(**kwargs) + nb = math.ceil(k / block) + codes = np.empty((8, nb, packed.shape[-1] * 2), dtype=np.float32) + codes[..., 0::2] = packed & 0x0F + codes[..., 1::2] = packed >> 4 + zeros = np.full((8, nb), 8, dtype=np.float32) + if has_zero_points: + zeros[:, 0::2] = packed_zeros & 0x0F + zeros[:, 1::2] = (packed_zeros >> 4)[:, : nb // 2] + weights = ((codes - zeros[..., None]) * scales[..., None]).reshape(8, -1)[:, :k] + reference = feeds["A"] @ weights.T + + self._inline(inlined) + + np.testing.assert_allclose(_run(inlined, feeds), reference, rtol=1e-5, atol=1e-5) + @pytest.mark.parametrize("dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16]) @pytest.mark.parametrize("keep_empty_slot", [False, True]) @pytest.mark.parametrize("num_blocks", [1, 2, 3]) diff --git a/src/mobius/integrations/diffusers/_builder.py b/src/mobius/integrations/diffusers/_builder.py index bdd5a8532..e85fcad16 100644 --- a/src/mobius/integrations/diffusers/_builder.py +++ b/src/mobius/integrations/diffusers/_builder.py @@ -23,6 +23,7 @@ import torch import tqdm +from mobius._build_context import validate_matmul_nbits_accuracy_level from mobius._builder import build_from_module, resolve_dtype from mobius._model_package import ModelPackage from mobius._optimizations import fold_initializers_after_weights @@ -486,6 +487,7 @@ def build_diffusers_pipeline( unet_loras: dict | None = None, components: set[str] | None = None, execution_provider: str = "default", + matmul_nbits_accuracy_level: int | None = None, workflow_config: HierarchicalAudioWorkflowConfig | None = None, ) -> ModelPackage: """Build ONNX models for all supported components in a diffusers pipeline. @@ -510,6 +512,8 @@ def build_diffusers_pipeline( components: Optional component-name allowlist. Non-neural pipeline metadata is still retained so a single-component export preserves its contract. execution_provider: Target execution provider for EP-aware graph optimization. + matmul_nbits_accuracy_level: Optional accuracy level (0 through 4) for + every emitted ``MatMulNBits``. ``None`` preserves EP defaults. workflow_config: Optional typed, model-agnostic workflow description for pipelines whose ONNX GenAI execution cannot be derived from the graphs alone (currently hierarchical audio). When omitted, mobius supplies the @@ -526,6 +530,7 @@ def build_diffusers_pipeline( Raises: ValueError: If the model does not have a ``model_index.json``. """ + validate_matmul_nbits_accuracy_level(matmul_nbits_accuracy_level) _init_diffusers_class_map() pipeline_index = _load_diffusers_pipeline_index(model_id, revision=revision) @@ -606,6 +611,7 @@ def build_diffusers_pipeline( config, task_name, execution_provider=execution_provider, + matmul_nbits_accuracy_level=matmul_nbits_accuracy_level, ) # Flatten sub-package into the top-level package diff --git a/src/mobius/integrations/diffusers/_builder_test.py b/src/mobius/integrations/diffusers/_builder_test.py index 24425ecc3..3589f15f2 100644 --- a/src/mobius/integrations/diffusers/_builder_test.py +++ b/src/mobius/integrations/diffusers/_builder_test.py @@ -721,6 +721,19 @@ def _mock_build_for_vae(self, mock_load_index, mock_load_config, mock_build_from model = ir.Model(graph, ir_version=10) mock_build_from_module.return_value = ModelPackage({"model": model}) + @pytest.mark.parametrize("level", [None, 0, 4]) + @patch("mobius.integrations.diffusers._builder.build_from_module") + @patch("mobius.integrations.diffusers._builder._load_diffusers_component_config") + @patch("mobius.integrations.diffusers._builder._load_diffusers_pipeline_index") + def test_forwards_accuracy_override( + self, mock_load_index, mock_load_config, mock_build_from_module, level + ): + self._mock_build_for_vae(mock_load_index, mock_load_config, mock_build_from_module) + build_diffusers_pipeline( + "fake/vae-model", load_weights=False, matmul_nbits_accuracy_level=level + ) + assert mock_build_from_module.call_args.kwargs["matmul_nbits_accuracy_level"] == level + @patch("mobius.integrations.diffusers._builder.build_from_module") @patch( "mobius.integrations.diffusers._builder._load_diffusers_component_config", diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index 3dbe632f6..007b4e264 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -13,6 +13,7 @@ import onnx_ir as ir from onnxscript import nn +from mobius._build_context import validate_matmul_nbits_accuracy_level from mobius._builder import build_from_module, resolve_dtype from mobius._component_quantization import ( attach_hf_component_sources, @@ -335,6 +336,7 @@ def build_transformers_model( keep_quantized: bool = True, trust_remote_code: bool = False, execution_provider: str = "default", + matmul_nbits_accuracy_level: int | None = None, trace_optimization: bool = False, text_only: bool = False, fp8_kv_cache: bool = False, @@ -359,7 +361,14 @@ def build_transformers_model( ``keep_quantized`` preserves supported compressed-tensors checkpoints in their native block-weight representation. Set it to ``False`` only to request explicit dense reconstruction. + + ``matmul_nbits_accuracy_level`` optionally overrides the accuracy level of + every exported ``MatMulNBits`` (0 through 4), independently of weight bit + width. ``None`` preserves existing EP defaults. It does not affect + ``GatherBlockQuantized`` or requantize weights. PersonaPlex and RE-USE + bespoke builders do not support this option. """ + validate_matmul_nbits_accuracy_level(matmul_nbits_accuracy_level) if input_sampling_rate is not None and bwe_sampling_rate is not None: raise ValueError("input_sampling_rate and bwe_sampling_rate are mutually exclusive") @@ -376,6 +385,7 @@ def build_transformers_model( "output_layer_indices": output_layer_indices is not None, "trace_optimization": trace_optimization, "dequantize": not keep_quantized, + "matmul_nbits_accuracy_level": matmul_nbits_accuracy_level is not None, "text_only": text_only, "fp8_kv_cache": fp8_kv_cache, "kv_cache_scales": kv_cache_scales is not None, @@ -462,6 +472,7 @@ def build_transformers_model( ): raise ValueError("RE-USE checkpoints only support task='speech-enhancement'.") unsupported = { + "matmul_nbits_accuracy_level": matmul_nbits_accuracy_level is not None, "output_layer_indices": output_layer_indices is not None, "text_only": text_only, "fp8_kv_cache": fp8_kv_cache, @@ -508,6 +519,7 @@ def build_transformers_model( dtype=dtype, load_weights=load_weights, execution_provider=execution_provider, + matmul_nbits_accuracy_level=matmul_nbits_accuracy_level, ) if input_sampling_rate is not None or bwe_sampling_rate is not None: @@ -664,6 +676,7 @@ def build_transformers_model( config, task, execution_provider=execution_provider, + matmul_nbits_accuracy_level=matmul_nbits_accuracy_level, trace_optimization=trace_optimization, fp8_kv_cache=fp8_kv_cache, kv_cache_scales=kv_cache_scales, diff --git a/src/mobius/integrations/transformers/_builder_test.py b/src/mobius/integrations/transformers/_builder_test.py index 4f02da31a..6cf05524a 100644 --- a/src/mobius/integrations/transformers/_builder_test.py +++ b/src/mobius/integrations/transformers/_builder_test.py @@ -38,6 +38,68 @@ def build_fp8_streaming_plan(self, *_args): raise AssertionError("mock streaming function should own planning") +@pytest.mark.parametrize("level", [None, 0, 4]) +def test_build_forwards_matmul_nbits_accuracy_override(monkeypatch, level): + from mobius import build + + config = make_config() + hf_config = SimpleNamespace(model_type="llama") + package = ModelPackage( + {"model": ir.Model(ir.Graph([], [], nodes=[], name="model"), ir_version=10)} + ) + core_build = mock.Mock(return_value=package) + monkeypatch.setattr( + transformers_builder, "_load_transformers_config", lambda *a, **k: (hf_config, False) + ) + monkeypatch.setattr( + transformers_builder, + "_resolve_module_class", + lambda *a, **k: (_DummyModule, "text-generation", "llama"), + ) + monkeypatch.setattr(_config_resolver, "_config_from_hf", lambda *a, **k: config) + monkeypatch.setattr(transformers_builder, "build_from_module", core_build) + assert ( + build("local-model", load_weights=False, matmul_nbits_accuracy_level=level) is package + ) + assert core_build.call_args.kwargs["matmul_nbits_accuracy_level"] == level + + +@pytest.mark.parametrize("level", [-1, 5, True, False, 4.0, "4"]) +def test_build_rejects_invalid_accuracy_override_before_loading(monkeypatch, level): + loader = mock.Mock(side_effect=AssertionError("must validate before model loading")) + monkeypatch.setattr(transformers_builder, "_load_transformers_config", loader) + with pytest.raises(ValueError, match="matmul_nbits_accuracy_level"): + transformers_builder.build_transformers_model( + "local-model", matmul_nbits_accuracy_level=level + ) + loader.assert_not_called() + + +@pytest.mark.parametrize("level", [0, 4]) +def test_personaplex_rejects_accuracy_override(monkeypatch, level): + from mobius.integrations import _moshi + + monkeypatch.setattr(_moshi, "_is_personaplex_checkpoint", lambda *a, **k: True) + with pytest.raises(ValueError, match="matmul_nbits_accuracy_level"): + transformers_builder.build_transformers_model( + "local-personaplex", matmul_nbits_accuracy_level=level + ) + + +@pytest.mark.parametrize("level", [0, 4]) +def test_reuse_rejects_accuracy_override(monkeypatch, level): + from mobius.models import reuse + + monkeypatch.setattr( + transformers_builder, "_load_transformers_config", lambda *a, **k: (None, False) + ) + monkeypatch.setattr(reuse, "_is_reuse_checkpoint", lambda *a, **k: True) + with pytest.raises(ValueError, match="matmul_nbits_accuracy_level"): + transformers_builder.build_transformers_model( + "local-reuse", matmul_nbits_accuracy_level=level + ) + + def _native_gptoss_config(*, dtype=ir.DataType.FLOAT16): return make_config( model_type="gpt_oss", @@ -1364,7 +1426,8 @@ def test_glm_full_attention_rejects_non_glm_model_type(monkeypatch) -> None: transformers_builder.build_transformers_model("fake/model", glm_full_attention=True) -def test_build_threads_revision_to_diffusers_fallback(monkeypatch) -> None: +@pytest.mark.parametrize("accuracy_level", [None, 0, 4]) +def test_build_threads_revision_to_diffusers_fallback(monkeypatch, accuracy_level) -> None: import transformers monkeypatch.setattr( @@ -1392,6 +1455,7 @@ def fake_build_diffusers(*args, **kwargs): "fake/diffusers", revision="pinned-revision", load_weights=False, + matmul_nbits_accuracy_level=accuracy_level, ) assert result is expected @@ -1403,6 +1467,7 @@ def fake_build_diffusers(*args, **kwargs): "dtype": None, "load_weights": False, "execution_provider": "default", + "matmul_nbits_accuracy_level": accuracy_level, }, ) ] diff --git a/src/mobius/models/gemma4_test.py b/src/mobius/models/gemma4_test.py index e82e6397f..4e5053e46 100644 --- a/src/mobius/models/gemma4_test.py +++ b/src/mobius/models/gemma4_test.py @@ -561,6 +561,42 @@ def test_each_component_emits_its_own_layout(self): assert self._layouts(package["audio_encoder"].graph) == {(2, 16)} assert self._layouts(package["embedding"].graph) == {(8, 16)} + def test_decoder_per_channel_keeps_other_component_layouts(self): + from mobius import build_from_module + + config = self._config() + decoder = dataclasses.replace( + config.component_quantization["decoder"], + group_size=-1, + quantize_lm_head=True, + overrides={"lm_head": QuantizationOverride(bits=8)}, + ) + config = dataclasses.replace( + config, + hidden_size=96, + intermediate_size=192, + quantization=decoder, + component_quantization={**config.component_quantization, "decoder": decoder}, + ) + + package = build_from_module(Gemma4Model(config), config, task="gemma4") + + projections = [ + node for node in package["decoder"].graph if node.op_type == "MatMulNBits" + ] + assert projections + assert sum(node.attributes["bits"].as_int() == 8 for node in projections) == 1 + for node in projections: + k = node.attributes["K"].as_int() + n = node.attributes["N"].as_int() + assert node.attributes["block_size"].as_int() == k + bits = node.attributes["bits"].as_int() + assert list(node.inputs[1].shape) == [n, 1, (k * bits + 7) // 8] + assert list(node.inputs[2].shape) == [n, 1] + assert self._layouts(package["vision_encoder"].graph) == {(8, 32)} + assert self._layouts(package["audio_encoder"].graph) == {(2, 16)} + assert self._layouts(package["embedding"].graph) == {(8, 16)} + def test_preprocess_preserves_raw_component_sidecars(self): model = Gemma4Model(self._config()) state_dict = { diff --git a/src/mobius/weights/_codecs.py b/src/mobius/weights/_codecs.py index 4d2b0ed9c..b90d58cea 100644 --- a/src/mobius/weights/_codecs.py +++ b/src/mobius/weights/_codecs.py @@ -46,6 +46,7 @@ def normalize( config: QuantizationConfig, *, kind: Literal["linear", "embedding"] = "linear", + in_features: int | None = None, ) -> dict[str, torch.Tensor]: """Convert one packed record to Mobius's canonical parameter layout.""" @@ -177,6 +178,7 @@ def normalize( config: QuantizationConfig, *, kind: Literal["linear", "embedding"] = "linear", + in_features: int | None = None, ) -> dict[str, torch.Tensor]: if not isinstance(record.storage, PackedWeight): raise TypeError(f"weight {record.name!r} is not packed") @@ -196,6 +198,8 @@ def normalize( tie_embeddings=False, embed_key=record.name if kind == "embedding" else "", qmoe_target_path=None, + in_features=in_features, + defer_per_channel_sidecars=False, ) if kind == "embedding" and self.method in {"gptq", "awq"}: # The affine packers produce [rows, blocks, bytes]; Gather consumes diff --git a/src/mobius/weights/_codecs_test.py b/src/mobius/weights/_codecs_test.py index 450f8db12..552982812 100644 --- a/src/mobius/weights/_codecs_test.py +++ b/src/mobius/weights/_codecs_test.py @@ -125,11 +125,13 @@ def test_compatibility_normalizer_uses_existing_packer(): @pytest.mark.parametrize(("method", "bits"), [("gptq", 4), ("gptq", 8), ("awq", 4)]) -def test_dotted_affine_embedding_normalizes_to_gather_bytes(method, bits): +@pytest.mark.parametrize("group_size", [-1, 16]) +@pytest.mark.parametrize("kind", ["linear", "embedding"]) +def test_dotted_affine_weights_normalize_to_component_layout(method, bits, group_size, kind): component = ComponentDescriptor("embedding", "embedding", "embedding") config = QuantizationConfig( bits=bits, - group_size=16, + group_size=group_size, quant_method=method, sym=True, quantize_embeddings=True, @@ -137,7 +139,9 @@ def test_dotted_affine_embedding_normalizes_to_gather_bytes(method, bits): qweight = ( torch.arange(bits * 8, dtype=torch.int32).reshape(bits, 8) * 0x010203 + 0x12345678 ) - scales = torch.linspace(0.125, 1.125, 16).reshape(2, 8) + block_size = 32 if group_size == -1 else group_size + groups = 32 // block_size + scales = torch.linspace(0.125, 1.125, groups * 8).reshape(groups, 8) codec = codec_registry.get(method) record = codec.group( component, @@ -145,7 +149,7 @@ def test_dotted_affine_embedding_normalizes_to_gather_bytes(method, bits): config, )["embedding.tokens.weight"] - normalized = codec.normalize(record, config, kind="embedding") + normalized = codec.normalize(record, config, kind=kind) expected = ( torch.stack([(qweight >> shift) & 255 for shift in (0, 8, 16, 24)], dim=-1) @@ -153,9 +157,14 @@ def test_dotted_affine_embedding_normalizes_to_gather_bytes(method, bits): .reshape(8, 32 * bits // 8) .to(torch.uint8) ) - assert set(normalized) == {"embedding.tokens.qweight", "embedding.tokens.scales"} - assert normalized["embedding.tokens.qweight"].dtype == torch.uint8 - torch.testing.assert_close(normalized["embedding.tokens.qweight"], expected) + weight_key = ( + "embedding.tokens.qweight" if kind == "embedding" else "embedding.tokens.weight" + ) + if kind == "linear": + expected = expected.reshape(8, groups, block_size * bits // 8) + assert set(normalized) == {weight_key, "embedding.tokens.scales"} + assert normalized[weight_key].dtype == torch.uint8 + torch.testing.assert_close(normalized[weight_key], expected) torch.testing.assert_close(normalized["embedding.tokens.scales"], scales.T)