From 6e71f62676a4c11a12d379a291ba231af030af1a Mon Sep 17 00:00:00 2001 From: titaiwang Date: Mon, 21 Sep 2026 22:48:15 +0000 Subject: [PATCH 1/8] Support mixed-width QMoE export Follow the ORT QMoE projection-specific bit-width contract for Olive checkpoints with INT2 gate/up experts and INT4 down experts. Preserve legacy uniform INT4 emission and add strict checkpoint, package, and rewrite regression coverage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang --- .../per-component-quantization-loading.md | 19 + src/mobius/_configs/_base.py | 2 + src/mobius/_configs/_quantization.py | 5 +- src/mobius/_weight_utils.py | 383 +++++++++++++++- src/mobius/components/_moe.py | 54 ++- src/mobius/components/_moe_test.py | 4 +- src/mobius/models/moe.py | 30 +- src/mobius/models/moe_test.py | 417 +++++++++++++++++- src/mobius/models/qwen35.py | 56 ++- src/mobius/models/qwen35_test.py | 80 +++- src/mobius/rewrite_rules/_qmoe_fusion.py | 17 + src/mobius/rewrite_rules/_qmoe_fusion_test.py | 19 + 12 files changed, 1055 insertions(+), 31 deletions(-) diff --git a/docs/design/per-component-quantization-loading.md b/docs/design/per-component-quantization-loading.md index 488ef3ad1..a231b0369 100644 --- a/docs/design/per-component-quantization-loading.md +++ b/docs/design/per-component-quantization-loading.md @@ -437,6 +437,25 @@ preprocessor can still pack QMoE expert banks and tie floating-point tables. A per-component `WeightBundle` adapter is a future migration target, not the currently implemented interface. +Native QMoE additionally has a projection-specific adapter for Olive fused +K-last routed experts. The initial mixed layout is deliberately narrow: +model-wide INT4, FC1 `gate_up_proj` overridden to INT2, FC2 `down_proj` +remaining INT4, and one common power-of-two group size of at least 16. Exact +and `re:` module overrides are resolved against the original Hugging Face +paths before FC1 and FC2 buffers are validated independently. The adapter +requires every routed layer and validates rank, dtype, expert count, packed +byte width, scale geometry, and optional zero-point geometry before binding. +Qwen3.5-VL keeps the decoder's full module plan +(``preserve_module_plan=True``) while constructing the split package: expert +overrides are resolved against authoritative ``model.language_model`` source +paths, while validation and binding use the decoder-prefixed checkpoint roots +created by the VL weight router. +It emits QMoE's FC-specific bit attributes (including FC3=FC1 for fused +SwiGLU). Uniform INT4 omits those attributes and retains the legacy payload +path. Mixed-width execution remains an ONNX Runtime `NOT_IMPLEMENTED` +boundary; Mobius tests this serialized graph contract without claiming +runtime numerical support. + Appropriate model-specific operations include: - HuggingFace-to-ONNX name alignment; diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 7078c9125..b34eeda16 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -651,6 +651,8 @@ class ArchitectureConfig(BaseModelConfig): num_experts_per_tok: int | None = None moe_intermediate_size: int | None = None shared_expert_intermediate_size: int | None = None + # Internal HF path aliases used to resolve routed-expert module overrides. + qmoe_source_paths: tuple[str, ...] = () norm_topk_prob: bool = True # When True, the decoder layer uses post-norm style (FlexOLMo): norms are applied # to sub-layer outputs instead of inputs, with an extra post_feedforward_layernorm. diff --git a/src/mobius/_configs/_quantization.py b/src/mobius/_configs/_quantization.py index f954c6b03..defb49ffc 100644 --- a/src/mobius/_configs/_quantization.py +++ b/src/mobius/_configs/_quantization.py @@ -100,8 +100,9 @@ class QuantizationConfig: # Full HuggingFace module paths (including descendants) or ``re:``-prefixed # full-match regexes that remain floating point inside this component. modules_to_not_convert: tuple[str, ...] = () - # The same path/regex matching rules apply to per-module overrides. - # Insertion order is significant: the first matching override wins. + # Generic component collapse treats a plain path as a subtree. Producer- + # specific adapters may apply stricter matching (Olive uses literal + # equality); insertion order remains significant in either case. overrides: dict[str, QuantizationOverride] = dataclasses.field(default_factory=dict) # Keep this field last: QuantizationConfig has historically supported # positional construction, so inserting a field earlier would silently diff --git a/src/mobius/_weight_utils.py b/src/mobius/_weight_utils.py index af8e40732..a1d377175 100644 --- a/src/mobius/_weight_utils.py +++ b/src/mobius/_weight_utils.py @@ -14,8 +14,10 @@ from __future__ import annotations +import dataclasses import logging import math +import re from collections.abc import Collection, Sequence import torch @@ -159,6 +161,138 @@ def supported_qmoe_quantization( return quantization +@dataclasses.dataclass(frozen=True) +class QMoEQuantizationLayout: + """Projection-specific integer-affine layout for native ``QMoE``.""" + + fallback: QuantizationConfig + fc1: QuantizationConfig + fc2: QuantizationConfig + + @property + def is_mixed_width(self) -> bool: + return self.fc1.bits != self.fc2.bits + + +def _resolve_qmoe_module( + quantization: QuantizationConfig, + source_module_names: tuple[str, ...], +) -> QuantizationConfig | None: + """Resolve Olive's QMoE override semantics for authoritative source names.""" + + def _matches_exclusion(pattern: str, module_name: str) -> bool: + if pattern.startswith("re:"): + return re.fullmatch(pattern[3:], module_name) is not None + return bool(pattern) and pattern in module_name + + if any( + _matches_exclusion(pattern, module_name) + for pattern in quantization.modules_to_not_convert + for module_name in source_module_names + ): + return None + for pattern, override in quantization.overrides.items(): + if pattern.startswith("re:"): + matches = any( + re.fullmatch(pattern[3:], module_name) is not None + for module_name in source_module_names + ) + else: + matches = pattern in source_module_names + if matches: + return override.apply(quantization) + return quantization + + +def resolve_qmoe_quantization( + quantization: QuantizationConfig | None, + source_moe_paths: tuple[str, ...], +) -> QMoEQuantizationLayout | None: + """Resolve and validate FC1/FC2 layouts for one native QMoE layer. + + Legacy uniform INT4 layouts retain the existing predicate and behavior. + The first mixed layout is intentionally narrower: Olive integer-affine + FC1 INT2 plus FC2 INT4, with a common QMoE block size and the model-wide + INT4 layout as the fallback for non-expert projections. + """ + if quantization is None: + return None + if not source_moe_paths: + uniform = supported_qmoe_quantization(quantization) + return ( + QMoEQuantizationLayout(uniform, uniform, uniform) if uniform is not None else None + ) + + fc1 = _resolve_qmoe_module( + quantization, + tuple(f"{path}.experts.gate_up_proj" for path in source_moe_paths), + ) + fc2 = _resolve_qmoe_module( + quantization, + tuple(f"{path}.experts.down_proj" for path in source_moe_paths), + ) + if fc1 is None and fc2 is None: + return None + if fc1 is None or fc2 is None: + raise ValueError( + "Native QMoE requires both routed expert projections to be quantized; " + f"resolved FC1={fc1 is not None}, FC2={fc2 is not None} " + f"for {source_moe_paths!r}." + ) + + fallback = supported_qmoe_quantization(quantization) + fc1_uniform = supported_qmoe_quantization(fc1) + fc2_uniform = supported_qmoe_quantization(fc2) + if fc1_uniform is not None and fc2_uniform is not None: + if fallback is None: + return None + if ( + fc1.group_size != fc2.group_size + or fc1.sym != fc2.sym + or fc1.float_zero_point != fc2.float_zero_point + ): + raise ValueError( + "Native QMoE uniform INT4 projections must share one layout; " + f"FC1=(group_size={fc1.group_size}, sym={fc1.sym}), " + f"FC2=(group_size={fc2.group_size}, sym={fc2.sym})." + ) + return QMoEQuantizationLayout(fallback, fc1, fc2) + + requested_mixed = ( + fc1.bits != fc2.bits or fc1.group_size != fc2.group_size or fc1.sym != fc2.sym + ) + if not requested_mixed: + return None + layouts = (fc1, fc2) + if any( + layout.weight_format is not QuantizedWeightFormat.INTEGER_AFFINE + or layout.float_zero_point + or layout.quant_method != "olive" + for layout in layouts + ): + raise ValueError( + "Mixed native QMoE requires Olive integer-affine weights with packed " + "uint8 zero points." + ) + if fallback is None or (fc1.bits, fc2.bits) != (2, 4): + raise ValueError( + "Unsupported mixed native QMoE expert widths: expected model-wide " + f"INT4 with FC1 INT2 and FC2 INT4, got fallback={quantization.bits}, " + f"FC1={fc1.bits}, FC2={fc2.bits}." + ) + if fc1.group_size != fc2.group_size: + raise ValueError( + "Mixed native QMoE requires one common FC1/FC2 group_size, got " + f"{fc1.group_size} and {fc2.group_size}." + ) + block_size = fc1.group_size + if block_size < 16 or (block_size & (block_size - 1)) != 0: + raise ValueError( + f"Mixed native QMoE requires a power-of-two group_size >= 16, got {block_size}." + ) + return QMoEQuantizationLayout(fallback, fc1, fc2) + + def merge_lora_weights( base_state_dict: dict[str, torch.Tensor], lora_state_dict: dict[str, torch.Tensor], @@ -850,12 +984,206 @@ def pack_qmoe_expert_weights( if source in key: key = key.replace(source, target) if flatten_blocks: - value = value.flatten(-2) + # Older packers retain separate packed block and byte axes. + # Current fused Olive tensors are already rank 3 and must + # not have their logical output and packed-K axes merged. + if value.ndim >= 4: + value = value.flatten(-2) break packed[key] = value return packed +def _mixed_olive_qmoe_roots( + quantization: QuantizationConfig, + *, + expected_moe_paths: tuple[str, ...], + source_moe_paths: tuple[str, ...], +) -> dict[str, QMoEQuantizationLayout]: + if len(expected_moe_paths) != len(source_moe_paths): + raise ValueError( + "QMoE canonical checkpoint roots and source module paths must align " + f"one-to-one, got {len(expected_moe_paths)} and {len(source_moe_paths)}." + ) + source_by_root = dict(zip(expected_moe_paths, source_moe_paths)) + layouts = { + root: resolve_qmoe_quantization(quantization, (source_path,)) + for root, source_path in source_by_root.items() + } + return { + root: layout + for root, layout in layouts.items() + if layout is not None and layout.is_mixed_width + } + + +def _preprocess_mixed_olive_qmoe_weights( + state_dict: dict[str, torch.Tensor], + quantization: QuantizationConfig, + layouts: dict[str, QMoEQuantizationLayout], + *, + qmoe_target_path: str, + num_experts: int, + hidden_size: int, + intermediate_size: int, + tie_embeddings: bool, + embed_key: str, + head_key: str, +) -> dict[str, torch.Tensor]: + """Validate and bind fused K-last Olive sidecars for mixed-width QMoE.""" + expert_keys = { + key + for key in state_dict + if qmoe_target_path in key and ".experts." in key and is_packed_quant_key(key) + } + expected_roots = set(layouts) + unknown = next( + ( + key + for key in expert_keys + if not any(key.startswith(f"{root}.experts.") for root in expected_roots) + ), + None, + ) + if unknown is not None: + raise ValueError( + "Mixed native QMoE only supports fused Olive K-last expert sidecars; " + f"cannot map {unknown!r}." + ) + + renamed_experts: dict[str, torch.Tensor] = {} + accepted_keys: set[str] = set() + for root, layout in layouts.items(): + for projection, projection_layout in ( + ("gate_up_proj", layout.fc1), + ("down_proj", layout.fc2), + ): + stem = f"{root}.experts.{projection}" + accepted_keys.update((f"{stem}_qweight", f"{stem}_scales")) + if not projection_layout.sym: + accepted_keys.add(f"{stem}_qzeros") + stale = expert_keys - accepted_keys + if stale: + raise ValueError( + "Mixed native QMoE found an unknown or stale packed expert sidecar " + f"under a recognized root: {min(stale)!r}." + ) + + for root, layout in layouts.items(): + block_size = layout.fc1.group_size + if hidden_size % block_size or intermediate_size % block_size: + raise ValueError( + f"Mixed QMoE dimensions for {root!r} must be divisible by " + f"group_size={block_size}: hidden_size={hidden_size}, " + f"intermediate_size={intermediate_size}." + ) + projections = { + "gate_up_proj": ( + layout.fc1, + (num_experts, 2 * intermediate_size), + hidden_size, + "fc1", + ), + "down_proj": ( + layout.fc2, + (num_experts, hidden_size), + intermediate_size, + "fc2", + ), + } + for projection, ( + projection_layout, + leading_shape, + logical_k, + label, + ) in projections.items(): + stem = f"{root}.experts.{projection}" + required = { + "qweight": f"{stem}_qweight", + "scales": f"{stem}_scales", + } + if not projection_layout.sym: + required["qzeros"] = f"{stem}_qzeros" + missing = [name for name, key in required.items() if key not in state_dict] + if missing: + raise ValueError( + f"Mixed QMoE {label.upper()} sidecars for {root!r} are incomplete; " + f"missing {missing!r}." + ) + unexpected_qzeros = projection_layout.sym and f"{stem}_qzeros" in state_dict + if unexpected_qzeros: + raise ValueError( + f"Mixed QMoE {label.upper()} is symmetric but checkpoint contains " + f"{stem + '_qzeros'!r}." + ) + + qweight = state_dict[required["qweight"]] + scales = state_dict[required["scales"]] + qzeros = state_dict[required["qzeros"]] if "qzeros" in required else None + expected_weight = ( + *leading_shape, + logical_k * projection_layout.bits // 8, + ) + expected_scales = (*leading_shape, logical_k // block_size) + expected_zeros = ( + *leading_shape, + math.ceil((logical_k // block_size) * projection_layout.bits / 8), + ) + tensors = ( + ("qweight", qweight, expected_weight), + ("scales", scales, expected_scales), + *((("qzeros", qzeros, expected_zeros),) if qzeros is not None else ()), + ) + for name, tensor, expected_shape in tensors: + assert tensor is not None + if tensor.ndim != 3: + raise ValueError( + f"Mixed QMoE {label.upper()} {name} for {root!r} must be " + f"rank 3, got rank {tensor.ndim}." + ) + if tuple(tensor.shape) != expected_shape: + raise ValueError( + f"Mixed QMoE {label.upper()} {name} shape for {root!r} must " + f"be {expected_shape}, got {tuple(tensor.shape)}." + ) + if qweight.dtype != torch.uint8: + raise ValueError( + f"Mixed QMoE {label.upper()} qweight must be uint8, got {qweight.dtype}." + ) + if scales.dtype not in { + torch.float16, + torch.bfloat16, + torch.float32, + }: + raise ValueError( + f"Mixed QMoE {label.upper()} scales must be float16, bfloat16, " + f"or float32, got {scales.dtype}." + ) + if qzeros is not None and qzeros.dtype != torch.uint8: + raise ValueError( + f"Mixed QMoE {label.upper()} qzeros must be uint8, got {qzeros.dtype}." + ) + + renamed_experts[f"{stem}.weight"] = qweight.contiguous() + renamed_experts[f"{stem}.scales"] = scales + if qzeros is not None: + renamed_experts[f"{stem}.zero_points"] = qzeros.contiguous() + + ordinary = {key: value for key, value in state_dict.items() if key not in expert_keys} + result = preprocess_olive_weights( + ordinary, + bits=quantization.bits, + group_size=quantization.group_size, + quantize_embeddings=quantization.quantize_embeddings, + quantize_lm_head=quantization.quantize_lm_head, + tie_word_embeddings=tie_embeddings or quantization.tie_word_embeddings, + embed_key=embed_key, + head_key=head_key, + ) + result.update(renamed_experts) + return pack_qmoe_expert_weights(result, target_moe_path=qmoe_target_path) + + def preprocess_gptq_weights( state_dict: dict[str, torch.Tensor], bits: int = 4, @@ -1157,6 +1485,11 @@ def preprocess_quantized_weights( head_key: str = "lm_head.weight", qmoe_target_path: str | None = None, qmoe_quant_methods: Collection[str] = ("gptq", "awq", "olive"), + qmoe_num_experts: int | None = None, + qmoe_hidden_size: int | None = None, + qmoe_intermediate_size: int | None = None, + qmoe_expected_moe_paths: tuple[str, ...] = (), + qmoe_source_moe_paths: tuple[str, ...] = (), reject_quantized_embeddings_lm_head: bool = False, defer_non_expert_sidecars: bool = False, ) -> dict[str, torch.Tensor]: @@ -1173,6 +1506,13 @@ def preprocess_quantized_weights( ``None`` means this model does not support or require QMoE handling. qmoe_quant_methods: Quantization methods supported by this caller when the config matches the native QMoE ABI. + qmoe_num_experts: Expected routed expert count for mixed layouts. + qmoe_hidden_size: Expected logical FC1 input dimension. + qmoe_intermediate_size: Expected logical FC2 input dimension. + qmoe_expected_moe_paths: Canonical target paths for every routed layer. + When any mixed expert sidecar is present, every path is required. + qmoe_source_moe_paths: Authoritative Hugging Face module paths aligned + with ``qmoe_expected_moe_paths`` for override resolution. reject_quantized_embeddings_lm_head: Reject packed embedding/head weights because the caller's graph requires float parameters. defer_non_expert_sidecars: Keep ordinary packed sidecars raw for the @@ -1187,11 +1527,21 @@ def preprocess_quantized_weights( :func:`preprocess_olive_weights` must handle quantized embedding/head tying internally before optional QMoE packing. """ - use_qmoe = ( + uniform_qmoe = ( supported_qmoe_quantization(quantization) is not None if qmoe_target_path is not None else False ) + mixed_layouts = ( + _mixed_olive_qmoe_roots( + quantization, + expected_moe_paths=qmoe_expected_moe_paths, + source_moe_paths=qmoe_source_moe_paths or qmoe_expected_moe_paths, + ) + if qmoe_target_path is not None and quantization is not None + else {} + ) + use_qmoe = uniform_qmoe or bool(mixed_layouts) if ( use_qmoe and quantization is not None @@ -1289,6 +1639,35 @@ def _is_packed_sidecar_of(key: str, float_key: str) -> bool: state_dict, bits=quantization.bits, group_size=quantization.group_size ) elif quantization is not None and quantization.quant_method == "olive": + if mixed_layouts: + assert qmoe_target_path is not None + geometry = ( + qmoe_num_experts, + qmoe_hidden_size, + qmoe_intermediate_size, + ) + if any(value is None for value in geometry): + raise ValueError( + "Mixed QMoE preprocessing requires num_experts, hidden_size, " + "and intermediate_size for strict sidecar validation." + ) + assert qmoe_num_experts is not None + assert qmoe_hidden_size is not None + assert qmoe_intermediate_size is not None + return_state_dict = _preprocess_mixed_olive_qmoe_weights( + state_dict, + quantization, + mixed_layouts, + qmoe_target_path=qmoe_target_path, + num_experts=qmoe_num_experts, + hidden_size=qmoe_hidden_size, + intermediate_size=qmoe_intermediate_size, + tie_embeddings=tie_embeddings, + embed_key=embed_key, + head_key=head_key, + ) + return_state_dict.update(deferred) + return return_state_dict olive_tie = tie_embeddings or quantization.tie_word_embeddings return_state_dict = preprocess_olive_weights( state_dict, diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py index 3e8f88d5e..0013aa15d 100644 --- a/src/mobius/components/_moe.py +++ b/src/mobius/components/_moe.py @@ -12,13 +12,22 @@ import onnx_ir as ir from onnxscript import OpBuilder, nn -from mobius._configs import ArchitectureConfig +from mobius._configs import ArchitectureConfig, QuantizationConfig from mobius._weight_utils import ( - supported_qmoe_quantization as _supported_qmoe_quantization, + QMoEQuantizationLayout, + resolve_qmoe_quantization, + supported_qmoe_quantization, ) from mobius.components._mlp import MLP +def _supported_qmoe_quantization( + quantization: QuantizationConfig | None, +) -> QuantizationConfig | None: + """Preserve the legacy predicate imported by model-specific MoE modules.""" + return supported_qmoe_quantization(quantization) + + def _interleave_gate_up_rows( op: OpBuilder, tensor: ir.Value, num_experts: int, fc1_out: int ) -> ir.Value: @@ -457,11 +466,15 @@ def __init__( if config.component_quantization is not None else config.quantization ) - self._qmoe_quantization = ( + quantization_layout = ( None if getattr(config, "disable_qmoe", False) - else _supported_qmoe_quantization(quantization) + else resolve_qmoe_quantization( + quantization, + config.qmoe_source_paths, + ) ) + self._qmoe_quantization: QMoEQuantizationLayout | None = quantization_layout # Clipped-SwiGLU attributes (QMoE's ``activation_alpha``/``activation_beta``/ # ``swiglu_limit``). Left ``None`` by default so existing callers get a # byte-identical QMoE call (the attributes are simply omitted, even though @@ -505,8 +518,9 @@ def _init_qmoe_parameters(self, expert_config: ArchitectureConfig) -> None: assert quantization is not None hidden_size = expert_config.hidden_size intermediate_size = expert_config.intermediate_size - block_size = quantization.group_size - bits = quantization.bits + block_size = quantization.fc1.group_size + fc1_bits = quantization.fc1.bits + fc2_bits = quantization.fc2.bits fc1_out = 2 * intermediate_size self._fc1_out = fc1_out if hidden_size % block_size or intermediate_size % block_size: @@ -515,34 +529,36 @@ def _init_qmoe_parameters(self, expert_config: ArchitectureConfig) -> None: ) self.fc1_experts_weights = nn.Parameter( - [self.num_experts, fc1_out, hidden_size * bits // 8], + [self.num_experts, fc1_out, hidden_size * fc1_bits // 8], dtype=ir.DataType.UINT8, ) self.fc1_scales = nn.Parameter([self.num_experts, fc1_out, hidden_size // block_size]) self.fc2_experts_weights = nn.Parameter( - [self.num_experts, hidden_size, intermediate_size * bits // 8], + [self.num_experts, hidden_size, intermediate_size * fc2_bits // 8], dtype=ir.DataType.UINT8, ) self.fc2_scales = nn.Parameter( [self.num_experts, hidden_size, intermediate_size // block_size] ) - if quantization.sym: + if quantization.fc1.sym: self.fc1_experts_zero_points = None - self.fc2_experts_zero_points = None else: self.fc1_experts_zero_points = nn.Parameter( [ self.num_experts, fc1_out, - math.ceil((hidden_size // block_size) * bits / 8), + math.ceil((hidden_size // block_size) * fc1_bits / 8), ], dtype=ir.DataType.UINT8, ) + if quantization.fc2.sym: + self.fc2_experts_zero_points = None + else: self.fc2_experts_zero_points = nn.Parameter( [ self.num_experts, hidden_size, - math.ceil((intermediate_size // block_size) * bits / 8), + math.ceil((intermediate_size // block_size) * fc2_bits / 8), ], dtype=ir.DataType.UINT8, ) @@ -582,6 +598,15 @@ def _qmoe_forward(self, op: OpBuilder, hidden_states: ir.Value, *gate_args): activation_kwargs["activation_beta"] = self.activation_beta if self.swiglu_limit is not None: activation_kwargs["swiglu_limit"] = self.swiglu_limit + bit_width_kwargs = {} + if quantization.is_mixed_width: + bit_width_kwargs = { + "fc1_expert_weight_bits": quantization.fc1.bits, + "fc2_expert_weight_bits": quantization.fc2.bits, + # Fused SwiGLU stores FC3 in FC1. ORT requires the effective + # FC3 width to match FC1; omission would inherit fallback INT4. + "fc3_expert_weight_bits": quantization.fc1.bits, + } result = op.QMoE( hidden_states, router_probs, @@ -601,13 +626,14 @@ def _qmoe_forward(self, op: OpBuilder, hidden_states: ir.Value, *gate_args): activation_type="swiglu", normalize_routing_weights=int(normalize), k=self.top_k, - expert_weight_bits=quantization.bits, - block_size=quantization.group_size, + expert_weight_bits=quantization.fallback.bits, + block_size=quantization.fc1.group_size, swiglu_fusion=1, quant_type="int", weights_prepacked=0, _domain="com.microsoft", **activation_kwargs, + **bit_width_kwargs, ) if output_scale != 1.0: # noqa: RUF069 result = op.Mul(result, op.CastLike(output_scale, result)) diff --git a/src/mobius/components/_moe_test.py b/src/mobius/components/_moe_test.py index 405a61927..7c6508a82 100644 --- a/src/mobius/components/_moe_test.py +++ b/src/mobius/components/_moe_test.py @@ -20,12 +20,12 @@ pack_qmoe_expert_weights, preprocess_gptq_weights, preprocess_olive_weights, + supported_qmoe_quantization, ) from mobius.components._moe import ( MoELayer, SparseMixerGate, TopKGate, - _supported_qmoe_quantization, ) @@ -193,7 +193,7 @@ def test_int4_olive_moe_emits_expert_major_qmoe(self): def test_olive_group_size_not_power_of_two_falls_back_to_dense(self): """Non-pow2 block_size is unrunnable by CUDA QMoE -> dense fallback.""" quant = QuantizationConfig(bits=4, group_size=48, quant_method="olive", sym=False) - assert _supported_qmoe_quantization(quant) is None + assert supported_qmoe_quantization(quant) is None config = make_config( hidden_size=96, intermediate_size=48, diff --git a/src/mobius/models/moe.py b/src/mobius/models/moe.py index 243bf66c3..2cc50917e 100644 --- a/src/mobius/models/moe.py +++ b/src/mobius/models/moe.py @@ -38,6 +38,17 @@ from mobius.models.phi3 import split_fused_qkv +def _config_for_qmoe_layer( + config: ArchitectureConfig, + layer_idx: int, +) -> ArchitectureConfig: + """Copy config with the canonical HF source path for one routed MoE layer.""" + return dataclasses.replace( + config, + qmoe_source_paths=(f"model.layers.{layer_idx}.mlp",), + ) + + def _quantized_linear_class(config: ArchitectureConfig) -> type | None: """Return a QuantizedLinear factory when the checkpoint is quantized. @@ -93,6 +104,17 @@ def _preprocess_moe_weights(model: CausalLMModel, state_dict) -> dict: tie_embeddings=model.config.tie_word_embeddings, qmoe_target_path=".mlp", qmoe_quant_methods=("gptq", "awq", "olive"), + qmoe_num_experts=model.config.num_local_experts, + qmoe_hidden_size=model.config.hidden_size, + qmoe_intermediate_size=model.config.moe_intermediate_size, + qmoe_expected_moe_paths=tuple( + f"model.layers.{layer_idx}.mlp" + for layer_idx in range(model.config.num_hidden_layers) + ), + qmoe_source_moe_paths=tuple( + f"model.layers.{layer_idx}.mlp" + for layer_idx in range(model.config.num_hidden_layers) + ), defer_non_expert_sidecars=( model.config.component_quantization is not None or (quantization is not None and quantization.has_module_plan) @@ -251,8 +273,12 @@ def _make_gate() -> nn.Module: ) self.layers = nn.ModuleList( [ - _layer_class(config, gate=_make_gate(), norm_class=norm_class) - for _ in range(config.num_hidden_layers) + _layer_class( + _config_for_qmoe_layer(config, layer_idx), + gate=_make_gate(), + norm_class=norm_class, + ) + for layer_idx in range(config.num_hidden_layers) ] ) self.norm = norm_class(config.hidden_size, eps=config.rms_norm_eps) diff --git a/src/mobius/models/moe_test.py b/src/mobius/models/moe_test.py index 3674c7452..69cd53796 100644 --- a/src/mobius/models/moe_test.py +++ b/src/mobius/models/moe_test.py @@ -23,13 +23,19 @@ from __future__ import annotations +import dataclasses import math import onnx_ir as ir import pytest import torch -from mobius._configs import QuantizationConfig +from mobius._configs import ( + QuantizationConfig, + QuantizationOverride, + QuantizedWeightFormat, +) +from mobius._model_package import ModelPackage from mobius._testing import make_config from mobius.models.moe import MoECausalLMModel, _rename_moe_expert_weights @@ -357,3 +363,412 @@ def test_unquantized_model_uses_dense_expert_fallback(self): assert not any("fc1_experts_weights" in k for k in out) for name, tensor in out.items(): assert tuple(params[name].shape) == tuple(tensor.shape) + + +def _mixed_quantization( + *, + pattern: str = f"{_LAYER}mlp.experts.gate_up_proj", + fc1_sym: bool = True, + fc2_sym: bool = True, + fc1_group_size: int = _BLK, +) -> QuantizationConfig: + return QuantizationConfig( + bits=4, + group_size=_BLK, + quant_method="olive", + sym=fc2_sym, + overrides={ + pattern: QuantizationOverride( + bits=2, + group_size=fc1_group_size, + sym=fc1_sym, + ) + }, + ) + + +def _mixed_expert_state_dict( + *, + num_experts: int = _E, + hidden_size: int = _H, + intermediate_size: int = _INT, + block_size: int = _BLK, + fc1_sym: bool = True, + fc2_sym: bool = True, +) -> dict[str, torch.Tensor]: + prefix = f"{_LAYER}mlp.experts." + fc1_out = 2 * intermediate_size + state = { + prefix + "gate_up_proj_qweight": torch.arange( + num_experts * fc1_out * (hidden_size * 2 // 8), + dtype=torch.int64, + ) + .remainder(256) + .to(torch.uint8) + .reshape(num_experts, fc1_out, hidden_size * 2 // 8), + prefix + "gate_up_proj_scales": torch.rand( + num_experts, + fc1_out, + hidden_size // block_size, + dtype=torch.bfloat16, + ), + prefix + "down_proj_qweight": torch.arange( + num_experts * hidden_size * (intermediate_size * 4 // 8), + dtype=torch.int64, + ) + .remainder(256) + .to(torch.uint8) + .reshape(num_experts, hidden_size, intermediate_size * 4 // 8), + prefix + "down_proj_scales": torch.rand( + num_experts, + hidden_size, + intermediate_size // block_size, + dtype=torch.bfloat16, + ), + } + if not fc1_sym: + state[prefix + "gate_up_proj_qzeros"] = torch.randint( + 0, + 256, + ( + num_experts, + fc1_out, + math.ceil((hidden_size // block_size) * 2 / 8), + ), + dtype=torch.uint8, + ) + if not fc2_sym: + state[prefix + "down_proj_qzeros"] = torch.randint( + 0, + 256, + ( + num_experts, + hidden_size, + math.ceil((intermediate_size // block_size) * 4 / 8), + ), + dtype=torch.uint8, + ) + return state + + +class TestMixedWidthQMoEExport: + @pytest.mark.parametrize( + "pattern", + [ + f"{_LAYER}mlp.experts.gate_up_proj", + r"re:model\.layers\.\d+\.mlp\.experts\.gate_up_proj", + ], + ) + def test_exact_and_regex_fc1_override_emit_mixed_qmoe(self, pattern): + config = _moe_config(_mixed_quantization(pattern=pattern)) + model = MoECausalLMModel(config) + block = model.model.layers[0].mlp + + assert block.experts is None + assert block.fc1_experts_weights.shape == ir.Shape([_E, _FC1_OUT, _H * 2 // 8]) + assert block.fc2_experts_weights.shape == ir.Shape([_E, _H, _INT * 4 // 8]) + + from mobius._builder import build_from_module + + graph = build_from_module(model, config)["model"].graph + qmoe = next(node for node in graph if node.op_type == "QMoE") + assert qmoe.attributes["expert_weight_bits"].value == 4 + assert qmoe.attributes["fc1_expert_weight_bits"].value == 2 + assert qmoe.attributes["fc2_expert_weight_bits"].value == 4 + assert qmoe.attributes["fc3_expert_weight_bits"].value == 2 + + @pytest.mark.parametrize("exact_first", [True, False]) + def test_plain_parent_override_does_not_shadow_exact_regex(self, exact_first): + parent = (f"{_LAYER}mlp", QuantizationOverride(bits=8)) + regex = ( + r"re:model\.layers\.\d+\.mlp\.experts\.gate_up_proj", + QuantizationOverride(bits=2), + ) + overrides = dict((parent, regex) if exact_first else (regex, parent)) + quantization = QuantizationConfig( + bits=4, + group_size=_BLK, + quant_method="olive", + overrides=overrides, + ) + + block = MoECausalLMModel(_moe_config(quantization)).model.layers[0].mlp + + assert block._qmoe_quantization is not None + assert block._qmoe_quantization.fc1.bits == 2 + assert block._qmoe_quantization.fc2.bits == 4 + + def test_mixed_olive_payload_is_bound_byte_exactly(self): + model = MoECausalLMModel(_moe_config(_mixed_quantization())) + raw = _mixed_expert_state_dict() + + out = model.preprocess_weights(raw) + + prefix = f"{_LAYER}mlp." + torch.testing.assert_close( + out[prefix + "fc1_experts_weights"], + raw[f"{_LAYER}mlp.experts.gate_up_proj_qweight"], + rtol=0, + atol=0, + ) + torch.testing.assert_close( + out[prefix + "fc2_experts_weights"], + raw[f"{_LAYER}mlp.experts.down_proj_qweight"], + rtol=0, + atol=0, + ) + + @pytest.mark.parametrize( + ("fc1_sym", "fc2_sym"), + [(False, True), (True, False), (False, False)], + ) + def test_independent_zero_point_presence(self, fc1_sym, fc2_sym): + model = MoECausalLMModel( + _moe_config(_mixed_quantization(fc1_sym=fc1_sym, fc2_sym=fc2_sym)) + ) + out = model.preprocess_weights( + _mixed_expert_state_dict(fc1_sym=fc1_sym, fc2_sym=fc2_sym) + ) + prefix = f"{_LAYER}mlp." + + assert (prefix + "fc1_experts_zero_points" in out) is not fc1_sym + assert (prefix + "fc2_experts_zero_points" in out) is not fc2_sym + + def test_odd_zero_point_byte_ceiling(self): + hidden_size, intermediate_size, num_experts = 80, 48, 3 + quantization = _mixed_quantization(fc1_sym=False, fc2_sym=False) + config = make_config( + num_hidden_layers=1, + hidden_size=hidden_size, + intermediate_size=64, + moe_intermediate_size=intermediate_size, + num_local_experts=num_experts, + num_experts_per_tok=2, + num_attention_heads=5, + num_key_value_heads=1, + head_dim=16, + vocab_size=32, + quantization=quantization, + ) + model = MoECausalLMModel(config) + out = model.preprocess_weights( + _mixed_expert_state_dict( + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + fc1_sym=False, + fc2_sym=False, + ) + ) + prefix = f"{_LAYER}mlp." + + assert out[prefix + "fc1_experts_zero_points"].shape[-1] == 2 + assert out[prefix + "fc2_experts_zero_points"].shape[-1] == 2 + + @pytest.mark.parametrize( + ("mutation", "message"), + [ + ("missing_scale", "incomplete"), + ("wrong_rank", "rank 3"), + ("wrong_dtype", "must be uint8"), + ("wrong_scale_dtype", "must be float16, bfloat16, or float32"), + ("wrong_experts", "shape"), + ("wrong_packed_bytes", "shape"), + ("wrong_scale_geometry", "shape"), + ], + ) + def test_invalid_mixed_sidecars_fail_closed(self, mutation, message): + state = _mixed_expert_state_dict() + gate_weight = f"{_LAYER}mlp.experts.gate_up_proj_qweight" + gate_scales = f"{_LAYER}mlp.experts.gate_up_proj_scales" + if mutation == "missing_scale": + del state[gate_scales] + elif mutation == "wrong_rank": + state[gate_weight] = state[gate_weight].flatten(1) + elif mutation == "wrong_dtype": + state[gate_weight] = state[gate_weight].to(torch.int8) + elif mutation == "wrong_scale_dtype": + state[gate_scales] = state[gate_scales].to(torch.float64) + elif mutation == "wrong_experts": + state[gate_weight] = state[gate_weight][:-1] + elif mutation == "wrong_packed_bytes": + state[gate_weight] = state[gate_weight][..., :-1] + elif mutation == "wrong_scale_geometry": + state[gate_scales] = state[gate_scales][..., :-1] + + model = MoECausalLMModel(_moe_config(_mixed_quantization())) + with pytest.raises(ValueError, match=message): + model.preprocess_weights(state) + + def test_unknown_packed_sidecar_under_recognized_root_fails_closed(self): + state = _mixed_expert_state_dict() + state[f"{_LAYER}mlp.experts.gate_proj_qweight"] = torch.zeros( + _E, _INT, _H // 4, dtype=torch.uint8 + ) + model = MoECausalLMModel(_moe_config(_mixed_quantization())) + + with pytest.raises(ValueError, match="unknown or stale packed expert sidecar"): + model.preprocess_weights(state) + + def test_unknown_packed_sidecar_root_fails_closed(self): + state = _mixed_expert_state_dict() + state["model.layers.9.mlp.experts.gate_up_proj_qweight"] = torch.zeros( + _E, 2 * _INT, _H // 4, dtype=torch.uint8 + ) + model = MoECausalLMModel(_moe_config(_mixed_quantization())) + + with pytest.raises(ValueError, match="only supports fused Olive K-last"): + model.preprocess_weights(state) + + def test_olive_substring_exclusion_prevents_partial_qmoe(self): + quantization = dataclasses.replace( + _mixed_quantization(), + modules_to_not_convert=("experts.gate_up_proj",), + ) + + with pytest.raises(ValueError, match="both routed expert projections"): + MoECausalLMModel(_moe_config(quantization)) + + @pytest.mark.parametrize( + ("mutation", "message"), + [ + ("missing", "incomplete"), + ("shape", "qzeros shape"), + ("dtype", "qzeros must be uint8"), + ], + ) + def test_invalid_zero_point_sidecars_fail_closed(self, mutation, message): + quantization = _mixed_quantization(fc1_sym=False) + state = _mixed_expert_state_dict(fc1_sym=False) + key = f"{_LAYER}mlp.experts.gate_up_proj_qzeros" + if mutation == "missing": + del state[key] + elif mutation == "shape": + state[key] = torch.cat((state[key], state[key]), dim=-1) + elif mutation == "dtype": + state[key] = state[key].to(torch.int8) + + model = MoECausalLMModel(_moe_config(quantization)) + with pytest.raises(ValueError, match=message): + model.preprocess_weights(state) + + def test_mismatched_group_size_fails_at_construction(self): + with pytest.raises(ValueError, match="common FC1/FC2 group_size"): + MoECausalLMModel(_moe_config(_mixed_quantization(fc1_group_size=32))) + + def test_uniform_int4_omits_projection_specific_attributes(self): + config = _moe_config(_quantization()) + model = MoECausalLMModel(config) + + from mobius._builder import build_from_module + + graph = build_from_module(model, config)["model"].graph + qmoe = next(node for node in graph if node.op_type == "QMoE") + assert "fc1_expert_weight_bits" not in qmoe.attributes + assert "fc2_expert_weight_bits" not in qmoe.attributes + assert "fc3_expert_weight_bits" not in qmoe.attributes + + def test_uniform_int4_experts_do_not_bypass_unsupported_global_fallback(self): + quantization = QuantizationConfig( + bits=2, + group_size=_BLK, + quant_method="olive", + overrides={ + f"{_LAYER}mlp.experts.gate_up_proj": QuantizationOverride(bits=4), + f"{_LAYER}mlp.experts.down_proj": QuantizationOverride(bits=4), + }, + ) + + block = MoECausalLMModel(_moe_config(quantization)).model.layers[0].mlp + + assert block._qmoe_quantization is None + assert block.experts is not None + + def test_model_package_roundtrip_preserves_mixed_attrs_and_bytes(self, tmp_path): + from mobius._builder import build_from_module + + config = _moe_config(_mixed_quantization()) + module = MoECausalLMModel(config) + package = build_from_module(module, config) + processed = module.preprocess_weights(_mixed_expert_state_dict()) + package.apply_weights(processed, fold_constants=False) + names = ( + f"{_LAYER}mlp.fc1_experts_weights", + f"{_LAYER}mlp.fc2_experts_weights", + ) + expected = { + name: package["model"].graph.initializers[name].const_value.numpy().tobytes() + for name in names + } + + package.save(str(tmp_path), check_weights=False, progress_bar=False) + reloaded = ModelPackage.load(str(tmp_path))["model"] + qmoe = next(node for node in reloaded.graph if node.op_type == "QMoE") + + assert qmoe.attributes["fc1_expert_weight_bits"].value == 2 + assert qmoe.attributes["fc2_expert_weight_bits"].value == 4 + assert qmoe.attributes["fc3_expert_weight_bits"].value == 2 + for name, payload in expected.items(): + assert reloaded.graph.initializers[name].const_value.numpy().tobytes() == payload + + @pytest.mark.parametrize( + ("quantization", "message"), + [ + ( + QuantizationConfig( + bits=4, + group_size=_BLK, + quant_method="gptq", + overrides={ + f"{_LAYER}mlp.experts.gate_up_proj": QuantizationOverride(bits=2) + }, + ), + "requires Olive integer-affine", + ), + ( + QuantizationConfig( + bits=4, + group_size=_BLK, + quant_method="olive", + weight_format=QuantizedWeightFormat.MXFP4, + overrides={ + f"{_LAYER}mlp.experts.gate_up_proj": QuantizationOverride(bits=2) + }, + ), + "requires Olive integer-affine", + ), + ( + QuantizationConfig( + bits=4, + group_size=_BLK, + quant_method="olive", + overrides={ + f"{_LAYER}mlp.experts.gate_up_proj": QuantizationOverride(bits=8) + }, + ), + "expected model-wide INT4 with FC1 INT2 and FC2 INT4", + ), + ], + ) + def test_unsupported_mixed_layout_fails_at_construction(self, quantization, message): + with pytest.raises(ValueError, match=message): + MoECausalLMModel(_moe_config(quantization)) + + def test_missing_entire_routed_layer_fails_closed(self): + quantization = _mixed_quantization( + pattern=r"re:model\.layers\.\d+\.mlp\.experts\.gate_up_proj" + ) + config = dataclasses.replace( + _moe_config(quantization), + num_hidden_layers=2, + ) + model = MoECausalLMModel(config) + + with pytest.raises(ValueError, match=r"model\.layers\.1\.mlp.*incomplete"): + model.preprocess_weights(_mixed_expert_state_dict()) + + def test_missing_all_routed_expert_sidecars_fails_closed(self): + model = MoECausalLMModel(_moe_config(_mixed_quantization())) + + with pytest.raises(ValueError, match=r"model\.layers\.0\.mlp.*incomplete"): + model.preprocess_weights({}) diff --git a/src/mobius/models/qwen35.py b/src/mobius/models/qwen35.py index 72cd58d12..64297e9a4 100644 --- a/src/mobius/models/qwen35.py +++ b/src/mobius/models/qwen35.py @@ -96,11 +96,17 @@ def _linear_factory(config: ArchitectureConfig) -> type | None: def _decoder_component_config( config: ArchitectureConfig, source_paths: tuple[str, ...], + *, + preserve_module_plan: bool = False, ) -> ArchitectureConfig: """Use the effective decoder layout while constructing its expert parameters.""" if config.component_quantization is None: return config - quantization = config.quantization_for_source_paths("decoder", source_paths) + quantization = ( + config.quantization_for("decoder") + if preserve_module_plan + else config.quantization_for_source_paths("decoder", source_paths) + ) return dataclasses.replace(config, quantization=quantization, component_quantization=None) @@ -421,7 +427,14 @@ class Qwen35MoEDecoderLayer(Qwen35DecoderLayer): def __init__(self, config: ArchitectureConfig, layer_idx: int): super().__init__(config, layer_idx) - self.mlp = Qwen35MoEBlock(config, linear_class=_linear_factory(config)) + layer_config = dataclasses.replace( + config, + qmoe_source_paths=(f"model.language_model.layers.{layer_idx}.mlp",), + ) + self.mlp = Qwen35MoEBlock( + layer_config, + linear_class=_linear_factory(config), + ) class Qwen35MoETextModel(nn.Module): @@ -586,6 +599,17 @@ def preprocess_weights( tie_embeddings=effective_tie_word_embeddings(self.config), qmoe_target_path=".mlp", qmoe_quant_methods=("gptq", "awq", "olive"), + qmoe_num_experts=self.config.num_local_experts, + qmoe_hidden_size=self.config.hidden_size, + qmoe_intermediate_size=self.config.moe_intermediate_size, + qmoe_expected_moe_paths=tuple( + f"model.layers.{layer_idx}.mlp" + for layer_idx in range(self.config.num_hidden_layers) + ), + qmoe_source_moe_paths=tuple( + f"model.language_model.layers.{layer_idx}.mlp" + for layer_idx in range(self.config.num_hidden_layers) + ), defer_non_expert_sidecars=( self.config.component_quantization is not None or (quantization is not None and quantization.has_module_plan) @@ -899,6 +923,7 @@ def __init__(self, config: ArchitectureConfig): decoder_config = _decoder_component_config( config, self.HF_COMPONENT_SOURCES["decoder"], + preserve_module_plan=True, ) self.decoder = Qwen35MoEVLDecoderModel(decoder_config) self.vision_encoder = Qwen3VLVisionEncoderModel(config) @@ -927,10 +952,7 @@ def preprocess_weights( """ quantization = self.config.quantization decoder_quantization = ( - self.config.quantization_for_source_paths( - "decoder", - self.HF_COMPONENT_SOURCES["decoder"], - ) + self.config.quantization_for("decoder") if self.config.component_quantization is not None else quantization ) @@ -1006,6 +1028,17 @@ def preprocess_weights( head_key="decoder.lm_head.weight", qmoe_target_path=".mlp", qmoe_quant_methods=("olive",), + qmoe_num_experts=self.config.num_local_experts, + qmoe_hidden_size=self.config.hidden_size, + qmoe_intermediate_size=self.config.moe_intermediate_size, + qmoe_expected_moe_paths=tuple( + f"decoder.model.layers.{layer_idx}.mlp" + for layer_idx in range(self.config.num_hidden_layers) + ), + qmoe_source_moe_paths=tuple( + f"model.language_model.layers.{layer_idx}.mlp" + for layer_idx in range(self.config.num_hidden_layers) + ), reject_quantized_embeddings_lm_head=True, defer_non_expert_sidecars=True, ) @@ -1019,6 +1052,17 @@ def preprocess_weights( head_key="decoder.lm_head.weight", qmoe_target_path=".mlp", qmoe_quant_methods=("olive",), + qmoe_num_experts=self.config.num_local_experts, + qmoe_hidden_size=self.config.hidden_size, + qmoe_intermediate_size=self.config.moe_intermediate_size, + qmoe_expected_moe_paths=tuple( + f"decoder.model.layers.{layer_idx}.mlp" + for layer_idx in range(self.config.num_hidden_layers) + ), + qmoe_source_moe_paths=tuple( + f"model.language_model.layers.{layer_idx}.mlp" + for layer_idx in range(self.config.num_hidden_layers) + ), reject_quantized_embeddings_lm_head=True, ) if tie: diff --git a/src/mobius/models/qwen35_test.py b/src/mobius/models/qwen35_test.py index ca02b77c7..552d86814 100644 --- a/src/mobius/models/qwen35_test.py +++ b/src/mobius/models/qwen35_test.py @@ -362,6 +362,27 @@ def _olive_expert_state_dict() -> dict[str, torch.Tensor]: } +def _mixed_olive_expert_state_dict() -> dict[str, torch.Tensor]: + state = _olive_expert_state_dict() + prefix = "model.language_model.layers.0.mlp.experts." + state[prefix + "gate_up_proj_qweight"] = state[prefix + "gate_up_proj_qweight"][ + ..., : _H * 2 // 8 + ].contiguous() + return state + + +def _mixed_qwen35_quantization( + pattern: str = "model.language_model.layers.0.mlp.experts.gate_up_proj", +) -> QuantizationConfig: + return QuantizationConfig( + bits=4, + group_size=_BLK, + quant_method="olive", + sym=False, + overrides={pattern: QuantizationOverride(bits=2)}, + ) + + def _moe_vl_config( quantization: QuantizationConfig | None, *, tie_word_embeddings: bool = False ) -> object: @@ -395,6 +416,43 @@ def _moe_vl_config( class TestQwen35MoEQMoEExport: + def test_language_model_source_prefix_resolves_mixed_layout(self): + model = Qwen35MoECausalLMModel(_moe_config(_mixed_qwen35_quantization())) + block = model.model.layers[0].mlp + + assert block.experts is None + assert block.fc1_experts_weights.shape[-1] == _H * 2 // 8 + assert block.fc2_experts_weights.shape[-1] == _INT * 4 // 8 + + out = model.preprocess_weights(_mixed_olive_expert_state_dict()) + prefix = "model.layers.0.mlp." + assert out[prefix + "fc1_experts_weights"].shape[-1] == _H * 2 // 8 + assert out[prefix + "fc2_experts_weights"].shape[-1] == _INT * 4 // 8 + + def test_language_model_regex_resolves_mixed_layout_and_sidecars(self): + quantization = _mixed_qwen35_quantization( + r"re:model\.language_model\.layers\.\d+\.mlp\.experts\.gate_up_proj" + ) + model = Qwen35MoECausalLMModel(_moe_config(quantization)) + + out = model.preprocess_weights(_mixed_olive_expert_state_dict()) + block = model.model.layers[0].mlp + + assert block._qmoe_quantization is not None + assert block._qmoe_quantization.fc1.bits == 2 + assert out["model.layers.0.mlp.fc1_experts_weights"].shape[-1] == _H * 2 // 8 + + def test_qwen3_only_regex_does_not_match_qwen35_source_alias(self): + quantization = _mixed_qwen35_quantization( + r"re:model\.layers\.\d+\.mlp\.experts\.gate_up_proj" + ) + + block = Qwen35MoECausalLMModel(_moe_config(quantization)).model.layers[0].mlp + + assert block._qmoe_quantization is not None + assert not block._qmoe_quantization.is_mixed_width + assert block._qmoe_quantization.fc1.bits == 4 + def test_moe_block_uses_qmoe_when_quantized(self): model = _moe_config( QuantizationConfig(bits=4, group_size=_BLK, quant_method="olive", sym=False) @@ -481,6 +539,22 @@ def test_unsupported_qmoe_quantization_with_packed_experts_raises(self): class TestQwen35MoEVL3ModelQMoEExport: + def test_component_plan_preserves_mixed_expert_overrides(self): + quantization = _mixed_qwen35_quantization() + config = dataclasses.replace( + _moe_vl_config(None), + component_quantization={"decoder": quantization}, + ) + model = Qwen35MoEVL3ModelCausalLMModel(config) + block = model.decoder.model.layers[0].mlp + + assert block.experts is None + assert block.fc1_experts_weights.shape[-1] == _H * 2 // 8 + result = model.preprocess_weights(_mixed_olive_expert_state_dict()) + assert ( + result["decoder.model.layers.0.mlp.fc1_experts_weights"].shape[-1] == _H * 2 // 8 + ) + def test_plan_only_quantization_keeps_graph_and_qmoe_weights_aligned(self): quantization = QuantizationConfig( bits=4, @@ -1018,7 +1092,9 @@ def test_moe_vl_decoder_override_sizes_fused_qmoe_parameters(self): bits=4, group_size=16, quant_method="olive", - overrides={"model.language_model": QuantizationOverride(bits=4, group_size=32)}, + overrides={ + r"re:model\.language_model\..*": QuantizationOverride(bits=4, group_size=32) + }, ) config = dataclasses.replace( _moe_vl_config(quantization), @@ -1030,7 +1106,7 @@ def test_moe_vl_decoder_override_sizes_fused_qmoe_parameters(self): block = model.decoder.model.layers[0].mlp assert block._qmoe_quantization is not None - assert block._qmoe_quantization.group_size == 32 + assert block._qmoe_quantization.fc1.group_size == 32 assert tuple(block.fc1_scales.shape) == (_E, 64, 1) assert tuple(block.fc2_scales.shape) == (_E, _H, 1) diff --git a/src/mobius/rewrite_rules/_qmoe_fusion.py b/src/mobius/rewrite_rules/_qmoe_fusion.py index 21d1c6cd0..19f24114e 100644 --- a/src/mobius/rewrite_rules/_qmoe_fusion.py +++ b/src/mobius/rewrite_rules/_qmoe_fusion.py @@ -311,6 +311,16 @@ def _qmoe_abi_supported(bits: int, block_size: int) -> bool: return block_size >= 16 and (block_size & (block_size - 1)) == 0 +def _has_uniform_expert_geometry(layer: _DenseMoELayer) -> bool: + """Whether every gate/up/down projection uses one legacy INT4 layout.""" + geometries = { + _expert_geometry(node) + for projections in layer.experts.values() + for node in (projections.gate, projections.up, projections.down) + } + return len(geometries) == 1 + + def _pack_projection(nodes: list[ir.Node], slot: int) -> np.ndarray: """Stack ``flatten(-2)`` of one ``MatMulNBits`` weight input across experts.""" stacked = [ @@ -486,6 +496,13 @@ def fuse_dense_moe_to_qmoe(model: ir.Model) -> int: layer.topk.name, ) continue + if not _has_uniform_expert_geometry(layer): + logger.warning( + "skipping MoE layer at %s: mixed expert projection layouts are " + "not supported by dense-MoE rewrite fusion", + layer.topk.name, + ) + continue down = layer.experts[sorted(layer.experts)[0]].down bits, block_size, _ = _expert_geometry(down) if not _qmoe_abi_supported(bits, block_size): diff --git a/src/mobius/rewrite_rules/_qmoe_fusion_test.py b/src/mobius/rewrite_rules/_qmoe_fusion_test.py index 65ceff1ff..a304e3ff2 100644 --- a/src/mobius/rewrite_rules/_qmoe_fusion_test.py +++ b/src/mobius/rewrite_rules/_qmoe_fusion_test.py @@ -564,5 +564,24 @@ def test_unsupported_geometry_keeps_dense_fallback() -> None: assert _count(graph, "TopK") == 1 +def test_mixed_projection_widths_keep_dense_fallback() -> None: + """The legacy rewrite must not pack INT2 FC1 bytes as uniform INT4 QMoE.""" + model, _, _ = _build_dense_graph() + graph = model.graph + for node in graph: + if ( + node.op_type == "MatMulNBits" + and node.name + and (".gate_proj" in node.name or ".up_proj" in node.name) + ): + node.attributes["bits"] = ir.AttrInt64("bits", 2) + + fused = fuse_dense_moe_to_qmoe(model) + + assert fused == 0 + assert _count(graph, "QMoE") == 0 + assert _count(graph, "TopK") == 1 + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__, "-v"])) From f4dc1792b20216747f0f483960d87fe24d071ed0 Mon Sep 17 00:00:00 2001 From: titaiwang Date: Tue, 22 Sep 2026 00:08:05 +0000 Subject: [PATCH 2/8] Fix per-layer QMoE preprocessing Resolve every routed layer against its effective Olive layout so mixed and legacy uniform QMoE layers can coexist. Fail closed when packed sidecars target dense fallbacks and honor projection-specific packing geometry. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang --- src/mobius/_weight_utils.py | 92 +++++++++++++++++++++++-------- src/mobius/models/moe_test.py | 101 ++++++++++++++++++++++++++++++++-- 2 files changed, 165 insertions(+), 28 deletions(-) diff --git a/src/mobius/_weight_utils.py b/src/mobius/_weight_utils.py index a1d377175..f0cf41c48 100644 --- a/src/mobius/_weight_utils.py +++ b/src/mobius/_weight_utils.py @@ -173,6 +173,25 @@ class QMoEQuantizationLayout: def is_mixed_width(self) -> bool: return self.fc1.bits != self.fc2.bits + @property + def requires_projection_preprocessing(self) -> bool: + fallback_layout = ( + self.fallback.bits, + self.fallback.group_size, + self.fallback.sym, + self.fallback.float_zero_point, + ) + return any( + ( + projection.bits, + projection.group_size, + projection.sym, + projection.float_zero_point, + ) + != fallback_layout + for projection in (self.fc1, self.fc2) + ) + def _resolve_qmoe_module( quantization: QuantizationConfig, @@ -994,30 +1013,25 @@ def pack_qmoe_expert_weights( return packed -def _mixed_olive_qmoe_roots( +def _resolve_olive_qmoe_layouts( quantization: QuantizationConfig, *, expected_moe_paths: tuple[str, ...], source_moe_paths: tuple[str, ...], -) -> dict[str, QMoEQuantizationLayout]: +) -> dict[str, QMoEQuantizationLayout | None]: if len(expected_moe_paths) != len(source_moe_paths): raise ValueError( "QMoE canonical checkpoint roots and source module paths must align " f"one-to-one, got {len(expected_moe_paths)} and {len(source_moe_paths)}." ) source_by_root = dict(zip(expected_moe_paths, source_moe_paths)) - layouts = { + return { root: resolve_qmoe_quantization(quantization, (source_path,)) for root, source_path in source_by_root.items() } - return { - root: layout - for root, layout in layouts.items() - if layout is not None and layout.is_mixed_width - } -def _preprocess_mixed_olive_qmoe_weights( +def _preprocess_olive_qmoe_weights( state_dict: dict[str, torch.Tensor], quantization: QuantizationConfig, layouts: dict[str, QMoEQuantizationLayout], @@ -1051,6 +1065,11 @@ def _preprocess_mixed_olive_qmoe_weights( f"cannot map {unknown!r}." ) + projection_layouts = { + root: layout + for root, layout in layouts.items() + if layout.requires_projection_preprocessing + } renamed_experts: dict[str, torch.Tensor] = {} accepted_keys: set[str] = set() for root, layout in layouts.items(): @@ -1069,7 +1088,7 @@ def _preprocess_mixed_olive_qmoe_weights( f"under a recognized root: {min(stale)!r}." ) - for root, layout in layouts.items(): + for root, layout in projection_layouts.items(): block_size = layout.fc1.group_size if hidden_size % block_size or intermediate_size % block_size: raise ValueError( @@ -1169,7 +1188,14 @@ def _preprocess_mixed_olive_qmoe_weights( if qzeros is not None: renamed_experts[f"{stem}.zero_points"] = qzeros.contiguous() - ordinary = {key: value for key, value in state_dict.items() if key not in expert_keys} + projection_expert_keys = { + key + for key in expert_keys + if any(key.startswith(f"{root}.experts.") for root in projection_layouts) + } + ordinary = { + key: value for key, value in state_dict.items() if key not in projection_expert_keys + } result = preprocess_olive_weights( ordinary, bits=quantization.bits, @@ -1532,8 +1558,8 @@ def preprocess_quantized_weights( if qmoe_target_path is not None else False ) - mixed_layouts = ( - _mixed_olive_qmoe_roots( + resolved_qmoe_layouts = ( + _resolve_olive_qmoe_layouts( quantization, expected_moe_paths=qmoe_expected_moe_paths, source_moe_paths=qmoe_source_moe_paths or qmoe_expected_moe_paths, @@ -1541,7 +1567,25 @@ def preprocess_quantized_weights( if qmoe_target_path is not None and quantization is not None else {} ) - use_qmoe = uniform_qmoe or bool(mixed_layouts) + packed_expert_keys = [ + key + for key in state_dict + if qmoe_target_path is not None + and qmoe_target_path in key + and ".experts." in key + and is_packed_quant_key(key) + ] + unsupported_roots = { + root for root, layout in resolved_qmoe_layouts.items() if layout is None + } + if unsupported_roots and packed_expert_keys: + raise ValueError( + "Quantized MoE expert weights were found for routed layers whose resolved " + "quantization layout does not match the native QMoE ABI: " + f"{sorted(unsupported_roots)!r}." + ) + layer_qmoe = bool(resolved_qmoe_layouts) and not unsupported_roots + use_qmoe = uniform_qmoe if not resolved_qmoe_layouts else layer_qmoe if ( use_qmoe and quantization is not None @@ -1561,11 +1605,6 @@ def preprocess_quantized_weights( ) if qmoe_target_path is not None and not use_qmoe: - packed_expert_keys = [ - key - for key in state_dict - if qmoe_target_path in key and ".experts." in key and is_packed_quant_key(key) - ] if packed_expert_keys: raise ValueError( "Quantized MoE expert weights were found for this model " @@ -1639,7 +1678,12 @@ def _is_packed_sidecar_of(key: str, float_key: str) -> bool: state_dict, bits=quantization.bits, group_size=quantization.group_size ) elif quantization is not None and quantization.quant_method == "olive": - if mixed_layouts: + projection_layouts = { + root: layout + for root, layout in resolved_qmoe_layouts.items() + if layout is not None and layout.requires_projection_preprocessing + } + if projection_layouts: assert qmoe_target_path is not None geometry = ( qmoe_num_experts, @@ -1654,10 +1698,14 @@ def _is_packed_sidecar_of(key: str, float_key: str) -> bool: assert qmoe_num_experts is not None assert qmoe_hidden_size is not None assert qmoe_intermediate_size is not None - return_state_dict = _preprocess_mixed_olive_qmoe_weights( + return_state_dict = _preprocess_olive_qmoe_weights( state_dict, quantization, - mixed_layouts, + { + root: layout + for root, layout in resolved_qmoe_layouts.items() + if layout is not None + }, qmoe_target_path=qmoe_target_path, num_experts=qmoe_num_experts, hidden_size=qmoe_hidden_size, diff --git a/src/mobius/models/moe_test.py b/src/mobius/models/moe_test.py index 69cd53796..7033e9c66 100644 --- a/src/mobius/models/moe_test.py +++ b/src/mobius/models/moe_test.py @@ -59,9 +59,11 @@ def _quantization(*, sym: bool = True, **overrides) -> QuantizationConfig: ) -def _moe_config(quantization: QuantizationConfig | None) -> object: +def _moe_config( + quantization: QuantizationConfig | None, *, num_hidden_layers: int = 1 +) -> object: return make_config( - num_hidden_layers=1, + num_hidden_layers=num_hidden_layers, hidden_size=_H, intermediate_size=64, moe_intermediate_size=_INT, @@ -389,23 +391,25 @@ def _mixed_quantization( def _mixed_expert_state_dict( *, + layer_index: int = 0, num_experts: int = _E, hidden_size: int = _H, intermediate_size: int = _INT, block_size: int = _BLK, + fc1_bits: int = 2, fc1_sym: bool = True, fc2_sym: bool = True, ) -> dict[str, torch.Tensor]: - prefix = f"{_LAYER}mlp.experts." + prefix = f"model.layers.{layer_index}.mlp.experts." fc1_out = 2 * intermediate_size state = { prefix + "gate_up_proj_qweight": torch.arange( - num_experts * fc1_out * (hidden_size * 2 // 8), + num_experts * fc1_out * (hidden_size * fc1_bits // 8), dtype=torch.int64, ) .remainder(256) .to(torch.uint8) - .reshape(num_experts, fc1_out, hidden_size * 2 // 8), + .reshape(num_experts, fc1_out, hidden_size * fc1_bits // 8), prefix + "gate_up_proj_scales": torch.rand( num_experts, fc1_out, @@ -433,7 +437,7 @@ def _mixed_expert_state_dict( ( num_experts, fc1_out, - math.ceil((hidden_size // block_size) * 2 / 8), + math.ceil((hidden_size // block_size) * fc1_bits / 8), ), dtype=torch.uint8, ) @@ -518,6 +522,91 @@ def test_mixed_olive_payload_is_bound_byte_exactly(self): atol=0, ) + @pytest.mark.parametrize("mixed_layer", [0, 1]) + def test_mixed_and_legacy_uniform_layers_bind_with_per_layer_layouts(self, mixed_layer): + quantization = _mixed_quantization( + pattern=f"model.layers.{mixed_layer}.mlp.experts.gate_up_proj" + ) + model = MoECausalLMModel(_moe_config(quantization, num_hidden_layers=2)) + raw = {} + for layer_index in range(2): + raw.update( + _mixed_expert_state_dict( + layer_index=layer_index, + fc1_bits=2 if layer_index == mixed_layer else 4, + ) + ) + + out = model.preprocess_weights(raw) + + for layer_index in range(2): + source = f"model.layers.{layer_index}.mlp.experts." + target = f"model.layers.{layer_index}.mlp." + torch.testing.assert_close( + out[target + "fc1_experts_weights"], + raw[source + "gate_up_proj_qweight"], + rtol=0, + atol=0, + ) + torch.testing.assert_close( + out[target + "fc2_experts_weights"], + raw[source + "down_proj_qweight"], + rtol=0, + atol=0, + ) + + def test_uniform_per_layer_group_size_uses_resolved_layout(self): + quantization = QuantizationConfig( + bits=4, + group_size=_BLK, + quant_method="olive", + overrides={ + f"{_LAYER}mlp.experts.gate_up_proj": QuantizationOverride(group_size=32), + f"{_LAYER}mlp.experts.down_proj": QuantizationOverride(group_size=32), + }, + ) + model = MoECausalLMModel(_moe_config(quantization)) + raw = _mixed_expert_state_dict(block_size=32, fc1_bits=4) + + out = model.preprocess_weights(raw) + + prefix = f"{_LAYER}mlp." + torch.testing.assert_close( + out[prefix + "fc1_experts_weights"], + raw[f"{_LAYER}mlp.experts.gate_up_proj_qweight"], + rtol=0, + atol=0, + ) + torch.testing.assert_close( + out[prefix + "fc2_experts_weights"], + raw[f"{_LAYER}mlp.experts.down_proj_qweight"], + rtol=0, + atol=0, + ) + + @pytest.mark.parametrize("dense_mode", ["unsupported_uniform", "excluded"]) + def test_packed_sidecars_fail_when_layer_resolves_to_dense(self, dense_mode): + if dense_mode == "unsupported_uniform": + quantization = QuantizationConfig( + bits=4, + group_size=_BLK, + quant_method="olive", + overrides={ + f"{_LAYER}mlp.experts.gate_up_proj": QuantizationOverride(bits=8), + f"{_LAYER}mlp.experts.down_proj": QuantizationOverride(bits=8), + }, + ) + else: + quantization = dataclasses.replace( + _quantization(), + modules_to_not_convert=(f"{_LAYER}mlp.experts",), + ) + model = MoECausalLMModel(_moe_config(quantization)) + + assert model.model.layers[0].mlp.experts is not None + with pytest.raises(ValueError, match="resolved quantization layout"): + model.preprocess_weights(_mixed_expert_state_dict()) + @pytest.mark.parametrize( ("fc1_sym", "fc2_sym"), [(False, True), (True, False), (False, False)], From 2f8a591e579e8fad71504df7374165694dc641c8 Mon Sep 17 00:00:00 2001 From: titaiwang Date: Tue, 22 Sep 2026 20:10:12 +0000 Subject: [PATCH 3/8] Fail closed on unsupported QMoE plans Reject routed-expert exclusions and projection-specific GPTQ/AWQ layouts that cannot bind the current graph, while preserving redundant dense-fallback overrides. Validate every routed layer's Olive sidecars in mixed plans. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang --- src/mobius/_weight_utils.py | 79 ++++++++++++++++++++++---------- src/mobius/models/moe_test.py | 50 +++++++++++++++++--- src/mobius/models/qwen35_test.py | 9 ++++ 3 files changed, 107 insertions(+), 31 deletions(-) diff --git a/src/mobius/_weight_utils.py b/src/mobius/_weight_utils.py index f0cf41c48..534c75ebe 100644 --- a/src/mobius/_weight_utils.py +++ b/src/mobius/_weight_utils.py @@ -175,24 +175,31 @@ def is_mixed_width(self) -> bool: @property def requires_projection_preprocessing(self) -> bool: - fallback_layout = ( - self.fallback.bits, - self.fallback.group_size, - self.fallback.sym, - self.fallback.float_zero_point, - ) return any( - ( - projection.bits, - projection.group_size, - projection.sym, - projection.float_zero_point, - ) - != fallback_layout + not _same_qmoe_packing_layout(projection, self.fallback) for projection in (self.fc1, self.fc2) ) +def _same_qmoe_packing_layout( + left: QuantizationConfig, + right: QuantizationConfig, +) -> bool: + return ( + left.bits, + left.group_size, + left.sym, + left.float_zero_point, + left.weight_format, + ) == ( + right.bits, + right.group_size, + right.sym, + right.float_zero_point, + right.weight_format, + ) + + def _resolve_qmoe_module( quantization: QuantizationConfig, source_module_names: tuple[str, ...], @@ -233,6 +240,9 @@ def resolve_qmoe_quantization( The first mixed layout is intentionally narrower: Olive integer-affine FC1 INT2 plus FC2 INT4, with a common QMoE block size and the model-wide INT4 layout as the fallback for non-expert projections. + + Explicit per-layer plans must either resolve to a bindable native layout + or raise; they never silently select an incompatible dense fallback. """ if quantization is None: return None @@ -251,7 +261,11 @@ def resolve_qmoe_quantization( tuple(f"{path}.experts.down_proj" for path in source_moe_paths), ) if fc1 is None and fc2 is None: - return None + raise ValueError( + "Per-layer routed expert exclusions are not supported: the dense fallback " + "cannot bind fused expert tensors. Keep both expert projections quantized " + f"for {source_moe_paths!r}." + ) if fc1 is None or fc2 is None: raise ValueError( "Native QMoE requires both routed expert projections to be quantized; " @@ -264,7 +278,10 @@ def resolve_qmoe_quantization( fc2_uniform = supported_qmoe_quantization(fc2) if fc1_uniform is not None and fc2_uniform is not None: if fallback is None: - return None + raise ValueError( + "Per-layer routed expert overrides cannot enable native QMoE when " + "the model-wide fallback is unsupported." + ) if ( fc1.group_size != fc2.group_size or fc1.sym != fc2.sym @@ -275,12 +292,26 @@ def resolve_qmoe_quantization( f"FC1=(group_size={fc1.group_size}, sym={fc1.sym}), " f"FC2=(group_size={fc2.group_size}, sym={fc2.sym})." ) - return QMoEQuantizationLayout(fallback, fc1, fc2) + layout = QMoEQuantizationLayout(fallback, fc1, fc2) + if quantization.quant_method != "olive" and layout.requires_projection_preprocessing: + raise ValueError( + "Projection-specific QMoE layouts currently require Olive " + "preprocessing; GPTQ and AWQ only support the model-wide layout." + ) + return layout requested_mixed = ( fc1.bits != fc2.bits or fc1.group_size != fc2.group_size or fc1.sym != fc2.sym ) if not requested_mixed: + if not ( + _same_qmoe_packing_layout(fc1, quantization) + and _same_qmoe_packing_layout(fc2, quantization) + ): + raise ValueError( + "Unsupported per-layer routed expert override: the resolved layout " + "cannot bind the dense expert fallback." + ) return None layouts = (fc1, fc2) if any( @@ -1078,9 +1109,15 @@ def _preprocess_olive_qmoe_weights( ("down_proj", layout.fc2), ): stem = f"{root}.experts.{projection}" - accepted_keys.update((f"{stem}_qweight", f"{stem}_scales")) + required_keys = [f"{stem}_qweight", f"{stem}_scales"] if not projection_layout.sym: - accepted_keys.add(f"{stem}_qzeros") + required_keys.append(f"{stem}_qzeros") + accepted_keys.update(required_keys) + missing = [key for key in required_keys if key not in state_dict] + if missing: + raise ValueError( + f"QMoE sidecars for {root!r} are incomplete; missing {missing!r}." + ) stale = expert_keys - accepted_keys if stale: raise ValueError( @@ -1123,12 +1160,6 @@ def _preprocess_olive_qmoe_weights( } if not projection_layout.sym: required["qzeros"] = f"{stem}_qzeros" - missing = [name for name, key in required.items() if key not in state_dict] - if missing: - raise ValueError( - f"Mixed QMoE {label.upper()} sidecars for {root!r} are incomplete; " - f"missing {missing!r}." - ) unexpected_qzeros = projection_layout.sym and f"{stem}_qzeros" in state_dict if unexpected_qzeros: raise ValueError( diff --git a/src/mobius/models/moe_test.py b/src/mobius/models/moe_test.py index 7033e9c66..0a6388254 100644 --- a/src/mobius/models/moe_test.py +++ b/src/mobius/models/moe_test.py @@ -585,7 +585,7 @@ def test_uniform_per_layer_group_size_uses_resolved_layout(self): ) @pytest.mark.parametrize("dense_mode", ["unsupported_uniform", "excluded"]) - def test_packed_sidecars_fail_when_layer_resolves_to_dense(self, dense_mode): + def test_per_layer_expert_plan_rejects_unbindable_dense_fallback(self, dense_mode): if dense_mode == "unsupported_uniform": quantization = QuantizationConfig( bits=4, @@ -601,11 +601,8 @@ def test_packed_sidecars_fail_when_layer_resolves_to_dense(self, dense_mode): _quantization(), modules_to_not_convert=(f"{_LAYER}mlp.experts",), ) - model = MoECausalLMModel(_moe_config(quantization)) - - assert model.model.layers[0].mlp.experts is not None - with pytest.raises(ValueError, match="resolved quantization layout"): - model.preprocess_weights(_mixed_expert_state_dict()) + with pytest.raises(ValueError, match="routed expert"): + MoECausalLMModel(_moe_config(quantization)) @pytest.mark.parametrize( ("fc1_sym", "fc2_sym"), @@ -757,7 +754,7 @@ def test_uniform_int4_omits_projection_specific_attributes(self): assert "fc2_expert_weight_bits" not in qmoe.attributes assert "fc3_expert_weight_bits" not in qmoe.attributes - def test_uniform_int4_experts_do_not_bypass_unsupported_global_fallback(self): + def test_uniform_int4_expert_override_requires_supported_global_fallback(self): quantization = QuantizationConfig( bits=2, group_size=_BLK, @@ -768,11 +765,40 @@ def test_uniform_int4_experts_do_not_bypass_unsupported_global_fallback(self): }, ) + with pytest.raises(ValueError, match="model-wide fallback is unsupported"): + MoECausalLMModel(_moe_config(quantization)) + + def test_redundant_unsupported_override_keeps_dense_fallback(self): + quantization = QuantizationConfig( + bits=8, + group_size=_BLK, + quant_method="olive", + overrides={ + f"{_LAYER}mlp.experts.gate_up_proj": QuantizationOverride(bits=8), + f"{_LAYER}mlp.experts.down_proj": QuantizationOverride(bits=8), + }, + ) + block = MoECausalLMModel(_moe_config(quantization)).model.layers[0].mlp assert block._qmoe_quantization is None assert block.experts is not None + @pytest.mark.parametrize("quant_method", ["gptq", "awq"]) + def test_non_olive_projection_specific_layout_fails_at_construction(self, quant_method): + quantization = QuantizationConfig( + bits=4, + group_size=_BLK, + quant_method=quant_method, + overrides={ + f"{_LAYER}mlp.experts.gate_up_proj": QuantizationOverride(group_size=32), + f"{_LAYER}mlp.experts.down_proj": QuantizationOverride(group_size=32), + }, + ) + + with pytest.raises(ValueError, match="require Olive preprocessing"): + MoECausalLMModel(_moe_config(quantization)) + def test_model_package_roundtrip_preserves_mixed_attrs_and_bytes(self, tmp_path): from mobius._builder import build_from_module @@ -856,6 +882,16 @@ def test_missing_entire_routed_layer_fails_closed(self): with pytest.raises(ValueError, match=r"model\.layers\.1\.mlp.*incomplete"): model.preprocess_weights(_mixed_expert_state_dict()) + def test_missing_legacy_layer_in_mixed_plan_fails_closed(self): + config = dataclasses.replace( + _moe_config(_mixed_quantization()), + num_hidden_layers=2, + ) + model = MoECausalLMModel(config) + + with pytest.raises(ValueError, match=r"model\.layers\.1\.mlp.*incomplete"): + model.preprocess_weights(_mixed_expert_state_dict()) + def test_missing_all_routed_expert_sidecars_fails_closed(self): model = MoECausalLMModel(_moe_config(_mixed_quantization())) diff --git a/src/mobius/models/qwen35_test.py b/src/mobius/models/qwen35_test.py index 552d86814..fce0089d2 100644 --- a/src/mobius/models/qwen35_test.py +++ b/src/mobius/models/qwen35_test.py @@ -453,6 +453,15 @@ def test_qwen3_only_regex_does_not_match_qwen35_source_alias(self): assert not block._qmoe_quantization.is_mixed_width assert block._qmoe_quantization.fc1.bits == 4 + def test_excluded_routed_experts_fail_at_construction(self): + quantization = dataclasses.replace( + _mixed_qwen35_quantization(), + modules_to_not_convert=("experts.gate_up_proj", "experts.down_proj"), + ) + + with pytest.raises(ValueError, match="routed expert exclusions"): + Qwen35MoECausalLMModel(_moe_config(quantization)) + def test_moe_block_uses_qmoe_when_quantized(self): model = _moe_config( QuantizationConfig(bits=4, group_size=_BLK, quant_method="olive", sym=False) From 1a45f54c6bbd6a76d80c8fa20390b50c7c698750 Mon Sep 17 00:00:00 2001 From: titaiwang Date: Tue, 22 Sep 2026 20:49:45 +0000 Subject: [PATCH 4/8] Validate every mixed QMoE layer Apply strict fused-expert tensor validation to fallback-equivalent layers in mixed plans while preserving their legacy Olive conversion path. Allow uniform INT4 FC1 and FC2 projections to select zero-point inputs independently, and add regression coverage for malformed sibling layers. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang --- src/mobius/_weight_utils.py | 39 ++++++++++++------------- src/mobius/models/moe_test.py | 54 ++++++++++++++++++++++++++++++++++- 2 files changed, 72 insertions(+), 21 deletions(-) diff --git a/src/mobius/_weight_utils.py b/src/mobius/_weight_utils.py index 534c75ebe..46b0df97e 100644 --- a/src/mobius/_weight_utils.py +++ b/src/mobius/_weight_utils.py @@ -209,6 +209,7 @@ def _resolve_qmoe_module( def _matches_exclusion(pattern: str, module_name: str) -> bool: if pattern.startswith("re:"): return re.fullmatch(pattern[3:], module_name) is not None + # Olive plain exclusions use substring matching, unlike exact plain overrides. return bool(pattern) and pattern in module_name if any( @@ -282,15 +283,13 @@ def resolve_qmoe_quantization( "Per-layer routed expert overrides cannot enable native QMoE when " "the model-wide fallback is unsupported." ) - if ( - fc1.group_size != fc2.group_size - or fc1.sym != fc2.sym - or fc1.float_zero_point != fc2.float_zero_point - ): + if fc1.group_size != fc2.group_size or fc1.float_zero_point != fc2.float_zero_point: raise ValueError( - "Native QMoE uniform INT4 projections must share one layout; " - f"FC1=(group_size={fc1.group_size}, sym={fc1.sym}), " - f"FC2=(group_size={fc2.group_size}, sym={fc2.sym})." + "Native QMoE uniform INT4 projections must share one block layout; " + f"FC1=(group_size={fc1.group_size}, " + f"float_zero_point={fc1.float_zero_point}), " + f"FC2=(group_size={fc2.group_size}, " + f"float_zero_point={fc2.float_zero_point})." ) layout = QMoEQuantizationLayout(fallback, fc1, fc2) if quantization.quant_method != "olive" and layout.requires_projection_preprocessing: @@ -1112,6 +1111,11 @@ def _preprocess_olive_qmoe_weights( required_keys = [f"{stem}_qweight", f"{stem}_scales"] if not projection_layout.sym: required_keys.append(f"{stem}_qzeros") + elif f"{stem}_qzeros" in state_dict: + raise ValueError( + f"QMoE {projection} is symmetric but checkpoint contains " + f"{stem + '_qzeros'!r}." + ) accepted_keys.update(required_keys) missing = [key for key in required_keys if key not in state_dict] if missing: @@ -1125,7 +1129,7 @@ def _preprocess_olive_qmoe_weights( f"under a recognized root: {min(stale)!r}." ) - for root, layout in projection_layouts.items(): + for root, layout in layouts.items(): block_size = layout.fc1.group_size if hidden_size % block_size or intermediate_size % block_size: raise ValueError( @@ -1160,13 +1164,6 @@ def _preprocess_olive_qmoe_weights( } if not projection_layout.sym: required["qzeros"] = f"{stem}_qzeros" - unexpected_qzeros = projection_layout.sym and f"{stem}_qzeros" in state_dict - if unexpected_qzeros: - raise ValueError( - f"Mixed QMoE {label.upper()} is symmetric but checkpoint contains " - f"{stem + '_qzeros'!r}." - ) - qweight = state_dict[required["qweight"]] scales = state_dict[required["scales"]] qzeros = state_dict[required["qzeros"]] if "qzeros" in required else None @@ -1214,10 +1211,12 @@ def _preprocess_olive_qmoe_weights( f"Mixed QMoE {label.upper()} qzeros must be uint8, got {qzeros.dtype}." ) - renamed_experts[f"{stem}.weight"] = qweight.contiguous() - renamed_experts[f"{stem}.scales"] = scales - if qzeros is not None: - renamed_experts[f"{stem}.zero_points"] = qzeros.contiguous() + if root in projection_layouts: + # Legacy roots are validated here but renamed by preprocess_olive_weights. + renamed_experts[f"{stem}.weight"] = qweight.contiguous() + renamed_experts[f"{stem}.scales"] = scales + if qzeros is not None: + renamed_experts[f"{stem}.zero_points"] = qzeros.contiguous() projection_expert_keys = { key diff --git a/src/mobius/models/moe_test.py b/src/mobius/models/moe_test.py index 0a6388254..5758309f1 100644 --- a/src/mobius/models/moe_test.py +++ b/src/mobius/models/moe_test.py @@ -555,6 +555,31 @@ def test_mixed_and_legacy_uniform_layers_bind_with_per_layer_layouts(self, mixed atol=0, ) + @pytest.mark.parametrize( + ("mutation", "message"), + [ + ("wrong_rank", "rank 3"), + ("wrong_experts", "FC1 qweight shape"), + ("wrong_scale_dtype", "must be float16, bfloat16, or float32"), + ], + ) + def test_malformed_legacy_layer_in_mixed_plan_fails_closed(self, mutation, message): + model = MoECausalLMModel(_moe_config(_mixed_quantization(), num_hidden_layers=2)) + state = {} + state.update(_mixed_expert_state_dict()) + state.update(_mixed_expert_state_dict(layer_index=1, fc1_bits=4)) + weight = "model.layers.1.mlp.experts.gate_up_proj_qweight" + scales = "model.layers.1.mlp.experts.gate_up_proj_scales" + if mutation == "wrong_rank": + state[weight] = state[weight].flatten(1) + elif mutation == "wrong_experts": + state[weight] = state[weight][:-1] + elif mutation == "wrong_scale_dtype": + state[scales] = state[scales].to(torch.float64) + + with pytest.raises(ValueError, match=message): + model.preprocess_weights(state) + def test_uniform_per_layer_group_size_uses_resolved_layout(self): quantization = QuantizationConfig( bits=4, @@ -584,6 +609,23 @@ def test_uniform_per_layer_group_size_uses_resolved_layout(self): atol=0, ) + def test_uniform_int4_projections_allow_independent_symmetry(self): + quantization = QuantizationConfig( + bits=4, + group_size=_BLK, + quant_method="olive", + sym=True, + overrides={f"{_LAYER}mlp.experts.gate_up_proj": QuantizationOverride(sym=False)}, + ) + model = MoECausalLMModel(_moe_config(quantization)) + out = model.preprocess_weights( + _mixed_expert_state_dict(fc1_bits=4, fc1_sym=False, fc2_sym=True) + ) + prefix = f"{_LAYER}mlp." + + assert prefix + "fc1_experts_zero_points" in out + assert prefix + "fc2_experts_zero_points" not in out + @pytest.mark.parametrize("dense_mode", ["unsupported_uniform", "excluded"]) def test_per_layer_expert_plan_rejects_unbindable_dense_fallback(self, dense_mode): if dense_mode == "unsupported_uniform": @@ -658,7 +700,7 @@ def test_odd_zero_point_byte_ceiling(self): ("wrong_rank", "rank 3"), ("wrong_dtype", "must be uint8"), ("wrong_scale_dtype", "must be float16, bfloat16, or float32"), - ("wrong_experts", "shape"), + ("wrong_experts", "FC1 qweight shape"), ("wrong_packed_bytes", "shape"), ("wrong_scale_geometry", "shape"), ], @@ -696,6 +738,16 @@ def test_unknown_packed_sidecar_under_recognized_root_fails_closed(self): with pytest.raises(ValueError, match="unknown or stale packed expert sidecar"): model.preprocess_weights(state) + def test_symmetric_projection_rejects_zero_point_sidecar(self): + state = _mixed_expert_state_dict() + state[f"{_LAYER}mlp.experts.gate_up_proj_qzeros"] = torch.zeros( + _E, 2 * _INT, 1, dtype=torch.uint8 + ) + model = MoECausalLMModel(_moe_config(_mixed_quantization())) + + with pytest.raises(ValueError, match="is symmetric but checkpoint contains"): + model.preprocess_weights(state) + def test_unknown_packed_sidecar_root_fails_closed(self): state = _mixed_expert_state_dict() state["model.layers.9.mlp.experts.gate_up_proj_qweight"] = torch.zeros( From 5fdadb6067de19008dd0fe59d83f2b959573fbde Mon Sep 17 00:00:00 2001 From: titaiwang Date: Tue, 22 Sep 2026 21:20:18 +0000 Subject: [PATCH 5/8] Verify per-layer mixed QMoE bindings Strengthen the two-layer regression to associate each QMoE node with its FC1 initializer, validate projection-specific bit attributes on only the mixed layer, and compare both layers' bound FC1 and FC2 initializer bytes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang --- src/mobius/models/moe_test.py | 56 ++++++++++++++++++++++++++++++++++- 1 file changed, 55 insertions(+), 1 deletion(-) diff --git a/src/mobius/models/moe_test.py b/src/mobius/models/moe_test.py index 5758309f1..e835db6a7 100644 --- a/src/mobius/models/moe_test.py +++ b/src/mobius/models/moe_test.py @@ -524,10 +524,14 @@ def test_mixed_olive_payload_is_bound_byte_exactly(self): @pytest.mark.parametrize("mixed_layer", [0, 1]) def test_mixed_and_legacy_uniform_layers_bind_with_per_layer_layouts(self, mixed_layer): + from mobius._builder import build_from_module + quantization = _mixed_quantization( pattern=f"model.layers.{mixed_layer}.mlp.experts.gate_up_proj" ) - model = MoECausalLMModel(_moe_config(quantization, num_hidden_layers=2)) + config = _moe_config(quantization, num_hidden_layers=2) + model = MoECausalLMModel(config) + package = build_from_module(model, config) raw = {} for layer_index in range(2): raw.update( @@ -539,9 +543,46 @@ def test_mixed_and_legacy_uniform_layers_bind_with_per_layer_layouts(self, mixed out = model.preprocess_weights(raw) + def _fc1_initializer_name(qmoe: ir.Node) -> str: + pending = [qmoe.inputs[2]] + visited = set() + while pending: + value = pending.pop() + if value is None or id(value) in visited: + continue + visited.add(id(value)) + if value.is_initializer() and value.name.endswith("fc1_experts_weights"): + return value.name + if producer := value.producer(): + pending.extend(producer.inputs) + raise AssertionError(f"QMoE node {qmoe.name!r} has no FC1 initializer input") + + qmoe_by_fc1 = { + _fc1_initializer_name(node): node + for node in package["model"].graph + if node.op_type == "QMoE" + } + assert set(qmoe_by_fc1) == { + f"model.layers.{layer_index}.mlp.fc1_experts_weights" for layer_index in range(2) + } for layer_index in range(2): source = f"model.layers.{layer_index}.mlp.experts." target = f"model.layers.{layer_index}.mlp." + qmoe = qmoe_by_fc1[target + "fc1_experts_weights"] + assert qmoe.attributes["expert_weight_bits"].value == 4 + projection_attrs = { + name: attribute.value + for name, attribute in qmoe.attributes.items() + if name.startswith("fc") and name.endswith("_expert_weight_bits") + } + if layer_index == mixed_layer: + assert projection_attrs == { + "fc1_expert_weight_bits": 2, + "fc2_expert_weight_bits": 4, + "fc3_expert_weight_bits": 2, + } + else: + assert projection_attrs == {} torch.testing.assert_close( out[target + "fc1_experts_weights"], raw[source + "gate_up_proj_qweight"], @@ -555,6 +596,19 @@ def test_mixed_and_legacy_uniform_layers_bind_with_per_layer_layouts(self, mixed atol=0, ) + package.apply_weights(out, fold_constants=False) + for layer_index in range(2): + source = f"model.layers.{layer_index}.mlp.experts." + target = f"model.layers.{layer_index}.mlp." + expected = { + target + "fc1_experts_weights": raw[source + "gate_up_proj_qweight"], + target + "fc2_experts_weights": raw[source + "down_proj_qweight"], + } + for name, tensor in expected.items(): + initializer = package["model"].graph.initializers[name] + assert initializer.const_value is not None + assert initializer.const_value.numpy().tobytes() == tensor.numpy().tobytes() + @pytest.mark.parametrize( ("mutation", "message"), [ From a0a05d0879006677ac049c554075212bdbb98dbf Mon Sep 17 00:00:00 2001 From: titaiwang Date: Tue, 22 Sep 2026 21:49:41 +0000 Subject: [PATCH 6/8] Use producer-specific QMoE matching Delegate GPTQ and AWQ module rules to the generic subtree matcher while preserving Olive's literal override and substring exclusion semantics. Add parent, near-match, and exclusion regressions and document the producer-specific contract. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang --- .../per-component-quantization-loading.md | 10 +++-- src/mobius/_weight_utils.py | 4 +- src/mobius/models/moe_test.py | 39 +++++++++++++++++++ 3 files changed, 48 insertions(+), 5 deletions(-) diff --git a/docs/design/per-component-quantization-loading.md b/docs/design/per-component-quantization-loading.md index a231b0369..267796e7f 100644 --- a/docs/design/per-component-quantization-loading.md +++ b/docs/design/per-component-quantization-loading.md @@ -441,10 +441,12 @@ Native QMoE additionally has a projection-specific adapter for Olive fused K-last routed experts. The initial mixed layout is deliberately narrow: model-wide INT4, FC1 `gate_up_proj` overridden to INT2, FC2 `down_proj` remaining INT4, and one common power-of-two group size of at least 16. Exact -and `re:` module overrides are resolved against the original Hugging Face -paths before FC1 and FC2 buffers are validated independently. The adapter -requires every routed layer and validates rank, dtype, expert count, packed -byte width, scale geometry, and optional zero-point geometry before binding. +and `re:` Olive overrides are resolved against the original Hugging Face paths; +plain Olive exclusions retain their producer-defined substring matching. +GPTQ and AWQ instead use the generic plain-path subtree and `re:` full-match +rules. FC1 and FC2 buffers are then validated independently. The adapter +requires every routed layer and validates rank, dtype, expert count, packed byte +width, scale geometry, and optional zero-point geometry before binding. Qwen3.5-VL keeps the decoder's full module plan (``preserve_module_plan=True``) while constructing the split package: expert overrides are resolved against authoritative ``model.language_model`` source diff --git a/src/mobius/_weight_utils.py b/src/mobius/_weight_utils.py index 46b0df97e..5bf5f4eab 100644 --- a/src/mobius/_weight_utils.py +++ b/src/mobius/_weight_utils.py @@ -204,7 +204,9 @@ def _resolve_qmoe_module( quantization: QuantizationConfig, source_module_names: tuple[str, ...], ) -> QuantizationConfig | None: - """Resolve Olive's QMoE override semantics for authoritative source names.""" + """Resolve producer-specific QMoE rules for authoritative source names.""" + if quantization.quant_method != "olive": + return quantization.for_module(source_module_names) def _matches_exclusion(pattern: str, module_name: str) -> bool: if pattern.startswith("re:"): diff --git a/src/mobius/models/moe_test.py b/src/mobius/models/moe_test.py index e835db6a7..c54be0fe0 100644 --- a/src/mobius/models/moe_test.py +++ b/src/mobius/models/moe_test.py @@ -905,6 +905,45 @@ def test_non_olive_projection_specific_layout_fails_at_construction(self, quant_ with pytest.raises(ValueError, match="require Olive preprocessing"): MoECausalLMModel(_moe_config(quantization)) + @pytest.mark.parametrize("quant_method", ["gptq", "awq"]) + def test_non_olive_parent_override_uses_subtree_matching(self, quant_method): + quantization = QuantizationConfig( + bits=4, + group_size=_BLK, + quant_method=quant_method, + overrides={f"{_LAYER}mlp": QuantizationOverride(group_size=32)}, + ) + + with pytest.raises(ValueError, match="require Olive preprocessing"): + MoECausalLMModel(_moe_config(quantization)) + + @pytest.mark.parametrize("quant_method", ["gptq", "awq"]) + def test_non_olive_plain_override_does_not_use_substring_matching(self, quant_method): + quantization = QuantizationConfig( + bits=4, + group_size=_BLK, + quant_method=quant_method, + overrides={"model.layers.0.ml": QuantizationOverride(group_size=32)}, + ) + + block = MoECausalLMModel(_moe_config(quantization)).model.layers[0].mlp + + assert block._qmoe_quantization is not None + assert block._qmoe_quantization.fc1.group_size == _BLK + assert block._qmoe_quantization.fc2.group_size == _BLK + + @pytest.mark.parametrize("quant_method", ["gptq", "awq"]) + def test_non_olive_parent_exclusion_uses_subtree_matching(self, quant_method): + quantization = QuantizationConfig( + bits=4, + group_size=_BLK, + quant_method=quant_method, + modules_to_not_convert=(f"{_LAYER}mlp",), + ) + + with pytest.raises(ValueError, match="routed expert exclusions"): + MoECausalLMModel(_moe_config(quantization)) + def test_model_package_roundtrip_preserves_mixed_attrs_and_bytes(self, tmp_path): from mobius._builder import build_from_module From 7fc3e8e82459bae83c14f0762b11efcf83e13ac3 Mon Sep 17 00:00:00 2001 From: titaiwang Date: Thu, 24 Sep 2026 22:39:30 +0000 Subject: [PATCH 7/8] Extend mixed-width QMoE export to INT8 experts Accept Olive INT4-fallback expert layouts (4,8), (8,4), and (8,8), preserving per-projection packing and correct uniform INT8 node attributes. Cover packed bytes, malformed sidecars, adjacent layers, and package roundtrips. Preserve pinned GGUF import-route hashes when the new QMoE source aliases are empty; document that mixed INT4/INT8 execution still needs ORT runtime support. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang --- .../per-component-quantization-loading.md | 22 +- src/mobius/_weight_utils.py | 16 +- src/mobius/components/_moe.py | 8 +- src/mobius/integrations/gguf/_builder.py | 3 + .../integrations/gguf/_builder_core_test.py | 16 ++ .../integrations/gguf/_moe_cohort_test.py | 1 + src/mobius/models/moe_test.py | 239 +++++++++++++++++- 7 files changed, 283 insertions(+), 22 deletions(-) diff --git a/docs/design/per-component-quantization-loading.md b/docs/design/per-component-quantization-loading.md index 267796e7f..0cccba059 100644 --- a/docs/design/per-component-quantization-loading.md +++ b/docs/design/per-component-quantization-loading.md @@ -438,9 +438,9 @@ A per-component `WeightBundle` adapter is a future migration target, not the currently implemented interface. Native QMoE additionally has a projection-specific adapter for Olive fused -K-last routed experts. The initial mixed layout is deliberately narrow: -model-wide INT4, FC1 `gate_up_proj` overridden to INT2, FC2 `down_proj` -remaining INT4, and one common power-of-two group size of at least 16. Exact +K-last routed experts. Under a model-wide INT4 fallback, supported FC1/FC2 +expert widths are (2,4), (4,8), (8,4), and (8,8), with one common power-of-two +group size of at least 16. A model-wide INT8 fallback is not enabled. Exact and `re:` Olive overrides are resolved against the original Hugging Face paths; plain Olive exclusions retain their producer-defined substring matching. GPTQ and AWQ instead use the generic plain-path subtree and `re:` full-match @@ -452,11 +452,17 @@ Qwen3.5-VL keeps the decoder's full module plan overrides are resolved against authoritative ``model.language_model`` source paths, while validation and binding use the decoder-prefixed checkpoint roots created by the VL weight router. -It emits QMoE's FC-specific bit attributes (including FC3=FC1 for fused -SwiGLU). Uniform INT4 omits those attributes and retains the legacy payload -path. Mixed-width execution remains an ONNX Runtime `NOT_IMPLEMENTED` -boundary; Mobius tests this serialized graph contract without claiming -runtime numerical support. +It emits QMoE's FC-specific bit attributes for differing widths (including +FC3=FC1 for fused SwiGLU). Uniform INT4 retains the legacy payload path; +uniform 8/8 uses global `expert_weight_bits=8` without FC-specific attributes. +The new INT8-containing layouts (4/8, 8/4, 8/8) are graph-only exports: +this work does not establish CPU or CUDA execution for them. Legacy 2/4 may +execute on CUDA in ORT builds containing microsoft/onnxruntime#32743, subject +to that fallback's bounds: raw weights, 3-D scales, block size 16–256, +FP16/BF16, and its scratch-memory budget. This export test does not validate +runtime execution. The packed-decode path in #32761 is open (not merged) and +covers only the 2/4-bit set, not INT8. Do not silently replace unsupported QMoE +layouts with dense expert execution; runtime support is a separate follow-up. Appropriate model-specific operations include: diff --git a/src/mobius/_weight_utils.py b/src/mobius/_weight_utils.py index 5bf5f4eab..098877b7b 100644 --- a/src/mobius/_weight_utils.py +++ b/src/mobius/_weight_utils.py @@ -240,9 +240,9 @@ def resolve_qmoe_quantization( """Resolve and validate FC1/FC2 layouts for one native QMoE layer. Legacy uniform INT4 layouts retain the existing predicate and behavior. - The first mixed layout is intentionally narrower: Olive integer-affine - FC1 INT2 plus FC2 INT4, with a common QMoE block size and the model-wide - INT4 layout as the fallback for non-expert projections. + Projection overrides support Olive integer-affine FC1/FC2 widths (2, 4) + and (4, 8), (8, 4), or (8, 8), with a common QMoE block size and the + model-wide INT4 layout as the fallback for non-expert projections. Explicit per-layer plans must either resolve to a bindable native layout or raise; they never silently select an incompatible dense fallback. @@ -302,7 +302,10 @@ def resolve_qmoe_quantization( return layout requested_mixed = ( - fc1.bits != fc2.bits or fc1.group_size != fc2.group_size or fc1.sym != fc2.sym + fc1.bits != fc2.bits + or fc1.group_size != fc2.group_size + or fc1.sym != fc2.sym + or (fallback is not None and fc1.bits == fc2.bits == 8) ) if not requested_mixed: if not ( @@ -325,10 +328,11 @@ def resolve_qmoe_quantization( "Mixed native QMoE requires Olive integer-affine weights with packed " "uint8 zero points." ) - if fallback is None or (fc1.bits, fc2.bits) != (2, 4): + if fallback is None or (fc1.bits, fc2.bits) not in {(2, 4), (4, 8), (8, 4), (8, 8)}: raise ValueError( "Unsupported mixed native QMoE expert widths: expected model-wide " - f"INT4 with FC1 INT2 and FC2 INT4, got fallback={quantization.bits}, " + "INT4 with FC1/FC2 in {(2, 4), (4, 8), (8, 4), (8, 8)}, " + f"got fallback={quantization.bits}, " f"FC1={fc1.bits}, FC2={fc2.bits}." ) if fc1.group_size != fc2.group_size: diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py index 0013aa15d..b1d084a9d 100644 --- a/src/mobius/components/_moe.py +++ b/src/mobius/components/_moe.py @@ -626,7 +626,13 @@ def _qmoe_forward(self, op: OpBuilder, hidden_states: ir.Value, *gate_args): activation_type="swiglu", normalize_routing_weights=int(normalize), k=self.top_k, - expert_weight_bits=quantization.fallback.bits, + # A uniform override can use the global attribute directly. + # Differing widths retain the model-wide INT4 fallback and FC attrs. + expert_weight_bits=( + quantization.fallback.bits + if quantization.is_mixed_width + else quantization.fc1.bits + ), block_size=quantization.fc1.group_size, swiglu_fusion=1, quant_type="int", diff --git a/src/mobius/integrations/gguf/_builder.py b/src/mobius/integrations/gguf/_builder.py index d7d6a941a..834896b7c 100644 --- a/src/mobius/integrations/gguf/_builder.py +++ b/src/mobius/integrations/gguf/_builder.py @@ -6701,6 +6701,9 @@ def _logical_source_filename(reference: str | Path, resolved_path: str | Path) - def _graph_config_fields_for_fingerprint(config, gguf_arch: str) -> dict[str, object]: """Serialize only fields consumed by an architecture's imported graph.""" fields = asdict(config) + if not fields.get("qmoe_source_paths"): + # This field postdates pinned GGUF routes; its empty default must preserve their hashes. + fields.pop("qmoe_source_paths", None) if fields.get("component_quantization") is None: # This field postdates the pinned GGUF evidence routes. An absent # component plan preserves legacy graph behavior and fingerprint bytes; diff --git a/src/mobius/integrations/gguf/_builder_core_test.py b/src/mobius/integrations/gguf/_builder_core_test.py index 7fb3b243a..a754e3d7b 100644 --- a/src/mobius/integrations/gguf/_builder_core_test.py +++ b/src/mobius/integrations/gguf/_builder_core_test.py @@ -114,6 +114,22 @@ def test_existing_routes_ignore_absent_component_quantization() -> None: assert changed != legacy +def test_existing_routes_ignore_empty_qmoe_source_paths() -> None: + from mobius._testing import make_config + from mobius.integrations.gguf._builder import _serialize_route_graph_config + + config = make_config() + legacy = _serialize_route_graph_config(config, "llama") + changed = _serialize_route_graph_config( + dataclasses.replace(config, qmoe_source_paths=("model.layers.0.mlp",)), + "llama", + ) + + assert "qmoe_source_paths" not in json.loads(legacy) + assert json.loads(changed)["qmoe_source_paths"] == ["model.layers.0.mlp"] + assert changed != legacy + + class TestReuseGgufWeights: """Tests for mixed GGUF references plus converted ONNX sidecar weights.""" diff --git a/src/mobius/integrations/gguf/_moe_cohort_test.py b/src/mobius/integrations/gguf/_moe_cohort_test.py index 72231da5d..45cb02764 100644 --- a/src/mobius/integrations/gguf/_moe_cohort_test.py +++ b/src/mobius/integrations/gguf/_moe_cohort_test.py @@ -82,6 +82,7 @@ def test_new_cohort_fields_preserve_existing_route_fingerprint_bytes( "attention_clamp", "component_quantization", "moe_layer_frequency", + "qmoe_source_paths", "routing_weight_normalization_floor", ): legacy_fields.pop(field_name, None) diff --git a/src/mobius/models/moe_test.py b/src/mobius/models/moe_test.py index c54be0fe0..e691eeb80 100644 --- a/src/mobius/models/moe_test.py +++ b/src/mobius/models/moe_test.py @@ -397,6 +397,7 @@ def _mixed_expert_state_dict( intermediate_size: int = _INT, block_size: int = _BLK, fc1_bits: int = 2, + fc2_bits: int = 4, fc1_sym: bool = True, fc2_sym: bool = True, ) -> dict[str, torch.Tensor]: @@ -417,12 +418,12 @@ def _mixed_expert_state_dict( dtype=torch.bfloat16, ), prefix + "down_proj_qweight": torch.arange( - num_experts * hidden_size * (intermediate_size * 4 // 8), + num_experts * hidden_size * (intermediate_size * fc2_bits // 8), dtype=torch.int64, ) .remainder(256) .to(torch.uint8) - .reshape(num_experts, hidden_size, intermediate_size * 4 // 8), + .reshape(num_experts, hidden_size, intermediate_size * fc2_bits // 8), prefix + "down_proj_scales": torch.rand( num_experts, hidden_size, @@ -448,7 +449,7 @@ def _mixed_expert_state_dict( ( num_experts, hidden_size, - math.ceil((intermediate_size // block_size) * 4 / 8), + math.ceil((intermediate_size // block_size) * fc2_bits / 8), ), dtype=torch.uint8, ) @@ -456,6 +457,229 @@ def _mixed_expert_state_dict( class TestMixedWidthQMoEExport: + @pytest.mark.parametrize(("fc1_bits", "fc2_bits"), [(4, 8), (8, 4), (8, 8)]) + @pytest.mark.parametrize("asymmetric", [False, True]) + def test_int8_expert_layout_attrs_geometry_and_binding( + self, fc1_bits, fc2_bits, asymmetric + ): + from mobius._builder import build_from_module + + overrides = { + f"{_LAYER}mlp.experts.{projection}": QuantizationOverride( + bits=bits, sym=not asymmetric + ) + for projection, bits in (("gate_up_proj", fc1_bits), ("down_proj", fc2_bits)) + } + config = _moe_config(_quantization(overrides=overrides)) + model = MoECausalLMModel(config) + block = model.model.layers[0].mlp + assert block.experts is None + raw = _mixed_expert_state_dict( + fc1_bits=fc1_bits, + fc2_bits=fc2_bits, + fc1_sym=not asymmetric, + fc2_sym=not asymmetric, + ) + out = model.preprocess_weights(raw) + prefix = f"{_LAYER}mlp." + for projection, label, bits, rows, k in ( + ("gate_up_proj", "fc1", fc1_bits, _FC1_OUT, _H), + ("down_proj", "fc2", fc2_bits, _H, _INT), + ): + target = prefix + label + "_experts_weights" + source = prefix + "experts." + projection + "_qweight" + assert tuple(getattr(block, label + "_experts_weights").shape) == ( + _E, + rows, + k * bits // 8, + ) + assert out[target].numpy().tobytes() == raw[source].numpy().tobytes() + if asymmetric: + zeros = prefix + label + "_experts_zero_points" + assert out[zeros].shape == (_E, rows, math.ceil((k // _BLK) * bits / 8)) + assert ( + out[zeros].numpy().tobytes() + == raw[prefix + "experts." + projection + "_qzeros"].numpy().tobytes() + ) + + package = build_from_module(model, config) + package.apply_weights(out, fold_constants=False) + qmoe = next(node for node in package["model"].graph if node.op_type == "QMoE") + assert qmoe.attributes["expert_weight_bits"].value == ( + 8 if fc1_bits == fc2_bits else 4 + ) + attrs = { + name: attribute.value + for name, attribute in qmoe.attributes.items() + if name.startswith("fc") and name.endswith("_expert_weight_bits") + } + assert attrs == ( + {} + if fc1_bits == fc2_bits + else { + "fc1_expert_weight_bits": fc1_bits, + "fc2_expert_weight_bits": fc2_bits, + "fc3_expert_weight_bits": fc1_bits, + } + ) + for label in ("fc1", "fc2"): + name = prefix + label + "_experts_weights" + assert package["model"].graph.initializers[name].const_value.numpy().tobytes() == ( + out[name].numpy().tobytes() + ) + + def test_uniform_int8_roundtrip_preserves_global_width_and_weights(self, tmp_path): + from mobius._builder import build_from_module + + config = _moe_config( + _quantization( + overrides={ + f"{_LAYER}mlp.experts.{projection}": QuantizationOverride(bits=8) + for projection in ("gate_up_proj", "down_proj") + } + ) + ) + module = MoECausalLMModel(config) + processed = module.preprocess_weights(_mixed_expert_state_dict(fc1_bits=8, fc2_bits=8)) + package = build_from_module(module, config) + package.apply_weights(processed, fold_constants=False) + package.save(str(tmp_path), check_weights=False, progress_bar=False) + graph = ModelPackage.load(str(tmp_path))["model"].graph + qmoe = next(node for node in graph if node.op_type == "QMoE") + + assert qmoe.attributes["expert_weight_bits"].value == 8 + assert not any( + name in qmoe.attributes + for name in ( + "fc1_expert_weight_bits", + "fc2_expert_weight_bits", + "fc3_expert_weight_bits", + ) + ) + for projection in ("fc1", "fc2"): + name = f"{_LAYER}mlp.{projection}_experts_weights" + assert graph.initializers[name].const_value.numpy().tobytes() == ( + processed[name].numpy().tobytes() + ) + + def test_adjacent_int8_layouts_save_load(self, tmp_path): + from mobius._builder import build_from_module + + config = _moe_config( + _quantization( + overrides={ + f"model.layers.{layer}.mlp.experts.{proj}": QuantizationOverride(bits=8) + for layer, proj in ((0, "down_proj"), (1, "gate_up_proj")) + } + ), + num_hidden_layers=2, + ) + model = MoECausalLMModel(config) + raw = {} + for layer, bits in enumerate(((4, 8), (8, 4))): + raw.update( + _mixed_expert_state_dict(layer_index=layer, fc1_bits=bits[0], fc2_bits=bits[1]) + ) + out = model.preprocess_weights(raw) + package = build_from_module(model, config) + package.apply_weights(out, fold_constants=False) + package.save(str(tmp_path), check_weights=False, progress_bar=False) + loaded = ModelPackage.load(str(tmp_path))["model"].graph + assert sorted( + ( + node.attributes["fc1_expert_weight_bits"].value, + node.attributes["fc2_expert_weight_bits"].value, + node.attributes["fc3_expert_weight_bits"].value, + ) + for node in loaded + if node.op_type == "QMoE" + ) == [(4, 8, 4), (8, 4, 8)] + for layer, bits in enumerate(((4, 8), (8, 4))): + for proj, label, width in ( + ("gate_up_proj", "fc1", bits[0]), + ("down_proj", "fc2", bits[1]), + ): + prefix = f"model.layers.{layer}.mlp." + name = prefix + label + "_experts_weights" + expected = raw[prefix + "experts." + proj + "_qweight"] + assert loaded.initializers[name].const_value.numpy().tobytes() == ( + expected.numpy().tobytes() + ) + assert out[name].shape[-1] == (_H if label == "fc1" else _INT) * width // 8 + + def test_int8_missing_layer_and_bad_sidecar_fail_closed(self): + config = _moe_config( + _quantization( + overrides={ + r"re:model\.layers\.\d+\.mlp\.experts\.down_proj": QuantizationOverride( + bits=8 + ) + } + ), + num_hidden_layers=2, + ) + model = MoECausalLMModel(config) + with pytest.raises(ValueError, match=r"model\.layers\.1\.mlp.*incomplete"): + model.preprocess_weights(_mixed_expert_state_dict(fc1_bits=4, fc2_bits=8)) + state = _mixed_expert_state_dict(fc1_bits=4, fc2_bits=8) + state[f"{_LAYER}mlp.experts.down_proj_qweight"] = state[ + f"{_LAYER}mlp.experts.down_proj_qweight" + ][..., :-1] + state.update(_mixed_expert_state_dict(layer_index=1, fc1_bits=4, fc2_bits=8)) + with pytest.raises(ValueError, match="FC2 qweight shape"): + model.preprocess_weights(state) + + def test_olive_int8_plain_override_requires_exact_source_path(self): + quantization = _quantization( + overrides={"model.layers.0.mlp.experts.down": QuantizationOverride(bits=8)} + ) + block = MoECausalLMModel(_moe_config(quantization)).model.layers[0].mlp + assert block._qmoe_quantization.fc2.bits == 4 + + @pytest.mark.parametrize("mutation", ["missing", "shape", "dtype"]) + def test_int8_asymmetric_zero_points_fail_closed(self, mutation): + quantization = _quantization( + overrides={ + f"{_LAYER}mlp.experts.down_proj": QuantizationOverride(bits=8, sym=False) + } + ) + state = _mixed_expert_state_dict(fc1_bits=4, fc2_bits=8, fc2_sym=False) + key = f"{_LAYER}mlp.experts.down_proj_qzeros" + if mutation == "missing": + del state[key] + elif mutation == "shape": + state[key] = state[key][..., :-1] + else: + state[key] = state[key].to(torch.int8) + with pytest.raises( + ValueError, + match={ + "missing": "incomplete", + "shape": "FC2 qzeros shape", + "dtype": "FC2 qzeros must be uint8", + }[mutation], + ): + MoECausalLMModel(_moe_config(quantization)).preprocess_weights(state) + + def test_int8_mismatched_group_size_rejected(self): + quantization = _quantization( + overrides={ + f"{_LAYER}mlp.experts.down_proj": QuantizationOverride(bits=8, group_size=32) + } + ) + with pytest.raises(ValueError, match="common FC1/FC2 group_size"): + MoECausalLMModel(_moe_config(quantization)) + + def test_int8_non_olive_override_rejected(self): + quantization = QuantizationConfig( + bits=4, + group_size=_BLK, + quant_method="gptq", + overrides={f"{_LAYER}mlp.experts.down_proj": QuantizationOverride(bits=8)}, + ) + with pytest.raises(ValueError, match="requires Olive integer-affine"): + MoECausalLMModel(_moe_config(quantization)) + @pytest.mark.parametrize( "pattern", [ @@ -688,8 +912,8 @@ def test_per_layer_expert_plan_rejects_unbindable_dense_fallback(self, dense_mod group_size=_BLK, quant_method="olive", overrides={ - f"{_LAYER}mlp.experts.gate_up_proj": QuantizationOverride(bits=8), - f"{_LAYER}mlp.experts.down_proj": QuantizationOverride(bits=8), + f"{_LAYER}mlp.experts.gate_up_proj": QuantizationOverride(bits=2), + f"{_LAYER}mlp.experts.down_proj": QuantizationOverride(bits=2), }, ) else: @@ -1003,10 +1227,11 @@ def test_model_package_roundtrip_preserves_mixed_attrs_and_bytes(self, tmp_path) group_size=_BLK, quant_method="olive", overrides={ - f"{_LAYER}mlp.experts.gate_up_proj": QuantizationOverride(bits=8) + f"{_LAYER}mlp.experts.gate_up_proj": QuantizationOverride(bits=2), + f"{_LAYER}mlp.experts.down_proj": QuantizationOverride(bits=8), }, ), - "expected model-wide INT4 with FC1 INT2 and FC2 INT4", + "expected model-wide INT4", ), ], ) From 02a4843fd3174a97009f25adf3d2c7d64cba7ba4 Mon Sep 17 00:00:00 2001 From: titaiwang Date: Thu, 24 Sep 2026 22:52:15 +0000 Subject: [PATCH 8/8] Correct mixed INT8 QMoE CUDA runtime boundary ORT's merged correctness fallback covers INT4/INT8 and INT8/INT4 even though its PR summary only lists INT2 combinations. Distinguish upstream CUDA coverage from unverified Mobius fused-SwiGLU end-to-end execution and CPU support. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang --- .../per-component-quantization-loading.md | 20 +++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/docs/design/per-component-quantization-loading.md b/docs/design/per-component-quantization-loading.md index 0cccba059..5d7aba62a 100644 --- a/docs/design/per-component-quantization-loading.md +++ b/docs/design/per-component-quantization-loading.md @@ -455,14 +455,18 @@ created by the VL weight router. It emits QMoE's FC-specific bit attributes for differing widths (including FC3=FC1 for fused SwiGLU). Uniform INT4 retains the legacy payload path; uniform 8/8 uses global `expert_weight_bits=8` without FC-specific attributes. -The new INT8-containing layouts (4/8, 8/4, 8/8) are graph-only exports: -this work does not establish CPU or CUDA execution for them. Legacy 2/4 may -execute on CUDA in ORT builds containing microsoft/onnxruntime#32743, subject -to that fallback's bounds: raw weights, 3-D scales, block size 16–256, -FP16/BF16, and its scratch-memory budget. This export test does not validate -runtime execution. The packed-decode path in #32761 is open (not merged) and -covers only the 2/4-bit set, not INT8. Do not silently replace unsupported QMoE -layouts with dense expert execution; runtime support is a separate follow-up. +ORT's merged CUDA fallback in microsoft/onnxruntime#32743 executes integer +mixed-width QMoE, including tested (4,8) and (8,4) widths. It dequantizes +both expert banks into FP16/BF16 scratch, subject to canonical raw weights, +3-D blockwise scales, block size 16–256, divisible reduction dimensions, +and a scratch-memory limit. ORT tests cover these INT8 combinations without +fused SwiGLU; only (2,4) has a fused-SwiGLU test. Mobius validates graph +export and weight binding here, not end-to-end numerical execution of its +exported Qwen graphs on this ORT build. CPU mixed-width execution and uniform +8/8 runtime behavior have not been qualified by these tests. The packed +decode path in microsoft/onnxruntime#32761 is open (not merged) and covers +only the 2/4-bit set, not INT8. Do not silently replace an unsupported QMoE +configuration with the dense expert-loop exporter. Appropriate model-specific operations include: