diff --git a/src/python/py/models/README.md b/src/python/py/models/README.md index 8adca82f3b..9c6bed481f 100644 --- a/src/python/py/models/README.md +++ b/src/python/py/models/README.md @@ -140,6 +140,18 @@ python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o pa python builder.py -i path_to_local_folder_on_disk -o path_to_output_folder -p int4 -e execution_provider -c cache_dir_to_store_temp_files ``` +#### 2-bit (INT2) pre-quantized MoE + +For a model whose MoE experts have already been quantized to 2-bit with [AMD Quark](https://quark.docs.amd.com/) (group-wise `uint2`), pass `-p int2`. The builder consumes the pre-quantized expert weights, scales, and zero-points directly (no float re-quantization) and emits a fused `QMoE` op with `expert_weight_bits=2`. This path is currently exercised by the Gemma 4 MoE (`gemma-4-26B-A4B-it`) text model on the CPU execution provider. + +```bash +# From wheel: +python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o path_to_output_folder -p int2 -e cpu -c cache_dir_to_store_temp_files + +# From source: +python builder.py -i path_to_local_folder_on_disk -o path_to_output_folder -p int2 -e cpu -c cache_dir_to_store_temp_files +``` + ### GGUF Model This scenario is where your float16/float32 GGUF model is already on disk. diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py index 6974cc8319..ee0d5636d9 100644 --- a/src/python/py/models/builder.py +++ b/src/python/py/models/builder.py @@ -22,6 +22,8 @@ ErnieModel, Gemma2Model, Gemma3Model, + Gemma4MoEModel, + Gemma4Model, GemmaModel, GPTOSSModel, GraniteModel, @@ -242,7 +244,7 @@ def check_extra_options( # `moe_quant_type` is the single option that selects the MoE quantization scheme. It replaces the # older per-type flags (`use_8bits_moe``) so new schemes can be added without a new flag. - supported_moe_quant_types = {"int4", "int8", "mxfp4", "nvfp4"} + supported_moe_quant_types = {"int2", "uint2", "int4", "int8", "mxfp4", "nvfp4"} # Backward compatibility: `use_8bits_moe` is deprecated in favor of `moe_quant_type`. if "use_8bits_moe" in extra_options: @@ -395,7 +397,10 @@ def set_io_dtype(precision, execution_provider, extra_options) -> ir.DataType: """ Set the input/output precision of the ONNX model based on the provided precision and execution provider. """ - cpu_quant = precision in {"int4", "int8"} and execution_provider == "cpu" + # int2/int4/int8 weight-only quantization builds a float graph and quantizes the weights at save time. + # On the CPU EP the I/O stays FP32; on GPU/WebGPU it follows the usual FP16 default (int8 must not + # be forced to FP32 I/O everywhere). + cpu_quant = precision in {"int2", "int4", "int8"} and execution_provider == "cpu" fp32_webgpu = execution_provider == "webgpu" and extra_options.get("use_webgpu_fp32", False) bf16_cuda = precision == "int4" and execution_provider in {"cuda", "trt-rtx"} and extra_options.get("use_cuda_bf16", False) @@ -418,6 +423,11 @@ def set_onnx_dtype(precision: str, extra_options: dict[str, Any]) -> ir.DataType if precision == "int4": return ir.DataType.INT4 if extra_options.get("is_symmetric", True) else ir.DataType.UINT4 + if precision == "int2": + # 2-bit quantized weights are emitted as MatMulNBits (per-module bits=2); non-quantized + # ops follow io_dtype, so onnx_dtype tracks the FLOAT io_dtype used on CPU. + return ir.DataType.FLOAT + if precision == "int8": return ir.DataType.INT8 if extra_options.get("is_symmetric", True) else ir.DataType.UINT8 @@ -459,6 +469,10 @@ def create_model( # Set input/output precision of ONNX model io_dtype = set_io_dtype(precision, execution_provider, extra_options) onnx_dtype = set_onnx_dtype(precision, extra_options) + # int2 carries onnx_dtype FLOAT (2-bit weights are per-module MatMulNBits bits=2), so the + # onnx_dtype alone can't distinguish an int2 build from a genuine fp32 build. Thread the true + # precision string through so the builder can resolve the QMoE bit-width / op selection. + extra_options["precision"] = precision config_only = extra_options.get("config_only", False) # List architecture options in alphabetical order @@ -484,6 +498,24 @@ def create_model( onnx_model = Gemma3Model(config, io_dtype, onnx_dtype, execution_provider, cache_dir, extra_options) if not onnx_model.exclude_embeds: onnx_model.model_type = "gemma3_vl_text" + elif config.architectures[0] == "Gemma4ForConditionalGeneration": + text_config = config.text_config + for key in text_config: + if not hasattr(config, key): + setattr(config, key, getattr(text_config, key)) + print("WARNING: This model loses accuracy with float16 precision. It is recommended to set `--precision bf16` or `--precision int4 --extra_options use_cuda_bf16=true` by default.") + print("WARNING: This is only generating the text component of the model. The vision and audio components are not supported.") + onnx_model = Gemma4MoEModel(config, io_dtype, onnx_dtype, execution_provider, cache_dir, extra_options) + onnx_model.model_type = "gemma4_text" + elif config.architectures[0] == "Gemma4UnifiedForConditionalGeneration": + text_config = config.text_config + for key in text_config: + if not hasattr(config, key): + setattr(config, key, getattr(text_config, key)) + print("WARNING: This model loses accuracy with float16 precision. It is recommended to set `--precision bf16` or `--precision int4 --extra_options use_cuda_bf16=true` by default.") + print("WARNING: This is only generating the text component of the model. The vision and audio components are not supported.") + onnx_model = Gemma4Model(config, io_dtype, onnx_dtype, execution_provider, cache_dir, extra_options) + onnx_model.model_type = "gemma4_text" elif config.architectures[0] == "GptOssForCausalLM": print("WARNING: This model only supports symmetric quantization for `QMoE`.") if hasattr(config, "quantization_config") and config.quantization_config.get("quant_method") != "quark": @@ -620,7 +652,7 @@ def get_args(): "-p", "--precision", required=True, - choices=["int4", "int8", "bf16", "fp16", "fp32"], + choices=["int2", "int4", "int8", "bf16", "fp16", "fp32"], help="Precision of model", ) diff --git a/src/python/py/models/builders/__init__.py b/src/python/py/models/builders/__init__.py index 9544a51862..121cdddb9c 100644 --- a/src/python/py/models/builders/__init__.py +++ b/src/python/py/models/builders/__init__.py @@ -9,7 +9,7 @@ from .base import Model from .chatglm import ChatGLMModel from .ernie import ErnieModel -from .gemma import Gemma2Model, Gemma3Model, GemmaModel +from .gemma import Gemma2Model, Gemma3Model, Gemma4MoEModel, Gemma4Model, GemmaModel from .gptoss import GPTOSSModel from .granite import GraniteModel, GraniteMoEHybridModel from .hunyuan import HunyuanDenseV1Model @@ -52,6 +52,8 @@ "GPTOSSModel", "Gemma2Model", "Gemma3Model", + "Gemma4MoEModel", + "Gemma4Model", "GemmaModel", "GraniteMoEHybridModel", "GraniteModel", diff --git a/src/python/py/models/builders/base.py b/src/python/py/models/builders/base.py index 71c32035cc..d0925bdcaa 100644 --- a/src/python/py/models/builders/base.py +++ b/src/python/py/models/builders/base.py @@ -927,9 +927,13 @@ def make_kv_cache_scale(per_layer, scale_index, layer_id): def make_quant_config_init(self): self.quant_config = self.extra_options.get("_quant_config", None) if self.quant_config is None: + # int2 carries onnx_dtype FLOAT, so passing self.onnx_dtype would collapse it to "fp32" + # and lose the 2-bit signal. Prefer the true precision string threaded through + # extra_options (set by builder.create_model) when present. + quant_precision = self.extra_options.get("precision") or self.onnx_dtype self.quant_config = QuantConfig.from_extra_options( extra_options=self.extra_options, - precision=self.onnx_dtype, + precision=quant_precision, execution_provider=self.ep, ) elif not isinstance(self.quant_config, QuantConfig): @@ -951,7 +955,12 @@ def make_moe_init(self): # MXFP4 and NVFP4 both resolve to the "mx" kind; the QMoE op tells them apart by dtype name # ("mxfp4" -> op "fp4", "nvfp4" -> op "nvfp4"). Integer dtypes use the plain "int" QMoE path. - self.moe_attrs["moe_op_type"] = "QMoE" if moe_descriptor.is_quantized else "MoE" + # Both keys carry the resolved op: the shared emitters (make_moe_expert_initializers / + # make_moe_op) read "op_type", while the gemma parallel-FFN path reads "moe_op_type". Setting + # only one left quantized experts on the float "MoE" op, which rejects uint8 weights. + moe_op_type = "QMoE" if moe_descriptor.is_quantized else "MoE" + self.moe_attrs["op_type"] = moe_op_type + self.moe_attrs["moe_op_type"] = moe_op_type if moe_descriptor.kind == "mx": self.moe_attrs["qmoe_quant_type"] = "nvfp4" if moe_descriptor.name == "nvfp4" else "fp4" else: @@ -2057,6 +2066,30 @@ def make_matmul_fp8(self, matmul, basename, root_input, **kwargs): ) def make_matmul_op(self, matmul, basename, root_input, **kwargs): + # Factored online rotation (Quark rotation algo): rotate the activation before the + # quantized projection. Weights are stored in the rotated basis, so this is required. + # LoRA (below) is applied to the ORIGINAL (pre-rotation) activation, so capture it first. + original_root_input = root_input + if getattr(matmul, "input_prescale", None) is not None and getattr(self, "shared_input_rotations", None): + root_input = self.make_factored_rotation(matmul, basename, root_input, **kwargs) + + matmul_name = self.make_matmul_core(matmul, basename, root_input, **kwargs) + + # Bake the additive PEFT LoRA delta into the graph, if present (never for `logits`, + # since the excluded lm_head carries no adapter). + if ( + not kwargs.get("logits", False) + and getattr(matmul, "lora_A", None) is not None + and getattr(matmul, "lora_B", None) is not None + ): + return self.make_lora_add(matmul, basename, original_root_input, matmul_name, **kwargs) + return matmul_name + + def make_matmul_core(self, matmul, basename, root_input, **kwargs): + # 2-bit Quark weights are emitted as MatMulNBits even when the io/onnx dtype is + # float (ORT runs the float zero-point 2-bit MatMulNBits on CPU/MLAS). + if hasattr(matmul, "qweight") and matmul.qweight is not None and getattr(matmul, "bits", None) == 2: + return self.make_matmul_nbits(matmul, basename, root_input, **kwargs) if self.onnx_dtype in {ir.DataType.FLOAT16, ir.DataType.BFLOAT16, ir.DataType.FLOAT}: return self.make_matmul_float(matmul, basename, root_input, **kwargs) elif self.onnx_dtype in {ir.DataType.INT4, ir.DataType.UINT4, ir.DataType.INT8, ir.DataType.UINT8}: @@ -2172,6 +2205,62 @@ def make_matmul_block_quantized_nvfp4_weight( ) return basename + def make_factored_rotation(self, matmul, basename, root_input, **kwargs): + """Emit the factored online input rotation for a projection: + + x_rot = (x * input_prescale) @ shared_input_rotation_ + + `input_prescale` is a per-projection [in] vector (small, emitted per projection). + `shared_input_rotation_` is a single [in, in] matrix shared by every projection + of the same in_features and emitted only once (a huge size saving vs. inlining a copy + per projection). Returns the rotated activation tensor name to feed into MatMulNBits. + """ + seq_dim = kwargs.get("seq_dim", "sequence_length") + in_features = matmul.in_features + + # 1. Multiply by the per-projection input pre-scale vector. + prescale_init = basename[1:].replace("/", ".") + ".input_prescale" + self.make_initializer(matmul.input_prescale, prescale_init, to=self.io_dtype) + div_output = f"{basename}/input_prescale/output_0" + self.make_node("Mul", inputs=[root_input, prescale_init], outputs=[div_output], name=f"{basename}/input_prescale") + self.make_value(div_output, self.io_dtype, shape=["batch_size", seq_dim, in_features]) + + # 2. Multiply by the shared rotation matrix (single initializer per in_features). + rot_init = f"model.shared_input_rotation_{in_features}" + if rot_init not in self.shared_rotation_initializers: + self.make_initializer(self.shared_input_rotations[in_features], rot_init, to=self.io_dtype) + self.shared_rotation_initializers.add(rot_init) + rot_output = f"{basename}/shared_input_rotation/output_0" + self.make_node("MatMul", inputs=[div_output, rot_init], outputs=[rot_output], name=f"{basename}/shared_input_rotation") + self.make_value(rot_output, self.io_dtype, shape=["batch_size", seq_dim, in_features]) + return rot_output + + def make_lora_add(self, matmul, basename, root_input, matmul_name, **kwargs): + # Additive PEFT LoRA delta, baked into the graph: + # y = MatMulNBits(rotate(x)) + ((x @ lora_A^T) @ (scaling * lora_B)^T) + # lora_A is [rank, in_features], lora_B is [out_features, rank]; both act on the + # ORIGINAL (pre-rotation) activation, matching how the adapter was trained. The + # `scaling` factor (lora_alpha / r) is pre-baked into the lora_B initializer. + seq_dim = kwargs.get("seq_dim", "sequence_length") + rank = matmul.lora_A.shape[0] + out_features = matmul.lora_B.shape[0] + + lora_A_weight = f"{basename}/lora_A/weight" + self.make_initializer(matmul.lora_A.T, lora_A_weight, to=self.io_dtype) + lora_A_output = f"{basename}/lora_A/output_0" + self.make_node("MatMul", inputs=[root_input, lora_A_weight], outputs=[lora_A_output], name=f"{basename}/lora_A") + self.make_value(lora_A_output, self.io_dtype, shape=["batch_size", seq_dim, rank]) + + lora_B_weight = f"{basename}/lora_B/weight" + self.make_initializer((matmul.lora_B * matmul.lora_scaling).T, lora_B_weight, to=self.io_dtype) + lora_B_output = f"{basename}/lora_B/output_0" + self.make_node("MatMul", inputs=[lora_A_output, lora_B_weight], outputs=[lora_B_output], name=f"{basename}/lora_B") + self.make_value(lora_B_output, self.io_dtype, shape=["batch_size", seq_dim, out_features]) + + add_name = f"{basename}/lora/Add" + self.make_add(add_name, [f"{matmul_name}/output_0", lora_B_output], dtype=self.io_dtype, shape=["batch_size", seq_dim, out_features]) + return add_name + def make_matmul_float(self, matmul, name, root_input, **kwargs): weight = name[1:].replace("/", ".") + ".weight" self.make_initializer(matmul.weight.T, weight, to=self.io_dtype) @@ -5175,6 +5264,13 @@ def load_weights(self, input_path): intermediate_size=self.intermediate_size, num_layers=self.num_layers, ) + # Factored online rotation (Quark rotation algo): shared [in, in] rotation + # matrices, emitted once and referenced by every projection of that in_features. + self.shared_input_rotations = getattr(model, "shared_input_rotations", {}) or {} + self.shared_rotation_initializers = set() + if self.shared_input_rotations: + # Each projection applies its own input rotation, so build q/k/v separately. + self.attention_attrs["use_packed_matmul"] = False else: extra_kwargs = {"num_hidden_layers": self.num_layers} if "num_hidden_layers" in self.extra_options else {} @@ -5210,7 +5306,10 @@ def load_weights(self, input_path): **extra_kwargs, ) - if "adapter_path" in self.extra_options: + # The quantized (Quark) loader attaches its LoRA adapter internally from + # /lora_adapters/, and QuantModel is not an nn.Module, so the + # PeftModel wrapping only applies to the plain HF path. + if "adapter_path" in self.extra_options and self.quant_type is None: from peft import PeftModel model = PeftModel.from_pretrained( diff --git a/src/python/py/models/builders/gemma.py b/src/python/py/models/builders/gemma.py index 8f6bdd8a69..9de48fd784 100644 --- a/src/python/py/models/builders/gemma.py +++ b/src/python/py/models/builders/gemma.py @@ -3,8 +3,13 @@ # Licensed under the MIT License. See License.txt in the project root for # license information. # -------------------------------------------------------------------------- +import os + import numpy as np +import torch +from onnx_ir.tensor_adapters import to_torch_dtype +from .base import Model from .mistral import MistralModel @@ -144,3 +149,714 @@ def make_rotary_embedding_caches(self, **kwargs): "sin_cache_name", self.sin_cache_global_name if self.window_size == -1 else self.sin_cache_local_name ) return super().make_rotary_embedding_caches(cos_cache_name=cos_cache_name, sin_cache_name=sin_cache_name) + + +class Gemma4Model(Gemma3Model): + """Builder for the text decoder of Gemma4Unified (gemma4-12b-it). + + Differs from Gemma3 in several structural ways (see below). Only the text + component is built; vision/audio configs are ignored. + """ + + def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): + # Per-layer-type geometry. Sliding layers are the base profile; full + # ("global") layers swap in their own head_dim / kv_heads per layer. + self.layer_types = list(config.layer_types) + self.sliding_head_dim = config.head_dim + self.global_head_dim = config.global_head_dim + self.sliding_num_kv_heads = config.num_key_value_heads + self.global_num_kv_heads = config.num_global_key_value_heads + self.attention_k_eq_v = getattr(config, "attention_k_eq_v", False) + + # RoPE parameters differ per layer type (nested dict in the HF config). + rope_params = config.rope_parameters + self.global_rope_theta = rope_params["full_attention"]["rope_theta"] + self.global_partial_rotary_factor = rope_params["full_attention"]["partial_rotary_factor"] + self.local_rope_theta = rope_params["sliding_attention"]["rope_theta"] + + # Base __init__ reads config.head_dim (= sliding) for self.head_size and + # config.num_key_value_heads (= sliding) for self.num_kv_heads, which is + # the default (sliding) profile. It also builds the RoPE caches via the + # Gemma3 multi-cache path, so set rope_local_base_freq for that. + config.rope_local_base_freq = self.local_rope_theta + # Gemma2Model.__init__ derives its attention scale from + # query_pre_attn_scalar, which Gemma4 does not have (its Q/K are + # RMS-normed, so the scale is 1.0). Provide a placeholder so the parent + # chain runs; the scale is overridden to 1.0 below. + if not hasattr(config, "query_pre_attn_scalar"): + config.query_pre_attn_scalar = config.head_dim + super().__init__(config, io_dtype, onnx_dtype, ep, cache_dir, extra_options) + + # Gemma4 RMSNorm uses the weight directly (no `1 + weight` offset that + # Gemma1/2/3 apply). This also governs q_norm/k_norm weights. + self.layernorm_attrs["add_offset"] = 0 + + # Q/K are RMS-normed, so HF sets attention scaling to 1.0 (not + # 1/sqrt(head_dim)). + self.attention_attrs["scale"] = 1.0 + + # Partial rotary on the global layers requires the standalone + # RotaryEmbedding op (GQA's fused rope has no rotary_embedding_dim + # attribute). Force external RoPE and keep position_ids as an input. + self.attention_attrs["use_rope_in_attn"] = False + if "position_ids" not in self.input_names: + self.input_names["position_ids"] = "position_ids" + self.input_types["position_ids"] = self.input_types.get("attention_mask") + self.input_shapes["position_ids"] = ["batch_size", "sequence_length"] + + # Per-layer residual output multipliers (`layer_scalar`), read from the + # weights in make_layer. + self.layer_scalars = {} + + def make_rope_init(self, config): + # Gemma4 stores rope_parameters as a per-layer-type nested dict + # ({"full_attention": {...}, "sliding_attention": {...}}) rather than the + # flat {"rope_type": ...} form the base reader expects. The per-type theta + # / partial_rotary_factor are applied directly in __init__ and via the + # Gemma3 multi-cache path, so skip the base initializer. + return + + def load_weights(self, input_path): + # The checkpoint is a full multimodal Gemma4Unified model whose text + # weights live under the `model.language_model.` prefix. Loading the full + # ConditionalGeneration model (as the base does) also pulls in the vision + # and audio towers, which are out of scope and roughly double the memory. + # Instead, build the text-only CausalLM and load just the remapped text + # weights. + if self.quant_type is not None or input_path.endswith(".gguf"): + return super().load_weights(input_path) + + import glob + + from safetensors import safe_open + from transformers import AutoConfig + from transformers.models.gemma4_unified import Gemma4UnifiedForCausalLM + + config = AutoConfig.from_pretrained( + self.model_name_or_path, token=self.hf_token, trust_remote_code=self.hf_remote + ) + text_config = config.text_config + text_config.num_hidden_layers = self.num_layers + text_config.layer_types = text_config.layer_types[: self.num_layers] + + with torch.device("meta"): + model = Gemma4UnifiedForCausalLM(text_config) + + # Remap `model.language_model.*` -> `model.*` and drop layers beyond the + # (possibly truncated) layer count. + prefix = "model.language_model." + state_dict = {} + for shard in sorted(glob.glob(os.path.join(self.model_name_or_path, "*.safetensors"))): + with safe_open(shard, framework="pt") as f: + for key in f.keys(): + if not key.startswith(prefix): + continue + new_key = key[len(prefix) :] + if new_key.startswith("layers."): + if int(new_key.split(".")[1]) >= self.num_layers: + continue + state_dict["model." + new_key] = f.get_tensor(key) + + if getattr(text_config, "tie_word_embeddings", False): + state_dict["lm_head.weight"] = state_dict["model.embed_tokens.weight"] + + missing, unexpected = model.load_state_dict(state_dict, strict=False, assign=True) + if missing: + raise ValueError(f"Missing weights while loading Gemma4 text model: {missing}") + if unexpected: + raise ValueError(f"Unexpected weights while loading Gemma4 text model: {unexpected}") + + return model + + def is_local(self, layer_id): + return self.layer_types[layer_id] == "sliding_attention" + + def layer_head_dim(self, layer_id): + return self.sliding_head_dim if self.is_local(layer_id) else self.global_head_dim + + def layer_num_kv_heads(self, layer_id): + return self.sliding_num_kv_heads if self.is_local(layer_id) else self.global_num_kv_heads + + def make_key_value_cache_shape(self, layer_id, shape): + # Emit concrete kv_heads (dim 1) and head_dim (dim 3) per layer so the + # runtime's DefaultKeyValueCache detects the per-layer variation. + shape = super().make_key_value_cache_shape(layer_id, shape) + return [shape[0], self.layer_num_kv_heads(layer_id), shape[2], self.layer_head_dim(layer_id)] + + def make_attention(self, layer_id, attention, root_input, **kwargs): + # Swap in this layer's geometry (head_dim, kv_heads, q/kv sizes) around + # the base implementation, restoring afterward. Window handling is done + # by Gemma2Model.make_attention (super) via is_local. + original = (self.head_size, self.num_kv_heads, self.q_size, self.kv_size) + self.head_size = self.layer_head_dim(layer_id) + self.num_kv_heads = self.layer_num_kv_heads(layer_id) + self.q_size = self.num_attn_heads * self.head_size + self.kv_size = self.num_kv_heads * self.head_size + # RoPE uses a full split-half rotation (rotary_embedding_dim=0) for both + # layer types. The global layers' partial rotary is baked into the global + # cache as a zero-padded NoPE tail (see make_proportional_rope_caches), so + # the op rotates the full head_dim with that pre-zeroed cache — matching + # HF's proportional RoPE. The external RotaryEmbedding op needs position_ids. + super().make_attention( + layer_id, attention, root_input, position_ids=self.input_names["position_ids"], **kwargs + ) + self.head_size, self.num_kv_heads, self.q_size, self.kv_size = original + + def make_attention_input_proj(self, layer_id, attention, root_input, **kwargs): + if self.attention_k_eq_v and not self.is_local(layer_id): + # Full-attention layers share the K projection as V (no v_proj). + attention.v_proj = attention.k_proj + super().make_attention_input_proj(layer_id, attention, root_input, **kwargs) + # Insert the scaleless value RMSNorm (v_norm) on the V path. + self.make_v_norm(layer_id) + + def make_v_norm(self, layer_id): + # Scaleless SimplifiedLayerNorm (weight = ones) applied per-head on V, + # matching HF's Gemma4UnifiedRMSNorm(with_scale=False) on value_states. + head_size = self.head_size + kv_size = self.kv_size + + reshape_1_name = f"/model/layers.{layer_id}/attn/v_norm/Reshape_1" + reshape_1_inputs = [self.attention_attrs["v_path"], f"/model/constants/INT64/[0, -1, {head_size}]"] + self.make_reshape( + reshape_1_name, + reshape_1_inputs, + dtype=self.io_dtype, + shape=["batch_size", "sequence_length * num_key_value_heads", head_size], + ) + + weight_name = f"model.layers.{layer_id}.attn.v_norm.layernorm.weight" + self.make_initializer(torch.ones(head_size), weight_name, to=self.io_dtype) + + layernorm_name = f"/model/layers.{layer_id}/attn/v_norm/SimplifiedLayerNormalization" + layernorm_output = f"{layernorm_name}/output_0" + self.make_node( + "SimplifiedLayerNormalization", + inputs=[f"{reshape_1_name}/output_0", weight_name], + outputs=[layernorm_output], + name=layernorm_name, + epsilon=self.layernorm_attrs["epsilon"], + axis=-1, + stash_type=1, + ) + self.make_value( + layernorm_output, + dtype=self.io_dtype, + shape=["batch_size", "sequence_length * num_key_value_heads", head_size], + ) + + reshape_2_name = f"/model/layers.{layer_id}/attn/v_norm/Reshape_2" + reshape_2_inputs = [layernorm_output, f"/model/constants/INT64/[0, -1, {kv_size}]"] + self.make_reshape( + reshape_2_name, + reshape_2_inputs, + dtype=self.io_dtype, + shape=["batch_size", "sequence_length", kv_size], + ) + self.attention_attrs["v_path"] = f"{reshape_2_name}/output_0" + + def make_rotary_embedding_multi_cache(self): + # Build the global cache with the proportional RoPE variant and the + # local cache with default RoPE (theta = local_rope_theta on the sliding + # head_dim). Overrides Gemma3's default-both-caches behavior. + self.cos_cache_global_name, self.sin_cache_global_name = "cos_cache_global", "sin_cache_global" + cos_global, sin_global = self.make_proportional_rope_caches() + self.make_initializer(cos_global, self.cos_cache_global_name, to=self.io_dtype) + self.make_initializer(sin_global, self.sin_cache_global_name, to=self.io_dtype) + + # Local (sliding) cache: default RoPE on the sliding head_dim. + original = (self.head_size, self.rope_attrs["theta"], self.rope_attrs["partial_rotary_factor"]) + self.head_size = self.sliding_head_dim + self.rope_attrs["theta"] = self.local_rope_theta + self.rope_attrs["partial_rotary_factor"] = 1.0 + self.rope_attrs["create_caches"] = True + self.cos_cache_local_name, self.sin_cache_local_name = "cos_cache_local", "sin_cache_local" + # Deliberately reach the base (Model) implementation, skipping Gemma3's + # multi-cache override, to build a single default RoPE cache. Called + # explicitly (rather than super()) because super() would resolve to + # Gemma3Model.make_rotary_embedding_caches and change behavior. + Model.make_rotary_embedding_caches( + self, cos_cache_name=self.cos_cache_local_name, sin_cache_name=self.sin_cache_local_name + ) + self.head_size, self.rope_attrs["theta"], self.rope_attrs["partial_rotary_factor"] = original + + def make_proportional_rope_caches(self): + # Replicates transformers' _compute_proportional_rope_parameters: + # partial rotary applied on the *global* head_dim, with a zero-padded + # NoPE tail so the emitted rotary_embedding_dim spans the full head_dim. + head_dim = self.global_head_dim + base = self.global_rope_theta + rope_angles = int(self.global_partial_rotary_factor * head_dim // 2) + inv_freq_rotated = 1.0 / ( + base ** (torch.arange(0, 2 * rope_angles, 2, dtype=torch.int64).float() / head_dim) + ) + nope_angles = head_dim // 2 - rope_angles + if nope_angles > 0: + inv_freq = torch.cat((inv_freq_rotated, torch.zeros(nope_angles, dtype=torch.float32)), dim=0) + else: + inv_freq = inv_freq_rotated + + t = torch.arange(self.context_length, dtype=torch.int64).float() + freqs = torch.outer(t, inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + cos_cache, sin_cache = emb.cos(), emb.sin() + cos_cache = cos_cache.squeeze().to(to_torch_dtype(self.io_dtype)) + sin_cache = sin_cache.squeeze().to(to_torch_dtype(self.io_dtype)) + # Halve to (M, head_dim/2) as the RotaryEmbedding kernel expects. The + # NoPE tail contributes zero-frequency (cos=1, sin=0) entries. + cos_cache = cos_cache[:, : (head_dim // 2)] + sin_cache = sin_cache[:, : (head_dim // 2)] + return cos_cache, sin_cache + + def make_layer(self, layer_id, layer): + super().make_layer(layer_id, layer) + + # Apply the per-layer residual multiplier (`layer_scalar`) to the whole + # layer output, matching HF's `hidden_states *= self.layer_scalar`. + # + # In the SkipLayerNorm design the layer output is carried as + # (root_input + skip_input), summed inside the next SkipLayerNorm (or the + # final norm). To scale the sum, scale both operands. `last_layernorm` is + # set by the super() call for the final layer; the final norm consumes the + # same (root_input, skip_input) pair, so scaling both is correct there too. + # + # Gemma2/3 keep the residual (root_input, via SkipLayerNorm output_3) in + # fp32 while skip_input is io_dtype, so each Mul must use its operand's own + # recorded dtype (and a matching scalar constant). + scalar = float(layer.layer_scalar.item()) + + for suffix, attr in (("skip", "skip_input"), ("root", "root_input")): + operand = self.layernorm_attrs[attr] + operand_dtype = self.values[operand].dtype + mul_name = f"/model/layers.{layer_id}/layer_scalar/Mul_{suffix}" + self.make_mul( + mul_name, + [operand, f"/model/constants/{self.to_str_dtype(operand_dtype)}/{scalar}"], + dtype=operand_dtype, + shape=["batch_size", "sequence_length", self.hidden_size], + ) + self.layernorm_attrs[attr] = f"{mul_name}/output_0" + + +class Gemma4MoEModel(Gemma4Model): + """Builder for the text decoder of Gemma4 MoE (gemma-4-26B-A4B-it). + + Inherits the entire attention / RoPE / per-layer-KV / layer_scalar stack from + the dense `Gemma4Model`. The only structural difference is the FFN: every layer + runs a dense MLP AND a top-k expert MoE block in parallel, then sums them: + + residual = hidden # after attention + post_attn norm + h1 = post_ffn_norm_1(mlp(pre_ffn_norm(residual))) # dense branch + h2 = post_ffn_norm_2(experts(pre_ffn_norm_2(residual))) # MoE branch + h = post_ffn_norm(h1 + h2) + hidden = residual + h + hidden *= layer_scalar + + Only `make_mlp` is overridden: the parent's LayerNorm/residual threading (and + layer_scalar) then work unchanged, with `make_mlp` producing the combined + dense+MoE contribution as the FFN block's `skip_input`. + """ + + def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): + # The base MoE init (base.Model) reads num_local_experts / num_experts_per_tok. + if not hasattr(config, "num_local_experts"): + config.num_local_experts = config.num_experts + if not hasattr(config, "num_experts_per_tok"): + config.num_experts_per_tok = config.top_k_experts + + # Base __init__ prefers config.moe_intermediate_size when setting self.intermediate_size, + # but Gemma4 keeps a parallel DENSE MLP alongside the experts and needs self.intermediate_size + # to stay the dense size. Hide moe_intermediate_size across super().__init__ so base sets the + # dense value (config.intermediate_size) directly, then restore it and track the expert size + # separately. (Assigning self.intermediate_size after super() would overwrite the base value.) + moe_intermediate_size = config.moe_intermediate_size + del config.moe_intermediate_size + super().__init__(config, io_dtype, onnx_dtype, ep, cache_dir, extra_options) + config.moe_intermediate_size = moe_intermediate_size + self.moe_intermediate_size = moe_intermediate_size + + # Gemma4 experts use GeGLU (gelu_pytorch_tanh(gate) * up). This maps to the fused QMoE + # op's "geglu" activation with swiglu_fusion=1 (interleaved gate|up), which applies the + # gelu-tanh gate on the gate half. HF renormalizes the selected top-k weights to sum to 1. + self.moe_attrs["activation_type"] = "geglu" + # Float re-quant path interleaves gate|up (swiglu_fusion=1). The pre-quantized Quark + # path re-fuses experts as gate|up CONCAT (matching the factored checkpoint's layout), + # which the CUDA op consumes with swiglu_fusion=2 (fused, non-interleaved). The CPU QMoE + # kernel only supports the interleaved layout (swiglu_fusion=1), so for the CPU quark path + # we interleave the fc1 expert rows at build time (see make_moe_quark) and use fusion=1. + if self.quant_type == "quark": + self.moe_attrs["swiglu_fusion"] = 1 if self.ep == "cpu" else 2 + else: + self.moe_attrs["swiglu_fusion"] = 1 + self.moe_attrs["normalize_routing_weights"] = True + + # The pre-quantized Quark experts are group-wise (asymmetric, per-group scales/zero + # points). Emit block_size so the QMoE op interprets the 3D block-wise scales/zero + # points ([E, out, in/group_size]) instead of treating them as per-row. + if self.quant_type == "quark": + quant_config = self.quant_attrs["config"] + group_size = quant_config["global_quant_config"]["weight"]["group_size"] + self.moe_attrs["block_size"] = group_size + + # MoE layers emit MoE/QMoE ops instead of dense /mlp/ MatMuls for the experts, but + # the parallel DENSE mlp is still a normal MatMul path — keep its mixed-precision + # overrides. (No pruning needed here, unlike pure-MoE models.) + + def load_weights(self, input_path): + # Same text-only remap as the dense model, but the MoE checkpoint is a + # `Gemma4ForConditionalGeneration` whose text tower is `Gemma4ForCausalLM`. + if self.quant_type is not None or input_path.endswith(".gguf"): + # Deliberately skip Gemma4Model.load_weights (the dense text-only + # remap) and use the generic base loader for quantized/GGUF inputs. + # Called explicitly rather than via super() because super() would + # resolve to Gemma4Model.load_weights and change behavior. + return Model.load_weights(self, input_path) + + import glob + + from safetensors import safe_open + from transformers import AutoConfig + from transformers.models.gemma4 import Gemma4ForCausalLM + + config = AutoConfig.from_pretrained( + self.model_name_or_path, token=self.hf_token, trust_remote_code=self.hf_remote + ) + text_config = config.text_config + text_config.num_hidden_layers = self.num_layers + text_config.layer_types = text_config.layer_types[: self.num_layers] + + with torch.device("meta"): + model = Gemma4ForCausalLM(text_config) + + prefix = "model.language_model." + state_dict = {} + for shard in sorted(glob.glob(os.path.join(self.model_name_or_path, "*.safetensors"))): + with safe_open(shard, framework="pt") as f: + for key in f.keys(): + if not key.startswith(prefix): + continue + new_key = key[len(prefix) :] + if new_key.startswith("layers."): + if int(new_key.split(".")[1]) >= self.num_layers: + continue + state_dict["model." + new_key] = f.get_tensor(key) + + if getattr(text_config, "tie_word_embeddings", False): + state_dict["lm_head.weight"] = state_dict["model.embed_tokens.weight"] + + missing, unexpected = model.load_state_dict(state_dict, strict=False, assign=True) + if missing: + raise ValueError(f"Missing weights while loading Gemma4 MoE text model: {missing}") + if unexpected: + raise ValueError(f"Unexpected weights while loading Gemma4 MoE text model: {unexpected}") + + return model + + def make_layer(self, layer_id, layer): + # Stash the full layer so make_mlp can reach the router/experts/extra norms + # (the parent only passes layer.mlp to make_mlp). + self._current_layer = layer + super().make_layer(layer_id, layer) + self._current_layer = None + + def make_gemma4_rmsnorm(self, name, root_input, weight, with_scale=True): + # Standalone Gemma4 RMSNorm (SimplifiedLayerNormalization, fp32 accumulation) + # on a [B, S, H] tensor. `weight` is the HF parameter (or None if scaleless). + weight_name = f"model.{name}.weight" + if with_scale: + self.make_initializer(weight + self.layernorm_attrs["add_offset"], weight_name, to=self.io_dtype) + else: + self.make_initializer(torch.ones(self.hidden_size), weight_name, to=self.io_dtype) + # Node paths use '/' separators (initializer names keep the '.'-joined form). + ln_name = f"/model/{name.replace('.', '/')}/SimplifiedLayerNormalization" + output = f"{ln_name}/output_0" + self.make_node( + "SimplifiedLayerNormalization", + inputs=[root_input, weight_name], + outputs=[output], + name=ln_name, + epsilon=self.layernorm_attrs["epsilon"], + axis=-1, + stash_type=1, + ) + self.make_value(output, self.io_dtype, shape=["batch_size", "sequence_length", self.hidden_size]) + return output + + def make_mlp(self, layer_id, mlp, root_input): + # Build the parallel dense-MLP + MoE FFN. `root_input` is + # pre_feedforward_layernorm(residual); the pre-FFN residual itself is carried + # in layernorm_attrs["root_input"] (output_3 of the pre-FFN SkipLayerNorm). + layer = self._current_layer + residual = self.layernorm_attrs["root_input"] + + # --- Dense branch: mlp(root_input) -> post_feedforward_layernorm_1 --- + super().make_mlp(layer_id, mlp, root_input) + dense_out = self.layernorm_attrs["skip_input"] + h1 = self.make_gemma4_rmsnorm( + f"layers.{layer_id}.post_feedforward_layernorm_1", dense_out, layer.post_feedforward_layernorm_1.weight + ) + + # --- MoE branch on the pre-FFN residual --- + expert_input = self.make_gemma4_rmsnorm( + f"layers.{layer_id}.pre_feedforward_layernorm_2", residual, layer.pre_feedforward_layernorm_2.weight + ) + moe_out = self.make_gemma4_moe(layer_id, layer, router_input=residual, expert_input=expert_input) + h2 = self.make_gemma4_rmsnorm( + f"layers.{layer_id}.post_feedforward_layernorm_2", moe_out, layer.post_feedforward_layernorm_2.weight + ) + + # --- Combine dense + MoE; the sum is the FFN block contribution --- + combine_name = f"/model/layers.{layer_id}/ffn_combine/Add" + self.make_add( + combine_name, [h1, h2], dtype=self.io_dtype, shape=["batch_size", "sequence_length", self.hidden_size] + ) + self.layernorm_attrs["skip_input"] = f"{combine_name}/output_0" + + def make_gemma4_moe(self, layer_id, layer, router_input, expert_input): + # Bespoke Gemma4 MoE builder (distinct signature from base Model.make_moe, + # which it does not override): the dense+MoE parallel FFN is driven from + # make_mlp, so this takes the full layer plus separate router/expert inputs. + # Pre-quantized Quark experts take a dedicated path (weights are already 2-bit; no + # float re-quantization). The float path below re-quantizes HF float experts. + if self.quant_type == "quark": + return self.make_moe_quark(layer_id, layer, router_input, expert_input) + + # Router pre-projection (built as explicit nodes) feeds raw logits to the fused + # MoE/QMoE op; the op runs the experts on `expert_input` and does topk+softmax + # internally. Returns the MoE op output tensor name. + basename = f"/model/layers.{layer_id}/moe" + op_type = self.moe_attrs["moe_op_type"] + num_experts = self.moe_attrs["num_experts"] + router = layer.router + + # scaleless RMSNorm(router_input) * router.scale * hidden^-0.5 + router_norm = self.make_gemma4_rmsnorm( + f"layers.{layer_id}.moe.router.norm", router_input, None, with_scale=False + ) + scale_vec = router.scale * (self.hidden_size**-0.5) + scale_name = f"model.layers.{layer_id}.moe.router.scale" + self.make_initializer(scale_vec, scale_name, to=self.io_dtype) + scale_mul_name = f"{basename}/router/scale/Mul" + self.make_mul( + scale_mul_name, + [router_norm, scale_name], + dtype=self.io_dtype, + shape=["batch_size", "sequence_length", self.hidden_size], + ) + + # Router projection -> logits, reshaped to [tokens, num_experts] + router_matmul_name = self.make_matmul(router.proj, f"{basename}/router/MatMul", f"{scale_mul_name}/output_0") + router_reshape_name = f"{basename}/router/Reshape" + self.make_reshape( + router_reshape_name, + [f"{router_matmul_name}/output_0", f"/model/constants/INT64/{[-1, num_experts]}"], + dtype=self.io_dtype, + shape=["batch_size * sequence_length", num_experts], + ) + + # Expert weights. HF: experts.gate_up_proj [E, 2*inter, hidden] (gate|up concat), + # experts.down_proj [E, hidden, inter]. Fold per_expert_scale into down_proj (a pure + # per-expert output constant, exactly equivalent to scaling each expert's output). + moe_weight_type = f"{'q' if op_type == 'QMoE' else ''}weight" + gate_up_proj_weight = f"model.layers.{layer_id}.moe.experts.gate_up_proj.{moe_weight_type}" + gate_up_proj_scales = f"model.layers.{layer_id}.moe.experts.gate_up_proj.scales" + gate_up_proj_bias = f"model.layers.{layer_id}.moe.experts.gate_up_proj.bias" + down_proj_weight = f"model.layers.{layer_id}.moe.experts.down_proj.{moe_weight_type}" + down_proj_scales = f"model.layers.{layer_id}.moe.experts.down_proj.scales" + down_proj_bias = f"model.layers.{layer_id}.moe.experts.down_proj.bias" + + raw_gate_up = layer.experts.gate_up_proj + half = raw_gate_up.shape[1] // 2 + # Interleave [gate|up] -> [g0,u0,g1,u1,...] for swiglu_fusion=1. + interleaved = torch.stack([raw_gate_up[:, :half, :], raw_gate_up[:, half:, :]], dim=2).reshape_as(raw_gate_up) + down_scaled = layer.experts.down_proj * layer.router.per_expert_scale.reshape(-1, 1, 1) + + if op_type == "MoE": + self.make_initializer(interleaved, gate_up_proj_weight, to=self.io_dtype) + self.make_initializer(down_scaled, down_proj_weight, to=self.io_dtype) + else: + gate_up_qw, gate_up_sc, down_qw, down_sc = [], [], [], [] + for i in range(num_experts): + qw1, sc1 = self.make_qmoe_weights(interleaved[i]) + gate_up_qw.append(qw1) + gate_up_sc.append(sc1) + qw2, sc2 = self.make_qmoe_weights(down_scaled[i]) + down_qw.append(qw2) + down_sc.append(sc2) + self.make_initializer(torch.stack(gate_up_qw, dim=0).to(torch.uint8), gate_up_proj_weight) + self.make_initializer(torch.stack(down_qw, dim=0).to(torch.uint8), down_proj_weight) + self.make_initializer(torch.stack(gate_up_sc, dim=0), gate_up_proj_scales, to=self.io_dtype) + self.make_initializer(torch.stack(down_sc, dim=0), down_proj_scales, to=self.io_dtype) + + # Experts have no bias; the op still expects the (empty) bias inputs. + self.make_initializer( + torch.zeros(num_experts, 2 * self.moe_intermediate_size), gate_up_proj_bias, to=self.io_dtype + ) + self.make_initializer(torch.zeros(num_experts, self.hidden_size), down_proj_bias, to=self.io_dtype) + + moe_name = f"{basename}/{op_type}" + self.make_moe_op( + moe_name, + root_input=expert_input, + router_probs=f"{router_reshape_name}/output_0", + weight1=gate_up_proj_weight, + scales1=gate_up_proj_scales if op_type == "QMoE" else "", + bias1=gate_up_proj_bias, + weight2=down_proj_weight, + scales2=down_proj_scales if op_type == "QMoE" else "", + bias2=down_proj_bias, + ) + return f"{moe_name}/output_0" + + def make_moe_quark(self, layer_id, layer, router_input, expert_input): + """QMoE from pre-quantized Quark uint2 experts (2-bit, split gate/up re-fused offline). + + The QuarkModel loader has already re-fused each layer's split experts into + `experts.fc1_weights/fc1_scales/fc1_zero_points` (gate|up CONCAT, [E, 2*inter, hidden/pack]) + and `experts.fc2_*` ([E, hidden, inter/pack]), with float zero_points. Here we: + - build the router pre-projection subgraph (scaleless RMSNorm * scale * hidden^-0.5 -> proj), + - fold router.per_expert_scale into fc2 scales (pure per-expert output constant), + - emit the shared gate/up input transform (prescale + rotation) ONCE before QMoE + (the experts are stored in the rotated+prescaled domain; down is plain), + - emit the QMoE op with weights/scales/zero_points. + """ + basename = f"/model/layers.{layer_id}/moe" + num_experts = self.moe_attrs["num_experts"] + experts = layer.mlp.experts + router = layer.router + + # --- Router pre-projection: scaleless RMSNorm(router_input) * router.scale * hidden^-0.5 --- + router_norm = self.make_gemma4_rmsnorm( + f"layers.{layer_id}.moe.router.norm", router_input, None, with_scale=False + ) + scale_vec = router.scale * (self.hidden_size**-0.5) + scale_name = f"model.layers.{layer_id}.moe.router.scale" + self.make_initializer(scale_vec, scale_name, to=self.io_dtype) + scale_mul_name = f"{basename}/router/scale/Mul" + self.make_mul( + scale_mul_name, + [router_norm, scale_name], + dtype=self.io_dtype, + shape=["batch_size", "sequence_length", self.hidden_size], + ) + router_matmul_name = self.make_matmul(router.proj, f"{basename}/router/MatMul", f"{scale_mul_name}/output_0") + router_reshape_name = f"{basename}/router/Reshape" + self.make_reshape( + router_reshape_name, + [f"{router_matmul_name}/output_0", f"/model/constants/INT64/{[-1, num_experts]}"], + dtype=self.io_dtype, + shape=["batch_size * sequence_length", num_experts], + ) + + # --- Shared gate/up input transform: x_rot = (x * input_prescale) @ shared_input_rotation --- + # Only the factored LoRA-KD checkpoint stores experts in a prescaled+rotated domain; it + # carries a per-input `input_prescale` (byte-identical across experts, gate==up). Plain + # Quark/AWQ experts have no prescale/rotation, so the MoE input feeds QMoE directly. + expert0 = experts[sorted(experts.keys())[0]] + has_input_transform = expert0.gate_proj.input_prescale is not None + if has_input_transform: + prescale_name = f"model.layers.{layer_id}.moe.experts.input_prescale" + self.make_initializer(expert0.gate_proj.input_prescale, prescale_name, to=self.io_dtype) + prescale_mul_name = f"{basename}/experts/input_prescale/Mul" + self.make_mul( + prescale_mul_name, + [expert_input, prescale_name], + dtype=self.io_dtype, + shape=["batch_size", "sequence_length", self.hidden_size], + ) + rot_init = f"model.shared_input_rotation_{self.hidden_size}" + if rot_init not in self.shared_rotation_initializers: + self.make_initializer(self.shared_input_rotations[self.hidden_size], rot_init, to=self.io_dtype) + self.shared_rotation_initializers.add(rot_init) + rot_matmul_name = f"{basename}/experts/shared_input_rotation/MatMul" + self.make_node( + "MatMul", + inputs=[f"{prescale_mul_name}/output_0", rot_init], + outputs=[f"{rot_matmul_name}/output_0"], + name=rot_matmul_name, + ) + self.make_value( + f"{rot_matmul_name}/output_0", self.io_dtype, shape=["batch_size", "sequence_length", self.hidden_size] + ) + moe_input = f"{rot_matmul_name}/output_0" + else: + moe_input = expert_input + + # --- Fold router.per_expert_scale into fc2 (down) scales: pure per-expert output constant --- + per_expert = layer.router.per_expert_scale.to(experts.fc2_scales.dtype).reshape(-1, 1, 1) + fc2_scales = experts.fc2_scales * per_expert + + # --- Emit expert weight / scale / zero-point initializers --- + gate_up_proj_weight = f"model.layers.{layer_id}.moe.experts.gate_up_proj.qweight" + gate_up_proj_scales = f"model.layers.{layer_id}.moe.experts.gate_up_proj.scales" + gate_up_proj_zero_points = f"model.layers.{layer_id}.moe.experts.gate_up_proj.zero_points" + gate_up_proj_bias = f"model.layers.{layer_id}.moe.experts.gate_up_proj.bias" + down_proj_weight = f"model.layers.{layer_id}.moe.experts.down_proj.qweight" + down_proj_scales = f"model.layers.{layer_id}.moe.experts.down_proj.scales" + down_proj_zero_points = f"model.layers.{layer_id}.moe.experts.down_proj.zero_points" + down_proj_bias = f"model.layers.{layer_id}.moe.experts.down_proj.bias" + + # The CPU QMoE kernel only supports the interleaved gate|up layout (swiglu_fusion=1). + # The 2-bit factored loader re-fuses fc1 as [gate(inter), up(inter)] CONCAT along the + # output dim, so for the CPU path we reorder those rows to interleaved + # ([gate0,up0,gate1,up1,...]) so the op's interleaved activation reads the correct + # gate/up pairs. fc1 rows are independent (quantization packs the input dim), so this is + # a pure row permutation on weights/scales/zero_points. Plain Quark/AWQ int4 experts are + # already emitted interleaved by `combine_and_repack_gate_up`, so no reorder is needed. + is_int2 = int(self.moe_attrs["expert_weight_bits"]) == 2 + fc1_weights = experts.fc1_weights + fc1_scales = experts.fc1_scales + fc1_zero_points = experts.fc1_zero_points + interleave_fc1 = self.quant_type == "quark" and self.ep == "cpu" and is_int2 + if interleave_fc1: + inter = self.moe_intermediate_size + + def _concat_to_interleaved(t): + # dim 1 is [gate(inter), up(inter)] -> [gate0,up0,gate1,up1,...] + gate, up = t[:, :inter], t[:, inter:] + return torch.stack((gate, up), dim=2).reshape(t.shape[0], 2 * inter, *t.shape[2:]) + + fc1_weights = _concat_to_interleaved(fc1_weights) + fc1_scales = _concat_to_interleaved(fc1_scales) + fc1_zero_points = _concat_to_interleaved(fc1_zero_points) + + self.make_initializer(fc1_weights, gate_up_proj_weight) + self.make_initializer(experts.fc2_weights, down_proj_weight) + self.make_initializer(fc1_scales, gate_up_proj_scales, to=self.io_dtype) + self.make_initializer(fc2_scales, down_proj_scales, to=self.io_dtype) + + # Experts have no bias; the op still expects the (empty) bias inputs. + self.make_initializer( + torch.zeros(num_experts, 2 * self.moe_intermediate_size), gate_up_proj_bias, to=self.io_dtype + ) + self.make_initializer(torch.zeros(num_experts, self.hidden_size), down_proj_bias, to=self.io_dtype) + + # zero_points: the Quark uint2 export is symmetric with a constant zp of 1.5 + # (codes {0,1,2,3} -> {-1.5,-0.5,0.5,1.5}*scale) stored as FLOAT per-group zero-points. + # On CUDA the GeGLU QMoE op reconstructs the -1.5*scale bias internally from scales when + # zp is omitted (bits==2, no zp input), so we emit NO zero_points tensor there. CPU keeps + # the float zp inputs as-is. Plain Quark/AWQ int4 experts carry INTEGER (uint8) per-group + # zero-points, which are emitted without an io_dtype cast. trt-rtx never supports ZP inputs. + zp_is_float = torch.is_floating_point(fc1_zero_points) + omit_zero_points = self.ep == "trt-rtx" or (self.ep == "cuda" and is_int2) + use_zero_points = not omit_zero_points + if use_zero_points: + zp_to = self.io_dtype if zp_is_float else None + self.make_initializer(fc1_zero_points, gate_up_proj_zero_points, to=zp_to) + self.make_initializer(experts.fc2_zero_points, down_proj_zero_points, to=zp_to) + + moe_name = f"{basename}/QMoE" + self.make_moe_op( + moe_name, + root_input=moe_input, + router_probs=f"{router_reshape_name}/output_0", + weight1=gate_up_proj_weight, + scales1=gate_up_proj_scales, + bias1=gate_up_proj_bias, + weight2=down_proj_weight, + scales2=down_proj_scales, + bias2=down_proj_bias, + zero_points1=gate_up_proj_zero_points if use_zero_points else "", + zero_points2=down_proj_zero_points if use_zero_points else "", + ) + return f"{moe_name}/output_0" diff --git a/src/python/py/models/loaders/base.py b/src/python/py/models/loaders/base.py index 3cc4b6ce10..38516ab244 100644 --- a/src/python/py/models/loaders/base.py +++ b/src/python/py/models/loaders/base.py @@ -16,6 +16,7 @@ import os import re +import json import torch from safetensors.torch import load_file @@ -33,6 +34,8 @@ def __init__(self): self.out_features = 0 self.bits = None self.group_size_value = None + # Factored-rotation per-input prescale (Quark rotation algo); None if absent. + self.input_prescale = None @property def group_size(self): @@ -198,6 +201,18 @@ def __str__(self): return "\n".join(lines) +class QuantizedRouter: + """MoE router: bias-free projection + scaleless-RMSNorm scale + per-expert output scale. + + Used by architectures (e.g. Gemma4 MoE) whose router pre-projection is more than a + plain Linear. `proj` is a TensorModule so the builder can emit it via `make_matmul`. + """ + def __init__(self): + self.proj = TensorModule() + self.scale = None + self.per_expert_scale = None + + class QuantizedMLP: def __init__(self): self.gate_proj = QuantizedTensorModule() @@ -218,6 +233,13 @@ def __init__(self, layer_id): self.post_attention_layernorm = TensorModule() self.pre_feedforward_layernorm = TensorModule() self.post_feedforward_layernorm = TensorModule() + # Extra Gemma4-MoE norms (parallel dense + MoE FFN); unused by other archs. + self.pre_feedforward_layernorm_2 = TensorModule() + self.post_feedforward_layernorm_1 = TensorModule() + self.post_feedforward_layernorm_2 = TensorModule() + # Gemma4-MoE per-layer residual multiplier + router; experts live under mlp.experts. + self.layer_scalar = None + self.router = QuantizedRouter() self.mlp = QuantizedMLP() def is_empty(self): @@ -245,6 +267,13 @@ def __init__( self.lm_head = lm_head if lm_head is not None else TensorModule() self.layers = {} if load_weights else [] self.num_layers = num_layers + self.q_size = q_size + self.kv_size = kv_size + self.intermediate_size = intermediate_size + # Factored online rotation (Quark rotation algo): x_rot = (x / input_prescale) @ shared_input_rotation_. + # `shared_input_rotations` maps in_features -> [in, in] rotation matrix (shared across all layers). + self.input_path = input_path + self.shared_input_rotations = {} if not load_weights: return @@ -298,6 +327,10 @@ def __init__( # transformer.rotary_pos_emb.inv_freq in ChatGLM3. # Skip rotary embedding weights since they can be re-calculated when looping through the model continue + elif name.startswith("shared_input_rotation_"): + # Model-level shared input-rotation matrix keyed by in_features (Quark factored rotation). + in_features = int(name.rsplit("_", 1)[1]) + self.shared_input_rotations[in_features] = tensor else: if name.startswith("transformer.encoder"): # Chatglm3, e.g., transformer.encoder.layers.0.input_layernorm.weight @@ -525,8 +558,8 @@ def __init__( tensor_map["self_attn.v_proj.qweight"] = tensor[q_size + kv_size :, :] else: # AWQ/GPTQ/Quark: (in_features, out_features), split on dim=1 - q_dim = q_size // (32 // local_bits) if quant_type in {"awq", "quark"} else q_size - kv_dim = kv_size // (32 // local_bits) if quant_type in {"awq", "quark"} else kv_size + q_dim = q_size // self._packed_out_factor(tensor, local_bits) if quant_type in {"awq", "quark"} else q_size + kv_dim = kv_size // self._packed_out_factor(tensor, local_bits) if quant_type in {"awq", "quark"} else kv_size tensor_map["self_attn.q_proj.qweight"] = tensor[:, :q_dim] tensor_map["self_attn.k_proj.qweight"] = tensor[:, q_dim : q_dim + kv_dim] tensor_map["self_attn.v_proj.qweight"] = tensor[:, q_dim + kv_dim :] @@ -570,10 +603,14 @@ def __init__( else: # AWQ/GPTQ/Quark: int32 packing, split on dim=1 q_dim = ( - q_size // (32 // local_bits) if quant_type in {"awq", "gptq", "quark"} else q_size + q_size // self._packed_out_factor(tensor, local_bits) + if quant_type in {"awq", "gptq", "quark"} + else q_size ) kv_dim = ( - kv_size // (32 // local_bits) if quant_type in {"awq", "gptq", "quark"} else kv_size + kv_size // self._packed_out_factor(tensor, local_bits) + if quant_type in {"awq", "gptq", "quark"} + else kv_size ) tensor_map["self_attn.q_proj.qzeros"] = tensor[:, :q_dim] tensor_map["self_attn.k_proj.qzeros"] = tensor[:, q_dim : q_dim + kv_dim] @@ -612,7 +649,7 @@ def __init__( else: # AWQ/GPTQ/Quark: (in_features, out_features), split on dim=1 intermediate_dim = ( - intermediate_size // (32 // local_bits) + intermediate_size // self._packed_out_factor(tensor, local_bits) if quant_type in {"awq", "quark"} else intermediate_size ) @@ -654,7 +691,7 @@ def __init__( else: # AWQ/GPTQ/Quark: int32 packing, split on dim=1 intermediate_dim = ( - intermediate_size // (32 // local_bits) + intermediate_size // self._packed_out_factor(tensor, local_bits) if quant_type in {"awq", "gptq", "quark"} else intermediate_size ) @@ -691,6 +728,51 @@ def __init__( elif bool(re.match(r"^model.layers\.\d+\.self_attn\.sinks$", name)): # model.layers.layer_id.self_attn.sinks tensor_map["self_attn.sinks"] = tensor + elif bool(re.match(r"^model.layers\.\d+\.(self_attn.qkv_proj|self_attention.query_key_value)\.input_prescale$", name)): + # Factored-rotation per-input scale, shared by the q/k/v splits (same input). + tensor_map["self_attn.q_proj.input_prescale"] = tensor + tensor_map["self_attn.k_proj.input_prescale"] = tensor + tensor_map["self_attn.v_proj.input_prescale"] = tensor + elif bool(re.match(r"^model.layers\.\d+\.self_attn\.(q_proj|k_proj|v_proj|o_proj)\.input_prescale$", name)): + tensor_map["self_attn." + name.split(".")[-2] + ".input_prescale"] = tensor + elif bool(re.match(r"^model.layers\.\d+\.mlp.(gate_up_proj|dense_h_to_4h)\.input_prescale$", name)): + # Shared by the gate/up splits (same input). + tensor_map["mlp.gate_proj.input_prescale"] = tensor + tensor_map["mlp.up_proj.input_prescale"] = tensor + elif bool(re.match(r"^model.layers\.\d+\.mlp.(gate_proj|up_proj|down_proj|dense_4h_to_h)\.input_prescale$", name)): + leaf = name.split(".")[-2] + leaf = "down_proj" if leaf == "dense_4h_to_h" else leaf + tensor_map["mlp." + leaf + ".input_prescale"] = tensor + # --- Gemma4 MoE: experts live directly under the layer (not `mlp.experts`) --- + elif bool(re.match(r"^model\.layers\.\d+\.experts\.\d+\.(gate_proj|up_proj|gate_up_proj|down_proj)\.(weight|bias|qweight|scales|qzeros|weight_scale|weight_zero_point|g_idx)$", name)): + # model.layers.layer_id.experts.expert_id.proj_type.param_type + split_name = name.split(".") + expert_id = int(split_name[4]) + proj_type = split_name[-2] + param_type = split_name[-1] + module.mlp.experts.set_weight_data(expert_id, proj_type, param_type, tensor, local_bits, local_group_size) + elif bool(re.match(r"^model\.layers\.\d+\.experts\.\d+\.(gate_proj|up_proj|down_proj)\.input_prescale$", name)): + # Per-expert factored-rotation input scale. Byte-identical across all experts and + # gate==up within a layer, so store one shared prescale on the router for the MoE input. + split_name = name.split(".") + expert_id = int(split_name[4]) + proj_type = split_name[-2] + module.mlp.experts.set_weight_data(expert_id, proj_type, "input_prescale", tensor, local_bits, local_group_size) + elif bool(re.match(r"^model\.layers\.\d+\.router\.proj\.weight$", name)): + # Router projection is a plain (non-quantized) bf16 Linear -> float MatMul. + tensor_map["router.proj.weight"] = tensor + elif bool(re.match(r"^model\.layers\.\d+\.router\.scale$", name)): + tensor_map["router.scale"] = tensor + elif bool(re.match(r"^model\.layers\.\d+\.router\.per_expert_scale$", name)): + tensor_map["router.per_expert_scale"] = tensor + elif bool(re.match(r"^model\.layers\.\d+\.layer_scalar$", name)): + tensor_map["layer_scalar"] = tensor + elif bool(re.match(r"^model\.layers\.\d+\.pre_feedforward_layernorm_2\.weight$", name)): + tensor_map["pre_feedforward_layernorm_2.weight"] = tensor + elif bool(re.match(r"^model\.layers\.\d+\.post_feedforward_layernorm_1\.weight$", name)): + tensor_map["post_feedforward_layernorm_1.weight"] = tensor + elif bool(re.match(r"^model\.layers\.\d+\.post_feedforward_layernorm_2\.weight$", name)): + tensor_map["post_feedforward_layernorm_2.weight"] = tensor else: raise NotImplementedError(f"{name} in your quantized model is not recognized.") @@ -724,9 +806,20 @@ def __init__( # Set properties of each layer based on quantization type self.set_properties() + # Bake additive PEFT LoRA adapters into the projections (if present). + self._load_lora_adapters() + def normalize_weight_name(self, name): """Normalize a checkpoint tensor key to the shared model structure.""" - if name.startswith(("model.visual.", "model.vision.", "visual.")): + if name.startswith(( + "model.visual.", + "model.vision.", + "visual.", + "model.vision_tower.", + "model.embed_vision.", + "model.audio_tower.", + "model.embed_audio.", + )): return None if name.startswith("model.language_model."): name = "model." + name[len("model.language_model.") :] @@ -734,6 +827,97 @@ def normalize_weight_name(self, name): name = name.replace(old, new) return name + def _load_lora_adapters(self): + """Attach PEFT LoRA adapters (baked additively into the graph) if present. + + The adapter lives at ``/lora_adapters/adapter_model.safetensors`` + with keys ``base_model.model.model.layers.{i}.{proj}.lora_A.weight`` [r, in] + and ``.lora_B.weight`` [out, r] for proj in {qkv_proj, o_proj, gate_up_proj, + down_proj}. ``lora_A`` is shared across split projections that share an input + (q/k/v share the qkv input, gate/up share the gate_up input); ``lora_B`` is + split along its output dimension. The builder emits the runtime delta + ``(lora_B @ lora_A @ x) * scaling`` added to the quantized projection output. + """ + adapter_dir = os.path.join(self.input_path, "lora_adapters") + adapter_path = os.path.join(adapter_dir, "adapter_model.safetensors") + if not os.path.exists(adapter_path): + return + + scaling = 1.0 + config_path = os.path.join(adapter_dir, "adapter_config.json") + if os.path.exists(config_path): + with open(config_path) as config_file: + adapter_config = json.load(config_file) + rank = adapter_config.get("r") + alpha = adapter_config.get("lora_alpha") + if rank: + scaling = alpha / (rank ** 0.5) if adapter_config.get("use_rslora", False) else alpha / rank + + weights = load_file(adapter_path) + layers_by_id = {layer.layer_id: layer for layer in self.layers} + + for name, tensor in weights.items(): + # Normalize the VLM prefix (Gemma4 stores the text tower under `language_model.`) + # so the same regex works for both flat-LLM and VLM adapters. + norm = name.replace("model.language_model.", "model.") + # Expert LoRA has no slot in the fused QMoE op and is intentionally dropped + # (unbakeable rank-64 residual; the dense/attention LoRA carries the recovery). + if re.match(r"^base_model\.model\.model\.layers\.\d+\.experts\.\d+\.", norm): + continue + match = re.match( + r"^base_model\.model\.model\.layers\.(\d+)\." + r"(self_attn\.qkv_proj|self_attn\.q_proj|self_attn\.k_proj|self_attn\.v_proj|self_attn\.o_proj|" + r"mlp\.gate_up_proj|mlp\.gate_proj|mlp\.up_proj|mlp\.down_proj)\." + r"(lora_A|lora_B)\.weight$", + norm, + ) + if match is None: + raise NotImplementedError(f"{name} in the LoRA adapter is not recognized.") + layer_id, proj, ab = int(match.group(1)), match.group(2), match.group(3) + layer = layers_by_id.get(layer_id) + if layer is None: + continue + + # Map the adapter projection to the builder's split modules. Fused adapter names + # (qkv_proj / gate_up_proj) fan out to multiple split targets; already-split names + # map 1:1. + if proj == "self_attn.qkv_proj": + targets = [layer.self_attn.q_proj, layer.self_attn.k_proj, layer.self_attn.v_proj] + out_splits = [self.q_size, self.kv_size, self.kv_size] + elif proj == "self_attn.q_proj": + targets, out_splits = [layer.self_attn.q_proj], None + elif proj == "self_attn.k_proj": + targets, out_splits = [layer.self_attn.k_proj], None + elif proj == "self_attn.v_proj": + targets, out_splits = [layer.self_attn.v_proj], None + elif proj == "self_attn.o_proj": + targets, out_splits = [layer.self_attn.o_proj], None + elif proj == "mlp.gate_up_proj": + targets = [layer.mlp.gate_proj, layer.mlp.up_proj] + out_splits = [self.intermediate_size, self.intermediate_size] + elif proj == "mlp.gate_proj": + targets, out_splits = [layer.mlp.gate_proj], None + elif proj == "mlp.up_proj": + targets, out_splits = [layer.mlp.up_proj], None + else: # mlp.down_proj + targets, out_splits = [layer.mlp.down_proj], None + + if ab == "lora_A": + # Shared input projection: every split target uses the same lora_A. + for target in targets: + target.lora_A = tensor + target.lora_scaling = scaling + elif out_splits is None: + targets[0].lora_B = tensor + targets[0].lora_scaling = scaling + else: + # lora_B is split along its output dimension across the targets. + start = 0 + for target, size in zip(targets, out_splits): + target.lora_B = tensor[start : start + size, :] + target.lora_scaling = scaling + start += size + def assign_lm_head_tensors(self, lm_head_tensors): """Assign collected lm_head tensors in a defined order so that weight/bias are always processed before quantization parameters (scales, qzeros, etc.), @@ -876,6 +1060,18 @@ def modules(self): """ return [self.embedding] + self.layers + [self.final_norm, self.lm_head] + @staticmethod + def _packed_out_factor(tensor, bits): + """Number of logical output channels stored per element along the packed + (column) axis: 8//bits for uint8 packing (Quark native uint2/uint4), + 32//bits for int32 packing (AWQ/GPTQ style), and 1 for unpacked floating + tensors (per-group float scales / float zero-points).""" + if tensor.dtype == torch.uint8: + return 8 // bits + if tensor.dtype == torch.int32: + return 32 // bits + return 1 + def unpack(self, module): """ Unpack `qzeros` and `qweight` to standard format diff --git a/src/python/py/models/loaders/quark.py b/src/python/py/models/loaders/quark.py index 5a73a448e6..3483d20e1e 100644 --- a/src/python/py/models/loaders/quark.py +++ b/src/python/py/models/loaders/quark.py @@ -15,10 +15,13 @@ class QuarkModel(QuantizedModel): (".weight_quantizer.zero_point", ".weight_zero_point"), ) + # uint2/int2 pack 4 codes per byte (MSB-first); uint4/int4 pack 2 codes per byte. + _DTYPE_BITS = {"uint4": 4, "int4": 4, "uint2": 2, "int2": 2} + def __init__(self, quant_type, input_path, quant_attrs, q_size, kv_size, intermediate_size, num_layers): self.global_quant_config = quant_attrs["config"]["global_quant_config"]["weight"] global_dtype = self.global_quant_config["dtype"] - if global_dtype not in {"uint4", "int4"}: + if global_dtype not in self._DTYPE_BITS: raise ValueError(f"Unexpected dtype: {global_dtype}.") super().__init__( quant_type, @@ -29,19 +32,34 @@ def __init__(self, quant_type, input_path, quant_attrs, q_size, kv_size, interme intermediate_size, num_layers, global_group_size=self.global_quant_config["group_size"], - global_bits=4, + global_bits=self._DTYPE_BITS[global_dtype], ) self.repack_quantized_tensors(clear_g_idx=True) def set_quantized_tensor_properties(self, module): + # bf16 outlier experts (Gemma4 uint2) carry a full-precision weight with no scales here; + # they are re-quantized later in repack_experts and their in/out_features aren't needed. + if module.scales is None: + return module.out_features = module.scales.shape[1] module.in_features = module.qweight.shape[0] self.set_g_idx(module) def repack_experts(self, experts): """ - Unpacks weights from pre-quantized Quark experts + Unpacks weights from pre-quantized Quark experts, then repacks them into the + format expected by the QMoE operator. """ + # Pre-quantized 2-bit split experts (Gemma4 factored uint2 checkpoint): consume the + # group-wise uint2 weights/scales/float zero-points directly and re-fuse gate|up as a + # CONCAT block. The per-expert prescale and per_expert_scale are folded in the builder's + # make_moe_quark, so the native QMoE attrs (finalize_packed_experts) are not used here. + first = next(iter(experts.values())) + is_2bit = first.gate_proj is not None and first.gate_proj.bits == 2 + if is_2bit: + self.unpack_repack_uint2_experts(experts) + return + for expert in experts.values(): # Process gate_proj if expert.gate_proj.qweight is not None: @@ -71,13 +89,133 @@ def repack_experts(self, experts): self.unpack_qweight_quark(expert.down_proj) expert.down_proj.g_idx = None - """ - Repacks weights from pre-quantized Quark experts - into the format expected by the QMoE operator. - """ self.repack_qmoe_weights(experts) self.finalize_packed_experts(experts) + def unpack_repack_uint2_experts(self, experts): + """Unpack + re-fuse pre-quantized 2-bit (uint2) Gemma4 split experts. + + Each expert projection is unpacked into int codes with per-group float scales/zeros + (bf16 outlier experts are re-quantized into the same rotated+prescaled uint2 domain), + then gate|up are concatenated and repacked into the ORT QMoE tensors. + """ + # gate/up prescale is byte-identical across experts and gate==up; find any set one so + # bf16 outliers (which lack their own prescale) can be transformed into the rotated domain. + shared_prescale = None + for expert in experts.values(): + if expert.gate_proj.input_prescale is not None: + shared_prescale = expert.gate_proj.input_prescale + break + + for expert in experts.values(): + self.unpack_expert_proj(expert.gate_proj, is_gate_up=True, prescale=shared_prescale) + self.unpack_expert_proj(expert.up_proj, is_gate_up=True, prescale=shared_prescale) + self.unpack_expert_proj(expert.gate_up_proj, is_gate_up=True, prescale=shared_prescale) + self.unpack_expert_proj(expert.down_proj, is_gate_up=False, prescale=None) + + self.combine_and_repack_split_experts_concat(experts) + + def unpack_expert_proj(self, proj, is_gate_up, prescale): + """Unpack one expert projection into int codes + per-group scales/zeros. + + For 2-bit Quark experts the weight is uint8 ``[in, out/4]`` packed MSB-first + along the output axis and the zero_point is a per-group FLOAT tensor + ``[n_groups, out]`` (dequant = (code - zero_point) * scale). We unpack to int + codes ``[in, out]`` (matching the 4-bit `unpack_qweight_quark` output) and keep + the float zeros as-is. + + A few experts are stored as UNQUANTIZED bf16 outliers (full-precision `[out, in]`, + no scales/zeros/prescale). The fused QMoE op requires all experts uniform, so those + are re-quantized here into the same rotated+prescaled (gate/up) or plain (down) uint2 + domain as their siblings. Non-2-bit checkpoints use the 4-bit AWQ path. + """ + if proj.qweight is None: + return + if proj.bits == 2 and proj.qweight.dtype == torch.uint8 and proj.scales is not None: + # Normal packed uint2: uint8 [in, out/4] MSB-first along out -> int codes [in, out]. + proj.qweight = self._unpack_uint2_msb_first(proj.qweight) + # Float zero points are already per-group [n_groups, out]; keep them. + elif proj.bits == 2: + # bf16 outlier expert: re-quantize into the same uint2 domain as the siblings. + self._requantize_float_expert(proj, is_gate_up, prescale) + else: + self.unpack_qzeros(proj) + self.pack_zeros_ort_format(proj) + self.unpack_qweight_quark(proj) + proj.g_idx = None + + def _requantize_float_expert(self, proj, is_gate_up, prescale): + """Re-quantize a full-precision (bf16) outlier expert projection to uint2. + + gate/up outliers are stored in the plain domain but the op applies a shared + prescale+rotation to the MoE input, so transform into the rotated domain first: + W_rot[in,out] = R^T @ (W_orig^T / prescale) (matches op x_rot=(x*prescale)@R) + down is plain (no transform). Then per-group asymmetric uint2 quant along `in`. + """ + W = proj.qweight.to(torch.float32) # stored as [out, in] (standard Linear) + WT = W.t().contiguous() # [in, out] + if is_gate_up: + R = self.shared_input_rotations[WT.shape[0]].to(torch.float32) # [in, in] + WT = WT / prescale.to(torch.float32).unsqueeze(1) # scale rows by 1/prescale + WT = R.t() @ WT # [in, out] + gs = self.global_group_size + in_dim, out_dim = WT.shape + ng = in_dim // gs + w = WT.reshape(ng, gs, out_dim) + wmin = w.min(dim=1).values # [ng, out] + wmax = w.max(dim=1).values + scale = (wmax - wmin) / 3.0 # 2-bit -> 4 levels (0..3) + scale = torch.where(scale == 0, torch.ones_like(scale), scale) + zp = -wmin / scale # float zp; dequant = (code - zp) * scale + codes = torch.clamp(torch.round(w / scale.unsqueeze(1) + zp.unsqueeze(1)), 0, 3).to(torch.uint8) + proj.qweight = codes.reshape(in_dim, out_dim) + proj.scales = scale.to(torch.float16) + proj.qzeros = zp.to(torch.float16) + proj.group_size = gs + + def combine_and_repack_split_experts_concat(self, experts): + """Re-fuse per-expert split uint2 gate/up/down into ORT QMoE tensors (CONCAT layout). + + Input per expert (int codes after `_unpack_uint2_msb_first`, before transpose): + gate/up.qweight [in=hidden, out=inter]; down.qweight [in=inter, out=hidden] + *.scales / *.qzeros are per-group FLOAT [n_groups, out]. + Output (native Gemma4 fused orientation, gate block then up block): + fc1_weights (E, 2*inter, hidden/pack) packed uint8; fc1_scales/zp (E, 2*inter, n_groups_h) + fc2_weights (E, hidden, inter/pack) packed uint8; fc2_scales/zp (E, hidden, n_groups_i) + per_expert_scale is folded into fc2_scales later in the builder's make_moe_quark. + """ + fc1_w, fc1_s, fc1_z, fc2_w, fc2_s, fc2_z = [], [], [], [], [], [] + for expert_id in sorted(experts.keys()): + e = experts[expert_id] + # [in, out] -> [out, in] + gate_w = e.gate_proj.qweight.T.contiguous() + up_w = e.up_proj.qweight.T.contiguous() + down_w = e.down_proj.qweight.T.contiguous() + # concat gate|up along out: [2*inter, hidden] + fc1_codes = torch.cat([gate_w, up_w], dim=0) + fc1_w.append(self.repack_qweight(fc1_codes, bits=2)) + fc2_w.append(self.repack_qweight(down_w, bits=2)) + # scales / zeros: [n_groups, out] -> [out, n_groups], concat gate|up along out + fc1_s.append(torch.cat([e.gate_proj.scales.T, e.up_proj.scales.T], dim=0)) + fc1_z.append(torch.cat([e.gate_proj.qzeros.T, e.up_proj.qzeros.T], dim=0)) + fc2_s.append(e.down_proj.scales.T) + fc2_z.append(e.down_proj.qzeros.T) + + experts.fc1_weights = torch.stack(fc1_w, dim=0) + experts.fc1_scales = torch.stack(fc1_s, dim=0).to(torch.float16) + experts.fc1_zero_points = torch.stack(fc1_z, dim=0).to(torch.float16) + experts.fc2_weights = torch.stack(fc2_w, dim=0) + experts.fc2_scales = torch.stack(fc2_s, dim=0).to(torch.float16) + experts.fc2_zero_points = torch.stack(fc2_z, dim=0).to(torch.float16) + + @staticmethod + def _unpack_uint2_msb_first(packed): + """Unpack uint8 ``[rows, cols/4]`` (4x 2-bit codes per byte, MSB-first along + columns) into int codes ``[rows, cols]``.""" + shifts = torch.tensor([6, 4, 2, 0], dtype=torch.int32, device=packed.device) + codes = (packed.to(torch.int32).unsqueeze(-1) >> shifts.view(1, 1, -1)) & 0x3 + return codes.reshape(packed.shape[0], -1) + def finalize_packed_experts(self, experts): first_expert = experts[min(experts.keys())] first_projection = ( @@ -246,55 +384,70 @@ def combine_and_repack_gate_up(self, experts): def repack_qweight(self, weights, bits) -> torch.Tensor: """ - Repacks unpacked uint8 weights (representing 4-bit values) into a packed uint8 tensor. - This mirrors the packing logic from builder.py's _symmetric_blockwise_quantize. + Repacks unpacked uint8 int-code weights into the packed uint8 tensor that + MatMulNBits / QMoE consume. Codes are packed low-order-first along the last + dim: 2 codes/byte for 4-bit, 4 codes/byte for 2-bit. Mirrors the packing in + builder.py's `_symmetric_blockwise_quantize`. """ - if bits != 4: - raise NotImplementedError("This repacking function is specifically for 4-bit weights.") + if bits not in (2, 4): + raise NotImplementedError("This repacking function supports 2-bit or 4-bit weights only.") - quantized_flat = weights.cpu() + quantized_flat = weights.cpu().to(torch.uint8) original_shape = quantized_flat.shape - quantized_uint4 = quantized_flat.to(torch.uint8) + pack = 8 // bits # codes per byte (4-bit -> 2, 2-bit -> 4) + mask = (1 << bits) - 1 # 0xF for 4-bit, 0x3 for 2-bit packed_shape = list(original_shape) - packed_shape[-1] = (original_shape[-1] + 1) // 2 + packed_shape[-1] = (original_shape[-1] + pack - 1) // pack packed_weight = torch.zeros(packed_shape, dtype=torch.uint8, device=quantized_flat.device) - # Pack two 4-bit values per byte - for i in range(0, quantized_uint4.shape[-1], 2): - val1 = quantized_uint4[..., i] - if i + 1 < quantized_uint4.shape[-1]: - val2 = quantized_uint4[..., i + 1] - packed_val = (val1 & 0xF) | ((val2 & 0xF) << 4) - else: - # Odd number of values - pack only lower 4 bits - packed_val = val1 & 0xF - packed_weight[..., i // 2] = packed_val + # Pack `pack` codes per byte, low-order code in the least-significant bits. + n = quantized_flat.shape[-1] + for out_idx in range(packed_shape[-1]): + byte = torch.zeros(original_shape[:-1], dtype=torch.uint8, device=quantized_flat.device) + for j in range(pack): + src = out_idx * pack + j + if src >= n: + break + byte = byte | ((quantized_flat[..., src] & mask) << (bits * j)) + packed_weight[..., out_idx] = byte return packed_weight def get_layer_bits(self, layer_name): name = layer_name.split(".")[0] - if name in self.quant_attrs["config"]["layer_quant_config"]: - layer_quant_config = self.quant_attrs["config"]["layer_quant_config"][name]["weight"] - local_dtype = layer_quant_config["dtype"] - - dtype_bits_maps = { - "uint4": 4, - "int4": 4, - } - if local_dtype not in dtype_bits_maps: + layer_quant_config = self.quant_attrs["config"].get("layer_quant_config", {}) + if name in layer_quant_config: + local_dtype = layer_quant_config[name]["weight"]["dtype"] + if local_dtype not in self._DTYPE_BITS: raise ValueError(f"Unexpected dtype: {local_dtype}.") - return dtype_bits_maps[local_dtype] + return self._DTYPE_BITS[local_dtype] return self.global_bits def get_layer_group_size(self, layer_name): name = layer_name.split(".")[0] - if name in self.quant_attrs["config"]["layer_quant_config"]: - layer_quant_config = self.quant_attrs["config"]["layer_quant_config"][name]["weight"] - return layer_quant_config["group_size"] + layer_quant_config = self.quant_attrs["config"].get("layer_quant_config", {}) + if name in layer_quant_config: + return layer_quant_config[name]["weight"]["group_size"] return self.global_group_size + def unpack(self, module): + """ + Unpack a Quark ``QuantizedTensorModule`` into the standard int-code layout. + + Standard Quark native uint2 stores weights as uint8 ``[in, out/4]`` packed + MSB-first along the output axis, with per-group float ``scales`` and a float + ``zero_point`` (e.g. a constant 1.5). These are unpacked + dequantized here so + the shared ``repack``/``pack_ort_format`` pipeline can emit MatMulNBits-ready + tensors (same packing as the 4-bit path). Other bit-widths use the base + (int32-packed, AWQ-reorder) path. + """ + if module.bits == 2: + module.qweight = self._unpack_uint2_msb_first(module.qweight) + self.dequant_weight(module) + else: + super().unpack(module) + def unpack_qweight(self, module): """ Unpack `qweight` to standard format diff --git a/src/python/py/models/quantization/quant_config.py b/src/python/py/models/quantization/quant_config.py index 8a787e5ea0..4fa1b7ba21 100644 --- a/src/python/py/models/quantization/quant_config.py +++ b/src/python/py/models/quantization/quant_config.py @@ -60,6 +60,8 @@ def is_quantized(self) -> bool: "uint8": DtypeDescriptor("uint8", "int", 8, signed=False), "int4": DtypeDescriptor("int4", "int", 4, signed=True), "uint4": DtypeDescriptor("uint4", "int", 4, signed=False), + "int2": DtypeDescriptor("int2", "int", 2, signed=True), + "uint2": DtypeDescriptor("uint2", "int", 2, signed=False), "mxfp4": DtypeDescriptor("mxfp4", "mx", 4, block_size=32), "nvfp4": DtypeDescriptor("nvfp4", "mx", 4, block_size=16), "none": DtypeDescriptor("none", "float", 0), # explicit "do not quantize this target" @@ -290,7 +292,9 @@ def to_dict(self) -> dict[str, Any]: } # ``--precision`` -> weights.type. Float precisions do not quantize weights. +# int2 quantizes dense weights to 2-bit MatMulNBits (asymmetric uint2, matching the Quark export). _PRECISION_TO_WEIGHTS_TYPE = { + "int2": "uint2", "int4": "int4", "int8": "int8", "fp16": "none", @@ -429,6 +433,9 @@ def from_extra_options( # Match the model precision unless the MoE target is configured independently. if precision == "int8" or extra_options.get("use_8bits_moe", False): moe_quant_type = "int8" + elif precision == "int2": + # int2 quantizes MoE experts to 2-bit (asymmetric uint2 QMoE), matching the dense weights. + moe_quant_type = "uint2" elif precision in IO_DTYPES: moe_quant_type = "none" else: diff --git a/test/python/builder/test_qmoe_weights.py b/test/python/builder/test_qmoe_weights.py index 3a47e4ecc1..629ffdf11e 100644 --- a/test/python/builder/test_qmoe_weights.py +++ b/test/python/builder/test_qmoe_weights.py @@ -200,6 +200,8 @@ def _load_builder_cli_module(monkeypatch): "ErnieModel", "Gemma2Model", "Gemma3Model", + "Gemma4MoEModel", + "Gemma4Model", "GemmaModel", "GPTOSSModel", "GraniteMoEHybridModel", diff --git a/test/python/models/test_gemma4_builder.py b/test/python/models/test_gemma4_builder.py new file mode 100644 index 0000000000..5bd8036af6 --- /dev/null +++ b/test/python/models/test_gemma4_builder.py @@ -0,0 +1,184 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +"""Unit tests for the Gemma4Model builder (gemma4-12b-it text component). + +These exercise the builder's structural deviations from Gemma3 without loading +weights or running the full pipeline: + - RMSNorm uses the weight directly (add_offset == 0), unlike Gemma1/2/3. + - Attention scale is 1.0 (Q/K are RMS-normed), not 1/sqrt(head_dim). + - Per-layer geometry: sliding layers use head_dim/kv_heads distinct from the + full ("global") layers, emitted as concrete KV-cache I/O shapes. + - Proportional RoPE on the global layers (partial rotary on the global + head_dim with a zero-padded NoPE tail) matches the HF reference. + - Local (sliding) RoPE uses the sliding theta on the sliding head_dim. + +Run with: + python -m pytest test/python/models/test_gemma4_builder.py -v --test_models +""" + +from __future__ import annotations + +import os +import sys +import types + +import numpy as np +import onnx_ir as ir +import torch + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..", "src", "python", "py")) + +from models.builders.gemma import Gemma4Model + +# 6-layer synthetic config: 5 sliding + 1 full, mirroring the gemma4 pattern. +LAYER_TYPES = ["sliding_attention"] * 5 + ["full_attention"] + ["sliding_attention"] * 4 + + +def _make_minimal_config(**overrides): + layer_types = overrides.pop("layer_types", LAYER_TYPES) + cfg = types.SimpleNamespace( + architectures=["Gemma4UnifiedForConditionalGeneration"], + model_type="gemma4_unified_text", + hidden_size=256, + num_attention_heads=8, + num_key_value_heads=4, # sliding KV heads + num_global_key_value_heads=1, # full KV heads + num_hidden_layers=len(layer_types), + intermediate_size=512, + vocab_size=128, + max_position_embeddings=64, + hidden_activation="gelu_pytorch_tanh", + head_dim=32, # sliding head dim + global_head_dim=64, # full head dim + sliding_window=16, + rms_norm_eps=1e-6, + layer_types=layer_types, + attention_k_eq_v=True, + final_logit_softcapping=30.0, + tie_word_embeddings=True, + rope_parameters={ + "full_attention": { + "partial_rotary_factor": 0.25, + "rope_theta": 1000000.0, + "rope_type": "proportional", + }, + "sliding_attention": {"rope_theta": 10000.0, "rope_type": "default"}, + }, + _name_or_path="", + ) + for k, v in overrides.items(): + setattr(cfg, k, v) + return cfg + + +def _hf_proportional_cos_sin(head_dim, base, partial_rotary_factor, cache_length): + """HF _compute_proportional_rope_parameters reference (cos/sin, halved).""" + rope_angles = int(partial_rotary_factor * head_dim // 2) + inv_freq_rotated = 1.0 / (base ** (torch.arange(0, 2 * rope_angles, 2, dtype=torch.int64).float() / head_dim)) + nope_angles = head_dim // 2 - rope_angles + if nope_angles > 0: + inv_freq = torch.cat((inv_freq_rotated, torch.zeros(nope_angles, dtype=torch.float32)), dim=0) + else: + inv_freq = inv_freq_rotated + t = torch.arange(cache_length, dtype=torch.int64).float() + freqs = torch.outer(t, inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos()[:, : head_dim // 2].numpy() + sin = emb.sin()[:, : head_dim // 2].numpy() + return cos, sin + + +def _hf_default_cos_sin(head_dim, base, cache_length): + """HF default RoPE reference (cos/sin, halved).""" + inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.int64).float() / head_dim)) + t = torch.arange(cache_length, dtype=torch.int64).float() + freqs = torch.outer(t, inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos()[:, : head_dim // 2].numpy() + sin = emb.sin()[:, : head_dim // 2].numpy() + return cos, sin + + +class TestGemma4Model: + def _make_model(self, **overrides): + cfg = _make_minimal_config(**overrides) + return Gemma4Model( + cfg, + io_dtype=ir.DataType.FLOAT, + onnx_dtype=ir.DataType.FLOAT, + ep="cpu", + cache_dir=None, + extra_options={}, + ) + + def test_layernorm_no_offset(self): + """Gemma4 RMSNorm uses the weight directly (add_offset == 0).""" + m = self._make_model() + assert m.layernorm_attrs["add_offset"] == 0 + + def test_attention_scale_is_one(self): + """Q/K are RMS-normed, so the attention scale is 1.0.""" + m = self._make_model() + assert m.attention_attrs["scale"] == 1.0 + + def test_is_local_matches_layer_types(self): + """is_local reflects the config's per-layer attention type.""" + m = self._make_model() + for i, lt in enumerate(LAYER_TYPES): + assert m.is_local(i) == (lt == "sliding_attention") + + def test_per_layer_kv_cache_shape(self): + """Sliding vs full layers emit concrete, distinct KV-cache shapes.""" + m = self._make_model() + template = ["batch_size", m.num_kv_heads, "past_sequence_length", "kv_cache_dim"] + sliding = m.make_key_value_cache_shape(0, list(template)) + full = m.make_key_value_cache_shape(5, list(template)) + # [batch, kv_heads, seq, head_dim] + assert sliding[1] == 4 and sliding[3] == 32 + assert full[1] == 1 and full[3] == 64 + + def test_external_rope_and_position_ids(self): + """Partial rotary needs the standalone RotaryEmbedding op + position_ids.""" + m = self._make_model() + assert m.attention_attrs["use_rope_in_attn"] is False + assert "position_ids" in m.input_names + + def test_proportional_rope_cache_parity(self): + """Global cache matches HF proportional RoPE (partial rotary + NoPE tail).""" + m = self._make_model() + builder_cos = m.values["cos_cache_global"].const_value.numpy() + builder_sin = m.values["sin_cache_global"].const_value.numpy() + hf_cos, hf_sin = _hf_proportional_cos_sin( + head_dim=64, base=1000000.0, partial_rotary_factor=0.25, cache_length=64 + ) + np.testing.assert_allclose(builder_cos, hf_cos, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(builder_sin, hf_sin, rtol=1e-5, atol=1e-5) + + def test_proportional_rope_nope_tail_is_identity(self): + """The NoPE tail (zero frequencies) yields cos=1, sin=0 (no rotation).""" + m = self._make_model() + cos = m.values["cos_cache_global"].const_value.numpy() + sin = m.values["sin_cache_global"].const_value.numpy() + rope_angles = int(0.25 * 64 // 2) # = 8 + assert np.allclose(cos[:, rope_angles:], 1.0) + assert np.allclose(sin[:, rope_angles:], 0.0) + + def test_local_rope_cache_parity(self): + """Local cache uses the sliding theta (1e4) on the sliding head_dim.""" + m = self._make_model() + builder_cos = m.values["cos_cache_local"].const_value.numpy() + builder_sin = m.values["sin_cache_local"].const_value.numpy() + hf_cos, hf_sin = _hf_default_cos_sin(head_dim=32, base=10000.0, cache_length=64) + np.testing.assert_allclose(builder_cos, hf_cos, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(builder_sin, hf_sin, rtol=1e-5, atol=1e-5) + + def test_global_and_local_caches_differ(self): + """Global (proportional, 1e6) and local (default, 1e4) caches must differ.""" + m = self._make_model() + # Compare on the overlapping rotary columns only (local head_dim/2 = 16). + gcos = m.values["cos_cache_global"].const_value.numpy() + lcos = m.values["cos_cache_local"].const_value.numpy() + assert not np.allclose(gcos[:, : lcos.shape[1]], lcos, rtol=1e-3, atol=1e-3) diff --git a/test/python/models/test_gemma4moe_builder.py b/test/python/models/test_gemma4moe_builder.py new file mode 100644 index 0000000000..94df2bf7c4 --- /dev/null +++ b/test/python/models/test_gemma4moe_builder.py @@ -0,0 +1,361 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +"""Unit tests for the Gemma4MoEModel builder (gemma-4-26B-A4B-it text component). + +Gemma4MoEModel subclasses the dense Gemma4Model and adds only the MoE FFN, so +these tests focus on the MoE-specific structure without loading weights: + - moe_attrs: experts/top_k mapped from config, QMoE under int4, geglu + + swiglu_fusion=1 + normalize_routing_weights (fused-path activation choice). + - Each layer builds a parallel dense-MLP + MoE FFN combined by an Add. + - The router pre-projection subgraph (scaleless RMSNorm -> scale -> proj) feeds + logits to the fused MoE op. + - per_expert_scale is folded into the expert down_proj initializer offline. + +Run with: + python -m pytest test/python/models/test_gemma4moe_builder.py -v --test_models +""" + +from __future__ import annotations + +import os +import sys +import types + +import numpy as np +import onnx_ir as ir +import torch + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..", "src", "python", "py")) + +from models.builders.gemma import Gemma4MoEModel + +LAYER_TYPES = ["sliding_attention", "full_attention"] + + +def _make_minimal_config(**overrides): + layer_types = overrides.pop("layer_types", LAYER_TYPES) + cfg = types.SimpleNamespace( + architectures=["Gemma4ForConditionalGeneration"], + model_type="gemma4_text", + hidden_size=64, + num_attention_heads=8, + num_key_value_heads=4, + num_global_key_value_heads=2, + num_hidden_layers=len(layer_types), + intermediate_size=128, # dense MLP intermediate + moe_intermediate_size=32, # MoE expert intermediate + num_experts=8, + top_k_experts=2, + vocab_size=128, + max_position_embeddings=128, + hidden_activation="gelu_pytorch_tanh", + head_dim=16, + global_head_dim=32, + sliding_window=32, + rms_norm_eps=1e-6, + layer_types=layer_types, + attention_k_eq_v=True, + final_logit_softcapping=30.0, + tie_word_embeddings=True, + enable_moe_block=True, + rope_parameters={ + "full_attention": {"partial_rotary_factor": 0.25, "rope_theta": 1000000.0, "rope_type": "proportional"}, + "sliding_attention": {"rope_theta": 10000.0, "rope_type": "default"}, + }, + _name_or_path="", + ) + for k, v in overrides.items(): + setattr(cfg, k, v) + return cfg + + +def _make_model(onnx_dtype=ir.DataType.FLOAT, **overrides): + return Gemma4MoEModel( + _make_minimal_config(**overrides), + io_dtype=ir.DataType.FLOAT, + onnx_dtype=onnx_dtype, + ep="cpu", + cache_dir=None, + extra_options={}, + ) + + +def _mock_moe_layer(hidden, num_experts, moe_inter, dense_inter): + def lin(o, i): + return types.SimpleNamespace(weight=torch.randn(o, i), bias=None) + + def ln(): + return types.SimpleNamespace(weight=torch.randn(hidden)) + + return types.SimpleNamespace( + router=types.SimpleNamespace( + proj=lin(num_experts, hidden), + scale=torch.randn(hidden), + per_expert_scale=torch.rand(num_experts) + 0.5, + ), + experts=types.SimpleNamespace( + gate_up_proj=torch.randn(num_experts, 2 * moe_inter, hidden), + down_proj=torch.randn(num_experts, hidden, moe_inter), + ), + post_feedforward_layernorm_1=ln(), + post_feedforward_layernorm_2=ln(), + pre_feedforward_layernorm_2=ln(), + ), types.SimpleNamespace( + gate_proj=lin(dense_inter, hidden), up_proj=lin(dense_inter, hidden), down_proj=lin(hidden, dense_inter) + ) + + +def _mock_quark_moe_layer(hidden, num_experts, moe_inter, group_size): + """Mock a pre-quantized Quark uint2 MoE layer, mirroring what the QuarkModel loader + produces: `layer.mlp.experts` is a container exposing the re-fused fc1/fc2 packed + tensors plus per-expert `gate_proj.input_prescale`, and `layer.router` carries the + scale / per_expert_scale. Weights are packed 4 codes/uint8 (2-bit); scales and float + zero-points are per-group ([E, out, in/group_size]).""" + blocks_in = hidden // group_size # fc1 packs along hidden + blocks_in_fc2 = moe_inter // group_size # fc2 packs along moe_inter + fc1_out = 2 * moe_inter # gate|up concat + packed_hidden = hidden // 4 # 4 uint2 codes per uint8 + packed_inter = moe_inter // 4 + + class _Experts: + def __init__(self): + self.fc1_weights = torch.randint(0, 256, (num_experts, fc1_out, packed_hidden), dtype=torch.uint8) + self.fc1_scales = torch.rand(num_experts, fc1_out, blocks_in) + 0.1 + self.fc1_zero_points = torch.full((num_experts, fc1_out, blocks_in), 1.5) + self.fc2_weights = torch.randint(0, 256, (num_experts, hidden, packed_inter), dtype=torch.uint8) + self.fc2_scales = torch.rand(num_experts, hidden, blocks_in_fc2) + 0.1 + self.fc2_zero_points = torch.full((num_experts, hidden, blocks_in_fc2), 1.5) + prescale = types.SimpleNamespace(input_prescale=torch.rand(hidden) + 0.5) + self._e = {i: types.SimpleNamespace(gate_proj=prescale) for i in range(num_experts)} + + def keys(self): + return self._e.keys() + + def __getitem__(self, i): + return self._e[i] + + return types.SimpleNamespace( + mlp=types.SimpleNamespace(experts=_Experts()), + router=types.SimpleNamespace( + proj=types.SimpleNamespace(weight=torch.randn(num_experts, hidden), bias=None), + scale=torch.randn(hidden), + per_expert_scale=torch.rand(num_experts) + 0.5, + ), + ) + + +class TestGemma4MoEQuarkPath: + """Pre-quantized Quark uint2 (2-bit) expert path: make_moe_quark.""" + + def _make_quark_model(self, ep="cpu", group_size=16): + m = _make_model(onnx_dtype=ir.DataType.FLOAT) + # The QuarkModel loader sets these; seed them so make_moe_quark can run standalone. + m.quant_type = "quark" + m.ep = ep + m.moe_attrs["op_type"] = "QMoE" + m.moe_attrs["expert_weight_bits"] = 2 + m.moe_attrs["swiglu_fusion"] = 1 if ep == "cpu" else 2 + m.moe_attrs["block_size"] = group_size + m.shared_input_rotations = {m.hidden_size: torch.randn(m.hidden_size, m.hidden_size)} + m.shared_rotation_initializers = set() + return m + + def _build_quark_layer(self, m, group_size=16): + hidden = m.hidden_size + num_experts = m.moe_attrs["num_experts"] + layer = _mock_quark_moe_layer(hidden, num_experts, m.moe_intermediate_size, group_size) + m.make_value("resid", ir.DataType.FLOAT, ["batch_size", "sequence_length", hidden]) + m.make_moe_quark(0, layer, "resid", "resid") + return layer + + def test_qmoe_node_emits_expert_weight_bits_2(self): + """The fused QMoE op carries expert_weight_bits=2 for the 2-bit Quark path.""" + m = self._make_quark_model(ep="cpu") + self._build_quark_layer(m) + qmoe = next(n for n in m.model.graph if n.op_type == "QMoE") + assert qmoe.attributes["expert_weight_bits"].as_int() == 2 + + def test_cpu_emits_float_zero_points(self): + """On CPU the float per-group zero-points are emitted as QMoE inputs (float io_dtype).""" + m = self._make_quark_model(ep="cpu") + self._build_quark_layer(m) + zp_name = "model.layers.0.moe.experts.gate_up_proj.zero_points" + assert zp_name in m.values + assert m.values[zp_name].const_value.dtype == ir.DataType.FLOAT + # And the QMoE node references it as an input (not the empty string). + qmoe = next(n for n in m.model.graph if n.op_type == "QMoE") + assert any(inp is not None and inp.name == zp_name for inp in qmoe.inputs) + + def test_cuda_omits_zero_points(self): + """On CUDA the symmetric -1.5*scale bias is reconstructed by the op, so zp is omitted.""" + m = self._make_quark_model(ep="cuda") + self._build_quark_layer(m) + assert "model.layers.0.moe.experts.gate_up_proj.zero_points" not in m.values + assert "model.layers.0.moe.experts.down_proj.zero_points" not in m.values + + +def _mock_quark_int4_moe_layer(hidden, num_experts, moe_inter, group_size): + """Mock a plain Quark/AWQ int4 MoE layer. Unlike the factored uint2 checkpoint, these + experts have NO input_prescale/rotation, are already emitted interleaved by + combine_and_repack_gate_up (no fc1 reorder), and carry INTEGER (uint8) per-group + zero-points. Weights pack 2 codes/uint8 (4-bit).""" + blocks_in = hidden // group_size + blocks_in_fc2 = moe_inter // group_size + fc1_out = 2 * moe_inter + packed_hidden = hidden // 2 # 2 uint4 codes per uint8 + packed_inter = moe_inter // 2 + + class _Experts: + def __init__(self): + self.fc1_weights = torch.randint(0, 256, (num_experts, fc1_out, packed_hidden), dtype=torch.uint8) + self.fc1_scales = torch.rand(num_experts, fc1_out, blocks_in) + 0.1 + self.fc1_zero_points = torch.randint(0, 16, (num_experts, fc1_out, blocks_in), dtype=torch.uint8) + self.fc2_weights = torch.randint(0, 256, (num_experts, hidden, packed_inter), dtype=torch.uint8) + self.fc2_scales = torch.rand(num_experts, hidden, blocks_in_fc2) + 0.1 + self.fc2_zero_points = torch.randint(0, 16, (num_experts, hidden, blocks_in_fc2), dtype=torch.uint8) + # Plain Quark/AWQ experts have no prescale. + noprescale = types.SimpleNamespace(input_prescale=None) + self._e = {i: types.SimpleNamespace(gate_proj=noprescale) for i in range(num_experts)} + + def keys(self): + return self._e.keys() + + def __getitem__(self, i): + return self._e[i] + + return types.SimpleNamespace( + mlp=types.SimpleNamespace(experts=_Experts()), + router=types.SimpleNamespace( + proj=types.SimpleNamespace(weight=torch.randn(num_experts, hidden), bias=None), + scale=torch.randn(hidden), + per_expert_scale=torch.rand(num_experts) + 0.5, + ), + ) + + +class TestGemma4MoEQuarkInt4Path: + """Plain Quark/AWQ int4 expert path: the shape-guarded generalization of make_moe_quark + (no prescale/rotation, no fc1 reorder, integer zero-points).""" + + def _make_quark_model(self, ep="cpu", group_size=16): + m = _make_model(onnx_dtype=ir.DataType.FLOAT) + m.quant_type = "quark" + m.ep = ep + m.moe_attrs["op_type"] = "QMoE" + m.moe_attrs["expert_weight_bits"] = 4 + m.moe_attrs["swiglu_fusion"] = 1 if ep == "cpu" else 2 + m.moe_attrs["block_size"] = group_size + m.shared_input_rotations = {m.hidden_size: torch.randn(m.hidden_size, m.hidden_size)} + m.shared_rotation_initializers = set() + return m + + def _build_quark_layer(self, m, group_size=16): + hidden = m.hidden_size + num_experts = m.moe_attrs["num_experts"] + layer = _mock_quark_int4_moe_layer(hidden, num_experts, m.moe_intermediate_size, group_size) + m.make_value("resid", ir.DataType.FLOAT, ["batch_size", "sequence_length", hidden]) + m.make_moe_quark(0, layer, "resid", "resid") + return layer + + def test_qmoe_node_emits_expert_weight_bits_4(self): + """The fused QMoE op carries expert_weight_bits=4 for the int4 Quark path.""" + m = self._make_quark_model(ep="cpu") + self._build_quark_layer(m) + qmoe = next(n for n in m.model.graph if n.op_type == "QMoE") + assert qmoe.attributes["expert_weight_bits"].as_int() == 4 + + def test_no_input_prescale_or_rotation(self): + """Without input_prescale the expert input feeds QMoE directly: no prescale Mul / rotation MatMul.""" + m = self._make_quark_model(ep="cpu") + self._build_quark_layer(m) + assert not any("input_prescale" in name for name in m.values) + assert not any("shared_input_rotation" in name for name in m.values) + + def test_integer_zero_points_not_cast_to_float(self): + """Plain int4 experts carry uint8 zero-points, emitted at native dtype (no io_dtype cast).""" + m = self._make_quark_model(ep="cpu") + self._build_quark_layer(m) + zp_name = "model.layers.0.moe.experts.gate_up_proj.zero_points" + assert zp_name in m.values + assert m.values[zp_name].const_value.dtype == ir.DataType.UINT8 + + def test_cuda_int4_keeps_zero_points(self): + """The zp-omit path is 2-bit-only; int4 keeps its zero-points even on CUDA.""" + m = self._make_quark_model(ep="cuda") + self._build_quark_layer(m) + assert "model.layers.0.moe.experts.gate_up_proj.zero_points" in m.values + + +class TestGemma4MoEModel: + def test_moe_attrs_mapping(self): + """num_experts/top_k come from config; activation/fusion/normalize set for fused path.""" + m = _make_model() + assert m.moe_attrs["num_experts"] == 8 + assert m.moe_attrs["top_k"] == 2 + assert m.moe_attrs["activation_type"] == "geglu" + assert m.moe_attrs["swiglu_fusion"] == 1 + assert m.moe_attrs["normalize_routing_weights"] is True + + def test_qmoe_under_int4(self): + """int4 build selects the fused QMoE op; float build uses MoE.""" + assert _make_model(onnx_dtype=ir.DataType.INT4).moe_attrs["moe_op_type"] == "QMoE" + assert _make_model(onnx_dtype=ir.DataType.FLOAT).moe_attrs["moe_op_type"] == "MoE" + + def test_dense_and_moe_intermediate_sizes(self): + """Dense MLP uses intermediate_size; MoE uses moe_intermediate_size.""" + m = _make_model() + assert m.intermediate_size == 128 + assert m.moe_intermediate_size == 32 + + def test_inherits_dense_attention(self): + """MoE model inherits the dense Gemma4 attention config (scale 1.0, no offset).""" + m = _make_model() + assert m.attention_attrs["scale"] == 1.0 + assert m.layernorm_attrs["add_offset"] == 0 + assert m.is_local(0) is True and m.is_local(1) is False + + def _build_one_layer(self, m): + hidden = m.hidden_size + moe_layer, dense_mlp = _mock_moe_layer( + hidden, m.moe_attrs["num_experts"], m.moe_intermediate_size, m.intermediate_size + ) + # Seed layernorm_attrs as if the pre-FFN SkipLayerNorm had run. + m.make_value("resid", ir.DataType.FLOAT, ["batch_size", "sequence_length", hidden]) + m.make_value("normed", ir.DataType.FLOAT, ["batch_size", "sequence_length", hidden]) + m.layernorm_attrs["root_input"] = "resid" + m.layernorm_attrs["output_0"] = "normed" + m.layernorm_attrs["skip_input"] = "normed" + m._current_layer = moe_layer + m.make_mlp(0, dense_mlp, "normed") + return moe_layer + + def test_parallel_dense_moe_combine(self): + """make_mlp builds both branches, a MoE op, and combines them with an Add.""" + m = _make_model() + self._build_one_layer(m) + names = {n.name for n in m.model.graph} + assert any("/model/layers.0/moe/MoE" in n for n in names) + assert any("/model/layers.0/ffn_combine/Add" in n for n in names) + assert any("/model/layers.0/moe/router/MatMul" in n for n in names) + # The FFN block contribution (skip_input) is the combine output. + assert m.layernorm_attrs["skip_input"] == "/model/layers.0/ffn_combine/Add/output_0" + + def test_router_preprojection_subgraph(self): + """Router path: scaleless RMSNorm -> * (scale * hidden^-0.5) -> proj -> reshape.""" + m = _make_model() + self._build_one_layer(m) + names = {n.name for n in m.model.graph} + assert any("moe/router/norm/SimplifiedLayerNormalization" in n for n in names) + assert any("moe/router/scale/Mul" in n for n in names) + assert any("moe/router/Reshape" in n for n in names) + + def test_per_expert_scale_folded_into_down_proj(self): + """down_proj initializer must equal raw_down_proj * per_expert_scale (folded offline).""" + m = _make_model() + layer = self._build_one_layer(m) + expected = (layer.experts.down_proj * layer.router.per_expert_scale.reshape(-1, 1, 1)).numpy() + init = m.values["model.layers.0.moe.experts.down_proj.weight"].const_value.numpy() + np.testing.assert_allclose(init, expected, rtol=1e-5, atol=1e-5)