diff --git a/docs/design/per-component-quantization-loading.md b/docs/design/per-component-quantization-loading.md index 488ef3ad1..5d7aba62a 100644 --- a/docs/design/per-component-quantization-loading.md +++ b/docs/design/per-component-quantization-loading.md @@ -437,6 +437,37 @@ 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. 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 +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 +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 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. +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: - 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..098877b7b 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,193 @@ 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 + + @property + def requires_projection_preprocessing(self) -> bool: + return any( + 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, ...], +) -> QuantizationConfig | None: + """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:"): + 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( + _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. + 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. + """ + 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: + 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; " + 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: + 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.float_zero_point != fc2.float_zero_point: + raise ValueError( + "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: + 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 + or (fallback is not None and fc1.bits == fc2.bits == 8) + ) + 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( + 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) not in {(2, 4), (4, 8), (8, 4), (8, 8)}: + raise ValueError( + "Unsupported mixed native QMoE expert widths: expected model-wide " + "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: + 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 +1039,213 @@ 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 _resolve_olive_qmoe_layouts( + quantization: QuantizationConfig, + *, + expected_moe_paths: tuple[str, ...], + source_moe_paths: tuple[str, ...], +) -> 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)) + return { + root: resolve_qmoe_quantization(quantization, (source_path,)) + for root, source_path in source_by_root.items() + } + + +def _preprocess_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}." + ) + + 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(): + for projection, projection_layout in ( + ("gate_up_proj", layout.fc1), + ("down_proj", layout.fc2), + ): + stem = f"{root}.experts.{projection}" + 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: + raise ValueError( + f"QMoE sidecars for {root!r} are incomplete; missing {missing!r}." + ) + 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" + 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}." + ) + + 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 + 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, + 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 +1547,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 +1568,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 +1589,39 @@ 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 ) + 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, + ) + if qmoe_target_path is not None and quantization is not None + else {} + ) + 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 @@ -1211,11 +1641,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 " @@ -1289,6 +1714,44 @@ 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": + 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, + 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_olive_qmoe_weights( + state_dict, + quantization, + { + 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, + 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..b1d084a9d 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,20 @@ 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, + # 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", 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/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.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..e691eeb80 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 @@ -53,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, @@ -357,3 +365,905 @@ 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( + *, + layer_index: int = 0, + num_experts: int = _E, + hidden_size: int = _H, + 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]: + 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 * fc1_bits // 8), + dtype=torch.int64, + ) + .remainder(256) + .to(torch.uint8) + .reshape(num_experts, fc1_out, hidden_size * fc1_bits // 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 * fc2_bits // 8), + dtype=torch.int64, + ) + .remainder(256) + .to(torch.uint8) + .reshape(num_experts, hidden_size, intermediate_size * fc2_bits // 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) * fc1_bits / 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) * fc2_bits / 8), + ), + dtype=torch.uint8, + ) + return state + + +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", + [ + 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("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" + ) + 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( + _mixed_expert_state_dict( + layer_index=layer_index, + fc1_bits=2 if layer_index == mixed_layer else 4, + ) + ) + + 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"], + rtol=0, + atol=0, + ) + torch.testing.assert_close( + out[target + "fc2_experts_weights"], + raw[source + "down_proj_qweight"], + rtol=0, + 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"), + [ + ("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, + 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, + ) + + 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": + quantization = QuantizationConfig( + bits=4, + group_size=_BLK, + quant_method="olive", + overrides={ + f"{_LAYER}mlp.experts.gate_up_proj": QuantizationOverride(bits=2), + f"{_LAYER}mlp.experts.down_proj": QuantizationOverride(bits=2), + }, + ) + else: + quantization = dataclasses.replace( + _quantization(), + modules_to_not_convert=(f"{_LAYER}mlp.experts",), + ) + with pytest.raises(ValueError, match="routed expert"): + MoECausalLMModel(_moe_config(quantization)) + + @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", "FC1 qweight 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_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( + _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_expert_override_requires_supported_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), + }, + ) + + 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)) + + @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 + + 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=2), + f"{_LAYER}mlp.experts.down_proj": QuantizationOverride(bits=8), + }, + ), + "expected model-wide 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_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())) + + 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..fce0089d2 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,52 @@ 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_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) @@ -481,6 +548,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 +1101,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 +1115,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"]))