From 30e818219df68e9480d114e10883fa417d8c5cff Mon Sep 17 00:00:00 2001 From: Shreyas8612 Date: Wed, 16 Sep 2026 19:13:31 +0100 Subject: [PATCH 1/4] Add fused Qwen3-MoE experts, BF16 router and phase-aware attention Qwen3-30B-A3B stores its routed experts in two fused 3-D Parameters, which the per-Linear quantisation path cannot reach. Add expert modules with independent prefill/decode weight banks, a BF16 router kept as a safety island with FP32 softmax and top-k, an MX router-precision ablation lane, phase-aware RMSNorm and decoder-layer residual rounding, and an import-time guard pinning the fused transformers 5.5.0 expert ABI. Attention and MLP are rebased onto the shared Llama phase mixins, adding the MXFP rotate variant and requiring identical K/V cache precision. --- src/chop/nn/quantized/modules/__init__.py | 24 + src/chop/nn/quantized/modules/phase_config.py | 8 + .../quantized/modules/qwen3_moe/__init__.py | 20 + .../quantized/modules/qwen3_moe/attention.py | 384 +++++++-------- .../nn/quantized/modules/qwen3_moe/compat.py | 40 ++ .../modules/qwen3_moe/decoder_layer.py | 86 ++++ .../nn/quantized/modules/qwen3_moe/experts.py | 445 ++++++++++++++++++ .../nn/quantized/modules/qwen3_moe/mlp.py | 36 +- .../quantized/modules/qwen3_moe/rms_norm.py | 59 +++ .../nn/quantized/modules/qwen3_moe/router.py | 100 ++++ .../modules/qwen3_moe/router_precision.py | 325 +++++++++++++ 11 files changed, 1281 insertions(+), 246 deletions(-) create mode 100644 src/chop/nn/quantized/modules/qwen3_moe/compat.py create mode 100644 src/chop/nn/quantized/modules/qwen3_moe/decoder_layer.py create mode 100644 src/chop/nn/quantized/modules/qwen3_moe/experts.py create mode 100644 src/chop/nn/quantized/modules/qwen3_moe/rms_norm.py create mode 100644 src/chop/nn/quantized/modules/qwen3_moe/router.py create mode 100644 src/chop/nn/quantized/modules/qwen3_moe/router_precision.py diff --git a/src/chop/nn/quantized/modules/__init__.py b/src/chop/nn/quantized/modules/__init__.py index c91466532..3a162c071 100644 --- a/src/chop/nn/quantized/modules/__init__.py +++ b/src/chop/nn/quantized/modules/__init__.py @@ -40,10 +40,22 @@ from .qwen3_moe import ( Qwen3MoeAttentionMXFP, + Qwen3MoeAttentionMXFPRotate, Qwen3MoeAttentionMXInt, Qwen3MoeAttentionMXIntRotate, Qwen3MoeMLPMXFP, Qwen3MoeMLPMXInt, + Qwen3MoeExpertsMXFP, + Qwen3MoeExpertsMXInt, + Qwen3MoeSparseMoeBlockBF16Router, + Qwen3MoeSparseMoeBlockMinifloat, + Qwen3MoeTopKRouterBF16, + Qwen3MoeTopKRouterMX, + Qwen3MoeTopKRouterMXInt8, + Qwen3MoeTopKRouterE4M3, + Qwen3MoeTopKRouterE5M2, + Qwen3MoeRMSNormMinifloat, + Qwen3MoeDecoderLayerMinifloat, ) from .glm4_moe import ( @@ -399,11 +411,23 @@ } quantized_qwen3_moe_module_map = { + "qwen3_moe_decoder_layer_minifloat": Qwen3MoeDecoderLayerMinifloat, "qwen3_moe_self_attention_mxfp": Qwen3MoeAttentionMXFP, + "qwen3_moe_self_attention_mxfp_rotate": Qwen3MoeAttentionMXFPRotate, "qwen3_moe_self_attention_mxint": Qwen3MoeAttentionMXInt, "qwen3_moe_self_attention_mxint_rotate": Qwen3MoeAttentionMXIntRotate, "qwen3_moe_mlp_mxfp": Qwen3MoeMLPMXFP, "qwen3_moe_mlp_mxint": Qwen3MoeMLPMXInt, + "qwen3_moe_experts_mxfp": Qwen3MoeExpertsMXFP, + "qwen3_moe_experts_mxint": Qwen3MoeExpertsMXInt, + "qwen3_moe_sparse_block_bf16_router": Qwen3MoeSparseMoeBlockBF16Router, + "qwen3_moe_sparse_block_minifloat": Qwen3MoeSparseMoeBlockMinifloat, + "qwen3_moe_router_bf16": Qwen3MoeTopKRouterBF16, + "qwen3_moe_router_mx": Qwen3MoeTopKRouterMX, + "qwen3_moe_router_mxint8": Qwen3MoeTopKRouterMXInt8, + "qwen3_moe_router_e4m3": Qwen3MoeTopKRouterE4M3, + "qwen3_moe_router_e5m2": Qwen3MoeTopKRouterE5M2, + "qwen3_moe_rms_norm_minifloat": Qwen3MoeRMSNormMinifloat, } quantized_glm4_moe_module_map = { diff --git a/src/chop/nn/quantized/modules/phase_config.py b/src/chop/nn/quantized/modules/phase_config.py index 82cdb313b..7a987783d 100644 --- a/src/chop/nn/quantized/modules/phase_config.py +++ b/src/chop/nn/quantized/modules/phase_config.py @@ -57,6 +57,14 @@ DECODE_FP_WEIGHT_ATTR = "_mase_decode_weight_fp" DECODE_FP_BIAS_ATTR = "_mase_decode_bias_fp" +# Qwen3-MoE stores all routed experts in two fused Parameters rather than +# child Linear modules. These attributes mirror the linear handoff contract +# for the fused gate/up and down tensors. +GPTQ_DECODE_EXPERT_GATE_UP_ATTR = "_mase_gptq_gate_up_proj_decode" +GPTQ_DECODE_EXPERT_DOWN_ATTR = "_mase_gptq_down_proj_decode" +DECODE_FP_EXPERT_GATE_UP_ATTR = "_mase_decode_gate_up_proj_fp" +DECODE_FP_EXPERT_DOWN_ATTR = "_mase_decode_down_proj_fp" + _VALID_DECODE_POLICIES = ("fp_only", "quantized") _VALID_KV_HANDOFFS = ("decode_format", "fp") diff --git a/src/chop/nn/quantized/modules/qwen3_moe/__init__.py b/src/chop/nn/quantized/modules/qwen3_moe/__init__.py index a43b6f3e7..d126b789d 100644 --- a/src/chop/nn/quantized/modules/qwen3_moe/__init__.py +++ b/src/chop/nn/quantized/modules/qwen3_moe/__init__.py @@ -1,6 +1,26 @@ +from .compat import SUPPORTED_TRANSFORMERS_VERSION, require_qwen3_moe_fused_abi from .attention import ( Qwen3MoeAttentionMXFP, + Qwen3MoeAttentionMXFPRotate, Qwen3MoeAttentionMXInt, Qwen3MoeAttentionMXIntRotate, ) from .mlp import Qwen3MoeMLPMXFP, Qwen3MoeMLPMXInt +from .experts import Qwen3MoeExpertsMXFP, Qwen3MoeExpertsMXInt +from .router import ( + Qwen3MoeSparseMoeBlockBF16Router, + Qwen3MoeSparseMoeBlockMinifloat, + Qwen3MoeTopKRouterBF16, +) +from .router_precision import ( + Qwen3MoeTopKRouterE4M3, + Qwen3MoeTopKRouterE5M2, + Qwen3MoeTopKRouterMX, + Qwen3MoeTopKRouterMXInt8, + ROUTER_MX_FORMATS, + ROUTER_PRECISION_SCHEMA, + router_decode_config, + router_phase_config, +) +from .rms_norm import Qwen3MoeRMSNormMinifloat +from .decoder_layer import Qwen3MoeDecoderLayerMinifloat diff --git a/src/chop/nn/quantized/modules/qwen3_moe/attention.py b/src/chop/nn/quantized/modules/qwen3_moe/attention.py index cec46a0ae..235de2e8b 100644 --- a/src/chop/nn/quantized/modules/qwen3_moe/attention.py +++ b/src/chop/nn/quantized/modules/qwen3_moe/attention.py @@ -1,283 +1,233 @@ -from typing import Optional, Tuple +"""Phase-aware Qwen3-MoE attention quantisation.""" -import torch -from torch import Tensor, nn, LongTensor +from __future__ import annotations + +from typing import Optional + +from torch import LongTensor, Tensor +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS from transformers.models.qwen3_moe.modeling_qwen3_moe import ( - apply_rotary_pos_emb, Cache, Qwen3MoeAttention, + apply_rotary_pos_emb, + eager_attention_forward as _hf_eager_attention_forward, ) -from chop.nn.quantized.functional.rope import rope_minifloat -from chop.nn.quantized.functional.softmax import softmax_minifloat -from chop.nn.quantized.functional.kvcache import ( - kv_cache_mxfp, - kv_cache_mxint, - kv_cache_mxint_rotate, -) from chop.nn.quantized.functional.attention import ( eager_attention_forward_mxfp as _eager_attention_forward_mxfp, + eager_attention_forward_mxfp_rotate as _eager_attention_forward_mxfp_rotate, eager_attention_forward_mxint as _eager_attention_forward_mxint, eager_attention_forward_mxint_rotate as _eager_attention_forward_mxint_rotate, ) +from chop.nn.quantized.functional.kvcache import kv_cache_mx +from chop.nn.quantized.functional.rope import rope_minifloat +from chop.nn.quantized.modules.llama.attention import _PhaseAwareAttentionMixin +from chop.nn.quantized.modules.qwen3.attention import _apply_qk_norm + + +def _hf_attention_dispatch( + module, + query_states, + key_states, + value_states, + attention_mask, + **kwargs, +): + interface = ALL_ATTENTION_FUNCTIONS.get_interface( + module.config._attn_implementation, _hf_eager_attention_forward + ) + return interface( + module, + query_states, + key_states, + value_states, + attention_mask, + **kwargs, + ) -class Qwen3MoeAttentionMXFP(Qwen3MoeAttention): - """MXFP-quantized Qwen3MoeAttention with RoPE, KV cache, and softmax quantization.""" +class _PhaseAwareQwen3MoeAttention(_PhaseAwareAttentionMixin, Qwen3MoeAttention): + _mx_format = "mxint" + _rotate = False - def __init__(self, config, layer_idx, q_config: dict = None): + def __init__(self, config, layer_idx, q_config: dict | None = None): super().__init__(config, layer_idx) - q_config = q_config or {} - self.qk_config = q_config.get("qk_matmul", {}) - self.av_config = q_config.get("av_matmul", {}) - self.rope_config = q_config.get("rope", {}) - self.softmax_config = q_config.get("softmax", {}) - self.kv_cache_config = q_config.get("kv_cache", {}) - self.qk_bypass = self.qk_config.get("bypass", False) - self.av_bypass = self.av_config.get("bypass", False) - self.rope_bypass = self.rope_config.get("bypass", False) - self.softmax_bypass = self.softmax_config.get("bypass", False) - self.kv_cache_bypass = self.kv_cache_config.get("bypass", False) + self._init_phase_attention_config(q_config) + for phase_cfg in self._phase_stage_cfgs.values(): + kv_cfg = phase_cfg["kv_cache_config"] + if "key" in kv_cfg or "value" in kv_cfg: + if ( + set(kv_cfg) != {"key", "value"} + or kv_cfg["key"] != kv_cfg["value"] + ): + raise ValueError( + "Qwen3-MoE requires identical K/V cache precision" + ) + if self._rotate: + decode_cfg = self._phase_stage_cfgs["decode"] + self.qk_use_rotate = decode_cfg["qk_config"].get("rotate", True) + self.av_use_rotate = decode_cfg["av_config"].get("rotate", True) + self.kv_cache_use_rotate = decode_cfg["kv_cache_config"].get( + "rotate", True + ) - def forward( + def _quantized_attention( self, - hidden_states: Tensor, - position_embeddings: Tuple[Tensor, Tensor], - attention_mask: Optional[Tensor], - past_key_values: Optional[Cache] = None, - cache_position: Optional[LongTensor] = None, + query_states, + key_states, + value_states, + attention_mask, + stage, **kwargs, - ) -> Tuple[Tensor, Optional[Tensor]]: - input_shape = hidden_states.shape[:-1] - hidden_shape = (*input_shape, -1, self.head_dim) - - query_states = self.q_norm( - self.q_proj(hidden_states).view(hidden_shape) - ).transpose(1, 2) - key_states = self.k_norm( - self.k_proj(hidden_states).view(hidden_shape) - ).transpose(1, 2) - value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) - - cos, sin = position_embeddings - if not self.rope_bypass: - query_states, key_states = rope_minifloat( - query_states, - key_states, - cos, - sin, - self.rope_config, + ): + common = dict( + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, + qk_bypass=stage["qk_bypass"], + qk_config=stage["qk_config"], + av_bypass=stage["av_bypass"], + av_config=stage["av_config"], + softmax_bypass=stage["softmax_bypass"], + softmax_config=stage["softmax_config"], + **kwargs, + ) + if self._mx_format == "mxfp": + function = ( + _eager_attention_forward_mxfp_rotate + if self._rotate + else _eager_attention_forward_mxfp ) else: - query_states, key_states = apply_rotary_pos_emb( - query_states, - key_states, - cos, - sin, + function = ( + _eager_attention_forward_mxint_rotate + if self._rotate + else _eager_attention_forward_mxint ) - - if past_key_values is not None: - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - if not self.kv_cache_bypass: - key_states, value_states = kv_cache_mxfp( - key_states, - value_states, - self.kv_cache_config, - ) - key_states, value_states = past_key_values.update( - key_states, - value_states, - self.layer_idx, - cache_kwargs, - ) - - attn_output, attn_weights = _eager_attention_forward_mxfp( + if self._rotate: + common["qk_use_rotate"] = self.qk_use_rotate + common["av_use_rotate"] = self.av_use_rotate + return function( self, query_states, key_states, value_states, attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - qk_bypass=self.qk_bypass, - qk_config=self.qk_config, - av_bypass=self.av_bypass, - av_config=self.av_config, - softmax_bypass=self.softmax_bypass, - softmax_config=self.softmax_config, - **kwargs, + **common, ) - attn_output = attn_output.reshape(*input_shape, -1).contiguous() - attn_output = self.o_proj(attn_output) - return attn_output, attn_weights - - -class Qwen3MoeAttentionMXInt(Qwen3MoeAttention): - """MXInt-quantized Qwen3MoeAttention with RoPE, KV cache, and softmax quantization.""" - - def __init__(self, config, layer_idx, q_config: dict = None): - super().__init__(config, layer_idx) - q_config = q_config or {} - self.qk_config = q_config.get("qk_matmul", {}) - self.av_config = q_config.get("av_matmul", {}) - self.rope_config = q_config.get("rope", {}) - self.softmax_config = q_config.get("softmax", {}) - self.kv_cache_config = q_config.get("kv_cache", {}) - self.qk_bypass = self.qk_config.get("bypass", False) - self.av_bypass = self.av_config.get("bypass", False) - self.rope_bypass = self.rope_config.get("bypass", False) - self.softmax_bypass = self.softmax_config.get("bypass", False) - self.kv_cache_bypass = self.kv_cache_config.get("bypass", False) - def forward( self, hidden_states: Tensor, - position_embeddings: Tuple[Tensor, Tensor], + position_embeddings: tuple[Tensor, Tensor], attention_mask: Optional[Tensor], past_key_values: Optional[Cache] = None, cache_position: Optional[LongTensor] = None, **kwargs, - ) -> Tuple[Tensor, Optional[Tensor]]: + ): + past_key_values = kwargs.pop("past_key_value", past_key_values) + stage = self._active_stage_cfgs() input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) - query_states = self.q_norm( - self.q_proj(hidden_states).view(hidden_shape) - ).transpose(1, 2) - key_states = self.k_norm( - self.k_proj(hidden_states).view(hidden_shape) - ).transpose(1, 2) + query_states = self.q_proj(hidden_states).view(hidden_shape) + key_states = self.k_proj(hidden_states).view(hidden_shape) + query_states, key_states = _apply_qk_norm( + self, query_states, key_states, stage + ) + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings - if not self.rope_bypass: - query_states, key_states = rope_minifloat( - query_states, - key_states, - cos, - sin, - self.rope_config, + if stage["rope_bypass"]: + query_states, key_states = apply_rotary_pos_emb( + query_states, key_states, cos, sin ) else: - query_states, key_states = apply_rotary_pos_emb( + query_states, key_states = rope_minifloat( query_states, key_states, cos, sin, + stage["rope_config"], ) - if past_key_values is not None: - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - if not self.kv_cache_bypass: - key_states, value_states = kv_cache_mxint( - key_states, - value_states, - self.kv_cache_config, - ) - key_states, value_states = past_key_values.update( - key_states, - value_states, - self.layer_idx, - cache_kwargs, + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + + def cache_quantizer(key, value, config): + return kv_cache_mx( + key, + value, + config, + rotate=self._rotate and self.kv_cache_use_rotate, ) - attn_output, attn_weights = _eager_attention_forward_mxint( - self, - query_states, + key_states, value_states = self._apply_kv_cache_phase_aware( key_states, value_states, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - qk_bypass=self.qk_bypass, - qk_config=self.qk_config, - av_bypass=self.av_bypass, - av_config=self.av_config, - softmax_bypass=self.softmax_bypass, - softmax_config=self.softmax_config, - **kwargs, + past_key_values, + cache_kwargs, + stage, + cache_quantizer, ) - attn_output = attn_output.reshape(*input_shape, -1).contiguous() - attn_output = self.o_proj(attn_output) - return attn_output, attn_weights - - -class Qwen3MoeAttentionMXIntRotate(Qwen3MoeAttentionMXInt): - """Qwen3MoeAttentionMXInt with online Hadamard rotation around the - activation quantizers (KV-cache, Q@K, A@V). See - ``Qwen3AttentionMXIntRotate`` for the same pattern in dense Qwen3.""" - - def forward( - self, - hidden_states: Tensor, - position_embeddings: Tuple[Tensor, Tensor], - attention_mask: Optional[Tensor], - past_key_values: Optional[Cache] = None, - cache_position: Optional[LongTensor] = None, - **kwargs, - ) -> Tuple[Tensor, Optional[Tensor]]: - input_shape = hidden_states.shape[:-1] - hidden_shape = (*input_shape, -1, self.head_dim) - - query_states = self.q_norm( - self.q_proj(hidden_states).view(hidden_shape) - ).transpose(1, 2) - key_states = self.k_norm( - self.k_proj(hidden_states).view(hidden_shape) - ).transpose(1, 2) - value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) - - cos, sin = position_embeddings - if not self.rope_bypass: - query_states, key_states = rope_minifloat( + if stage["qk_bypass"] and stage["av_bypass"] and stage["softmax_bypass"]: + attn_output, attn_weights = _hf_attention_dispatch( + self, query_states, key_states, - cos, - sin, - self.rope_config, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, + **kwargs, ) else: - query_states, key_states = apply_rotary_pos_emb( - query_states, - key_states, - cos, - sin, - ) - - if past_key_values is not None: - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - if not self.kv_cache_bypass: - key_states, value_states = kv_cache_mxint_rotate( - key_states, - value_states, - self.kv_cache_config, + if self.config._attn_implementation != "eager": + raise ValueError( + "quantized Qwen3-MoE attention requires " + "_attn_implementation='eager'" ) - key_states, value_states = past_key_values.update( + attn_output, attn_weights = self._quantized_attention( + query_states, key_states, value_states, - self.layer_idx, - cache_kwargs, + attention_mask, + stage, + **kwargs, ) - attn_output, attn_weights = _eager_attention_forward_mxint_rotate( - self, - query_states, - key_states, - value_states, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - qk_bypass=self.qk_bypass, - qk_config=self.qk_config, - av_bypass=self.av_bypass, - av_config=self.av_config, - softmax_bypass=self.softmax_bypass, - softmax_config=self.softmax_config, - **kwargs, + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + return self.o_proj(attn_output), attn_weights + + @classmethod + def from_self( + cls, attention: Qwen3MoeAttention, q_config: dict | None = None + ): + new = cls( + config=attention.config, + layer_idx=attention.layer_idx, + q_config=q_config, ) + parameter = next(attention.parameters()) + new = new.to(device=parameter.device, dtype=parameter.dtype) + new.load_state_dict(attention.state_dict(), strict=True) + return new - attn_output = attn_output.reshape(*input_shape, -1).contiguous() - attn_output = self.o_proj(attn_output) - return attn_output, attn_weights + +class Qwen3MoeAttentionMXFP(_PhaseAwareQwen3MoeAttention): + _mx_format = "mxfp" + + +class Qwen3MoeAttentionMXInt(_PhaseAwareQwen3MoeAttention): + _mx_format = "mxint" + + +class Qwen3MoeAttentionMXFPRotate(Qwen3MoeAttentionMXFP): + _rotate = True + + +class Qwen3MoeAttentionMXIntRotate(Qwen3MoeAttentionMXInt): + _rotate = True diff --git a/src/chop/nn/quantized/modules/qwen3_moe/compat.py b/src/chop/nn/quantized/modules/qwen3_moe/compat.py new file mode 100644 index 000000000..f537ff83b --- /dev/null +++ b/src/chop/nn/quantized/modules/qwen3_moe/compat.py @@ -0,0 +1,40 @@ +"""Transformers ABI guard for fused Qwen3-MoE modules.""" + +from __future__ import annotations + +import transformers + + +SUPPORTED_TRANSFORMERS_VERSION = "5.5.0" + + +def require_qwen3_moe_fused_abi() -> None: + """Reject releases whose Qwen3-MoE weight layout is incompatible.""" + + if transformers.__version__ != SUPPORTED_TRANSFORMERS_VERSION: + raise ImportError( + "MASE fused Qwen3-MoE quantization requires " + f"transformers=={SUPPORTED_TRANSFORMERS_VERSION}; found " + f"{transformers.__version__}. Transformers 4.51 uses a " + "ModuleList expert ABI and is not compatible." + ) + try: + from transformers.models.qwen3_moe.modeling_qwen3_moe import ( + Qwen3MoeExperts, + Qwen3MoeTopKRouter, + ) + except ImportError as exc: + raise ImportError( + "transformers==5.5.0 must expose fused Qwen3MoeExperts and " + "Qwen3MoeTopKRouter" + ) from exc + if not all(hasattr(Qwen3MoeExperts, name) for name in ("forward",)): + raise ImportError("Qwen3MoeExperts does not expose the expected fused ABI") + if not hasattr(Qwen3MoeTopKRouter, "forward"): + raise ImportError("Qwen3MoeTopKRouter does not expose the expected router ABI") + + +require_qwen3_moe_fused_abi() + + +__all__ = ["SUPPORTED_TRANSFORMERS_VERSION", "require_qwen3_moe_fused_abi"] diff --git a/src/chop/nn/quantized/modules/qwen3_moe/decoder_layer.py b/src/chop/nn/quantized/modules/qwen3_moe/decoder_layer.py new file mode 100644 index 000000000..eef48d189 --- /dev/null +++ b/src/chop/nn/quantized/modules/qwen3_moe/decoder_layer.py @@ -0,0 +1,86 @@ +"""Qwen3-MoE decoder layer with phase-aware residual rounding.""" + +from __future__ import annotations + +from torch import Tensor +from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeDecoderLayer + +from chop.nn.quantized.functional.vector import VectorRoundingPolicy +from chop.nn.quantized.modules.phase_config import ( + normalize_phase_q_config, + resolve_module_phase_config, +) +from chop.nn.quantized.modules.phase_context import get_runtime_phase + + +class Qwen3MoeDecoderLayerMinifloat(Qwen3MoeDecoderLayer): + """Apply the active FP_SETTING policy at both residual boundaries.""" + + def __init__(self, config, layer_idx: int, q_config: dict | None = None): + super().__init__(config, layer_idx) + self.q_config = q_config or {} + self.phase_q_config = normalize_phase_q_config(q_config) + self.decode_policy = self.phase_q_config["decode_policy"] + self._phase_policies = { + phase: VectorRoundingPolicy.from_config( + resolve_module_phase_config(self.phase_q_config, phase) + ) + for phase in ("prefill", "decode") + } + self.bypass = not self._phase_policies["decode"].enabled + + def forward( + self, + hidden_states: Tensor, + attention_mask: Tensor | None = None, + position_ids: Tensor | None = None, + past_key_values=None, + use_cache: bool | None = False, + cache_position: Tensor | None = None, + position_embeddings: tuple[Tensor, Tensor] | None = None, + **kwargs, + ) -> Tensor: + policy = self._phase_policies[get_runtime_phase()] + if policy.enabled: + hidden_states = policy.round(hidden_states) + + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = ( + policy.residual_add(residual, hidden_states) + if policy.enabled + else residual + hidden_states + ) + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + return ( + policy.residual_add(residual, hidden_states) + if policy.enabled + else residual + hidden_states + ) + + @classmethod + def from_self( + cls, layer: Qwen3MoeDecoderLayer, q_config: dict | None = None + ): + new = cls( + config=layer.self_attn.config, + layer_idx=layer.self_attn.layer_idx, + q_config=q_config, + ) + parameter = next(layer.parameters()) + new = new.to(device=parameter.device, dtype=parameter.dtype) + new.load_state_dict(layer.state_dict(), strict=True) + return new diff --git a/src/chop/nn/quantized/modules/qwen3_moe/experts.py b/src/chop/nn/quantized/modules/qwen3_moe/experts.py new file mode 100644 index 000000000..f770838d4 --- /dev/null +++ b/src/chop/nn/quantized/modules/qwen3_moe/experts.py @@ -0,0 +1,445 @@ +"""Phase-aware quantisation for fused Qwen3-MoE expert tensors.""" + +from __future__ import annotations + +import os + +import torch +import torch.nn.functional as F +from torch import Tensor, nn +from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeExperts + +from chop.nn.quantizers import mxfp_quantizer, mxint_quantizer +from chop.nn.quantized.functional.matrix import plena_matrix_product +from chop.nn.quantized.functional.vector import VectorRoundingPolicy +from chop.nn.quantized.modules.phase_config import ( + DECODE_FP_EXPERT_DOWN_ATTR, + DECODE_FP_EXPERT_GATE_UP_ATTR, + GPTQ_DECODE_EXPERT_DOWN_ATTR, + GPTQ_DECODE_EXPERT_GATE_UP_ATTR, + normalize_phase_q_config, + resolve_module_phase_config, +) +from chop.nn.quantized.modules.phase_context import get_runtime_phase + + +class _PhaseAwareQwen3MoeExperts(Qwen3MoeExperts): + """Qwen3 routed experts with independent prefill and decode weight banks.""" + + _WEIGHT_CFG_KEYS: tuple[str, ...] = () + + def __init__(self, config, q_config: dict | None = None): + super().__init__(config) + self._init_phase_state(q_config) + self._register_phase_banks() + + def _init_phase_state(self, q_config: dict | None) -> None: + self.q_config = q_config or {} + self.phase_q_config = normalize_phase_q_config(q_config) + self.decode_policy = self.phase_q_config["decode_policy"] + self.prefill_config = resolve_module_phase_config( + self.phase_q_config, "prefill" + ) + self.decode_config = resolve_module_phase_config( + self.phase_q_config, "decode" + ) + self.bypass = self.prefill_config.get("bypass", False) + self.gptq = self.prefill_config.get("gptq", False) + self.decode_gptq = self.decode_config.get("gptq", False) + self.shared_phase_banks = self.prefill_config == self.decode_config + self._validate_weight_configs() + + def _register_phase_banks(self) -> None: + self.register_buffer( + "_decode_gate_up_q", torch.empty(0), persistent=False + ) + self.register_buffer("_decode_down_q", torch.empty(0), persistent=False) + self.register_buffer( + "_decode_gate_up_fp", torch.empty(0), persistent=False + ) + self.register_buffer("_decode_down_fp", torch.empty(0), persistent=False) + self._decode_bank_collapsed = False + self._decode_bank_sealed = False + self._weight_quantization_events = 0 + + def _validate_weight_configs(self) -> None: + for phase, cfg in ( + ("prefill", self.prefill_config), + ("decode", self.decode_config), + ): + if cfg.get("bypass", False) or cfg.get("gptq", False): + continue + present = [key for key in self._WEIGHT_CFG_KEYS if cfg.get(key) is not None] + if present and len(present) != len(self._WEIGHT_CFG_KEYS): + missing = [ + key for key in self._WEIGHT_CFG_KEYS if cfg.get(key) is None + ] + raise ValueError( + f"{type(self).__name__}: incomplete expert weight config in " + f"the {phase} bucket; missing {missing}" + ) + + def _weight_config_present(self, cfg: dict) -> bool: + return all(cfg.get(key) is not None for key in self._WEIGHT_CFG_KEYS) + + def _quantize_weight(self, weight: Tensor, cfg: dict) -> Tensor: + raise NotImplementedError + + def _quantize_activation(self, value: Tensor, cfg: dict) -> Tensor: + block_size = cfg.get("data_in_block_size") + if block_size is None: + return value + if cfg.get("data_in_width") is not None: + return mxint_quantizer( + value, + block_size=block_size, + element_bits=cfg["data_in_width"], + block_dim=-1, + ) + if ( + cfg.get("data_in_exponent_width") is not None + and cfg.get("data_in_frac_width") is not None + ): + return mxfp_quantizer( + value, + block_size=block_size, + element_exp_bits=cfg["data_in_exponent_width"], + element_frac_bits=cfg["data_in_frac_width"], + block_dim=-1, + ) + return value + + @torch.no_grad() + def _quantize_expert_bank(self, weight: Tensor, cfg: dict) -> Tensor: + """Quantise one expert at a time to bound peak temporary storage.""" + + self._weight_quantization_events += 1 + output = torch.empty_like(weight) + route = os.environ.get("MASE_PHASE_BANK_DEVICE") or None + for expert_idx in range(weight.shape[0]): + source = weight[expert_idx] + quantized = self._quantize_weight( + source.to(route) if route and str(source.device) != route else source, + cfg, + ) + output[expert_idx].copy_(quantized.to(output.device)) + return output + + @torch.no_grad() + def _build_phase_weight_banks(self) -> None: + if self._decode_bank_sealed: + raise RuntimeError("a sealed decode expert bank cannot be rebuilt") + + device = self.gate_up_proj.device + self._decode_gate_up_q = torch.empty(0, device=device) + self._decode_down_q = torch.empty(0, device=device) + self._decode_gate_up_fp = torch.empty(0, device=device) + self._decode_down_fp = torch.empty(0, device=device) + + prefill_quantizes = ( + not self.bypass + and not self.gptq + and self._weight_config_present(self.prefill_config) + ) + if not self.shared_phase_banks: + decode_cfg = self.decode_config + if ( + not decode_cfg.get("bypass", False) + and not self.decode_gptq + and self._weight_config_present(decode_cfg) + ): + self._decode_gate_up_q = self._quantize_expert_bank( + self.gate_up_proj.data, decode_cfg + ) + self._decode_down_q = self._quantize_expert_bank( + self.down_proj.data, decode_cfg + ) + elif decode_cfg.get("bypass", False) and prefill_quantizes: + self._decode_gate_up_fp = self.gate_up_proj.detach().clone() + self._decode_down_fp = self.down_proj.detach().clone() + + if prefill_quantizes: + self.gate_up_proj.data.copy_( + self._quantize_expert_bank( + self.gate_up_proj.data, self.prefill_config + ) + ) + self.down_proj.data.copy_( + self._quantize_expert_bank(self.down_proj.data, self.prefill_config) + ) + + @torch.no_grad() + def adopt_decode_gptq_expert_weights( + self, gate_up_proj: Tensor, down_proj: Tensor + ) -> None: + self._decode_gate_up_q = gate_up_proj.detach().to( + device=self.gate_up_proj.device, + dtype=self.gate_up_proj.dtype, + copy=True, + ) + self._decode_down_q = down_proj.detach().to( + device=self.down_proj.device, + dtype=self.down_proj.dtype, + copy=True, + ) + + @torch.no_grad() + def adopt_decode_fp_expert_weights( + self, gate_up_proj: Tensor, down_proj: Tensor + ) -> None: + self._decode_gate_up_fp = gate_up_proj.detach().to( + device=self.gate_up_proj.device, + dtype=self.gate_up_proj.dtype, + copy=True, + ) + self._decode_down_fp = down_proj.detach().to( + device=self.down_proj.device, + dtype=self.down_proj.dtype, + copy=True, + ) + + @classmethod + def from_self(cls, experts: Qwen3MoeExperts, q_config: dict | None = None): + """Create a wrapper without allocating another pair of expert Parameters.""" + + new = cls.__new__(cls) + nn.Module.__init__(new) + new.num_experts = experts.num_experts + new.hidden_dim = experts.hidden_dim + new.intermediate_dim = experts.intermediate_dim + new.gate_up_proj = experts.gate_up_proj + new.down_proj = experts.down_proj + new.act_fn = experts.act_fn + new._init_phase_state(q_config) + new._register_phase_banks() + new._build_phase_weight_banks() + + gate_up_gptq = getattr(experts, GPTQ_DECODE_EXPERT_GATE_UP_ATTR, None) + down_gptq = getattr(experts, GPTQ_DECODE_EXPERT_DOWN_ATTR, None) + if gate_up_gptq is not None or down_gptq is not None: + if gate_up_gptq is None or down_gptq is None: + raise RuntimeError("incomplete Qwen3-MoE GPTQ decode bank") + new.adopt_decode_gptq_expert_weights(gate_up_gptq, down_gptq) + + gate_up_fp = getattr(experts, DECODE_FP_EXPERT_GATE_UP_ATTR, None) + down_fp = getattr(experts, DECODE_FP_EXPERT_DOWN_ATTR, None) + if gate_up_fp is not None or down_fp is not None: + if gate_up_fp is None or down_fp is None: + raise RuntimeError("incomplete Qwen3-MoE FP decode snapshot") + new.adopt_decode_fp_expert_weights(gate_up_fp, down_fp) + return new + + @torch.no_grad() + def collapse_to_decode_bank(self) -> bool: + if self._decode_gate_up_q.numel() == 0: + return False + if self._decode_down_q.numel() == 0: + raise RuntimeError("incomplete Qwen3-MoE decode bank") + self.gate_up_proj.data.copy_(self._decode_gate_up_q) + self.down_proj.data.copy_(self._decode_down_q) + device = self.gate_up_proj.device + self._decode_gate_up_q = torch.empty(0, device=device) + self._decode_down_q = torch.empty(0, device=device) + self._decode_bank_collapsed = True + return True + + def seal_decode_weight_bank(self) -> int: + if not self._decode_bank_collapsed: + raise RuntimeError("decode expert bank was not collapsed") + if self._decode_gate_up_q.numel() or self._decode_down_q.numel(): + raise RuntimeError("decode expert bank still has an auxiliary copy") + self._decode_bank_sealed = True + return self._weight_quantization_events + + def _active_weights(self) -> tuple[Tensor, Tensor, dict]: + if get_runtime_phase() != "decode" or self.shared_phase_banks: + return self.gate_up_proj, self.down_proj, self.prefill_config + if self._decode_gate_up_q.numel(): + return self._decode_gate_up_q, self._decode_down_q, self.decode_config + if self._decode_gate_up_fp.numel(): + return self._decode_gate_up_fp, self._decode_down_fp, self.decode_config + return self.gate_up_proj, self.down_proj, self.decode_config + + def forward( + self, + hidden_states: Tensor, + top_k_index: Tensor, + top_k_weights: Tensor, + ) -> Tensor: + gate_up_bank, down_bank, cfg = self._active_weights() + if cfg.get("bypass", False): + # Calling the HF implementation would read self Parameters instead + # of an FP decode snapshot, so retain the shared routed loop here. + quantize = False + else: + quantize = True + + final_hidden_states = torch.zeros_like(hidden_states) + with torch.no_grad(): + expert_mask = F.one_hot(top_k_index, num_classes=self.num_experts) + expert_mask = expert_mask.permute(2, 1, 0) + expert_hit = torch.greater( + expert_mask.sum(dim=(-1, -2)), 0 + ).nonzero() + + vector_token = cfg.get("output_format") or cfg.get("vector_format") + vector_policy = ( + VectorRoundingPolicy.from_token(vector_token) + if isinstance(vector_token, str) + else VectorRoundingPolicy.disabled() + ) + if self._use_gathered_dispatch(top_k_index): + return self._forward_gathered( + hidden_states, + top_k_index, + top_k_weights, + gate_up_bank, + down_bank, + cfg, + quantize, + vector_policy, + ) + for expert_hit_idx in expert_hit: + expert_idx = int(expert_hit_idx[0].item()) + if expert_idx >= self.num_experts: + continue + top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) + current_state = hidden_states[token_idx] + if quantize: + current_state = self._quantize_activation(current_state, cfg) + gate_up = plena_matrix_product( + current_state, + gate_up_bank[expert_idx].transpose(-1, -2), + cfg, + ) + gate, up = gate_up.chunk(2, dim=-1) + current_hidden = ( + vector_policy.silu_gate(gate, up) + if vector_policy.enabled + else self.act_fn(gate) * up + ) + if quantize: + current_hidden = self._quantize_activation(current_hidden, cfg) + current_hidden = plena_matrix_product( + current_hidden, + down_bank[expert_idx].transpose(-1, -2), + cfg, + ) + current_hidden = current_hidden * top_k_weights[ + token_idx, top_k_pos, None + ] + final_hidden_states.index_add_( + 0, token_idx, current_hidden.to(final_hidden_states.dtype) + ) + return final_hidden_states + + # ------------------------------------------------------------------ + # Gathered dispatch for decode-sized token counts. + # + # The per-expert loop above issues ~70 small kernels per hit expert, and a + # q_len=1 microbatch of 16 hits ~100 of the 128 experts in every layer, so + # one decode step costs hundreds of thousands of launches. The gathered + # path evaluates every (token, expert) assignment at once: the activation + # quantiser works per row along ``block_dim=-1`` and ``plena_matrix_product`` + # partitions the reduction dimension per row, so both are numerically the + # same operations applied to stacked rows. The only semantic difference is + # the combine: the loop accumulates expert partials into the BF16 output in + # expert-index order, the gathered path sums the weighted partials of each + # token in FP32 and rounds once. ``MASE_MOE_EXPERT_DISPATCH=loop`` restores + # the original loop; ``gather`` forces the gathered path; the default uses + # the gathered path below ``MASE_MOE_GATHER_MAX_ASSIGNMENTS`` assignments + # (prefill-sized token counts keep the loop, whose per-expert matmuls are + # already large). + # ------------------------------------------------------------------ + _GATHER_CHUNK = 512 + + def _use_gathered_dispatch(self, top_k_index: Tensor) -> bool: + mode = os.environ.get("MASE_MOE_EXPERT_DISPATCH", "auto") + if mode == "loop": + return False + if mode == "gather": + return True + limit = int(os.environ.get("MASE_MOE_GATHER_MAX_ASSIGNMENTS", "4096")) + return top_k_index.numel() <= limit + + def _forward_gathered( + self, + hidden_states: Tensor, + top_k_index: Tensor, + top_k_weights: Tensor, + gate_up_bank: Tensor, + down_bank: Tensor, + cfg: dict, + quantize: bool, + vector_policy: VectorRoundingPolicy, + ) -> Tensor: + tokens, top_k = top_k_index.shape + assignment_expert = top_k_index.reshape(-1) + assignment_token = torch.arange( + tokens, device=hidden_states.device + ).repeat_interleave(top_k) + assignment_weight = top_k_weights.reshape(-1) + accumulator = torch.zeros( + hidden_states.shape, dtype=torch.float32, device=hidden_states.device + ) + total = assignment_expert.numel() + for start in range(0, total, self._GATHER_CHUNK): + stop = min(total, start + self._GATHER_CHUNK) + experts = assignment_expert[start:stop] + token_idx = assignment_token[start:stop] + current_state = hidden_states[token_idx] + if quantize: + current_state = self._quantize_activation(current_state, cfg) + gate_up = plena_matrix_product( + current_state.unsqueeze(1), + gate_up_bank[experts].transpose(-1, -2), + cfg, + ).squeeze(1) + gate, up = gate_up.chunk(2, dim=-1) + current_hidden = ( + vector_policy.silu_gate(gate, up) + if vector_policy.enabled + else self.act_fn(gate) * up + ) + if quantize: + current_hidden = self._quantize_activation(current_hidden, cfg) + current_hidden = plena_matrix_product( + current_hidden.unsqueeze(1), + down_bank[experts].transpose(-1, -2), + cfg, + ).squeeze(1) + current_hidden = current_hidden.to(torch.float32) * assignment_weight[ + start:stop, None + ].to(torch.float32) + accumulator.index_add_(0, token_idx, current_hidden) + return accumulator.to(hidden_states.dtype) + + +class Qwen3MoeExpertsMXFP(_PhaseAwareQwen3MoeExperts): + _WEIGHT_CFG_KEYS = ( + "weight_block_size", + "weight_exponent_width", + "weight_frac_width", + ) + + def _quantize_weight(self, weight: Tensor, cfg: dict) -> Tensor: + return mxfp_quantizer( + weight, + block_size=cfg["weight_block_size"], + element_exp_bits=cfg["weight_exponent_width"], + element_frac_bits=cfg["weight_frac_width"], + block_dim=1, + ) + + +class Qwen3MoeExpertsMXInt(_PhaseAwareQwen3MoeExperts): + _WEIGHT_CFG_KEYS = ("weight_block_size", "weight_width") + + def _quantize_weight(self, weight: Tensor, cfg: dict) -> Tensor: + return mxint_quantizer( + weight, + block_size=cfg["weight_block_size"], + element_bits=cfg["weight_width"], + block_dim=1, + quantile_search=cfg.get("clip_search", False), + ) diff --git a/src/chop/nn/quantized/modules/qwen3_moe/mlp.py b/src/chop/nn/quantized/modules/qwen3_moe/mlp.py index 2bda0bfd1..637066ea1 100644 --- a/src/chop/nn/quantized/modules/qwen3_moe/mlp.py +++ b/src/chop/nn/quantized/modules/qwen3_moe/mlp.py @@ -1,41 +1,19 @@ -from torch import Tensor - -from chop.nn.quantized.functional.silu import silu_minifloat +from chop.nn.quantized.modules.llama.mlp import _PhaseAwareGLUMLPMixin from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeMLP -class Qwen3MoeMLPMXFP(Qwen3MoeMLP): - """MXFP-quantized Qwen3MoeMLP (individual expert). SiLU uses minifloat quantization.""" - +class _PhaseAwareQwen3MoeMLP(_PhaseAwareGLUMLPMixin, Qwen3MoeMLP): def __init__( self, config, intermediate_size=None, layer_idx=None, q_config: dict = None ): super().__init__(config, intermediate_size) - self.layer_idx = layer_idx - self.q_config = q_config or {} - self.bypass = self.q_config.get("bypass", False) + self._init_phase_mlp_config(layer_idx, q_config) - def forward(self, x: Tensor) -> Tensor: - if self.bypass: - return super().forward(x) - x = silu_minifloat(self.gate_proj(x), self.q_config) * self.up_proj(x) - return self.down_proj(x) +class Qwen3MoeMLPMXFP(_PhaseAwareQwen3MoeMLP): + """Phase-aware vector rounding for a dense Qwen3-MoE MLP layer.""" -class Qwen3MoeMLPMXInt(Qwen3MoeMLP): - """MXInt-quantized Qwen3MoeMLP (individual expert). SiLU uses minifloat quantization.""" - - def __init__( - self, config, intermediate_size=None, layer_idx=None, q_config: dict = None - ): - super().__init__(config, intermediate_size) - self.layer_idx = layer_idx - self.q_config = q_config or {} - self.bypass = self.q_config.get("bypass", False) - def forward(self, x: Tensor) -> Tensor: - if self.bypass: - return super().forward(x) - x = silu_minifloat(self.gate_proj(x), self.q_config) * self.up_proj(x) - return self.down_proj(x) +class Qwen3MoeMLPMXInt(_PhaseAwareQwen3MoeMLP): + """Phase-aware vector rounding for a dense Qwen3-MoE MLP layer.""" diff --git a/src/chop/nn/quantized/modules/qwen3_moe/rms_norm.py b/src/chop/nn/quantized/modules/qwen3_moe/rms_norm.py new file mode 100644 index 000000000..b9a915b37 --- /dev/null +++ b/src/chop/nn/quantized/modules/qwen3_moe/rms_norm.py @@ -0,0 +1,59 @@ +"""Phase-aware vector rounding for Qwen3-MoE RMSNorm.""" + +from chop.nn.quantized.modules.llama.rms_norm import build_vector_phase_policies +from chop.nn.quantized.modules.phase_config import ( + normalize_phase_q_config, + resolve_module_phase_config, +) +from chop.nn.quantized.modules.phase_context import get_runtime_phase +from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeRMSNorm + + +class Qwen3MoeRMSNormMinifloat(Qwen3MoeRMSNorm): + def __init__( + self, + config=None, + layer_idx=None, + q_config: dict | None = None, + *, + hidden_size: int | None = None, + eps: float | None = None, + ): + hidden_size = hidden_size if hidden_size is not None else config.hidden_size + eps = eps if eps is not None else config.rms_norm_eps + super().__init__(hidden_size=hidden_size, eps=eps) + self.layer_idx = layer_idx + self.q_config = q_config or {} + self.phase_q_config = normalize_phase_q_config(q_config) + self.decode_policy = self.phase_q_config["decode_policy"] + self._phase_policies = { + phase: build_vector_phase_policies( + resolve_module_phase_config(self.phase_q_config, phase) + ) + for phase in ("prefill", "decode") + } + self.bypass = resolve_module_phase_config( + self.phase_q_config, "decode" + ).get("bypass", False) + + def forward(self, hidden_states): + weight_policy, input_policy = self._phase_policies[get_runtime_phase()] + if not weight_policy.enabled and not input_policy.enabled: + return super().forward(hidden_states) + return input_policy.rms_norm( + hidden_states, + self.weight, + self.variance_epsilon, + weight_policy=weight_policy, + ) + + @classmethod + def from_self(cls, norm: Qwen3MoeRMSNorm, q_config: dict | None = None): + new = cls( + q_config=q_config, + hidden_size=norm.weight.numel(), + eps=norm.variance_epsilon, + ) + new = new.to(device=norm.weight.device, dtype=norm.weight.dtype) + new.load_state_dict(norm.state_dict(), strict=True) + return new diff --git a/src/chop/nn/quantized/modules/qwen3_moe/router.py b/src/chop/nn/quantized/modules/qwen3_moe/router.py new file mode 100644 index 000000000..3c914bfd0 --- /dev/null +++ b/src/chop/nn/quantized/modules/qwen3_moe/router.py @@ -0,0 +1,100 @@ +"""BF16 router safety island for Qwen3-MoE.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import nn +from transformers.models.qwen3_moe.modeling_qwen3_moe import ( + Qwen3MoeSparseMoeBlock, + Qwen3MoeTopKRouter, +) + +from chop.nn.quantized.functional.vector import VectorRoundingPolicy +from chop.nn.quantized.modules.phase_config import ( + normalize_phase_q_config, + resolve_module_phase_config, +) +from chop.nn.quantized.modules.phase_context import get_runtime_phase + + +class Qwen3MoeTopKRouterBF16(Qwen3MoeTopKRouter): + """Keep router GEMM in BF16 and its probability/top-k path in FP32.""" + + def forward(self, hidden_states): + hidden_states = hidden_states.reshape(-1, self.hidden_dim) + router_logits = F.linear( + hidden_states.to(torch.bfloat16), self.weight.to(torch.bfloat16) + ) + router_probs = F.softmax(router_logits, dtype=torch.float32, dim=-1) + router_top_value, router_indices = torch.topk( + router_probs, self.top_k, dim=-1 + ) + if self.norm_topk_prob: + router_top_value = router_top_value / router_top_value.sum( + dim=-1, keepdim=True + ) + return router_probs, router_top_value, router_indices + + @classmethod + def from_self(cls, router: Qwen3MoeTopKRouter, q_config: dict | None = None): + del q_config + new = cls.__new__(cls) + nn.Module.__init__(new) + new.top_k = router.top_k + new.num_experts = router.num_experts + new.norm_topk_prob = router.norm_topk_prob + new.hidden_dim = router.hidden_dim + # Keep one resident BF16 copy; per-forward weight conversion would be + # prohibitive in token-by-token decode. + new.weight = nn.Parameter( + router.weight.detach().to(torch.bfloat16), + requires_grad=router.weight.requires_grad, + ) + return new + + +class Qwen3MoeSparseMoeBlockBF16Router(Qwen3MoeSparseMoeBlock): + """Sparse block retaining HF routing semantics with a BF16 router.""" + + @classmethod + def from_self( + cls, block: Qwen3MoeSparseMoeBlock, q_config: dict | None = None + ): + new = cls.__new__(cls) + nn.Module.__init__(new) + new.experts = block.experts + new.gate = Qwen3MoeTopKRouterBF16.from_self(block.gate, q_config=q_config) + return new + + +class Qwen3MoeSparseMoeBlockMinifloat(Qwen3MoeSparseMoeBlockBF16Router): + """BF16 routing with phase-aware rounding at routed combine.""" + + def _init_phase_config(self, q_config: dict | None) -> None: + self.q_config = q_config or {} + self.phase_q_config = normalize_phase_q_config(q_config) + self.decode_policy = self.phase_q_config["decode_policy"] + self._phase_policies = { + phase: VectorRoundingPolicy.from_config( + resolve_module_phase_config(self.phase_q_config, phase) + ) + for phase in ("prefill", "decode") + } + self.bypass = not self._phase_policies["decode"].enabled + + def forward(self, hidden_states): + output = super().forward(hidden_states) + policy = self._phase_policies[get_runtime_phase()] + return policy.round(output) if policy.enabled else output + + @classmethod + def from_self( + cls, block: Qwen3MoeSparseMoeBlock, q_config: dict | None = None + ): + new = cls.__new__(cls) + nn.Module.__init__(new) + new.experts = block.experts + new.gate = Qwen3MoeTopKRouterBF16.from_self(block.gate, q_config=q_config) + new._init_phase_config(q_config) + return new diff --git a/src/chop/nn/quantized/modules/qwen3_moe/router_precision.py b/src/chop/nn/quantized/modules/qwen3_moe/router_precision.py new file mode 100644 index 000000000..ac609a9f7 --- /dev/null +++ b/src/chop/nn/quantized/modules/qwen3_moe/router_precision.py @@ -0,0 +1,325 @@ +"""Phase-aware MX router ablations for fused Qwen3-MoE blocks. + +The canonical Qwen3-MoE path keeps routing in BF16. These modules form a +separate numerical-ablation lane: only the router matrix operands are +quantized during decode, while softmax, top-k selection, and probability +renormalization retain the Hugging Face FP32 semantics. +""" + +from __future__ import annotations + +from copy import deepcopy +from typing import Any + +import torch +import torch.nn.functional as F +from torch import Tensor, nn +from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeTopKRouter + +from chop.nn.quantized.functional.matrix import plena_matrix_product +from chop.nn.quantized.modules.phase_config import ( + normalize_phase_q_config, + resolve_module_phase_config, +) +from chop.nn.quantized.modules.phase_context import get_runtime_phase +from chop.nn.quantizers import mxfp_quantizer, mxint_quantizer + + +ROUTER_PRECISION_SCHEMA = "mase-qwen3-moe-router-mx/v1" +ROUTER_MX_FORMATS = ("MXINT8", "E4M3", "E5M2") +ROUTER_BLOCK_SIZE = 8 +ROUTER_OUTPUT_FORMAT = "BF16" + + +def _canonical_format(value: Any, *, label: str) -> str: + token = str(value).upper() + if token not in ROUTER_MX_FORMATS: + raise ValueError( + f"{label} must be one of {ROUTER_MX_FORMATS}; observed {value!r}" + ) + return token + + +def _format_bits(token: str) -> tuple[str, int, int | None]: + if token == "MXINT8": + return "mxint", 8, None + exponent, fraction = token.removeprefix("E").split("M", maxsplit=1) + return "mxfp", int(exponent), int(fraction) + + +def _quantize(value: Tensor, token: str, *, block_dim: int) -> Tensor: + family, first, second = _format_bits(token) + if family == "mxint": + return mxint_quantizer( + value, + block_size=ROUTER_BLOCK_SIZE, + element_bits=first, + block_dim=block_dim, + ) + assert second is not None + return mxfp_quantizer( + value, + block_size=ROUTER_BLOCK_SIZE, + element_exp_bits=first, + element_frac_bits=second, + block_dim=block_dim, + ) + + +def router_decode_config( + weight_format: str, + activation_format: str, + *, + matrix_mlen: int, +) -> dict[str, Any]: + """Return the sealed decode bucket for one router-only ablation.""" + + if ( + isinstance(matrix_mlen, bool) + or not isinstance(matrix_mlen, int) + or matrix_mlen <= 0 + or matrix_mlen % ROUTER_BLOCK_SIZE + ): + raise ValueError("matrix_mlen must be a positive multiple of 8") + return { + "router_weight_format": _canonical_format( + weight_format, label="router_weight_format" + ), + "router_activation_format": _canonical_format( + activation_format, label="router_activation_format" + ), + "router_block_size": ROUTER_BLOCK_SIZE, + "matrix_mlen": matrix_mlen, + "output_format": ROUTER_OUTPUT_FORMAT, + "matrix_partial_format": ROUTER_OUTPUT_FORMAT, + "matrix_partial_to_accumulator": "truncate_to_signed_fixed16_16", + "matrix_cross_partition_accumulation": "signed_fixed16_16_wrap", + "matrix_final_writeout": ROUTER_OUTPUT_FORMAT, + "router_logits_container": "BF16", + "router_probability_dtype": "FP32", + "router_selection": "exact_torch_topk_sorted", + "router_topk_renormalization_dtype": "FP32", + } + + +def router_phase_config( + weight_format: str, + activation_format: str, + *, + matrix_mlen: int, +) -> dict[str, Any]: + """Build the disaggregated prefill-BF16/decode-MX router config.""" + + return { + "decode_policy": "quantized", + "prefill": {"bypass": True, "router_dtype": "BF16"}, + "decode": router_decode_config( + weight_format, + activation_format, + matrix_mlen=matrix_mlen, + ), + } + + +class Qwen3MoeTopKRouterMX(Qwen3MoeTopKRouter): + """Decode-only MX router matrix followed by the exact FP32 route path.""" + + fixed_weight_format: str | None = None + fixed_activation_format: str | None = None + + def _init_phase_state(self, q_config: dict[str, Any] | None) -> None: + self.q_config = deepcopy(q_config or {}) + self.phase_q_config = normalize_phase_q_config(self.q_config) + self.decode_policy = self.phase_q_config["decode_policy"] + self.prefill_config = resolve_module_phase_config( + self.phase_q_config, "prefill" + ) + self.decode_config = resolve_module_phase_config( + self.phase_q_config, "decode" + ) + if not self.prefill_config.get("bypass", False): + raise ValueError("router MX ablation requires a BF16-bypassed prefill") + if self.decode_policy != "quantized" or self.decode_config.get( + "bypass", False + ): + raise ValueError("router MX ablation requires quantized decode") + + weight_token = self.fixed_weight_format or self.decode_config.get( + "router_weight_format" + ) + activation_token = self.fixed_activation_format or self.decode_config.get( + "router_activation_format" + ) + self.router_weight_format = _canonical_format( + weight_token, label="router_weight_format" + ) + self.router_activation_format = _canonical_format( + activation_token, label="router_activation_format" + ) + if self.fixed_weight_format is not None and self.decode_config.get( + "router_weight_format", self.fixed_weight_format + ) != self.fixed_weight_format: + raise ValueError("router variant weight format conflicts with its config") + if self.fixed_activation_format is not None and self.decode_config.get( + "router_activation_format", self.fixed_activation_format + ) != self.fixed_activation_format: + raise ValueError("router variant activation format conflicts with its config") + + if self.decode_config.get("router_block_size") != ROUTER_BLOCK_SIZE: + raise ValueError("router_block_size must equal the PLENA block size 8") + if self.decode_config.get("output_format") != ROUTER_OUTPUT_FORMAT: + raise ValueError("router matrix output format must be BF16") + matrix_mlen = self.decode_config.get("matrix_mlen") + if ( + isinstance(matrix_mlen, bool) + or not isinstance(matrix_mlen, int) + or matrix_mlen <= 0 + or matrix_mlen % ROUTER_BLOCK_SIZE + ): + raise ValueError("router matrix_mlen must be a positive multiple of 8") + for key, required in ( + ("matrix_partial_format", "BF16"), + ( + "matrix_partial_to_accumulator", + "truncate_to_signed_fixed16_16", + ), + ( + "matrix_cross_partition_accumulation", + "signed_fixed16_16_wrap", + ), + ("matrix_final_writeout", "BF16"), + ("router_logits_container", "BF16"), + ("router_probability_dtype", "FP32"), + ("router_selection", "exact_torch_topk_sorted"), + ("router_topk_renormalization_dtype", "FP32"), + ): + if self.decode_config.get(key) != required: + raise ValueError(f"router decode config requires {key}={required}") + self.matrix_mlen = matrix_mlen + self.bypass = False + + @torch.no_grad() + def _build_decode_weight_bank(self) -> None: + self._decode_weight_q = _quantize( + self.weight.detach(), self.router_weight_format, block_dim=1 + ) + + @property + def router_precision_contract(self) -> dict[str, Any]: + return { + "schema_version": ROUTER_PRECISION_SCHEMA, + "scope": "router_only_decode_ablation", + "prefill_router": "BF16", + "decode_router_weight_format": self.router_weight_format, + "decode_router_activation_format": self.router_activation_format, + "block_size": ROUTER_BLOCK_SIZE, + "matrix_mlen": self.matrix_mlen, + "matrix_arithmetic_chain": [ + "per_mlen_fp32_matmul", + "bf16_partial_rounding", + "truncate_partial_to_signed_fixed16_16", + "signed_fixed16_16_wrap_across_partitions", + "final_bf16_writeout", + ], + "router_logits_container": "BF16", + "softmax_dtype": "FP32", + "topk": "torch.topk_sorted_exact", + "renormalization_dtype": "FP32", + "publication_rankable": False, + "selection_eligible": False, + "hardware_evidence_present": False, + "compiler_evidence_present": False, + } + + def _route(self, logits_bf16: Tensor): + router_probs = F.softmax(logits_bf16, dtype=torch.float32, dim=-1) + router_top_value, router_indices = torch.topk( + router_probs, self.top_k, dim=-1, sorted=True + ) + if self.norm_topk_prob: + router_top_value = router_top_value / router_top_value.sum( + dim=-1, keepdim=True + ) + return router_probs, router_top_value, router_indices + + def forward(self, hidden_states): + hidden_states = hidden_states.reshape(-1, self.hidden_dim) + if get_runtime_phase() != "decode": + logits_bf16 = F.linear( + hidden_states.to(torch.bfloat16), self.weight + ).to(torch.bfloat16) + return self._route(logits_bf16) + + activation_q = _quantize( + hidden_states.to(torch.bfloat16), + self.router_activation_format, + block_dim=-1, + ) + # The matrix oracle rounds every MLEN partial to BF16 before the + # fixed16.16 bank, wraps cross-partition accumulation, and performs a + # final BF16 writeout. Probability and selection then execute in FP32. + logits_bf16 = plena_matrix_product( + activation_q, + self._decode_weight_q.transpose(-1, -2), + { + "matrix_mlen": self.matrix_mlen, + "output_format": ROUTER_OUTPUT_FORMAT, + }, + ).to(torch.bfloat16) + return self._route(logits_bf16) + + @classmethod + def from_self( + cls, + router: Qwen3MoeTopKRouter, + q_config: dict[str, Any] | None = None, + ): + new = cls.__new__(cls) + nn.Module.__init__(new) + new.top_k = router.top_k + new.num_experts = router.num_experts + new.norm_topk_prob = router.norm_topk_prob + new.hidden_dim = router.hidden_dim + new.weight = nn.Parameter( + router.weight.detach().to(torch.bfloat16), + requires_grad=router.weight.requires_grad, + ) + new._init_phase_state(q_config) + new.register_buffer("_decode_weight_q", torch.empty(0), persistent=False) + new._build_decode_weight_bank() + return new + + +class Qwen3MoeTopKRouterMXInt8(Qwen3MoeTopKRouterMX): + """Uniform MXINT8 router operand ablation.""" + + fixed_weight_format = "MXINT8" + fixed_activation_format = "MXINT8" + + +class Qwen3MoeTopKRouterE4M3(Qwen3MoeTopKRouterMX): + """Uniform MXFP E4M3 router operand ablation.""" + + fixed_weight_format = "E4M3" + fixed_activation_format = "E4M3" + + +class Qwen3MoeTopKRouterE5M2(Qwen3MoeTopKRouterMX): + """Uniform MXFP E5M2 router operand ablation.""" + + fixed_weight_format = "E5M2" + fixed_activation_format = "E5M2" + + +__all__ = [ + "Qwen3MoeTopKRouterE4M3", + "Qwen3MoeTopKRouterE5M2", + "Qwen3MoeTopKRouterMX", + "Qwen3MoeTopKRouterMXInt8", + "ROUTER_BLOCK_SIZE", + "ROUTER_MX_FORMATS", + "ROUTER_OUTPUT_FORMAT", + "ROUTER_PRECISION_SCHEMA", + "router_decode_config", + "router_phase_config", +] From 3fb8809681825f260be2b243fe000ff57408ab9f Mon Sep 17 00:00:00 2001 From: Shreyas8612 Date: Wed, 16 Sep 2026 19:13:31 +0100 Subject: [PATCH 2/4] Wire Qwen3-MoE modules into replacement, phase hooks and rotation Register Qwen3MoeDecoderLayer, Experts, RMSNorm, SparseMoeBlock and TopKRouter in the prefix and from_self maps so the modify helper can resolve them, and carry the fused expert decode banks across wholesale replacement via adopt_decode_{gptq,fp}_expert_weights. The decoder layer joins the runtime phase pre-hook list. Rotation search gains the MoE rotate attention classes and excludes up, gate and down projections on fused experts, which have no rotation lowerer; the scope is recorded in the results and cached decisions are rejected on mismatch. --- .../passes/module/module_modify_helper.py | 48 +++++++++++- .../module/transforms/quantize/quantize.py | 7 +- .../transforms/quantize/rotation_search.py | 74 +++++++++++++++++-- 3 files changed, 122 insertions(+), 7 deletions(-) diff --git a/src/chop/passes/module/module_modify_helper.py b/src/chop/passes/module/module_modify_helper.py index 580b70afa..92aa3642c 100644 --- a/src/chop/passes/module/module_modify_helper.py +++ b/src/chop/passes/module/module_modify_helper.py @@ -42,7 +42,12 @@ from transformers.models.qwen3_moe.modeling_qwen3_moe import ( Qwen3MoeAttention, + Qwen3MoeDecoderLayer, + Qwen3MoeExperts, Qwen3MoeMLP, + Qwen3MoeRMSNorm, + Qwen3MoeSparseMoeBlock, + Qwen3MoeTopKRouter, ) from transformers.models.glm4_moe.modeling_glm4_moe import ( @@ -56,8 +61,12 @@ LLaDALlamaBlock, ) from chop.nn.quantized.modules.phase_config import ( + DECODE_FP_EXPERT_DOWN_ATTR, + DECODE_FP_EXPERT_GATE_UP_ATTR, DECODE_FP_BIAS_ATTR, DECODE_FP_WEIGHT_ATTR, + GPTQ_DECODE_EXPERT_DOWN_ATTR, + GPTQ_DECODE_EXPERT_GATE_UP_ATTR, GPTQ_DECODE_WEIGHT_ATTR, ) @@ -94,8 +103,13 @@ } qwen3_moe_prefix_map = { + Qwen3MoeDecoderLayer: "qwen3_moe_decoder_layer", Qwen3MoeAttention: "qwen3_moe_self_attention", + Qwen3MoeExperts: "qwen3_moe_experts", Qwen3MoeMLP: "qwen3_moe_mlp", + Qwen3MoeRMSNorm: "qwen3_moe_rms_norm", + Qwen3MoeSparseMoeBlock: "qwen3_moe_sparse_block", + Qwen3MoeTopKRouter: "qwen3_moe_router", } glm4_moe_prefix_map = { @@ -124,6 +138,12 @@ Qwen3RMSNorm: "qwen3_rms_norm", LlamaDecoderLayer: "llama_decoder_layer", LlamaAttention: "llama_self_attention", + Qwen3MoeAttention: "qwen3_moe_self_attention", + Qwen3MoeDecoderLayer: "qwen3_moe_decoder_layer", + Qwen3MoeExperts: "qwen3_moe_experts", + Qwen3MoeRMSNorm: "qwen3_moe_rms_norm", + Qwen3MoeSparseMoeBlock: "qwen3_moe_sparse_block", + Qwen3MoeTopKRouter: "qwen3_moe_router", } @@ -199,6 +219,24 @@ def _transfer_phase_weight_banks(source_module, target_module): fp_weight, getattr(source_module, DECODE_FP_BIAS_ATTR, None) ) + gate_up_gptq = getattr(source_module, GPTQ_DECODE_EXPERT_GATE_UP_ATTR, None) + down_gptq = getattr(source_module, GPTQ_DECODE_EXPERT_DOWN_ATTR, None) + if gate_up_gptq is not None or down_gptq is not None: + if gate_up_gptq is None or down_gptq is None: + raise RuntimeError("incomplete Qwen3-MoE GPTQ decode bank") + adopt = getattr(target_module, "adopt_decode_gptq_expert_weights", None) + if callable(adopt): + adopt(gate_up_gptq, down_gptq) + + gate_up_fp = getattr(source_module, DECODE_FP_EXPERT_GATE_UP_ATTR, None) + down_fp = getattr(source_module, DECODE_FP_EXPERT_DOWN_ATTR, None) + if gate_up_fp is not None or down_fp is not None: + if gate_up_fp is None or down_fp is None: + raise RuntimeError("incomplete Qwen3-MoE FP decode snapshot") + adopt = getattr(target_module, "adopt_decode_fp_expert_weights", None) + if callable(adopt): + adopt(gate_up_fp, down_fp) + def _restash_child_weight_banks(source_module, target_module): """Carry per-child phase-bank stashes across a WHOLESALE module replacement. @@ -211,7 +249,15 @@ def _restash_child_weight_banks(source_module, target_module): so ``_transfer_phase_weight_banks`` finds them at that replacement. """ - stash_attrs = (GPTQ_DECODE_WEIGHT_ATTR, DECODE_FP_WEIGHT_ATTR, DECODE_FP_BIAS_ATTR) + stash_attrs = ( + GPTQ_DECODE_WEIGHT_ATTR, + DECODE_FP_WEIGHT_ATTR, + DECODE_FP_BIAS_ATTR, + GPTQ_DECODE_EXPERT_GATE_UP_ATTR, + GPTQ_DECODE_EXPERT_DOWN_ATTR, + DECODE_FP_EXPERT_GATE_UP_ATTR, + DECODE_FP_EXPERT_DOWN_ATTR, + ) target_children = dict(target_module.named_modules()) for name, src_child in source_module.named_modules(): if not name or name not in target_children: diff --git a/src/chop/passes/module/transforms/quantize/quantize.py b/src/chop/passes/module/transforms/quantize/quantize.py index f0ab43c66..cf04504cb 100644 --- a/src/chop/passes/module/transforms/quantize/quantize.py +++ b/src/chop/passes/module/transforms/quantize/quantize.py @@ -16,6 +16,7 @@ import torch from transformers.models.llama.modeling_llama import LlamaDecoderLayer from transformers.models.qwen3.modeling_qwen3 import Qwen3DecoderLayer +from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeDecoderLayer from chop.nn.quantized.modules import quantized_module_map from chop.nn.quantized.modules.phase_context import ( @@ -29,7 +30,11 @@ # Decoder-layer classes that get runtime phase pre-hooks. Extending phase # support to a new architecture only requires phase-aware modules plus an # entry here. -_PHASE_HOOKED_DECODER_LAYERS = (LlamaDecoderLayer, Qwen3DecoderLayer) +_PHASE_HOOKED_DECODER_LAYERS = ( + LlamaDecoderLayer, + Qwen3DecoderLayer, + Qwen3MoeDecoderLayer, +) def get_config(config: dict, name: str): diff --git a/src/chop/passes/module/transforms/quantize/rotation_search.py b/src/chop/passes/module/transforms/quantize/rotation_search.py index 700284fbf..4149c892f 100644 --- a/src/chop/passes/module/transforms/quantize/rotation_search.py +++ b/src/chop/passes/module/transforms/quantize/rotation_search.py @@ -9,7 +9,7 @@ class but with all per-stage rotate flags set to False — this gives the same numerics as plain mx{int,fp} while making per-trial flag-flipping cheap. 2. Computes baseline perplexity on a calibration loader. -3. For each of the 9 matmul types per decoder layer +3. For each supported matmul type per decoder layer (q_proj / k_proj / v_proj / o_proj / qk_matmul / av_matmul / up_proj / gate_proj / down_proj), toggles online Hadamard rotation ON for that type uniformly across all @@ -27,8 +27,10 @@ class but with all per-stage rotate flags set to False — this gives the attrs on the rotate attention class (Qwen3AttentionMXIntRotate / LlamaAttentionMXIntRotate / LlamaAttentionMXFPRotate). -Coverage today: MXInt rotate exists for Qwen3 + Llama; MXFP rotate is -Llama-only. Adding a new arch only requires (a) implementing +Coverage today includes MXInt and MXFP attention rotation for fused Qwen3-MoE. +Its fused expert tensors are deliberately outside rotation search; only the +four attention projections and three attention stages are eligible. Adding a +new architecture requires (a) implementing ``AttentionMX{Int,FP}Rotate`` with the per-stage flag attrs and (b) appending it to ``_ROTATE_ATTENTION_CLASSES`` below. Linear-side wiring is shared and needs no changes. @@ -56,6 +58,10 @@ class but with all per-stage rotate flags set to False — this gives the ) from chop.nn.quantized.modules.phase_context import force_runtime_phase from chop.nn.quantized.modules.qwen3.attention import Qwen3AttentionMXIntRotate +from chop.nn.quantized.modules.qwen3_moe.attention import ( + Qwen3MoeAttentionMXFPRotate, + Qwen3MoeAttentionMXIntRotate, +) from chop.nn.quantized.modules.llama.attention import ( LlamaAttentionMXIntRotate, LlamaAttentionMXFPRotate, @@ -79,6 +85,8 @@ class but with all per-stage rotate flags set to False — this gives the # All rotate attention classes the search can flip per-stage flags on. _ROTATE_ATTENTION_CLASSES = ( Qwen3AttentionMXIntRotate, + Qwen3MoeAttentionMXFPRotate, + Qwen3MoeAttentionMXIntRotate, LlamaAttentionMXIntRotate, LlamaAttentionMXFPRotate, ) @@ -102,6 +110,9 @@ class but with all per-stage rotate flags set to False — this gives the ) ATTENTION_MATMUL_TYPES = ("qk_matmul", "av_matmul", "kv_cache") ALL_MATMUL_TYPES = LINEAR_MATMUL_TYPES + ATTENTION_MATMUL_TYPES +_FUSED_EXPERT_MATMUL_TYPES = frozenset( + {"up_proj", "gate_proj", "down_proj"} +) # Map attention stage name -> instance attr on the rotate attention classes. _ATTN_STAGE_TO_FLAG = { @@ -111,6 +122,45 @@ class but with all per-stage rotate flags set to False — this gives the } +def _has_fused_qwen3_moe_experts(network: torch.nn.Module) -> bool: + return any( + hasattr(module, "gate_up_proj") + and hasattr(module, "down_proj") + and type(module).__name__.startswith("Qwen3MoeExperts") + for module in network.modules() + ) + + +def _resolve_rotation_scope( + network: torch.nn.Module, + requested: tuple[str, ...], + *, + explicitly_requested: bool, +) -> tuple[tuple[str, ...], dict]: + if not _has_fused_qwen3_moe_experts(network): + return requested, { + "architecture": "generic", + "excluded_matmul_types": [], + } + unsupported = tuple( + matmul for matmul in requested if matmul in _FUSED_EXPERT_MATMUL_TYPES + ) + if unsupported and explicitly_requested: + raise ValueError( + "selective rotation does not support fused Qwen3-MoE expert " + f"tensors: {unsupported}; search attention projections/stages only" + ) + eligible = tuple( + matmul for matmul in requested if matmul not in _FUSED_EXPERT_MATMUL_TYPES + ) + return eligible, { + "architecture": "qwen3_moe_fused", + "eligible_matmul_types": list(eligible), + "excluded_matmul_types": sorted(_FUSED_EXPERT_MATMUL_TYPES), + "excluded_reason": "fused expert tensors have no rotation lowerer", + } + + def _patch_base_args_for_rotate_class(base_args: dict) -> dict: """Force every attention selector to use the rotate class with all stage flags OFF — gives baseline (non-rotated) numerics while leaving the @@ -403,7 +453,8 @@ def rotation_search_transform_pass(network, pass_args): device (str, default "cuda:0"): Where the perplexity forwards run. The model is moved here. matmul_types (Iterable[str], default ALL_MATMUL_TYPES): - Subset of the 9 matmul types to search over. + Subset of the 10 generic matmul types to search over. Fused + Qwen3-MoE automatically narrows the default to 7 attention types. output_json (str | None): If given, write a JSON summary (winners, per-round history, final ppl) to this path. Also acts as the cache file when @@ -435,6 +486,7 @@ def rotation_search_transform_pass(network, pass_args): base_args = pass_args["base_quantize_args"] calib_loader = pass_args["calib_loader"] device = pass_args.get("device", "cuda:0") + explicit_matmul_types = "matmul_types" in pass_args matmul_types = tuple(pass_args.get("matmul_types", ALL_MATMUL_TYPES)) output_json = pass_args.get("output_json", None) improvement_eps = float(pass_args.get("improvement_eps", 0.0)) @@ -453,12 +505,22 @@ def rotation_search_transform_pass(network, pass_args): raise ValueError( f"Unknown matmul_type {t!r}; valid types: {ALL_MATMUL_TYPES}" ) + matmul_types, rotation_scope = _resolve_rotation_scope( + network, + matmul_types, + explicitly_requested=explicit_matmul_types, + ) # Cache check: if a saved decisions file exists, load and apply — skip # the whole search. Same spirit as GPTQ's auto_load_quantized_layers. cached = None if cache_winners and output_json and Path(output_json).exists(): cached = json.loads(Path(output_json).read_text()) + if cached.get("rotation_scope") != rotation_scope: + raise ValueError( + "cached rotation scope does not match this model; invalidate " + f"{output_json} and rerun calibration" + ) logger.info("=" * 64) logger.info("CACHE HIT — loading rotation decisions from %s", output_json) logger.info( @@ -514,6 +576,7 @@ def rotation_search_transform_pass(network, pass_args): results = dict(cached) results["from_cache"] = True results["matmul_types_searched"] = list(matmul_types) + results["rotation_scope"] = rotation_scope logger.info("-----Rotation Search (cached) Done-----") return network, results @@ -533,6 +596,7 @@ def rotation_search_transform_pass(network, pass_args): results["baseline_ppl"] = baseline_ppl results["improvement_eps"] = improvement_eps results["matmul_types_searched"] = list(matmul_types) + results["rotation_scope"] = rotation_scope if output_json: out_path = Path(output_json) @@ -566,7 +630,7 @@ def dispatch_rotation_search_block(network, base_pass_args: dict, rot_cfg: dict) calib_nsamples (default 32), calib_seqlen (default 1024), improvement_eps (default 0.0), - matmul_types (default = all 9), + matmul_types (default = all 10 generic types; 7 for fused Qwen3-MoE), cache_path (default = /rotation_decisions.json), cache_winners (default True — set False to force a re-search). From ec323470cd9526f502a3d36ea039d455e36726b1 Mon Sep 17 00:00:00 2001 From: Shreyas8612 Date: Wed, 16 Sep 2026 19:13:31 +0100 Subject: [PATCH 3/4] Extend GPTQ to fused Qwen3-MoE experts and sharded device maps Add per-expert GPTQ over the fused gate_up and down banks: a forward hook on the experts module collects each expert's routed token slice, quantises slice by slice and falls back to RTN below min_expert_calibration_hits, recording hits and coverage on the model. The FP snapshot and decode/prefill hand-off now cover the fused Parameters as well as Linears. A device_map_aware mode keeps layers where accelerate placed them with CPU-side activation buffers, the calibration catcher uses a private exception instead of ValueError, and clip-search activations are only retained when clip_search_y is set. --- src/chop/passes/module/transforms/gptq/run.py | 508 +++++++++++++++--- 1 file changed, 439 insertions(+), 69 deletions(-) diff --git a/src/chop/passes/module/transforms/gptq/run.py b/src/chop/passes/module/transforms/gptq/run.py index 688b5789b..91a8dfd99 100644 --- a/src/chop/passes/module/transforms/gptq/run.py +++ b/src/chop/passes/module/transforms/gptq/run.py @@ -10,22 +10,102 @@ import torch import torch.nn as nn +import torch.nn.functional as F from .gptq import GPTQ +from .quantize_dispatch import quantize_tensor from .utils import find_qlayers, cleanup_memory from .data_utils import get_loaders from .checkpoint import save_layer_checkpoint, auto_load_quantized_layers from chop.nn.quantized.modules.phase_config import ( + DECODE_FP_EXPERT_DOWN_ATTR, + DECODE_FP_EXPERT_GATE_UP_ATTR, DECODE_FP_BIAS_ATTR, DECODE_FP_WEIGHT_ATTR, + GPTQ_DECODE_EXPERT_DOWN_ATTR, + GPTQ_DECODE_EXPERT_GATE_UP_ATTR, GPTQ_DECODE_WEIGHT_ATTR, ) _VALID_GPTQ_PHASES = ("both", "decode", "prefill") +_DENSE_SEQUENTIAL = ( + ("self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj"), + ("self_attn.o_proj",), + ("mlp.up_proj", "mlp.gate_proj"), + ("mlp.down_proj",), +) + + +class _WeightView: + """Minimal GPTQ adapter for one expert's fused weight slice.""" + + def __init__(self, weight: torch.Tensor): + self.weight = weight + + +class _CalibrationCaptureComplete(RuntimeError): + pass + + +def _module_device(module: nn.Module, module_name: str) -> torch.device: + devices = {param.device for param in module.parameters(recurse=True)} + devices.update(buffer.device for buffer in module.buffers(recurse=True)) + devices = {device for device in devices if device.type != "meta"} + if not devices: + return torch.device("cpu") + if len(devices) != 1: + formatted = ", ".join(sorted(str(device) for device in devices)) + raise RuntimeError( + "GPTQ device_map_aware requires layer-level sharding, but " + f"{module_name} spans multiple devices: {formatted}" + ) + return next(iter(devices)) + + +def _move_to_device(value, device: torch.device): + if value is None: + return None + if torch.is_tensor(value): + return value.to(device) + if isinstance(value, tuple): + return tuple(_move_to_device(item, device) for item in value) + if isinstance(value, list): + return [_move_to_device(item, device) for item in value] + if isinstance(value, dict): + return {key: _move_to_device(item, device) for key, item in value.items()} + return value + + +def _run_layer_sample( + *, + layer, + sample, + outs, + sample_idx, + rope, + attention_mask, + position_ids, + layer_device, +): + x = sample.unsqueeze(0).to(layer_device) + local_mask = _move_to_device(attention_mask, layer_device) + local_positions = _move_to_device(position_ids, layer_device) + local_rope = rope.to(layer_device) + cos, sin = local_rope(x, local_positions) + output = layer( + x, + attention_mask=local_mask, + position_embeddings=(cos, sin), + ) + if isinstance(output, (tuple, list)): + output = output[0] + outs[sample_idx].copy_( + output.squeeze(0).to(device=outs.device, dtype=outs.dtype) + ) def _snapshot_fp_linear_weights(layers) -> None: - """Snapshot FP weights of every decoder linear before GPTQ mutates them. + """Snapshot FP decoder weights before GPTQ mutates them. Snapshots live on CPU (``_mase_fp_weight``) so large models don't hold a second GPU copy. Must run before checkpoint resume, which loads already @@ -38,6 +118,20 @@ def _snapshot_fp_linear_weights(layers) -> None: module, "_mase_fp_weight" ): module._mase_fp_weight = module.weight.detach().clone().cpu() + if ( + hasattr(module, "gate_up_proj") + and isinstance(module.gate_up_proj, nn.Parameter) + and module.gate_up_proj.ndim == 3 + and hasattr(module, "down_proj") + and isinstance(module.down_proj, nn.Parameter) + ): + if not hasattr(module, "_mase_fp_gate_up_proj"): + module._mase_fp_gate_up_proj = ( + module.gate_up_proj.detach().clone().cpu() + ) + module._mase_fp_down_proj = ( + module.down_proj.detach().clone().cpu() + ) def _finalize_gptq_phase(network, phase: str) -> None: @@ -57,33 +151,251 @@ def _finalize_gptq_phase(network, phase: str) -> None: return for layer in network.model.layers: for module in layer.modules(): - if not isinstance(module, nn.Linear): + if isinstance(module, nn.Linear): + fp_weight = getattr(module, "_mase_fp_weight", None) + if fp_weight is None: + continue + if phase == "decode": + setattr( + module, + GPTQ_DECODE_WEIGHT_ATTR, + module.weight.detach().clone().cpu(), + ) + module.weight.data.copy_( + fp_weight.to( + device=module.weight.device, dtype=module.weight.dtype + ) + ) + else: # prefill + setattr(module, DECODE_FP_WEIGHT_ATTR, fp_weight) + if module.bias is not None and not hasattr( + module, DECODE_FP_BIAS_ATTR + ): + setattr( + module, + DECODE_FP_BIAS_ATTR, + module.bias.detach().clone().cpu(), + ) + del module._mase_fp_weight continue - fp_weight = getattr(module, "_mase_fp_weight", None) - if fp_weight is None: + + fp_gate_up = getattr(module, "_mase_fp_gate_up_proj", None) + fp_down = getattr(module, "_mase_fp_down_proj", None) + if fp_gate_up is None and fp_down is None: continue + if fp_gate_up is None or fp_down is None: + raise RuntimeError("incomplete Qwen3-MoE FP expert snapshot") if phase == "decode": setattr( module, - GPTQ_DECODE_WEIGHT_ATTR, - module.weight.detach().clone().cpu(), + GPTQ_DECODE_EXPERT_GATE_UP_ATTR, + module.gate_up_proj.detach().clone().cpu(), ) - module.weight.data.copy_( - fp_weight.to( - device=module.weight.device, dtype=module.weight.dtype + setattr( + module, + GPTQ_DECODE_EXPERT_DOWN_ATTR, + module.down_proj.detach().clone().cpu(), + ) + module.gate_up_proj.data.copy_( + fp_gate_up.to( + device=module.gate_up_proj.device, + dtype=module.gate_up_proj.dtype, ) ) - else: # prefill - setattr(module, DECODE_FP_WEIGHT_ATTR, fp_weight) - if module.bias is not None and not hasattr( - module, DECODE_FP_BIAS_ATTR - ): - setattr( - module, - DECODE_FP_BIAS_ATTR, - module.bias.detach().clone().cpu(), + module.down_proj.data.copy_( + fp_down.to( + device=module.down_proj.device, + dtype=module.down_proj.dtype, ) - del module._mase_fp_weight + ) + else: + setattr(module, DECODE_FP_EXPERT_GATE_UP_ATTR, fp_gate_up) + setattr(module, DECODE_FP_EXPERT_DOWN_ATTR, fp_down) + del module._mase_fp_gate_up_proj + del module._mase_fp_down_proj + + +def _get_qwen3_moe_experts(layer): + experts = getattr(getattr(layer, "mlp", None), "experts", None) + if experts is None: + return None + gate_up = getattr(experts, "gate_up_proj", None) + down = getattr(experts, "down_proj", None) + if not isinstance(gate_up, nn.Parameter) or not isinstance(down, nn.Parameter): + return None + if gate_up.ndim != 3 or down.ndim != 3: + return None + return experts + + +def _collect_expert_inputs( + *, + layer, + experts, + inps, + outs, + nsamples, + rope, + attention_mask, + position_ids, + layer_device, + target, +): + if target not in {"gate_up", "down"}: + raise ValueError(f"unknown expert GPTQ target {target!r}") + chunks = [[] for _ in range(experts.num_experts)] + hits = [0 for _ in range(experts.num_experts)] + + def hook(_, inputs, _output): + hidden_states, top_k_index, _top_k_weights = inputs + with torch.no_grad(): + expert_mask = F.one_hot( + top_k_index, num_classes=experts.num_experts + ).permute(2, 1, 0) + for hit in torch.greater( + expert_mask.sum(dim=(-1, -2)), 0 + ).nonzero(): + expert_idx = int(hit[0].item()) + _, token_idx = torch.where(expert_mask[expert_idx]) + current = hidden_states[token_idx] + if current.numel() == 0: + continue + if target == "down": + gate, up = F.linear( + current, experts.gate_up_proj[expert_idx] + ).chunk(2, dim=-1) + current = experts.act_fn(gate) * up + chunks[expert_idx].append(current.detach().cpu()) + hits[expert_idx] += int(current.shape[0]) + + handle = experts.register_forward_hook(hook) + try: + for sample_idx in range(nsamples): + _run_layer_sample( + layer=layer, + sample=inps[sample_idx], + outs=outs, + sample_idx=sample_idx, + rope=rope, + attention_mask=attention_mask, + position_ids=position_ids, + layer_device=layer_device, + ) + finally: + handle.remove() + return chunks, hits + + +def _quantize_expert_slice( + *, + weight, + chunks, + fmt, + weight_config, + quantile_search, + clip_search_y, + cali_batch_size, + layer_name, +): + if not chunks: + return quantize_tensor( + weight.data, + block_dim=1, + fmt=fmt, + config=weight_config, + quantile_search=quantile_search, + ) + gptq = GPTQ(_WeightView(weight)) + for chunk in chunks: + gptq.add_batch(chunk.to(weight.device).unsqueeze(0).data, None) + activation = None + if clip_search_y: + activation = torch.cat(chunks, dim=0).unsqueeze(0).to(weight.device) + quantized = gptq.fasterquant( + activation=activation, + fmt=fmt, + weight_config=weight_config, + percdamp=0.01, + cali_batch_size=cali_batch_size, + layer_name=layer_name, + quant_search=quantile_search, + ) + if quantized.shape != weight.shape: + raise RuntimeError( + f"GPTQ returned {tuple(quantized.shape)} for {tuple(weight.shape)}" + ) + gptq.free() + return quantized + + +def _quantize_qwen3_moe_experts( + *, + layer, + layer_idx, + inps, + outs, + nsamples, + rope, + attention_mask, + position_ids, + fmt, + weight_config, + quantile_search, + clip_search_y, + cali_batch_size, + layer_device, + min_expert_calibration_hits, +): + experts = _get_qwen3_moe_experts(layer) + if experts is None: + return None + + coverage = {"layer": layer_idx, "gate_up": [], "down": []} + for target, weight_name in ( + ("gate_up", "gate_up_proj"), + ("down", "down_proj"), + ): + chunks, hits = _collect_expert_inputs( + layer=layer, + experts=experts, + inps=inps, + outs=outs, + nsamples=nsamples, + rope=rope, + attention_mask=attention_mask, + position_ids=position_ids, + layer_device=layer_device, + target=target, + ) + weight_bank = getattr(experts, weight_name) + for expert_idx in range(experts.num_experts): + n_hits = hits[expert_idx] + used_gptq = n_hits >= min_expert_calibration_hits + if not used_gptq: + logging.warning( + "layer %d expert %d %s has %d calibration hits; using RTN", + layer_idx, + expert_idx, + weight_name, + n_hits, + ) + quantized = _quantize_expert_slice( + weight=weight_bank[expert_idx], + chunks=chunks[expert_idx] if used_gptq else [], + fmt=fmt, + weight_config=weight_config, + quantile_search=quantile_search, + clip_search_y=clip_search_y, + cali_batch_size=cali_batch_size, + layer_name=( + f"layers{layer_idx}.mlp.experts.{expert_idx}.{weight_name}" + ), + ) + weight_bank.data[expert_idx].copy_(quantized) + coverage[target].append( + {"expert": expert_idx, "hits": n_hits, "gptq": used_gptq} + ) + return coverage @torch.no_grad() @@ -138,6 +450,12 @@ def run_gptq(network, gptq_config): checkpoint_dir = gptq_config.get("checkpoint_dir", None) hf_token = gptq_config.get("hf_token", None) max_layers = gptq_config.get("max_layers", None) + device_map_aware = bool(gptq_config.get("device_map_aware", False)) + min_expert_calibration_hits = int( + gptq_config.get("min_expert_calibration_hits", 1) + ) + if min_expert_calibration_hits < 1: + raise ValueError("min_expert_calibration_hits must be positive") if phase != "both": # FP snapshots must precede checkpoint resume: resume overwrites @@ -173,16 +491,26 @@ def run_gptq(network, gptq_config): layers = network.model.layers - # Move embedding + norm + rope to device - network.model.embed_tokens = network.model.embed_tokens.to(dev) - network.model.norm = network.model.norm.to(dev) - rope = network.model.rotary_emb.to(dev) - layers[0] = layers[0].to(dev) + if device_map_aware: + embed_device = _module_device( + network.model.embed_tokens, "model.embed_tokens" + ) + _module_device(layers[0], "model.layers.0") + rope = network.model.rotary_emb + else: + embed_device = torch.device(dev) + network.model.embed_tokens = network.model.embed_tokens.to(dev) + network.model.norm = network.model.norm.to(dev) + rope = network.model.rotary_emb.to(dev) + layers[0] = layers[0].to(dev) dtype = next(iter(network.parameters())).dtype + buffer_device = torch.device("cpu") if device_map_aware else torch.device(dev) inps = torch.zeros( - (nsamples, seqlen, network.config.hidden_size), dtype=dtype, device=dev + (nsamples, seqlen, network.config.hidden_size), + dtype=dtype, + device=buffer_device, ) cache = {"i": 0, "attention_mask": None} @@ -193,22 +521,31 @@ def __init__(self, module): def forward(self, inp, **kwargs): if cache["i"] >= nsamples: - raise ValueError - inps[cache["i"]] = inp + raise _CalibrationCaptureComplete + inps[cache["i"]].copy_( + inp.squeeze(0).detach().to(device=inps.device, dtype=inps.dtype) + ) cache["i"] += 1 - cache["attention_mask"] = kwargs["attention_mask"] - cache["position_ids"] = kwargs["position_ids"] - raise ValueError - - layers[0] = Catcher(layers[0]) - for batch in dataloader: - if cache["i"] >= nsamples: - break - try: - network(batch[0].to(dev)) - except ValueError: - pass - layers[0] = layers[0].module + cache["attention_mask"] = _move_to_device( + kwargs.get("attention_mask"), inps.device + ) + cache["position_ids"] = _move_to_device( + kwargs.get("position_ids"), inps.device + ) + raise _CalibrationCaptureComplete + + original_first_layer = layers[0] + layers[0] = Catcher(original_first_layer) + try: + for batch in dataloader: + if cache["i"] >= nsamples: + break + try: + network(batch[0].to(embed_device)) + except _CalibrationCaptureComplete: + pass + finally: + layers[0] = original_first_layer torch.cuda.empty_cache() collected = int(cache["i"]) @@ -226,13 +563,6 @@ def forward(self, inp, **kwargs): attention_mask = cache["attention_mask"] position_ids = cache["position_ids"] - sequential = [ - ["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj"], - ["self_attn.o_proj"], - ["mlp.up_proj", "mlp.gate_proj"], - ["mlp.down_proj"], - ] - end_layer = ( len(layers) if max_layers is None @@ -242,13 +572,21 @@ def forward(self, inp, **kwargs): f"GPTQ: quantizing layers {start_layer} to {end_layer - 1} (of {len(layers)} total)" ) + expert_coverage = [] for i in range(start_layer, end_layer): print(f"\nLayer {i}:", flush=True, end=" ") - layer = layers[i].to(dev) + if device_map_aware: + layer = layers[i] + layer_device = _module_device(layer, f"model.layers.{i}") + else: + layer = layers[i].to(dev) + layer_device = torch.device(dev) full = find_qlayers(layer, layers=[torch.nn.Linear]) - for names in sequential: - subset = {n: full[n] for n in names} + for names in _DENSE_SEQUENTIAL: + subset = {name: full[name] for name in names if name in full} + if not subset: + continue gptq = {} for name in subset: @@ -259,7 +597,8 @@ def forward(self, inp, **kwargs): def make_pre_hook(): def pre_hook(_, inp): - pre_act.append(inp[0]) + if clip_search_y: + pre_act.append(inp[0].detach().cpu()) return pre_hook @@ -272,25 +611,31 @@ def tmp(_, inp, out): handles = [] for name in subset: handles.append(subset[name].register_forward_hook(add_batch(name))) - handles.append(subset[name].register_forward_pre_hook(make_pre_hook())) + first_module = next(iter(subset.values())) + handles.append(first_module.register_forward_pre_hook(make_pre_hook())) for j in range(nsamples): - x = inps[j].unsqueeze(0) - cos, sin = rope(x, position_ids) - outs[j] = layer( - x, + _run_layer_sample( + layer=layer, + sample=inps[j], + outs=outs, + sample_idx=j, + rope=rope, attention_mask=attention_mask, - position_embeddings=(cos, sin), - )[0] + position_ids=position_ids, + layer_device=layer_device, + ) - pre_act = torch.cat(pre_act, dim=0) + activation = ( + torch.cat(pre_act, dim=0).to(layer_device) if pre_act else None + ) for h in handles: h.remove() for name in subset: quantized_w = gptq[name].fasterquant( - activation=pre_act if clip_search_y else None, + activation=activation, fmt=fmt, weight_config=weight_config, percdamp=0.01, @@ -304,19 +649,43 @@ def tmp(_, inp, out): gptq[name].layer.weight.data.copy_(quantized_w) gptq[name].free() + if getattr(network.config, "model_type", None) == "qwen3_moe": + coverage = _quantize_qwen3_moe_experts( + layer=layer, + layer_idx=i, + inps=inps, + outs=outs, + nsamples=nsamples, + rope=rope, + attention_mask=attention_mask, + position_ids=position_ids, + fmt=fmt, + weight_config=weight_config, + quantile_search=quantile_search, + clip_search_y=clip_search_y, + cali_batch_size=cali_batch_size, + layer_device=layer_device, + min_expert_calibration_hits=min_expert_calibration_hits, + ) + if coverage is not None: + expert_coverage.append(coverage) + # Forward pass with quantized weights to get inputs for next layer for j in range(nsamples): - x = inps[j].unsqueeze(0) - cos, sin = network.model.rotary_emb(x, position_ids) - outs[j] = layer( - x, + _run_layer_sample( + layer=layer, + sample=inps[j], + outs=outs, + sample_idx=j, + rope=rope, attention_mask=attention_mask, - position_embeddings=(cos, sin), - )[0] + position_ids=position_ids, + layer_device=layer_device, + ) - layers[i] = layer.cpu() - del layer - del gptq + if not device_map_aware: + layers[i] = layer.cpu() + del layer torch.cuda.empty_cache() inps, outs = outs, inps @@ -325,6 +694,7 @@ def tmp(_, inp, out): save_layer_checkpoint(network, i, checkpoint_dir) network.config.use_cache = use_cache + network._mase_gptq_expert_coverage = expert_coverage _finalize_gptq_phase(network, phase) cleanup_memory(verbos=True) logging.info("-----GPTQ Quantization Done-----\n") From 2b19da7a96ddcd09cc5f63b69a94502579ddad05 Mon Sep 17 00:00:00 2001 From: Shreyas8612 Date: Wed, 16 Sep 2026 19:13:31 +0100 Subject: [PATCH 4/4] Add CPU tests for Qwen3-MoE phase quantisation and the router lane Seventeen tests over a tiny Qwen3MoeConfig: the transformers ABI pin, fused expert prefill/decode banks and the GPTQ decode hand-off end to end, rotation scope exclusion of fused experts, phase-aware attention with tied K/V precision, module replacement plus decoder phase hooks and the BF16 router island. The router-precision file checks MXINT8, E4M3 and E5M2 against explicit matrix products, FP32 route agreement and rejection of unplanned formats. --- .../quantize/test_qwen3_moe_phase_quantize.py | 391 ++++++++++++++++++ .../test_qwen3_moe_router_precision.py | 238 +++++++++++ 2 files changed, 629 insertions(+) create mode 100644 test/passes/module/transforms/quantize/test_qwen3_moe_phase_quantize.py create mode 100644 test/passes/module/transforms/quantize/test_qwen3_moe_router_precision.py diff --git a/test/passes/module/transforms/quantize/test_qwen3_moe_phase_quantize.py b/test/passes/module/transforms/quantize/test_qwen3_moe_phase_quantize.py new file mode 100644 index 000000000..10249ca8b --- /dev/null +++ b/test/passes/module/transforms/quantize/test_qwen3_moe_phase_quantize.py @@ -0,0 +1,391 @@ +import torch +import pytest +import transformers +import chop.passes.module.transforms.gptq.run as gptq_run + +from transformers.models.qwen3_moe import Qwen3MoeConfig, Qwen3MoeForCausalLM +from transformers.models.qwen3_moe.modeling_qwen3_moe import ( + Qwen3MoeAttention, + Qwen3MoeExperts, + Qwen3MoeSparseMoeBlock, +) + +from chop.nn.quantized.modules.phase_config import ( + GPTQ_DECODE_EXPERT_DOWN_ATTR, + GPTQ_DECODE_EXPERT_GATE_UP_ATTR, +) +from chop.nn.quantized.modules.phase_context import force_runtime_phase +from chop.nn.quantized.functional.matrix import plena_matrix_product +from chop.nn.quantized.modules.linear import LinearMXInt +from chop.nn.quantized.modules.qwen3_moe import ( + SUPPORTED_TRANSFORMERS_VERSION, + Qwen3MoeAttentionMXInt, + Qwen3MoeAttentionMXIntRotate, + Qwen3MoeExpertsMXInt, + Qwen3MoeSparseMoeBlockBF16Router, + Qwen3MoeSparseMoeBlockMinifloat, +) +from chop.nn.quantized.modules.qwen3_moe.compat import require_qwen3_moe_fused_abi +from chop.passes.module.transforms.gptq.run import ( + _finalize_gptq_phase, + _snapshot_fp_linear_weights, +) +from chop.passes.module.transforms.quantize.quantize import ( + quantize_module_transform_pass, +) +from chop.passes.module.transforms.quantize.rotation_search import ( + ALL_MATMUL_TYPES, + _resolve_rotation_scope, +) +from chop.nn.quantizers import mxint_quantizer + + +def _config(): + config = Qwen3MoeConfig( + vocab_size=64, + hidden_size=16, + intermediate_size=32, + moe_intermediate_size=8, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=4, + num_experts=4, + num_experts_per_tok=2, + norm_topk_prob=True, + decoder_sparse_step=1, + mlp_only_layers=[], + max_position_embeddings=32, + ) + config._attn_implementation = "eager" + return config + + +def _expert_qconfig(): + return { + "prefill": {"bypass": True}, + "decode": { + "weight_block_size": 4, + "weight_width": 3, + "data_in_block_size": 4, + "data_in_width": 4, + }, + } + + +def test_qwen3_moe_transformers_fused_abi_is_exactly_pinned(): + assert transformers.__version__ == SUPPORTED_TRANSFORMERS_VERSION == "5.5.0" + require_qwen3_moe_fused_abi() + + +def test_untied_local_head_uses_decode_mx_matrix_and_bf16_logits(): + torch.manual_seed(20260725) + model = Qwen3MoeForCausalLM(_config()).to(torch.bfloat16).eval() + original = model.lm_head.weight.detach().clone() + head_decode = { + "weight_block_size": 8, + "weight_width": 4, + "data_in_block_size": 8, + "data_in_width": 8, + "output_format": "FP_E3M2", + "matrix_mlen": 8, + } + model, _ = quantize_module_transform_pass( + model, + { + "by": "regex_name", + r"lm_head$": { + "config": { + "name": "mxint", + "decode": head_decode, + } + }, + }, + ) + head = model.lm_head + hidden = torch.randn(2, 1, 16, dtype=torch.bfloat16) + + assert isinstance(head, LinearMXInt) + assert head._decode_weight_q.numel() == original.numel() + torch.testing.assert_close(head.weight, original, rtol=0.0, atol=0.0) + with force_runtime_phase("prefill"), torch.no_grad(): + prefill = head(hidden) + torch.testing.assert_close( + prefill, + torch.nn.functional.linear(hidden, original), + rtol=0.0, + atol=0.0, + ) + + activation_q = mxint_quantizer( + hidden, + block_size=8, + element_bits=8, + block_dim=-1, + ) + weight_q = mxint_quantizer( + original, + block_size=8, + element_bits=4, + block_dim=1, + ) + expected = plena_matrix_product( + activation_q, + weight_q.transpose(-1, -2), + head_decode, + ) + with force_runtime_phase("decode"), torch.no_grad(): + decode = head(hidden) + tied = head(torch.zeros_like(hidden)) + torch.testing.assert_close(decode, expected, rtol=0.0, atol=0.0) + assert decode.dtype == torch.bfloat16 + assert int(tied.argmax(dim=-1)[0, 0].item()) == 0 + + +def test_mase_matrix_oracle_mlen_partition_is_numerically_observable(): + """The sealed MLEN is part of accuracy, not only a timing parameter.""" + + lhs = torch.tensor( + [[0.69086325, -8.41099262, 1.58839154, 0.66403484]], + dtype=torch.float32, + ) + rhs = torch.tensor( + [[0.68662047], [1.20273066], [6.97725391], [6.73511314]], + dtype=torch.float32, + ) + narrow = plena_matrix_product( + lhs, rhs, {"output_format": "FP_E3M2", "matrix_mlen": 2} + ) + wide = plena_matrix_product( + lhs, rhs, {"output_format": "FP_E3M2", "matrix_mlen": 4} + ) + + torch.testing.assert_close(narrow, torch.tensor([[4.0]]), rtol=0.0, atol=0.0) + torch.testing.assert_close(wide, torch.tensor([[6.0]]), rtol=0.0, atol=0.0) + + +def test_selective_rotation_scope_excludes_fused_expert_tensors(): + model = Qwen3MoeForCausalLM(_config()) + eligible, scope = _resolve_rotation_scope( + model, ALL_MATMUL_TYPES, explicitly_requested=False + ) + + assert set(eligible) == { + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "qk_matmul", + "av_matmul", + "kv_cache", + } + assert scope["architecture"] == "qwen3_moe_fused" + assert scope["excluded_matmul_types"] == [ + "down_proj", + "gate_proj", + "up_proj", + ] + with pytest.raises(ValueError, match="does not support fused"): + _resolve_rotation_scope( + model, ("gate_proj",), explicitly_requested=True + ) + + +def test_fused_expert_banks_preserve_prefill_and_quantize_decode(): + torch.manual_seed(7) + experts = Qwen3MoeExperts(_config()) + with torch.no_grad(): + experts.gate_up_proj.normal_(0.0, 0.1) + experts.down_proj.normal_(0.0, 0.1) + hidden = torch.randn(6, 16) + indices = torch.tensor( + [[0, 1], [1, 2], [2, 3], [3, 0], [0, 2], [1, 3]] + ) + scores = torch.full((6, 2), 0.5) + reference = experts(hidden, indices, scores) + + quantized = Qwen3MoeExpertsMXInt.from_self(experts, _expert_qconfig()) + with force_runtime_phase("prefill"): + prefill = quantized(hidden, indices, scores) + with force_runtime_phase("decode"): + decode = quantized(hidden, indices, scores) + + torch.testing.assert_close(prefill, reference) + assert quantized._decode_gate_up_q.shape == experts.gate_up_proj.shape + assert not torch.allclose(decode, reference) + + +def test_fused_expert_gptq_decode_handoff_restores_fp_source(): + model = Qwen3MoeForCausalLM(_config()).eval() + experts = model.model.layers[0].mlp.experts + fp_gate_up = experts.gate_up_proj.detach().clone() + fp_down = experts.down_proj.detach().clone() + + _snapshot_fp_linear_weights(model.model.layers) + with torch.no_grad(): + experts.gate_up_proj.add_(0.25) + experts.down_proj.sub_(0.125) + gptq_gate_up = experts.gate_up_proj.detach().clone() + gptq_down = experts.down_proj.detach().clone() + _finalize_gptq_phase(model, "decode") + + torch.testing.assert_close(experts.gate_up_proj, fp_gate_up) + torch.testing.assert_close(experts.down_proj, fp_down) + torch.testing.assert_close( + getattr(experts, GPTQ_DECODE_EXPERT_GATE_UP_ATTR), gptq_gate_up + ) + torch.testing.assert_close( + getattr(experts, GPTQ_DECODE_EXPERT_DOWN_ATTR), gptq_down + ) + replacement = Qwen3MoeExpertsMXInt.from_self(experts, _expert_qconfig()) + torch.testing.assert_close(replacement._decode_gate_up_q, gptq_gate_up) + torch.testing.assert_close(replacement._decode_down_q, gptq_down) + + +def test_fused_expert_gptq_runs_end_to_end_on_cpu(monkeypatch): + torch.manual_seed(20260725) + model = Qwen3MoeForCausalLM(_config()).eval() + experts = model.model.layers[0].mlp.experts + fp_gate_up = experts.gate_up_proj.detach().clone() + fp_down = experts.down_proj.detach().clone() + samples = [] + for _ in range(4): + input_ids = torch.randint(0, model.config.vocab_size, (1, 8)) + samples.append((input_ids, input_ids.clone())) + monkeypatch.setattr(gptq_run, "get_loaders", lambda *args, **kwargs: samples) + + gptq_run.run_gptq( + model, + { + "model_name": "in-memory-tiny-qwen3-moe", + "device": "cpu", + "dataset": "in-memory", + "nsamples": 4, + "seqlen": 8, + "format": "mxint", + "weight_config": {"weight_block_size": 4, "weight_width": 4}, + "phase": "decode", + "quantile_search": False, + "clip_search_y": False, + "cali_batch_size": 2, + "min_expert_calibration_hits": 1, + }, + ) + + coverage = model._mase_gptq_expert_coverage + assert len(coverage) == 1 + assert len(coverage[0]["gate_up"]) == experts.num_experts + assert len(coverage[0]["down"]) == experts.num_experts + assert all(record["gptq"] for record in coverage[0]["gate_up"]) + assert all(record["gptq"] for record in coverage[0]["down"]) + torch.testing.assert_close(experts.gate_up_proj, fp_gate_up) + torch.testing.assert_close(experts.down_proj, fp_down) + assert not torch.equal( + getattr(experts, GPTQ_DECODE_EXPERT_GATE_UP_ATTR), fp_gate_up + ) + assert not torch.equal( + getattr(experts, GPTQ_DECODE_EXPERT_DOWN_ATTR), fp_down + ) + assert model.config.use_cache is True + with torch.no_grad(): + logits = model(input_ids=torch.tensor([[1, 2, 3]]), use_cache=False).logits + assert torch.isfinite(logits).all() + + +def test_sparse_block_router_is_bf16_with_fp32_probabilities(): + block = Qwen3MoeSparseMoeBlockBF16Router.from_self( + Qwen3MoeSparseMoeBlock(_config()) + ) + hidden = torch.randn(5, 16) + probabilities, scores, indices = block.gate(hidden) + assert block.gate.weight.dtype == torch.bfloat16 + assert probabilities.dtype == torch.float32 + assert scores.dtype == torch.float32 + assert indices.shape == (5, 2) + torch.testing.assert_close(scores.sum(dim=-1), torch.ones(5)) + + +def test_moe_attention_is_phase_aware_and_requires_tied_kv_precision(): + attention = Qwen3MoeAttention(_config(), layer_idx=0) + rotate = Qwen3MoeAttentionMXIntRotate.from_self( + attention, + { + "decode": { + "qk_matmul": {"bypass": True, "rotate": False}, + "av_matmul": {"bypass": True, "rotate": True}, + "kv_cache": {"bypass": True, "rotate": False}, + } + }, + ) + assert rotate.qk_use_rotate is False + assert rotate.av_use_rotate is True + assert rotate.kv_cache_use_rotate is False + + with pytest.raises(ValueError, match="identical K/V"): + Qwen3MoeAttentionMXInt.from_self( + attention, + { + "decode": { + "kv_cache": { + "key": {"data_in_block_size": 4, "data_in_width": 4}, + "value": {"data_in_block_size": 4, "data_in_width": 3}, + } + } + }, + ) + + +def test_quantize_pass_replaces_packed_experts_and_hooks_moe_decoder(): + model = Qwen3MoeForCausalLM(_config()).eval() + model, _ = quantize_module_transform_pass( + model, + { + "by": "regex_name", + r"model\.layers\.\d+\.mlp$": { + "config": {"name": "bf16_router"} + }, + r"model\.layers\.\d+\.mlp\.experts$": { + "config": {"name": "mxint", **_expert_qconfig()} + }, + }, + ) + layer = model.model.layers[0] + assert isinstance(layer.mlp, Qwen3MoeSparseMoeBlockBF16Router) + assert isinstance(layer.mlp.experts, Qwen3MoeExpertsMXInt) + assert layer._mase_phase_hook_installed + with torch.no_grad(): + logits = model(input_ids=torch.tensor([[1, 2, 3]]), use_cache=False).logits + assert logits.shape == (1, 3, 64) + assert torch.isfinite(logits).all() + + +def test_sparse_block_minifloat_keeps_router_safety_island(): + model = Qwen3MoeForCausalLM(_config()).eval() + expert_config = _expert_qconfig() + expert_config["decode"].update( + {"output_format": "FP_E3M2", "matrix_mlen": 8} + ) + model, _ = quantize_module_transform_pass( + model, + { + "by": "regex_name", + r"model\.layers\.\d+\.mlp$": { + "config": { + "name": "minifloat", + "decode": {"format": "FP_E3M2"}, + } + }, + r"model\.layers\.\d+\.mlp\.experts$": { + "config": {"name": "mxint", **expert_config} + }, + }, + ) + block = model.model.layers[0].mlp + assert isinstance(block, Qwen3MoeSparseMoeBlockMinifloat) + assert block.gate.weight.dtype == torch.bfloat16 + assert isinstance(block.experts, Qwen3MoeExpertsMXInt) + assert block.experts.decode_config["output_format"] == "FP_E3M2" + with force_runtime_phase("decode"), torch.no_grad(): + output = block(torch.randn(1, 3, 16)) + assert output.shape == (1, 3, 16) + assert torch.isfinite(output).all() diff --git a/test/passes/module/transforms/quantize/test_qwen3_moe_router_precision.py b/test/passes/module/transforms/quantize/test_qwen3_moe_router_precision.py new file mode 100644 index 000000000..a0cfec611 --- /dev/null +++ b/test/passes/module/transforms/quantize/test_qwen3_moe_router_precision.py @@ -0,0 +1,238 @@ +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F +from transformers.models.qwen3_moe import Qwen3MoeConfig, Qwen3MoeForCausalLM +from transformers.models.qwen3_moe.modeling_qwen3_moe import ( + Qwen3MoeSparseMoeBlock, +) + +from chop.nn.quantized.functional.matrix import plena_matrix_product +from chop.nn.quantized.modules.phase_context import force_runtime_phase +from chop.nn.quantized.modules.qwen3_moe import ( + Qwen3MoeTopKRouterBF16, + Qwen3MoeTopKRouterE4M3, + Qwen3MoeTopKRouterE5M2, + Qwen3MoeTopKRouterMX, + Qwen3MoeTopKRouterMXInt8, + router_phase_config, +) +from chop.nn.quantizers import mxfp_quantizer, mxint_quantizer +from chop.passes.module.transforms.quantize.quantize import ( + quantize_module_transform_pass, +) + + +def _config() -> Qwen3MoeConfig: + config = Qwen3MoeConfig( + vocab_size=64, + hidden_size=16, + intermediate_size=32, + moe_intermediate_size=8, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=4, + num_experts=4, + num_experts_per_tok=2, + norm_topk_prob=True, + decoder_sparse_step=1, + mlp_only_layers=[], + max_position_embeddings=32, + ) + config._attn_implementation = "eager" + return config + + +def _mlen_config() -> Qwen3MoeConfig: + config = Qwen3MoeConfig( + vocab_size=64, + hidden_size=2048, + intermediate_size=32, + moe_intermediate_size=8, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=512, + num_experts=4, + num_experts_per_tok=2, + norm_topk_prob=True, + decoder_sparse_step=1, + mlp_only_layers=[], + max_position_embeddings=32, + ) + config._attn_implementation = "eager" + return config + + +def _manual_quantize(value: torch.Tensor, token: str, block_dim: int): + if token == "MXINT8": + return mxint_quantizer( + value, block_size=8, element_bits=8, block_dim=block_dim + ) + exponent, fraction = token.removeprefix("E").split("M") + return mxfp_quantizer( + value, + block_size=8, + element_exp_bits=int(exponent), + element_frac_bits=int(fraction), + block_dim=block_dim, + ) + + +def _manual_route(logits: torch.Tensor, top_k: int): + probabilities = F.softmax(logits.to(torch.bfloat16), dtype=torch.float32, dim=-1) + scores, indices = torch.topk(probabilities, top_k, dim=-1, sorted=True) + scores = scores / scores.sum(dim=-1, keepdim=True) + return probabilities, scores, indices + + +@pytest.mark.parametrize( + ("token", "router_cls"), + ( + ("MXINT8", Qwen3MoeTopKRouterMXInt8), + ("E4M3", Qwen3MoeTopKRouterE4M3), + ("E5M2", Qwen3MoeTopKRouterE5M2), + ), +) +def test_router_mx_cpu_matches_explicit_matrix_and_fp32_route(token, router_cls): + torch.manual_seed(20260820) + source = Qwen3MoeSparseMoeBlock(_config()).gate + source.weight.data.copy_( + torch.linspace(-0.75, 0.75, source.weight.numel()).reshape_as(source.weight) + ) + config = router_phase_config(token, token, matrix_mlen=8) + router = router_cls.from_self(source, config) + hidden = torch.linspace(-1.0, 1.0, 5 * 16).reshape(5, 16) + + activation_q = _manual_quantize(hidden.to(torch.bfloat16), token, -1) + weight_q = _manual_quantize(source.weight.to(torch.bfloat16), token, 1) + logits = plena_matrix_product( + activation_q, + weight_q.transpose(-1, -2), + {"matrix_mlen": 8, "output_format": "BF16"}, + ).to(torch.bfloat16) + expected = _manual_route(logits, top_k=2) + + with force_runtime_phase("decode"), torch.no_grad(): + observed = router(hidden) + for actual, reference in zip(observed, expected): + torch.testing.assert_close(actual, reference, rtol=0.0, atol=0.0) + assert observed[0].dtype == torch.float32 + assert observed[1].dtype == torch.float32 + assert router.router_precision_contract["publication_rankable"] is False + assert router.router_precision_contract["selection_eligible"] is False + + +def test_router_mx_prefill_has_exact_bf16_baseline_ancestry(): + torch.manual_seed(9) + source = Qwen3MoeSparseMoeBlock(_config()).gate + baseline = Qwen3MoeTopKRouterBF16.from_self(source) + variant = Qwen3MoeTopKRouterMX.from_self( + source, + router_phase_config("E4M3", "E5M2", matrix_mlen=16), + ) + hidden = torch.randn(7, 16) + with force_runtime_phase("prefill"), torch.no_grad(): + baseline_output = baseline(hidden) + variant_output = variant(hidden) + for actual, reference in zip(variant_output, baseline_output): + torch.testing.assert_close(actual, reference, rtol=0.0, atol=0.0) + + +def test_router_matrix_oracle_exposes_mlen_1024_vs_2048_partition_boundary(): + torch.manual_seed(17) + source = Qwen3MoeSparseMoeBlock(_mlen_config()).gate + source.weight.data.normal_(mean=0.0, std=0.01) + hidden = torch.randn(5, 2048) * 0.1 + split = Qwen3MoeTopKRouterE4M3.from_self( + source, + router_phase_config("E4M3", "E4M3", matrix_mlen=1024), + ) + whole = Qwen3MoeTopKRouterE4M3.from_self( + source, + router_phase_config("E4M3", "E4M3", matrix_mlen=2048), + ) + with force_runtime_phase("decode"), torch.no_grad(): + split_probabilities, _, _ = split(hidden) + whole_probabilities, _, _ = whole(hidden) + assert not torch.equal(split_probabilities, whole_probabilities) + assert split.router_precision_contract["matrix_arithmetic_chain"] == [ + "per_mlen_fp32_matmul", + "bf16_partial_rounding", + "truncate_partial_to_signed_fixed16_16", + "signed_fixed16_16_wrap_across_partitions", + "final_bf16_writeout", + ] + + +def test_router_route_agreement_measurement_is_deterministic_on_cpu(): + torch.manual_seed(112) + source = Qwen3MoeSparseMoeBlock(_config()).gate + baseline = Qwen3MoeTopKRouterBF16.from_self(source) + variant = Qwen3MoeTopKRouterE5M2.from_self( + source, + router_phase_config("E5M2", "E5M2", matrix_mlen=8), + ) + hidden = torch.randn(13, 16) + + def measure(): + with force_runtime_phase("decode"), torch.no_grad(): + base_prob, _, base_idx = baseline(hidden) + test_prob, _, test_idx = variant(hidden) + ordered = (base_idx == test_idx).all(dim=-1) + set_equal = torch.tensor( + [ + torch.equal(torch.sort(left).values, torch.sort(right).values) + for left, right in zip(base_idx, test_idx) + ], + dtype=torch.bool, + ) + return { + "tokens": int(hidden.shape[0]), + "ordered_matches": int(ordered.sum()), + "set_matches": int(set_equal.sum()), + "probability_l1_sum": float((base_prob - test_prob).abs().sum()), + "probability_linf": float((base_prob - test_prob).abs().max()), + } + + first = measure() + second = measure() + assert first == second + assert 0 <= first["ordered_matches"] <= first["set_matches"] <= 13 + assert first["probability_l1_sum"] >= 0.0 + assert first["probability_linf"] >= 0.0 + + +def test_quantize_pass_installs_router_variant_without_touching_profile_space(): + model = Qwen3MoeForCausalLM(_config()).eval() + model, _ = quantize_module_transform_pass( + model, + { + "by": "regex_name", + r"model\.layers\.\d+\.mlp\.gate$": { + "config": { + "name": "e4m3", + **router_phase_config("E4M3", "E4M3", matrix_mlen=8), + } + }, + }, + ) + gate = model.model.layers[0].mlp.gate + assert isinstance(gate, Qwen3MoeTopKRouterE4M3) + assert gate.router_weight_format == gate.router_activation_format == "E4M3" + with force_runtime_phase("decode"), torch.no_grad(): + probability, score, index = gate(torch.randn(3, 16)) + assert probability.shape == (3, 4) + assert score.shape == index.shape == (3, 2) + + +@pytest.mark.parametrize("token", ("MXINT4", "E3M4", "BF16")) +def test_router_lane_rejects_unplanned_formats(token): + source = Qwen3MoeSparseMoeBlock(_config()).gate + with pytest.raises(ValueError, match="must be one of"): + Qwen3MoeTopKRouterMX.from_self( + source, + router_phase_config(token, "MXINT8", matrix_mlen=8), + )