diff --git a/.github/workflows/gpu_l4_golden_parity.yml b/.github/workflows/gpu_l4_golden_parity.yml index 28a95e1f..5f55e656 100644 --- a/.github/workflows/gpu_l4_golden_parity.yml +++ b/.github/workflows/gpu_l4_golden_parity.yml @@ -59,7 +59,8 @@ jobs: pip install -r requirements/ci/requirements.txt pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs pip install -e '.[testing,transformers]' - pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ onnxruntime-gpu onnxruntime-genai-cuda + # ORT 1.27+ PyPI wheels require CUDA 13; these runners and PyTorch use CUDA 12.8. + pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ 'onnxruntime-gpu<1.27' onnxruntime-genai-cuda - name: Check for golden test data id: check_golden diff --git a/.github/workflows/gpu_l5_generation_e2e.yml b/.github/workflows/gpu_l5_generation_e2e.yml index ba3f2b39..4c806517 100644 --- a/.github/workflows/gpu_l5_generation_e2e.yml +++ b/.github/workflows/gpu_l5_generation_e2e.yml @@ -59,7 +59,8 @@ jobs: pip install -r requirements/ci/requirements.txt pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs pip install -e '.[testing,transformers]' - pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ onnxruntime-gpu onnxruntime-genai-cuda + # ORT 1.27+ PyPI wheels require CUDA 13; these runners and PyTorch use CUDA 12.8. + pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ 'onnxruntime-gpu<1.27' onnxruntime-genai-cuda - name: Run L5 generation E2E tests env: diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index dc705dc1..9945da40 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -283,6 +283,7 @@ jobs: pip install -r requirements/ci/requirements.txt pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs pip install -e '.[testing]' + # ORT 1.27+ PyPI wheels require CUDA 13; these runners and PyTorch use CUDA 12.8. pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ "onnxruntime-gpu<1.27" onnxruntime-genai-cuda - name: Run fast integration tests env: diff --git a/.github/workflows/validation_examples_gpu.yml b/.github/workflows/validation_examples_gpu.yml index 062fdfed..13e9873c 100644 --- a/.github/workflows/validation_examples_gpu.yml +++ b/.github/workflows/validation_examples_gpu.yml @@ -83,7 +83,8 @@ jobs: pip install -r requirements/ci/requirements.txt pip install -e '.[testing,transformers]' pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs - pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ onnxruntime-gpu onnxruntime-genai-cuda + # ORT 1.27+ PyPI wheels require CUDA 13; these runners and PyTorch use CUDA 12.8. + pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ 'onnxruntime-gpu<1.27' onnxruntime-genai-cuda - name: Run ${{ matrix.name }} env: diff --git a/docs/research/glm52-export.md b/docs/research/glm52-export.md new file mode 100644 index 00000000..c225d1f3 --- /dev/null +++ b/docs/research/glm52-export.md @@ -0,0 +1,147 @@ +# GLM-5.2 export design + +## Verified architecture + +Sources: + +- [`zai-org/GLM-5.2` config](https://huggingface.co/zai-org/GLM-5.2/blob/main/config.json) +- Transformers `GlmMoeDsaForCausalLM` implementation +- llama.cpp `glm-dsa` GGUF metadata and tensor names +- [`unsloth/GLM-5.2-GGUF`](https://huggingface.co/unsloth/GLM-5.2-GGUF) +- [IndexShare paper, arXiv:2603.12201](https://arxiv.org/abs/2603.12201) +- [GLM-5 report, arXiv:2602.15763](https://arxiv.org/abs/2602.15763) + +The Hugging Face model type is `glm_moe_dsa`, with +`GlmMoeDsaForCausalLM` as its architecture. llama.cpp writes the exact +`general.architecture` value `glm-dsa`. + +| Field | GLM-5.2 value | +| --- | --- | +| Backbone layers | 78 | +| MTP layers | 1 additional NextN layer | +| Hidden / dense FFN size | 6144 / 12288 | +| Attention heads / KV heads | 64 / 64 | +| QK / NoPE / RoPE / V head dimensions | 256 / 192 / 64 / 256 | +| Q-LoRA / KV-LoRA rank | 2048 / 512 | +| Routed experts / experts per token | 256 / 8 | +| Routed / shared-expert FFN size | 2048 / 2048 | +| Shared experts | 1 | +| Dense prefix | first 3 layers | +| Router | float32 sigmoid, `noaux_tc`, normalized top-k, scale 2.5 | +| Norm / activation | pre-RMSNorm, epsilon 1e-5 / SiLU | +| RoPE | interleaved, theta 8,000,000 | +| Maximum context | 1,048,576 | +| Indexer | 32 heads, head dimension 128, top-k 2048 | + +GGUF reports `block_count=79` and `nextn_predict_layers=1`; therefore the +causal-LM backbone has 78 layers. GGUF also splits the MLA KV decompression +weight into `attn_k_b` and `attn_v_b`, while Transformers exposes one +`kv_b_proj`. + +## Implemented scope + +The exporter now preserves the 78-layer backbone, IndexShare DSA, and the +additional improved-MTP layer. Full indexers are emitted for layers 0, 1, 2, +6, 10, ... 74; each full layer's INT64 top-k result is reused by the following +three shared layers. The indexer implements the checkpoint equations in fp32: +Q-LoRA residual projection, LayerNorm+interleaved-RoPE key projection, ReLU +scores, learned signed head weights, causal bias, and dynamic +`min(2048, key_length)` TopK. + +Each full-indexer layer packs its 128-value index key beside the decompressed +MLA key in the standard `past_key_values.N.key` tensor. Shared layers retain +the normal MLA cache. This keeps the ORT GenAI key/value naming contract +without duplicating the index key across 64 attention heads. + +The MTP component exports: + +1. `enorm(next-token embedding)` and `hnorm(target hidden state)`; +2. concatenation and the `2H -> H` `eh_proj`; +3. the complete layer-78 DSA+MoE decoder block; and +4. `shared_head.norm`, producing `mtp_hidden` for the shared target LM head. + +Multi-component packages contain `model/model.onnx`, `mtp/model.onnx`, and an +`mtp_config.json` sidecar. Both runtime aliases (`onnx-genai` and `ort-genai`) +emit these artifacts. ORT GenAI does not yet orchestrate the sidecar itself. + +The importer: + +- derives all MLA, MoE, IndexShare, RoPE, and NextN fields from GGUF metadata; +- separates the GGUF MTP block from the 78-layer backbone and exports it as `mtp`; +- fuses split `attn_k_b` and `attn_v_b` tensors into `kv_b_proj`; +- maps `blk.N.indexer.*` and supported `blk.78.nextn.*` tensors; +- streams stacked routed-expert tensors one expert at a time instead of + materializing all 256 experts as one float tensor; and +- preserves supported 4-bit and 8-bit tensors as MatMulNBits, requantizing + mixed source types to the selected graph-wide layout when necessary. + +The generic portable MoE implementation evaluates every expert and masks the +outputs. It is correct for routing, but a fused sparse-MoE runtime kernel will +be required for practical GLM-5.2 execution. + +## Runtime follow-ups + +Portable ONNX preserves DSA selection numerics by converting the selected +indices to an additive mask for the standard `Attention` operator. No custom op +is required for correctness. Practical million-token performance still needs +a selected-token sparse-attention kernel so the runtime avoids evaluating +masked K/V positions. + +`index_share_for_mtp_iteration` also needs runtime orchestration. MTP step 0 can +return `topk_indices`, but accepted speculative positions make the MTP MLA and +indexer caches advance at different logical lengths. ORT GenAI currently has +no separate indexer-cache/control-state contract, so steps 1+ recompute the +index selection rather than silently applying an incorrect packed-cache reuse. + +Use `--glm-full-attention` to retain the prior dense MLA fallback. The fallback +disables the DSA-dependent MTP artifact. + +## Separate workstream: Unsloth dynamic quantization + +The six `UD-IQ1_M` GGUF shard headers were inspected. The dominant routed +expert tensors are below the currently targeted 4-bit floor: + +- routed gate/up projections: 76 `IQ1_M` and 74 `IQ2_XXS` tensors; +- routed down projections: 71 `IQ3_XXS` and 4 `IQ4_XS` tensors; +- shared experts: mainly `Q5_K` / `Q6_K`; +- MLA projections: mainly `Q8_0` / `Q5_K`; +- norms: F32; embedding and output: Q4_K. + +This increment can dequantize IQ sources and requantize them to a supported +Q4 MatMulNBits layout. That enables export but loses the intended sub-4-bit +size and may alter accuracy. Native preservation should use one or both of the +following paths. + +### Runtime custom-op requirement + +Add an ONNX Runtime GenAI execution-provider capability for a fused sparse-MoE +operator that: + +1. consumes top-k expert indices and routing weights without evaluating all + 256 experts; +2. supports the llama.cpp IQ1_M, IQ2_XXS, and IQ3_XXS block/codebook layouts + (plus IQ4_XS), including their scales/minimums; +3. fuses gate/up activation and down projection where profitable; +4. loads only selected expert blocks and supports CPU and target accelerator + implementations; and +5. advertises supported bit layouts through Mobius EP capabilities so export + selects native packed initializers, otherwise falling back to Q4 + MatMulNBits requantization. + +### Draft ONNX Runtime issue + +**Title:** Support sub-4-bit GGUF/IQ weights for MatMulNBits and sparse MoE + +**Body:** + +> GLM-5.2 dynamic GGUF distributions use IQ1_M and IQ2_XXS for most routed +> gate/up expert matrices and IQ3_XXS for most down matrices. Exporting these +> models to ONNX currently requires expansion or requantization to 4-bit, +> losing the distribution's size advantage and making a 256-expert model +> impractical. Please define supported 1/2/3-bit MatMulNBits representations, +> or provide an extensible packed/codebook-weight contract for the llama.cpp +> IQ formats, with CPU and execution-provider kernels. For MoE, an operator +> that accepts expert weights plus per-token top-k indices/weights should avoid +> evaluating or unpacking every expert. Required formats for the first target +> are IQ1_M, IQ2_XXS, IQ3_XXS, and IQ4_XS. Please also document zero-point, +> scale/codebook packing, block-size, alignment, and external-data constraints. diff --git a/export_deepseek_v2_tiny.py b/export_deepseek_v2_tiny.py new file mode 100644 index 00000000..68a27a7a --- /dev/null +++ b/export_deepseek_v2_tiny.py @@ -0,0 +1,74 @@ +"""Export the tiny synthetic DeepSeek-V2 MLA + MoE model for onnx-genai E2E.""" + +from __future__ import annotations + +import argparse +import os +import sys + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "tests")) + +import numpy as np +import onnx_ir as ir +from _test_configs import ALL_CAUSAL_LM_CONFIGS, _base_config + +from mobius._config_resolver import _default_task_for_model +from mobius._registry import registry +from mobius.integrations.onnx_genai import write_onnx_genai_config +from mobius.tasks import get_task + +ARTIFACTS_DIR = os.environ.get( + "MOBIUS_ARTIFACTS_DIR", os.path.join(os.path.dirname(__file__), "artifacts") +) + + +def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None: + for init in model.graph.initializers.values(): + if init.const_value is not None: + continue + shape = tuple(int(d) for d in init.shape) + if not shape: + continue + if init.dtype == ir.DataType.FLOAT: + data = rng.standard_normal(shape).astype(np.float32) * 0.02 + elif init.dtype == ir.DataType.FLOAT16: + data = (rng.standard_normal(shape) * 0.02).astype(np.float16) + elif init.dtype in (ir.DataType.INT64, ir.DataType.INT32): + npd = np.int64 if init.dtype == ir.DataType.INT64 else np.int32 + data = rng.integers(0, 10, size=shape).astype(npd) + else: + data = rng.standard_normal(shape).astype(np.float32) * 0.02 + init.const_value = ir.Tensor(data) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--output-dir", + default=os.path.join(ARTIFACTS_DIR, "deepseek-v2-tiny"), + ) + args = parser.parse_args() + + overrides = dict(next(ov for mt, ov, _ in ALL_CAUSAL_LM_CONFIGS if mt == "deepseek_v2")) + config = _base_config(**overrides) + config.dtype = ir.DataType.FLOAT + + module = registry.get("deepseek_v2")(config) + task = get_task(_default_task_for_model("deepseek_v2")) + pkg = task.build(module, config) + + rng = np.random.default_rng(0) + for model in pkg.values(): + _fill_random_weights(model, rng) + + os.makedirs(args.output_dir, exist_ok=True) + pkg.save(args.output_dir, external_data="onnx", check_weights=False) + artifacts = write_onnx_genai_config(pkg, args.output_dir, config=config) + for name, path in artifacts.items(): + print(f"{name}:", path) + print("Saved to", args.output_dir) + print("files:", sorted(os.listdir(args.output_dir))) + + +if __name__ == "__main__": + main() diff --git a/export_glm_tiny_qmoe.py b/export_glm_tiny_qmoe.py new file mode 100644 index 00000000..7f8a98a6 --- /dev/null +++ b/export_glm_tiny_qmoe.py @@ -0,0 +1,133 @@ +"""Export a tiny synthetic glm_moe_dsa model with a FUSED ``com.microsoft::QMoE``. + +Same tiny ``glm_moe_dsa`` config as ``export_glm_tiny_quant.py``, but sets +``config.fused_quantized_moe = True`` so the routed MoE experts are emitted as a +single fused ``com.microsoft::QMoE`` node (int4, block-32, expert-major layout) +instead of a per-expert unroll of ``com.microsoft::MatMulNBits``. + +Routing (GLM sigmoid + noaux_tc, n_group=1) is preserved exactly: the QMoE +kernel re-derives top-k selection from ``router_probs`` (the gate's +``scores_for_choice``), and the normalized+scaled combine weights are scattered +into the optional ``router_weights`` input with ``normalize_routing_weights=0``. + +Writes an onnx-genai-loadable artifact directory: + /model.onnx (+ external data), inference_metadata.yaml, tokenizer.json +""" + +from __future__ import annotations + +import argparse +import os +import shutil +import sys + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "tests")) + +import numpy as np +import onnx_ir as ir +from _test_configs import ALL_CAUSAL_LM_CONFIGS, _base_config + +from mobius._config_resolver import _default_task_for_model +from mobius._configs import QuantizationConfig +from mobius._registry import registry +from mobius.integrations.onnx_genai import write_onnx_genai_config +from mobius.tasks import get_task + +ARTIFACTS_DIR = os.environ.get( + "MOBIUS_ARTIFACTS_DIR", os.path.join(os.path.dirname(__file__), "artifacts") +) + + +def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None: + for init in model.graph.initializers.values(): + if init.const_value is not None: + continue + shape = tuple(int(d) for d in init.shape) + if not shape: + continue + if init.dtype == ir.DataType.FLOAT: + data = rng.standard_normal(shape).astype(np.float32) * 0.02 + elif init.dtype == ir.DataType.FLOAT16: + data = (rng.standard_normal(shape) * 0.02).astype(np.float16) + elif init.dtype == ir.DataType.UINT8: + # Packed int4/int8 quantized weights or bit-packed zero points. + data = rng.integers(0, 256, size=shape).astype(np.uint8) + elif init.dtype in (ir.DataType.INT64, ir.DataType.INT32): + npd = np.int64 if init.dtype == ir.DataType.INT64 else np.int32 + data = rng.integers(0, 10, size=shape).astype(npd) + else: + data = rng.standard_normal(shape).astype(np.float32) * 0.02 + init.const_value = ir.Tensor(data) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--output-dir", + default=os.path.join(ARTIFACTS_DIR, "glm-5.2-tiny-qmoe"), + ) + parser.add_argument( + "--fp32-output-dir", + default=os.path.join(ARTIFACTS_DIR, "glm-5.2-tiny"), + ) + args = parser.parse_args() + + overrides = dict(next(ov for mt, ov, _ in ALL_CAUSAL_LM_CONFIGS if mt == "glm_moe_dsa")) + config = _base_config(**overrides) + config.dtype = ir.DataType.FLOAT + # int4, block-32, symmetric -> fused QMoE (no zero points; the kernel + # defaults per-block zero-point to 1 << (bits - 1)). + config.quantization = QuantizationConfig( + bits=4, + group_size=32, + quant_method="gguf", + sym=True, + ) + # Emit routed experts as a single fused com.microsoft::QMoE node. + config.fused_quantized_moe = True + + model_type = "glm_moe_dsa" + module = registry.get(model_type)(config) + task = get_task(_default_task_for_model(model_type)) + pkg = task.build(module, config) + + rng = np.random.default_rng(0) + for model in pkg.values(): + _fill_random_weights(model, rng) + + os.makedirs(args.output_dir, exist_ok=True) + pkg.save(args.output_dir, external_data="onnx", check_weights=False) + artifacts = write_onnx_genai_config(pkg, args.output_dir, config=config) + for name, path in artifacts.items(): + print(f"{name}:", path) + + # Reuse tokenizer from the fp32 artifacts if present. + tok_src = os.path.join(args.fp32_output_dir, "tokenizer.json") + if os.path.exists(tok_src): + shutil.copy(tok_src, os.path.join(args.output_dir, "tokenizer.json")) + + for attr in ( + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "head_dim", + "hidden_size", + "vocab_size", + "max_position_embeddings", + "qk_nope_head_dim", + "qk_rope_head_dim", + "v_head_dim", + "kv_lora_rank", + "q_lora_rank", + "num_local_experts", + "num_experts_per_tok", + "moe_intermediate_size", + ): + print(f" {attr} =", getattr(config, attr, None)) + + print("Saved to", args.output_dir) + print("files:", sorted(os.listdir(args.output_dir))) + + +if __name__ == "__main__": + main() diff --git a/export_glm_tiny_quant.py b/export_glm_tiny_quant.py new file mode 100644 index 00000000..16b72b6c --- /dev/null +++ b/export_glm_tiny_quant.py @@ -0,0 +1,129 @@ +"""Export a quantized tiny synthetic glm_moe_dsa model for onnx-genai E2E. + +Same tiny ``glm_moe_dsa`` config as ``export_glm_tiny.py``, but with a +``QuantizationConfig`` attached so the linear projections (and per-expert +MoE MLPs) are emitted as ``com.microsoft::MatMulNBits`` (int4, block-32). + +This script leaves ``fused_quantized_moe`` disabled, so the GLM/DeepSeek MoE +path decomposes experts into per-expert MLPs and each expert projection becomes +MatMulNBits. For the fused ``com.microsoft::QMoE`` path, use +``export_glm_tiny_qmoe.py``. + +Writes an onnx-genai-loadable artifact directory: + /model.onnx (+ external data), inference_metadata.yaml, tokenizer.json +""" + +from __future__ import annotations + +import argparse +import os +import shutil +import sys + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "tests")) + +import numpy as np +import onnx_ir as ir +from _test_configs import ALL_CAUSAL_LM_CONFIGS, _base_config + +from mobius._config_resolver import _default_task_for_model +from mobius._configs import QuantizationConfig +from mobius._registry import registry +from mobius.integrations.onnx_genai import write_onnx_genai_config +from mobius.tasks import get_task + +ARTIFACTS_DIR = os.environ.get( + "MOBIUS_ARTIFACTS_DIR", os.path.join(os.path.dirname(__file__), "artifacts") +) + + +def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None: + for init in model.graph.initializers.values(): + if init.const_value is not None: + continue + shape = tuple(int(d) for d in init.shape) + if not shape: + continue + if init.dtype == ir.DataType.FLOAT: + data = rng.standard_normal(shape).astype(np.float32) * 0.02 + elif init.dtype == ir.DataType.FLOAT16: + data = (rng.standard_normal(shape) * 0.02).astype(np.float16) + elif init.dtype == ir.DataType.UINT8: + # Packed int4/int8 quantized weights or bit-packed zero points. + data = rng.integers(0, 256, size=shape).astype(np.uint8) + elif init.dtype in (ir.DataType.INT64, ir.DataType.INT32): + npd = np.int64 if init.dtype == ir.DataType.INT64 else np.int32 + data = rng.integers(0, 10, size=shape).astype(npd) + else: + data = rng.standard_normal(shape).astype(np.float32) * 0.02 + init.const_value = ir.Tensor(data) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--output-dir", + default=os.path.join(ARTIFACTS_DIR, "glm-5.2-tiny-q4"), + ) + parser.add_argument( + "--fp32-output-dir", + default=os.path.join(ARTIFACTS_DIR, "glm-5.2-tiny"), + ) + args = parser.parse_args() + + overrides = dict(next(ov for mt, ov, _ in ALL_CAUSAL_LM_CONFIGS if mt == "glm_moe_dsa")) + config = _base_config(**overrides) + config.dtype = ir.DataType.FLOAT + # int4, block-32, asymmetric (GGUF/AWQ-style) -> MatMulNBits with uint8 + # zero points. This matches the validated glm_moe_dsa quant unit test. + config.quantization = QuantizationConfig( + bits=4, + group_size=32, + quant_method="gguf", + sym=False, + ) + + model_type = "glm_moe_dsa" + module = registry.get(model_type)(config) + task = get_task(_default_task_for_model(model_type)) + pkg = task.build(module, config) + + rng = np.random.default_rng(0) + for model in pkg.values(): + _fill_random_weights(model, rng) + + os.makedirs(args.output_dir, exist_ok=True) + pkg.save(args.output_dir, external_data="onnx", check_weights=False) + artifacts = write_onnx_genai_config(pkg, args.output_dir, config=config) + for name, path in artifacts.items(): + print(f"{name}:", path) + + # Reuse tokenizer from the fp32 artifacts. + tok_src = os.path.join(args.fp32_output_dir, "tokenizer.json") + if os.path.exists(tok_src): + shutil.copy(tok_src, os.path.join(args.output_dir, "tokenizer.json")) + + for attr in ( + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "head_dim", + "hidden_size", + "vocab_size", + "max_position_embeddings", + "qk_nope_head_dim", + "qk_rope_head_dim", + "v_head_dim", + "kv_lora_rank", + "q_lora_rank", + "num_local_experts", + "moe_intermediate_size", + ): + print(f" {attr} =", getattr(config, attr, None)) + + print("Saved to", args.output_dir) + print("files:", sorted(os.listdir(args.output_dir))) + + +if __name__ == "__main__": + main() diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index f6693401..6952fce6 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -243,6 +243,12 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: config = _config_from_hf(hf_config, parent_config=parent_config) if dtype_override is not None: config = dataclasses.replace(config, dtype=dtype_override) + if args.glm_full_attention: + config = dataclasses.replace( + config, + use_dsa=False, + num_nextn_predict_layers=0, + ) if static_cache_params is not None: task = _resolve_static_cache_task(model_type) elif task is None: @@ -278,6 +284,11 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: trust_remote_code=trust_remote_code, execution_provider=execution_provider, text_only=args.text_only, + config_overrides=( + {"use_dsa": False, "num_nextn_predict_layers": 0} + if args.glm_full_attention + else None + ), ) _save_package(pkg, output_dir, args, optimize, component_filter) @@ -615,6 +626,11 @@ def main(argv: list[str] | None = None) -> None: action="store_true", help="Do not include weights in the output ONNX model.", ) + build_parser.add_argument( + "--glm-full-attention", + action="store_true", + help="Disable GLM-5.2 IndexShare DSA and export the dense MLA fallback.", + ) build_parser.add_argument( "--trust-remote-code", action="store_true", diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py index 4638ecd4..6b49d176 100644 --- a/src/mobius/_builder.py +++ b/src/mobius/_builder.py @@ -355,6 +355,7 @@ def build( execution_provider: str = "default", trace_optimization: bool = False, text_only: bool = False, + config_overrides: dict[str, object] | None = None, ) -> ModelPackage: """Build an ONNX :class:`ModelPackage` from a HuggingFace model ID. @@ -418,6 +419,8 @@ def build( Raises :class:`ValueError` if the resolved ``model_type`` has no text-only sibling. Currently supported for ``gemma4_unified`` (``google/gemma-4-12B``). + config_overrides: Optional dataclass field overrides applied after + HuggingFace config extraction and dtype resolution. Returns: A :class:`ModelPackage` containing the built model(s). @@ -569,6 +572,8 @@ def build( if dtype is not None: dtype = resolve_dtype(dtype) config = dataclasses.replace(config, dtype=dtype) + if config_overrides: + config = dataclasses.replace(config, **config_overrides) if output_layer_indices is not None: # Opt-in: emit additional `hidden_states.{k}` ONNX outputs for each diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 2ecbc655..c3fce95d 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -467,6 +467,11 @@ class ArchitectureConfig(BaseModelConfig): topk_method: str = "greedy" first_k_dense_replace: int = 0 n_shared_experts: int | None = None + mlp_layer_types: list[str] | None = None + # When True (and quantization is active), routed MoE experts are emitted as a + # single fused ``com.microsoft::QMoE`` op instead of a per-expert unroll of + # ``MatMulNBits``. Only wired for the GLM/DeepSeek MoE path today. + fused_quantized_moe: bool = False # Multi-head Latent Attention (MLA) config — DeepSeek-V2/V3 q_lora_rank: int | None = None @@ -476,12 +481,20 @@ class ArchitectureConfig(BaseModelConfig): v_head_dim: int | None = None rope_interleave: bool = False + # Deep Sparse Attention / IndexShare config — GLM-5.2 and DeepSeek-V4 CSA. + use_dsa: bool = True + index_topk: int | None = None + index_head_dim: int | None = None + index_n_heads: int | None = None + index_topk_freq: int = 4 + index_skip_topk_offset: int = 3 + indexer_rope_interleave: bool = True + indexer_types: list[str] | None = None + index_share_for_mtp_iteration: bool = False + # DeepSeek-V4 compressed sparse attention / Hyper-Connections. o_groups: int = 1 o_lora_rank: int | None = None - index_n_heads: int | None = None - index_head_dim: int | None = None - index_topk: int | None = None compress_ratios: list[int] | None = None compress_rope_theta: float | None = None hc_mult: int = 1 @@ -650,7 +663,7 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: config, "rope_interleave", (getattr(config, "qk_rope_head_dim", None) or 0) > 0 - or model_type in ("glm", "glm4", "glm4_moe", "chatglm"), + or model_type in ("glm", "glm4", "glm4_moe", "glm_moe_dsa", "chatglm"), ) if rope_config is not None: rope_config = dataclasses.replace(rope_config, rope_interleave=rope_interleave) @@ -881,18 +894,28 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: topk_method=getattr(config, "topk_method", "greedy"), first_k_dense_replace=getattr(config, "first_k_dense_replace", 0), n_shared_experts=getattr(config, "n_shared_experts", None), + mlp_layer_types=getattr(config, "mlp_layer_types", None), # Multi-head Latent Attention (MLA) q_lora_rank=getattr(config, "q_lora_rank", None), kv_lora_rank=getattr(config, "kv_lora_rank", None), qk_nope_head_dim=getattr(config, "qk_nope_head_dim", None), qk_rope_head_dim=getattr(config, "qk_rope_head_dim", None), v_head_dim=getattr(config, "v_head_dim", None), + # Deep Sparse Attention / IndexShare + use_dsa=getattr(config, "use_dsa", True), + index_topk=getattr(config, "index_topk", None), + index_head_dim=getattr(config, "index_head_dim", None), + index_n_heads=getattr(config, "index_n_heads", None), + index_topk_freq=getattr(config, "index_topk_freq", 4), + index_skip_topk_offset=getattr(config, "index_skip_topk_offset", 3), + indexer_rope_interleave=getattr(config, "indexer_rope_interleave", True), + indexer_types=getattr(config, "indexer_types", None), + index_share_for_mtp_iteration=getattr( + config, "index_share_for_mtp_iteration", False + ), # DeepSeek-V4 compressed sparse attention / Hyper-Connections o_groups=getattr(config, "o_groups", 1), o_lora_rank=getattr(config, "o_lora_rank", None), - index_n_heads=getattr(config, "index_n_heads", None), - index_head_dim=getattr(config, "index_head_dim", None), - index_topk=getattr(config, "index_topk", None), compress_ratios=_deepseek_v4_compress_ratios(config), compress_rope_theta=getattr(config, "compress_rope_theta", None), hc_mult=getattr(config, "hc_mult", 1), diff --git a/src/mobius/_configs_test.py b/src/mobius/_configs_test.py index da97a82e..410b7efd 100644 --- a/src/mobius/_configs_test.py +++ b/src/mobius/_configs_test.py @@ -1085,6 +1085,23 @@ def test_boa_token_id_none_when_absent(self): class TestActivationFallbacks: """Tests for hidden_act extraction fallbacks (ff_activation, gelu_activation).""" + def test_chatglm_silu_fallback(self): + """ChatGLM hardcodes SiLU without exposing an activation field.""" + + class FakeConfig: + model_type = "chatglm" + num_attention_heads = 8 + num_key_value_heads = 8 + num_hidden_layers = 2 + vocab_size = 1000 + hidden_size = 256 + intermediate_size = 512 + max_position_embeddings = 1024 + head_dim = 32 + + config = ArchitectureConfig.from_transformers(FakeConfig()) + assert config.hidden_act == "silu" + def test_ff_activation_fallback(self): """ff_activation is used when hidden_act is absent (XLNet pattern).""" diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index a9597d76..3dc32183 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -57,6 +57,7 @@ Glm4CausalLMModel, Glm4MoECausalLMModel, GlmCausalLMModel, + GlmMoeDsaCausalLMModel, GPTOSSCausalLMModel, GraniteCausalLMModel, GraniteMoECausalLMModel, @@ -526,6 +527,7 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: "ernie4_5_moe": ModelRegistration(Ernie45MoECausalLMModel), "flex_olmo": ModelRegistration(MoECausalLMModel), "glm4_moe": ModelRegistration(Glm4MoECausalLMModel), + "glm_moe_dsa": ModelRegistration(GlmMoeDsaCausalLMModel, task="glm-moe-dsa"), "granitemoe": ModelRegistration(GraniteMoECausalLMModel), "granitemoehybrid": ModelRegistration(GraniteMoeHybridCausalLMModel), "granitemoeshared": ModelRegistration(GraniteMoECausalLMModel), @@ -905,6 +907,7 @@ def _create_default_registry() -> ModelRegistry: "ernie4_5_moe": "baidu/ERNIE-4.5-21B-A3B-PT", "flex_olmo": "allenai/Flex-reddit-2x7B-1T", "glm4_moe": "zai-org/GLM-4.5-Air", + "glm_moe_dsa": "zai-org/GLM-5.2", "granitemoehybrid": "ibm-granite/granite-4.0-tiny-preview", "granitemoeshared": "ibm-research/moe-7b-1b-active-shared-experts", "qwen3_omni_moe": "Qwen/Qwen3-Omni-30B-A3B-Instruct", @@ -1235,6 +1238,7 @@ def _create_default_registry() -> ModelRegistry: "deepseek_v2_moe": "mla+moe", "deepseek_v3": "mla+moe", "deepseek_v4": "dense-csa-fallback+mtp+moe+hc", + "glm_moe_dsa": "mla+moe-indexshare-dsa+mtp", "phi3small": "blocksparse", "falcon_h1": "hybrid-ssm", "mamba": "ssm", diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index 4362e374..4fe47727 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -48,6 +48,7 @@ "MLP", "MLPMultiModalProjector", "MoELayer", + "FusedQuantizedMoE", "OffsetRMSNorm", "PatchEmbed", "PatchEmbedding", @@ -163,6 +164,7 @@ from mobius.components._mamba_block import MambaBlock as MambaBlock from mobius.components._mlp import FCMLP, MLP, FusedGateUpMLP, GatedMLP from mobius.components._moe import ( + FusedQuantizedMoE, MoELayer, SigmoidTopKGate, SoftmaxTopKGate, diff --git a/src/mobius/components/_deepseek_mla_test.py b/src/mobius/components/_deepseek_mla_test.py index 3946ec41..508d457e 100644 --- a/src/mobius/components/_deepseek_mla_test.py +++ b/src/mobius/components/_deepseek_mla_test.py @@ -5,8 +5,10 @@ from __future__ import annotations +import onnx_ir as ir import pytest +from mobius._builder import build_from_module from mobius._testing import ( count_op_type, create_test_builder, @@ -14,6 +16,7 @@ make_config, ) from mobius.components._deepseek_mla import DeepSeekMLA +from mobius.models.deepseek import DeepSeekV3CausalLMModel # DeepSeek MLA config: uses low-rank KV compression with separate # nope (non-positional) and rope (rotary) head dimensions. @@ -31,6 +34,29 @@ ) +@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16]) +def test_low_precision_export_matches_attention_auxiliary_dtypes(dtype): + config = _mla_config( + dtype=dtype, + intermediate_size=128, + num_hidden_layers=1, + num_key_value_heads=4, + head_dim=8, + first_k_dense_replace=1, + num_local_experts=None, + ) + graph = build_from_module(DeepSeekV3CausalLMModel(config), config)["model"].graph + + rotary_nodes = [node for node in graph if node.op_type == "RotaryEmbedding"] + assert len(rotary_nodes) == 2 + for node in rotary_nodes: + assert [value.dtype for value in node.inputs] == [dtype, dtype, dtype] + + attention = next(node for node in graph if node.op_type == "Attention") + assert attention.inputs[3].dtype == dtype + assert attention.inputs[3].dtype == attention.inputs[0].dtype + + def _mla_config(**overrides): """Create a test config for DeepSeek MLA.""" kw = {**_MLA_DEFAULTS, **overrides} diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py index fe1edc05..fd0b49df 100644 --- a/src/mobius/components/_moe.py +++ b/src/mobius/components/_moe.py @@ -7,13 +7,21 @@ import dataclasses import math +import re +from typing import TYPE_CHECKING import onnx_ir as ir +import torch from onnxscript import OpBuilder, nn +from mobius._build_context import ep_capabilities from mobius._configs import ArchitectureConfig, QuantizationConfig +from mobius._weight_utils import preprocess_awq_weights, preprocess_gptq_weights from mobius.components._mlp import MLP +if TYPE_CHECKING: + pass + class TopKGate(nn.Module): """Standard top-k expert routing gate. @@ -267,7 +275,13 @@ class MoELayer(nn.Module): and compatibility path, not the grouped-expert performance representation. """ - def __init__(self, config: ArchitectureConfig, gate: nn.Module | None = None): + def __init__( + self, + config: ArchitectureConfig, + gate: nn.Module | None = None, + linear_class: type | None = None, + enable_qmoe: bool = True, + ): super().__init__() assert config.num_local_experts is not None assert config.num_experts_per_tok is not None @@ -284,11 +298,17 @@ def __init__(self, config: ArchitectureConfig, gate: nn.Module | None = None): if config.moe_intermediate_size is not None else config ) - if self._qmoe_quantization is not None and hasattr(self.gate, "qmoe_routing"): + if ( + enable_qmoe + and self._qmoe_quantization is not None + and hasattr(self.gate, "qmoe_routing") + ): self.experts = None self._init_qmoe_parameters(expert_config) else: - self.experts = nn.ModuleList([MLP(expert_config) for _ in range(self.num_experts)]) + self.experts = nn.ModuleList( + [MLP(expert_config, linear_class=linear_class) for _ in range(self.num_experts)] + ) def _init_qmoe_parameters(self, expert_config: ArchitectureConfig) -> None: quantization = self._qmoe_quantization @@ -372,6 +392,7 @@ def _qmoe_forward(self, op: OpBuilder, hidden_states: ir.Value): result = op.Mul(result, op.CastLike(output_scale, result)) return result + def forward(self, op: OpBuilder, hidden_states: ir.Value): if self.experts is None: return self._qmoe_forward(op, hidden_states) @@ -395,6 +416,376 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): return result +class FusedQuantizedMoE(nn.Module): + """Routed MoE experts emitted as a single fused ``com.microsoft::QMoE`` op. + + Replaces the per-expert ``MatMulNBits`` unroll (:class:`MoELayer`) for + weight-only int-quantized MoE. All routed experts are packed into + expert-major integer weight tensors laid out exactly as the ORT contrib + ``QMoE`` kernel expects:: + + fc1_experts_weights: [E, 2*inter, hidden // pack_size] uint8 + fc1_scales: [E, 2*inter, hidden // block_size] float32 + fc2_experts_weights: [E, hidden, inter // pack_size] uint8 + fc2_scales: [E, hidden, inter // block_size] float32 + + where ``pack_size = 8 // bits``. ``fc1`` fuses the SwiGLU gate/up + projections in the **interleaved** layout ``[g_0, u_0, g_1, u_1, ...]`` and is + consumed with ``swiglu_fusion=1`` (the only SwiGLU layout the ORT CPU QMoE + kernel supports), so the kernel computes ``silu(gate) * up`` — matching + GLM/DeepSeek's SiLU-gated experts (``activation_alpha=1``, ``activation_beta=0``, + ``swiglu_limit=inf``, all ORT defaults). + + Routing is delegated to ``gate.route_for_qmoe`` (GLM sigmoid + noaux_tc / + DeepSeek softmax group-limited). The kernel re-derives top-k selection from + ``router_probs`` (the gate's ``scores_for_choice``); the exact combine + weights are scattered into a dense ``[rows, E]`` aggregation tensor and fed + through the optional ``router_weights`` input with + ``normalize_routing_weights=0``, so the fused op reproduces the per-expert + path's routing bit-for-bit. + + Symmetric quantization omits zero-point inputs and uses the kernel's implicit + midpoint. Asymmetric quantization supplies packed per-block zero-points. + """ + + _MICROSOFT_DOMAIN = "com.microsoft" + + def __init__( + self, + config: ArchitectureConfig, + gate: nn.Module, + ): + super().__init__() + assert config.num_local_experts is not None + assert config.num_experts_per_tok is not None + assert config.moe_intermediate_size is not None + assert config.quantization is not None + if not hasattr(gate, "route_for_qmoe"): + raise TypeError( + f"gate {type(gate).__name__} does not support the fused QMoE path " + "(missing route_for_qmoe); use MoELayer instead" + ) + + qc = config.quantization + bits = qc.bits + block_size = qc.group_size + if bits not in (1, 2, 4, 8): + raise ValueError(f"QMoE expert_weight_bits must be 1/2/4/8, got {bits}") + if block_size < 16 or (block_size & (block_size - 1)): + raise ValueError(f"QMoE block_size must be a power of 2 >= 16, got {block_size}") + + self.gate = gate + self._num_experts = config.num_local_experts + self._top_k = config.num_experts_per_tok + self._hidden = config.hidden_size + self._inter = config.moe_intermediate_size + self._bits = bits + self._block_size = block_size + + if self._hidden % block_size != 0: + raise ValueError( + f"hidden_size {self._hidden} must be divisible by block_size {block_size}" + ) + if self._inter % block_size != 0: + raise ValueError( + f"moe_intermediate_size {self._inter} must be divisible by " + f"block_size {block_size}" + ) + + pack_size = 8 // bits + e = self._num_experts + fc1_out = 2 * self._inter # interleaved [g_0, u_0, ...] for swiglu_fusion=1 + # fc1: [E, 2*inter, hidden] quantized along hidden (K) + self.fc1_experts_weights = nn.Parameter( + [e, fc1_out, self._hidden // pack_size], + dtype=ir.DataType.UINT8, + ) + self.fc1_scales = nn.Parameter( + [e, fc1_out, self._hidden // block_size], + dtype=ir.DataType.FLOAT, + ) + # fc2: [E, hidden, inter] quantized along inter (K) + self.fc2_experts_weights = nn.Parameter( + [e, self._hidden, self._inter // pack_size], + dtype=ir.DataType.UINT8, + ) + self.fc2_scales = nn.Parameter( + [e, self._hidden, self._inter // block_size], + dtype=ir.DataType.FLOAT, + ) + if qc.sym: + self.fc1_experts_zero_points = None + self.fc2_experts_zero_points = None + else: + self.fc1_experts_zero_points = nn.Parameter( + [ + e, + fc1_out, + math.ceil((self._hidden // block_size) * bits / 8), + ], + dtype=ir.DataType.UINT8, + ) + self.fc2_experts_zero_points = nn.Parameter( + [ + e, + self._hidden, + math.ceil((self._inter // block_size) * bits / 8), + ], + dtype=ir.DataType.UINT8, + ) + + def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: + if ep_capabilities().name == "cuda": + raise ValueError( + "FusedQuantizedMoE is disabled for CUDA because ORT CUDA QMoE " + "currently ignores router_weights (input 14), which is required " + "for GLM/DeepSeek group-limited routing. Use the decomposed " + "MatMulNBits export path for CUDA." + ) + hidden = self._hidden + + # QMoE requires 2-D router_probs, so flatten [B, S, H] -> [rows, H]. + orig_shape = op.Shape(hidden_states) + flat = op.Reshape(hidden_states, op.Constant(value_ints=[-1, hidden])) + # Router math must be float32, but QMoE activation input stays in the + # model dtype: CUDA QMoE kernels are registered for fp16/bf16 inputs. + flat_f32 = op.Cast(flat, to=1) + + scores_for_choice, routing_weights, selected_experts = self._route(op, flat_f32) + + # Dense [rows, E] aggregation weights: combine weight at each selected + # expert position, 0 elsewhere. QMoE reads these at its own top-k picks. + zeros = op.Mul(scores_for_choice, 0.0) + aggregation = op.ScatterElements(zeros, selected_experts, routing_weights, axis=-1) + + moe_out = op.QMoE( + flat, # 0: input + scores_for_choice, # 1: router_probs (selection logits) + self.fc1_experts_weights, # 2 + op.Cast(self.fc1_scales, to=1), # 3 + None, # 4: fc1_experts_bias + self.fc2_experts_weights, # 5 + op.Cast(self.fc2_scales, to=1), # 6 + None, # 7: fc2_experts_bias + None, # 8: fc3_experts_weights + None, # 9: fc3_scales + None, # 10: fc3_experts_bias + self.fc1_experts_zero_points, # 11 + self.fc2_experts_zero_points, # 12 + None, # 13: fc3_zero_points + aggregation, # 14: router_weights (explicit combine weights) + activation_type="swiglu", + k=self._top_k, + normalize_routing_weights=0, + swiglu_fusion=1, + expert_weight_bits=self._bits, + block_size=self._block_size, + quant_type="int", + _domain=self._MICROSOFT_DOMAIN, + ) + moe_out = op.CastLike(moe_out, hidden_states) + return op.Reshape(moe_out, orig_shape) + + def _route(self, op: OpBuilder, hidden_states: ir.Value): + """Invoke ``gate.route_for_qmoe`` under the gate's module scope. + + ``route_for_qmoe`` is a plain method (not ``forward``), so it never goes + through :meth:`nn.Module.__call__`. We replicate the parameter-realization + step here so the gate's ``weight`` / ``e_score_correction_bias`` are + registered as graph initializers under the ``...moe.gate`` scope (matching + the per-expert path) instead of dangling as unqualified names. + """ + builder = op.builder + module_name = self.gate._name or "gate" + class_name = type(self.gate).__qualname__ + builder.push_module(module_name, class_name) + try: + for param in self.gate._parameters.values(): + param._realize(builder) + return self.gate.route_for_qmoe(op, hidden_states) + finally: + builder.pop_module() + + +_PER_EXPERT_RE = re.compile( + r"^(?P.*\.mlp)\.experts\.(?P\d+)\." + r"(?Pgate_proj|up_proj|down_proj)\." + r"(?Pweight|scales|zero_points)$" +) +_FUSED_EXPERT_RE = re.compile( + r"^(?P.*\.mlp)\.experts\." + r"(?Pgate_up_proj|down_proj)\." + r"(?Pqweight|qzeros|g_idx|weight|scales|zero_points)$" +) + + +def pack_fused_quantized_moe_weights( + state_dict: dict[str, torch.Tensor], + config: ArchitectureConfig, +) -> dict[str, torch.Tensor]: + """Pack routed expert tensors into the expert-major QMoE initializer ABI. + + Accepts either fused 3-D GPTQ/AWQ checkpoint tensors or already-repacked + per-expert GGUF/MatMulNBits tensors. Gate/up rows are interleaved for + ``swiglu_fusion=1``; shared-expert tensors are left untouched. + """ + quantization = config.quantization + if quantization is None: + raise ValueError("Fused QMoE packing requires quantization settings") + + result: dict[str, torch.Tensor] = {} + fused: dict[str, dict[str, dict[str, torch.Tensor]]] = {} + per_expert: dict[str, dict[int, dict[str, dict[str, torch.Tensor]]]] = {} + for key, value in state_dict.items(): + match = _FUSED_EXPERT_RE.match(key) + if match is not None: + fused.setdefault(match["prefix"], {}).setdefault( + match["projection"], {} + )[match["kind"]] = value + continue + match = _PER_EXPERT_RE.match(key) + if match is not None: + expert = int(match["expert"]) + per_expert.setdefault(match["prefix"], {}).setdefault(expert, {}).setdefault( + match["projection"], {} + )[match["kind"]] = value + continue + result[key] = value + + prefixes = set(fused) | set(per_expert) + for prefix in prefixes: + if prefix in fused and prefix in per_expert: + raise ValueError(f"Mixed fused and per-expert weights under {prefix}") + if prefix in fused: + gate_up = _prepare_fused_projection( + fused[prefix].get("gate_up_proj"), + quantization.quant_method, + quantization.bits, + quantization.group_size, + ) + down = _prepare_fused_projection( + fused[prefix].get("down_proj"), + quantization.quant_method, + quantization.bits, + quantization.group_size, + ) + fc1 = { + kind: _interleave_gate_up(value) + for kind, value in gate_up.items() + } + fc2 = down + else: + experts = per_expert[prefix] + expected = set(range(config.num_local_experts or 0)) + if set(experts) != expected: + raise ValueError( + f"{prefix} has expert indices {sorted(experts)}, " + f"expected {sorted(expected)}" + ) + fc1, fc2 = _stack_per_expert_projections(experts) + + _store_qmoe_projection( + result, + f"{prefix}.moe.fc1", + fc1, + symmetric=quantization.sym, + ) + _store_qmoe_projection( + result, + f"{prefix}.moe.fc2", + fc2, + symmetric=quantization.sym, + ) + return result + + +def _prepare_fused_projection( + tensors: dict[str, torch.Tensor] | None, + quant_method: str, + bits: int, + block_size: int, +) -> dict[str, torch.Tensor]: + if tensors is None: + raise ValueError("Missing fused routed-expert projection") + if "qweight" in tensors: + preprocess = ( + preprocess_gptq_weights if quant_method == "gptq" else preprocess_awq_weights + ) + num_experts = tensors["qweight"].shape[0] + processed: dict[str, list[torch.Tensor]] = {} + for expert in range(num_experts): + expert_state = {} + for kind in {"qweight", "qzeros", "scales", "g_idx"} & tensors.keys(): + value = tensors[kind] + expert_state[f"projection.{kind}"] = ( + value if kind == "g_idx" and value.ndim == 1 else value[expert] + ) + for key, value in preprocess( + expert_state, bits=bits, group_size=block_size + ).items(): + processed.setdefault(key.rsplit(".", 1)[-1], []).append(value) + return {kind: torch.stack(values) for kind, values in processed.items()} + required = {"weight", "scales"} + if not required.issubset(tensors): + raise ValueError(f"Missing fused projection tensors: {required - tensors.keys()}") + return tensors + + +def _stack_per_expert_projections( + experts: dict[int, dict[str, dict[str, torch.Tensor]]], +) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + fc1: dict[str, list[torch.Tensor]] = {} + fc2: dict[str, list[torch.Tensor]] = {} + for expert in sorted(experts): + projections = experts[expert] + for name in ("gate_proj", "up_proj", "down_proj"): + if name not in projections: + raise ValueError(f"Expert {expert} is missing {name}") + kinds = set(projections["gate_proj"]) | set(projections["up_proj"]) + for kind in kinds: + if kind not in projections["gate_proj"] or kind not in projections["up_proj"]: + raise ValueError(f"Expert {expert} gate/up {kind} tensors are incomplete") + gate = projections["gate_proj"][kind] + up = projections["up_proj"][kind] + fc1.setdefault(kind, []).append( + torch.stack((gate, up), dim=1).flatten(0, 1) + ) + for kind, value in projections["down_proj"].items(): + fc2.setdefault(kind, []).append(value) + return ( + {kind: torch.stack(values) for kind, values in fc1.items()}, + {kind: torch.stack(values) for kind, values in fc2.items()}, + ) + + +def _interleave_gate_up(value: torch.Tensor) -> torch.Tensor: + if value.shape[1] % 2: + raise ValueError(f"gate_up projection has odd output size {value.shape[1]}") + intermediate = value.shape[1] // 2 + return value.reshape(value.shape[0], 2, intermediate, *value.shape[2:]).transpose( + 1, 2 + ).flatten(1, 2) + + +def _store_qmoe_projection( + result: dict[str, torch.Tensor], + stem: str, + tensors: dict[str, torch.Tensor], + *, + symmetric: bool, +) -> None: + weight = tensors["weight"] + if weight.ndim == 4: + weight = weight.flatten(-2) + result[f"{stem}_experts_weights"] = weight + result[f"{stem}_scales"] = tensors["scales"].float() + zero_points = tensors.get("zero_points") + if not symmetric: + if zero_points is None: + raise ValueError(f"Asymmetric QMoE projection {stem} requires zero-points") + result[f"{stem}_experts_zero_points"] = zero_points + + def _supported_qmoe_quantization( quantization: QuantizationConfig | None, ) -> QuantizationConfig | None: diff --git a/src/mobius/components/_moe_test.py b/src/mobius/components/_moe_test.py index 16c3f8df..a6faab6a 100644 --- a/src/mobius/components/_moe_test.py +++ b/src/mobius/components/_moe_test.py @@ -9,15 +9,22 @@ import pytest import torch -from mobius._builder import _cast_module_dtype +from mobius._builder import build_from_module from mobius._configs import QuantizationConfig from mobius._testing import ( + count_op_type, create_test_builder, create_test_input, make_config, ) -from mobius._weight_utils import pack_qmoe_expert_weights, preprocess_gptq_weights -from mobius.components._moe import MoELayer, SparseMixerGate, TopKGate +from mobius.components._moe import ( + FusedQuantizedMoE, + MoELayer, + SparseMixerGate, + TopKGate, + pack_fused_quantized_moe_weights, +) +from mobius.models.deepseek import DeepSeekMoEGate, DeepSeekV3CausalLMModel class TestTopKGate: @@ -70,6 +77,32 @@ def test_gate_top_k_1(self): assert graph.num_nodes() > 0 +@pytest.mark.parametrize("route_for_qmoe", [False, True]) +def test_deepseek_gate_casts_both_router_matmul_inputs_to_float32(route_for_qmoe): + config = make_config( + dtype=ir.DataType.FLOAT16, + num_local_experts=4, + num_experts_per_tok=2, + ) + gate = DeepSeekMoEGate(config) + builder, op, graph = create_test_builder() + hidden = create_test_input( + builder, + "hidden", + [1, 2, config.hidden_size], + dtype=ir.DataType.FLOAT16, + ) + + outputs = gate.route_for_qmoe(op, hidden) if route_for_qmoe else gate(op, hidden) + builder._adapt_outputs(list(outputs), "") + + router_matmul = next(node for node in graph if node.op_type == "MatMul") + assert [value.dtype for value in router_matmul.inputs] == [ + ir.DataType.FLOAT, + ir.DataType.FLOAT, + ] + + class TestMoELayer: def test_moe_layer_has_gate(self): config = make_config(num_local_experts=4, num_experts_per_tok=2) @@ -123,98 +156,276 @@ def test_moe_layer_forward_with_sparse_mixer_gate(self): builder._adapt_outputs([result], "") assert graph.num_nodes() > 0 - def test_int4_moe_emits_expert_major_qmoe(self): - config = make_config( - hidden_size=64, - intermediate_size=32, - moe_intermediate_size=32, - num_local_experts=64, - num_experts_per_tok=6, - quantization=QuantizationConfig( - bits=4, - group_size=32, - quant_method="gptq", - sym=False, - ), - ) - layer = MoELayer(config) - builder, op, graph = create_test_builder() - hidden = create_test_input(builder, "hidden", [1, 8, 64]) - result = layer(op, hidden) - builder._adapt_outputs([result], "") - qmoe_nodes = [node for node in graph if node.op_type == "QMoE"] - assert len(qmoe_nodes) == 1 - qmoe = qmoe_nodes[0] - assert qmoe.attributes["k"].value == 6 - assert qmoe.attributes["block_size"].value == 32 - assert qmoe.attributes["swiglu_fusion"].value == 2 - assert layer.fc1_experts_weights.shape == ir.Shape([64, 64, 32]) - assert layer.fc1_scales.shape == ir.Shape([64, 64, 2]) - assert layer.fc1_experts_zero_points.shape == ir.Shape([64, 64, 1]) - assert layer.fc2_experts_weights.shape == ir.Shape([64, 64, 16]) - assert layer.fc2_scales.shape == ir.Shape([64, 64, 1]) - assert layer.fc2_experts_zero_points.shape == ir.Shape([64, 64, 1]) - assert sum(node.op_type == "MatMul" for node in graph) == 1 # router only - - _cast_module_dtype(layer, ir.DataType.FLOAT16) - assert layer.gate.weight.dtype == ir.DataType.FLOAT16 - assert layer.fc1_scales.dtype == ir.DataType.FLOAT - assert layer.fc2_scales.dtype == ir.DataType.FLOAT - - def test_expert_major_packing_matches_static_64_expert_top6_reference(self): - """Packed QMoE math matches the existing loop-over-experts semantics.""" - torch.manual_seed(0) - num_experts, top_k = 64, 6 - hidden_size, intermediate_size, block_size = 32, 16, 16 - fc1_codes = torch.randint(0, 16, (num_experts, 2 * intermediate_size, hidden_size)) - fc2_codes = torch.randint(0, 16, (num_experts, hidden_size, intermediate_size)) - fc1_scales = torch.rand(num_experts, 2 * intermediate_size, hidden_size // block_size) - fc2_scales = torch.rand(num_experts, hidden_size, intermediate_size // block_size) - raw = { - "model.layers.1.mlp.experts.gate_up_proj.qweight": _to_gptq_qweight(fc1_codes), - "model.layers.1.mlp.experts.gate_up_proj.scales": fc1_scales.transpose(-1, -2), - "model.layers.1.mlp.experts.down_proj.qweight": _to_gptq_qweight(fc2_codes), - "model.layers.1.mlp.experts.down_proj.scales": fc2_scales.transpose(-1, -2), - } - packed = pack_qmoe_expert_weights( - preprocess_gptq_weights(raw, bits=4, group_size=block_size) +def test_fused_qmoe_wires_asymmetric_zero_points(): + config = make_config( + hidden_size=32, + intermediate_size=16, + moe_intermediate_size=16, + num_local_experts=4, + num_experts_per_tok=2, + quantization=QuantizationConfig( + bits=4, + group_size=16, + quant_method="gptq", + sym=False, + ), + ) + layer = FusedQuantizedMoE(config, DeepSeekMoEGate(config)) + builder, op, graph = create_test_builder() + hidden = create_test_input(builder, "hidden", [1, 2, 32]) + builder._adapt_outputs([layer(op, hidden)], "") + + qmoe = next(node for node in graph if node.op_type == "QMoE") + assert qmoe.inputs[11] is not None + assert qmoe.inputs[12] is not None + assert qmoe.inputs[3] is not None + assert qmoe.inputs[3].producer().op_type == "Cast" + assert qmoe.inputs[3].dtype == ir.DataType.FLOAT + assert layer.fc1_experts_zero_points.shape == ir.Shape([4, 32, 1]) + assert layer.fc2_experts_zero_points.shape == ir.Shape([4, 32, 1]) + + +def test_fused_qmoe_keeps_activation_input_model_dtype(): + config = make_config( + hidden_size=32, + intermediate_size=16, + moe_intermediate_size=16, + num_local_experts=4, + num_experts_per_tok=2, + quantization=QuantizationConfig( + bits=4, + group_size=16, + quant_method="gptq", + sym=True, + ), + ) + layer = FusedQuantizedMoE(config, DeepSeekMoEGate(config)) + builder, op, graph = create_test_builder() + hidden = create_test_input(builder, "hidden", [1, 2, 32], dtype=ir.DataType.FLOAT16) + builder._adapt_outputs([layer(op, hidden)], "") + + qmoe = next(node for node in graph if node.op_type == "QMoE") + assert qmoe.inputs[0].dtype == ir.DataType.FLOAT16 + assert qmoe.inputs[1].dtype == ir.DataType.FLOAT + + +def test_deepseek_fused_qmoe_graph_has_one_node_per_moe_layer(): + config = make_config( + hidden_size=32, + intermediate_size=32, + moe_intermediate_size=16, + num_hidden_layers=2, + first_k_dense_replace=1, + num_local_experts=4, + num_experts_per_tok=2, + n_shared_experts=1, + fused_quantized_moe=True, + quantization=QuantizationConfig( + bits=4, + group_size=16, + quant_method="gptq", + sym=True, + ), + ) + graph = build_from_module(DeepSeekV3CausalLMModel(config), config)["model"].graph + + assert count_op_type(graph, "QMoE") == 1 + assert not any( + value is not None + and value.name is not None + and ".moe.experts." in value.name + for node in graph + for value in node.inputs + ) + shared_prefix = "model.layers.1.mlp.shared_experts." + assert ( + sum( + node.op_type == "MatMulNBits" + and any( + value is not None + and value.name is not None + and shared_prefix in value.name + for value in node.inputs + ) + for node in graph ) - fc1_packed = packed["model.layers.1.mlp.moe.fc1_experts_weights"] - fc2_packed = packed["model.layers.1.mlp.moe.fc2_experts_weights"] - assert fc1_packed.shape == (num_experts, 2 * intermediate_size, hidden_size // 2) - assert fc2_packed.shape == (num_experts, hidden_size, intermediate_size // 2) - - packed_fc1 = _dequant_qmoe(fc1_packed, fc1_scales, block_size) - packed_fc2 = _dequant_qmoe(fc2_packed, fc2_scales, block_size) - static_fc1 = _dequant_codes(fc1_codes, fc1_scales, block_size) - static_fc2 = _dequant_codes(fc2_codes, fc2_scales, block_size) - - hidden = torch.randn(3, hidden_size) - router_probs = torch.softmax(torch.randn(3, num_experts), dim=-1) - selected_weights, selected_experts = router_probs.topk(top_k, dim=-1) - selected_weights /= selected_weights.sum(dim=-1, keepdim=True) - static = _static_moe( - hidden, - selected_weights, - selected_experts, - static_fc1, - static_fc2, - intermediate_size, + == 3 + ) + + +def test_deepseek_fused_qmoe_rejects_cuda(): + config = make_config( + hidden_size=32, + intermediate_size=32, + moe_intermediate_size=16, + num_hidden_layers=2, + first_k_dense_replace=1, + num_local_experts=4, + num_experts_per_tok=2, + n_shared_experts=1, + fused_quantized_moe=True, + quantization=QuantizationConfig( + bits=4, + group_size=16, + quant_method="gguf", + sym=True, + ), + ) + with pytest.raises(ValueError, match=r"CUDA.*ignores router_weights"): + build_from_module( + DeepSeekV3CausalLMModel(config), + config, + execution_provider="cuda", ) - fused_reference = _static_moe( - hidden, - selected_weights, - selected_experts, - packed_fc1, - packed_fc2, - intermediate_size, - ) - torch.testing.assert_close(fused_reference, static, atol=1e-5, rtol=1e-5) + + +def test_expert_major_packing_matches_static_64_expert_top6_reference(): + torch.manual_seed(0) + experts, top_k = 64, 6 + hidden_size, intermediate_size, block_size = 32, 16, 16 + gate_codes = torch.randint(0, 16, (experts, intermediate_size, hidden_size)) + up_codes = torch.randint(0, 16, (experts, intermediate_size, hidden_size)) + down_codes = torch.randint(0, 16, (experts, hidden_size, intermediate_size)) + gate_scales = torch.rand(experts, intermediate_size, hidden_size // block_size) + up_scales = torch.rand(experts, intermediate_size, hidden_size // block_size) + down_scales = torch.rand(experts, hidden_size, intermediate_size // block_size) + state_dict = {} + for expert in range(experts): + prefix = f"model.layers.0.mlp.experts.{expert}" + for name, codes, scales in ( + ("gate_proj", gate_codes, gate_scales), + ("up_proj", up_codes, up_scales), + ("down_proj", down_codes, down_scales), + ): + state_dict[f"{prefix}.{name}.weight"] = _pack_matmul_nbits( + codes[expert], block_size + ) + state_dict[f"{prefix}.{name}.scales"] = scales[expert] + + config = make_config( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + moe_intermediate_size=intermediate_size, + num_local_experts=experts, + num_experts_per_tok=top_k, + fused_quantized_moe=True, + quantization=QuantizationConfig( + bits=4, + group_size=block_size, + quant_method="gguf", + sym=True, + ), + ) + packed = DeepSeekV3CausalLMModel(config).preprocess_weights(state_dict) + fc1_weight = packed["model.layers.0.mlp.moe.fc1_experts_weights"] + fc1_scales = packed["model.layers.0.mlp.moe.fc1_scales"] + fc2_weight = packed["model.layers.0.mlp.moe.fc2_experts_weights"] + fc2_scales = packed["model.layers.0.mlp.moe.fc2_scales"] + assert fc1_weight.shape == (experts, 2 * intermediate_size, hidden_size // 2) + assert fc2_weight.shape == (experts, hidden_size, intermediate_size // 2) + + packed_fc1 = _dequant(fc1_weight, fc1_scales, block_size) + packed_fc2 = _dequant(fc2_weight, fc2_scales, block_size) + static_fc1 = torch.stack( + ( + _dequant_codes(gate_codes, gate_scales, block_size), + _dequant_codes(up_codes, up_scales, block_size), + ), + dim=2, + ).flatten(1, 2) + static_fc2 = _dequant_codes(down_codes, down_scales, block_size) + + hidden = torch.randn(3, hidden_size) + router_probs = torch.softmax(torch.randn(3, experts), dim=-1) + weights, selected = router_probs.topk(top_k, dim=-1) + weights /= weights.sum(dim=-1, keepdim=True) + expected = _static_moe( + hidden, weights, selected, static_fc1, static_fc2, intermediate_size + ) + actual = _static_moe( + hidden, weights, selected, packed_fc1, packed_fc2, intermediate_size + ) + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("shared_g_idx", [False, True]) +def test_fused_gptq_checkpoint_tensors_pack_to_qmoe_layout(caplog, shared_g_idx): + experts, hidden_size, intermediate_size, block_size = 4, 32, 16, 16 + gate_up_codes = torch.randint( + 0, 16, (experts, 2 * intermediate_size, hidden_size) + ) + down_codes = torch.randint(0, 16, (experts, hidden_size, intermediate_size)) + gate_up_scales = torch.rand( + experts, 2 * intermediate_size, hidden_size // block_size + ) + down_scales = torch.rand( + experts, hidden_size, intermediate_size // block_size + ) + gate_up_g_idx = torch.arange(hidden_size, dtype=torch.int32) // block_size + gate_up_g_idx[0] = 1 + down_g_idx = torch.arange(intermediate_size, dtype=torch.int32) // block_size + if not shared_g_idx: + gate_up_g_idx = gate_up_g_idx.repeat(experts, 1) + down_g_idx = down_g_idx.repeat(experts, 1) + state_dict = { + "model.layers.0.mlp.experts.gate_up_proj.qweight": _to_gptq_qweight( + gate_up_codes + ), + "model.layers.0.mlp.experts.gate_up_proj.scales": gate_up_scales.transpose( + -1, -2 + ), + "model.layers.0.mlp.experts.gate_up_proj.g_idx": gate_up_g_idx, + "model.layers.0.mlp.experts.down_proj.qweight": _to_gptq_qweight( + down_codes + ), + "model.layers.0.mlp.experts.down_proj.scales": down_scales.transpose(-1, -2), + "model.layers.0.mlp.experts.down_proj.g_idx": down_g_idx, + } + config = make_config( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + moe_intermediate_size=intermediate_size, + num_local_experts=experts, + num_experts_per_tok=2, + fused_quantized_moe=True, + quantization=QuantizationConfig( + bits=4, + group_size=block_size, + quant_method="gptq", + sym=True, + ), + ) + + packed = pack_fused_quantized_moe_weights(state_dict, config) + fc1 = packed["model.layers.0.mlp.moe.fc1_experts_weights"] + scales = packed["model.layers.0.mlp.moe.fc1_scales"] + expected_codes = gate_up_codes.reshape( + experts, 2, intermediate_size, hidden_size + ).transpose(1, 2).flatten(1, 2) + expected_scales = gate_up_scales.reshape( + experts, 2, intermediate_size, hidden_size // block_size + ).transpose(1, 2).flatten(1, 2) + + assert fc1.shape == (experts, 2 * intermediate_size, hidden_size // 2) + torch.testing.assert_close( + _dequant(fc1, scales, block_size), + _dequant_codes(expected_codes, expected_scales, block_size), + ) + assert "desc_act models with non-trivial g_idx" in caplog.text + assert not any(key.endswith(".g_idx") for key in packed) + + +def _pack_matmul_nbits(codes: torch.Tensor, block_size: int) -> torch.Tensor: + low = codes[..., 0::2].to(torch.uint8) + high = codes[..., 1::2].to(torch.uint8) + return (low | (high << 4)).reshape( + codes.shape[0], codes.shape[1] // block_size, block_size // 2 + ) def _to_gptq_qweight(codes: torch.Tensor) -> torch.Tensor: - """Pack output-major int4 codes into GPTQ's [E, K/8, N] int32 layout.""" + shifts = torch.arange(8, dtype=torch.int64) * 4 packed = torch.sum( codes.to(torch.int64).reshape(*codes.shape[:-1], -1, 8) << shifts, @@ -223,16 +434,20 @@ def _to_gptq_qweight(codes: torch.Tensor) -> torch.Tensor: return packed.transpose(-1, -2).contiguous() -def _dequant_codes(codes: torch.Tensor, scales: torch.Tensor, block_size: int) -> torch.Tensor: +def _dequant_codes( + codes: torch.Tensor, scales: torch.Tensor, block_size: int +) -> torch.Tensor: + blocks = codes.shape[-1] // block_size values = codes.reshape(*codes.shape[:-1], blocks, block_size).float() - 8.0 return (values * scales.unsqueeze(-1)).flatten(-2) -def _dequant_qmoe(packed: torch.Tensor, scales: torch.Tensor, block_size: int) -> torch.Tensor: - low = packed & 0x0F - high = packed >> 4 - codes = torch.stack((low, high), dim=-1).flatten(-2) +def _dequant( + packed: torch.Tensor, scales: torch.Tensor, block_size: int +) -> torch.Tensor: + codes = torch.stack((packed & 0x0F, packed >> 4), dim=-1).flatten(-2) + return _dequant_codes(codes, scales, block_size) @@ -245,13 +460,16 @@ def _static_moe( intermediate_size: int, ) -> torch.Tensor: result = torch.zeros_like(hidden) - for expert_idx in range(fc1.shape[0]): - projected = hidden @ fc1[expert_idx].T - activated = torch.nn.functional.silu(projected[:, :intermediate_size]) - activated *= projected[:, intermediate_size:] - expert_output = activated @ fc2[expert_idx].T + for expert in range(fc1.shape[0]): + projected = hidden @ fc1[expert].T + activated = torch.nn.functional.silu(projected[:, 0::2]) + activated *= projected[:, 1::2] + assert activated.shape[-1] == intermediate_size + expert_output = activated @ fc2[expert].T weight = ( - routing_weights * (selected_experts == expert_idx).to(routing_weights.dtype) + routing_weights + * (selected_experts == expert).to(routing_weights.dtype) + ).sum(dim=-1, keepdim=True) result += expert_output * weight return result diff --git a/src/mobius/integrations/gguf/_builder.py b/src/mobius/integrations/gguf/_builder.py index b5aed1db..839dcb38 100644 --- a/src/mobius/integrations/gguf/_builder.py +++ b/src/mobius/integrations/gguf/_builder.py @@ -466,7 +466,19 @@ def _normalize_gguf_weights( import torch result: dict[str, torch.Tensor] = {} + kv_b_parts: dict[str, dict[str, torch.Tensor]] = {} for key, value in state_dict.items(): + for part in ("k", "v"): + suffix = f".self_attn.kv_b_proj.{part}_proj.weight" + if key.endswith(suffix): + prefix = key[: -len(suffix)] + kv_b_parts.setdefault(prefix, {})[part] = value + break + else: + part = None + if part is not None: + continue + # Stacked expert weights [num_experts, out, in] → per-expert unpacked = False for proj in ("gate_proj", "up_proj", "down_proj"): @@ -495,6 +507,10 @@ def _normalize_gguf_weights( result[key[: -len(".bias")]] = value continue + if key.endswith(".gate.e_score_correction_bias.bias"): + result[key[: -len(".bias")]] = value + continue + # layer_scalar.weight → layer_scalar (Gemma4 per-layer output scale is an # nn.Parameter, not a module weight). GGUF stores it as # blk.{i}.layer_output_scale.weight, which the tensor mapping renames to @@ -505,6 +521,25 @@ def _normalize_gguf_weights( result[key] = value + for prefix, parts in kv_b_parts.items(): + if set(parts) != {"k", "v"}: + missing = {"k", "v"} - set(parts) + raise ValueError( + f"Incomplete split kv_b_proj for {prefix}: missing {sorted(missing)}" + ) + k_proj = parts["k"].transpose(1, 2) + v_proj = parts["v"] + if k_proj.shape[0] != v_proj.shape[0] or k_proj.shape[2] != v_proj.shape[2]: + raise ValueError( + f"Incompatible split kv_b_proj shapes for {prefix}: " + f"k={tuple(parts['k'].shape)}, v={tuple(v_proj.shape)}" + ) + fused = torch.cat((k_proj, v_proj), dim=1).reshape( + -1, + k_proj.shape[-1], + ) + result[f"{prefix}.self_attn.kv_b_proj.weight"] = fused + return result @@ -699,12 +734,11 @@ def _require_supported_requantization( block_size: int, tensor_name: str, ) -> None: - if bits != 4 or block_size != 32: + if bits not in (4, 8) or block_size != 32: raise ValueError( "keep_quantized MatMulNBits requantization currently supports only " - f"4-bit/block-32 targets; got bits={bits} block={block_size} " - f"for tensor {tensor_name}. Use keep_quantized=False or a " - "4-bit/block-32 target." + f"4-bit or 8-bit block-32 targets; got bits={bits} block={block_size} " + f"for tensor {tensor_name}. Use keep_quantized=False or a supported target." ) @@ -828,6 +862,7 @@ def _load_quantized_state_dict( preserve_native_blocks, repack_dequantized_tensor, repack_gguf_tensor, + repack_quant_params, ) from mobius.integrations.gguf._tencent_q1_0 import ( is_tencent_q1_0_layout, @@ -876,6 +911,7 @@ def _load_quantized_state_dict( target_bits = config.quantization.bits target_block_size = config.quantization.group_size target_symmetric = config.quantization.sym + split_kv_b: dict[str, dict[str, np.ndarray]] = {} for gguf_name, raw, qtype, np_shape in tqdm.tqdm( gguf_model.tensor_items_raw(), @@ -890,6 +926,116 @@ def _load_quantized_state_dict( # Determine the int value of the quant type for can_repack qtype_val = qtype.value if hasattr(qtype, "value") else qtype + # GGUF stores every routed expert projection in one 3D tensor. Repack + # one expert at a time to avoid materializing a multi-gigabyte float + # tensor for GLM-5.2's 256 experts. + expert_proj = next( + ( + proj + for proj in ("gate_proj", "up_proj", "down_proj") + if hf_name.endswith(f".mlp.experts.{proj}.weight") + ), + None, + ) + if expert_proj is not None and len(np_shape) == 3: + num_experts, n_out, k_in = (int(dim) for dim in np_shape) + raw_flat = raw.reshape(-1) + if raw_flat.size % num_experts != 0: + raise ValueError( + f"Cannot split GGUF expert tensor {gguf_name}: " + f"{raw_flat.size} blocks across {num_experts} experts" + ) + blocks_per_expert = raw_flat.size // num_experts + source_params = repack_quant_params(qtype_val) + for expert_idx in range(num_experts): + raw_expert = raw_flat[ + expert_idx * blocks_per_expert : (expert_idx + 1) * blocks_per_expert + ] + repacked = None + if source_params == (target_bits, target_block_size): + candidate = repack_gguf_tensor( + raw_expert.ravel().view(np.uint8), + qtype_val, + (n_out, k_in), + ) + target_has_zero_points = not target_symmetric + if (candidate.zero_points is not None) == target_has_zero_points: + repacked = candidate + if repacked is None: + _require_supported_requantization( + bits=target_bits, + block_size=target_block_size, + tensor_name=hf_name, + ) + values = gguf_model.dequantize_raw_tensor( + raw_expert, + qtype, + (n_out, k_in), + ) + repacked = repack_dequantized_tensor( + values, + bits=target_bits, + block_size=target_block_size, + symmetric=target_symmetric, + ) + n_requantized += 1 + + expert_name = hf_name.replace( + f".mlp.experts.{expert_proj}.weight", + f".mlp.experts.{expert_idx}.{expert_proj}.weight", + ) + expert_stem = expert_name[: -len(".weight")] + state_dict[expert_name] = torch.from_numpy(repacked.weight) + state_dict[f"{expert_stem}.scales"] = torch.from_numpy(repacked.scales) + if repacked.zero_points is not None: + state_dict[f"{expert_stem}.zero_points"] = torch.from_numpy( + repacked.zero_points + ) + n_repacked += 1 + continue + + # GLM-5.2 GGUF splits the MLA kv_b projection into per-head K and V + # tensors. Fuse them in float and immediately requantize the resulting + # 2D projection so only one layer's temporary arrays are resident. + kv_b_part = None + for part in ("k", "v"): + suffix = f".self_attn.kv_b_proj.{part}_proj.weight" + if hf_name.endswith(suffix): + kv_b_part = part + kv_b_prefix = hf_name[: -len(suffix)] + break + if kv_b_part is not None: + values = gguf_model.dequantize_raw_tensor(raw, qtype, np_shape) + parts = split_kv_b.setdefault(kv_b_prefix, {}) + parts[kv_b_part] = values + if set(parts) == {"k", "v"}: + k_proj = parts["k"].transpose(0, 2, 1) + v_proj = parts["v"] + fused = np.concatenate((k_proj, v_proj), axis=1).reshape( + -1, + k_proj.shape[-1], + ) + _require_supported_requantization( + bits=target_bits, + block_size=target_block_size, + tensor_name=f"{kv_b_prefix}.self_attn.kv_b_proj.weight", + ) + repacked = repack_dequantized_tensor( + fused, + bits=target_bits, + block_size=target_block_size, + symmetric=target_symmetric, + ) + stem = f"{kv_b_prefix}.self_attn.kv_b_proj" + state_dict[f"{stem}.weight"] = torch.from_numpy(repacked.weight) + state_dict[f"{stem}.scales"] = torch.from_numpy(repacked.scales) + if repacked.zero_points is not None: + state_dict[f"{stem}.zero_points"] = torch.from_numpy(repacked.zero_points) + del split_kv_b[kv_b_prefix] + n_repacked += 1 + n_requantized += 1 + continue + # Repack every target QuantizedLinear weight. Mixed GGUF presets # otherwise leave unsupported source types as full float matrices, # which cannot fit the graph's packed MatMulNBits initializer shape. @@ -1020,6 +1166,9 @@ def _load_quantized_state_dict( arr = dequantize(raw, qtype).reshape(np_shape) state_dict[hf_name] = torch.from_numpy(arr) + if split_kv_b: + raise ValueError(f"Incomplete split kv_b_proj tensors: {sorted(split_kv_b)}") + logger.info( "Loaded %d state_dict entries (%d GGUF tensors repacked for quantized ops, " "%d requantized from mixed source types)", diff --git a/src/mobius/integrations/gguf/_builder_test.py b/src/mobius/integrations/gguf/_builder_test.py index 8b6fbaed..47c31aa1 100644 --- a/src/mobius/integrations/gguf/_builder_test.py +++ b/src/mobius/integrations/gguf/_builder_test.py @@ -427,20 +427,18 @@ def test_untied_quantized_head_uses_q4_matmulnbits( assert list(model.graph.initializers["lm_head.scales"].shape) == [256, 2] assert "lm_head.weight_t" not in model.graph.initializers - def test_unsupported_requantization_target_has_clear_error( - self, q8_0_projection_q4_head_gguf: Path - ): - """Mixed targets outside 4-bit/block-32 fail before the Q4 repacker.""" + def test_q8_target_requantizes_mixed_q4_head(self, q8_0_projection_q4_head_gguf: Path): + """A mixed Q8 target requantizes its Q4 output head to block-32 Q8.""" from mobius.integrations.gguf import build_from_gguf - with pytest.raises( - ValueError, - match=( - "keep_quantized MatMulNBits requantization currently supports only " - r"4-bit/block-32 targets; got bits=8 block=32 for tensor lm_head\.weight" - ), - ): - build_from_gguf(q8_0_projection_q4_head_gguf, keep_quantized=True) + model = build_from_gguf(q8_0_projection_q4_head_gguf, keep_quantized=True)["model"] + head = next( + node + for node in model.graph + if node.op_type == "MatMulNBits" and node.outputs[0].name == "logits" + ) + assert head.attributes["bits"].value == 8 + assert head.attributes["block_size"].value == 32 def test_norms_are_float(self, q4_0_gguf: Path): """Norm weights remain float, not quantized.""" @@ -657,3 +655,189 @@ def test_dequantize_raw_tensor_matches_get_tensor(self, q4_0_gguf: Path): expected = model.get_tensor(name) actual = model.dequantize_raw_tensor(raw, qtype, shape) np.testing.assert_array_equal(actual, expected) + + +def test_normalize_glm_dsa_split_kv_b_and_router_bias(): + import torch + + from mobius.integrations.gguf._builder import _normalize_gguf_weights + + k_proj = torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(2, 3, 4) + v_proj = torch.arange(2 * 5 * 3, dtype=torch.float32).reshape(2, 5, 3) + result = _normalize_gguf_weights( + { + "model.layers.0.self_attn.kv_b_proj.k_proj.weight": k_proj, + "model.layers.0.self_attn.kv_b_proj.v_proj.weight": v_proj, + "model.layers.0.mlp.gate.e_score_correction_bias.bias": torch.zeros(2), + } + ) + + expected = torch.cat((k_proj.transpose(1, 2), v_proj), dim=1).reshape(18, 3) + assert torch.equal(result["model.layers.0.self_attn.kv_b_proj.weight"], expected) + assert "model.layers.0.mlp.gate.e_score_correction_bias" in result + + +def test_load_quantized_glm_experts_repacked_individually(monkeypatch): + from types import SimpleNamespace + + from gguf import GGMLQuantizationType + + from mobius.integrations.gguf import _repacker, _tencent_q1_0, _tensor_mapping + from mobius.integrations.gguf._builder import _load_quantized_state_dict + + repacked_inputs = [] + + def _repack(raw, qtype, shape): + repacked_inputs.append((raw.copy(), qtype, shape)) + return SimpleNamespace( + weight=np.full((2, 1, 16), len(repacked_inputs), dtype=np.uint8), + scales=np.ones((2, 1), dtype=np.float32), + zero_points=None, + ) + + monkeypatch.setattr(_tencent_q1_0, "is_tencent_q1_0_layout", lambda _model: False) + monkeypatch.setattr( + _tensor_mapping, + "map_gguf_to_hf_names", + lambda _name, _arch: "model.layers.0.mlp.experts.gate_proj.weight", + ) + monkeypatch.setattr(_repacker, "repack_quant_params", lambda _qtype: (4, 32)) + monkeypatch.setattr(_repacker, "repack_gguf_tensor", _repack) + + raw = np.arange(8, dtype=np.uint8) + model = SimpleNamespace( + _tensor_index={"experts": object()}, + tensor_items_raw=lambda: [("experts", raw, GGMLQuantizationType.Q4_0, (2, 2, 2))], + ) + module = SimpleNamespace(named_modules=list) + config = SimpleNamespace( + quantization=SimpleNamespace(bits=4, group_size=32, sym=True), + num_attention_heads=1, + num_key_value_heads=1, + model_type="glm_moe_dsa", + ) + + result = _load_quantized_state_dict(model, "glm-dsa", module, config) + + assert [call[2] for call in repacked_inputs] == [(2, 2), (2, 2)] + assert np.array_equal(repacked_inputs[0][0], raw[:4]) + assert np.array_equal(repacked_inputs[1][0], raw[4:]) + assert "model.layers.0.mlp.experts.0.gate_proj.weight" in result + assert "model.layers.0.mlp.experts.1.gate_proj.weight" in result + + +def test_load_quantized_experts_requantizes_when_zero_point_presence_differs( + monkeypatch, +): + from types import SimpleNamespace + + from gguf import GGMLQuantizationType + + from mobius.integrations.gguf import _repacker, _tencent_q1_0, _tensor_mapping + from mobius.integrations.gguf._builder import _load_quantized_state_dict + + requantized = [] + + def _repack(*_args, **_kwargs): + return SimpleNamespace( + weight=np.zeros((2, 1, 16), dtype=np.uint8), + scales=np.ones((2, 1), dtype=np.float32), + zero_points=None, + ) + + def _requantize(values, **_kwargs): + requantized.append(values) + return SimpleNamespace( + weight=np.zeros((2, 1, 16), dtype=np.uint8), + scales=np.ones((2, 1), dtype=np.float32), + zero_points=np.zeros((2, 1), dtype=np.uint8), + ) + + monkeypatch.setattr(_tencent_q1_0, "is_tencent_q1_0_layout", lambda _model: False) + monkeypatch.setattr( + _tensor_mapping, + "map_gguf_to_hf_names", + lambda _name, _arch: "model.layers.0.mlp.experts.gate_proj.weight", + ) + monkeypatch.setattr(_repacker, "repack_quant_params", lambda _qtype: (4, 32)) + monkeypatch.setattr(_repacker, "repack_gguf_tensor", _repack) + monkeypatch.setattr(_repacker, "repack_dequantized_tensor", _requantize) + + raw = np.arange(8, dtype=np.uint8) + model = SimpleNamespace( + _tensor_index={"experts": object()}, + tensor_items_raw=lambda: [ + ("experts", raw, GGMLQuantizationType.Q4_0, (2, 2, 2)) + ], + dequantize_raw_tensor=lambda *_args: np.ones((2, 2), dtype=np.float32), + ) + config = SimpleNamespace( + quantization=SimpleNamespace(bits=4, group_size=32, sym=False), + num_attention_heads=1, + num_key_value_heads=1, + model_type="glm_moe_dsa", + ) + + result = _load_quantized_state_dict(model, "glm-dsa", SimpleNamespace(named_modules=list), config) + + assert len(requantized) == 2 + assert "model.layers.0.mlp.experts.0.gate_proj.zero_points" in result + assert "model.layers.0.mlp.experts.1.gate_proj.zero_points" in result + + +def test_load_quantized_glm_fuses_split_kv_b(monkeypatch): + from types import SimpleNamespace + + from gguf import GGMLQuantizationType + + from mobius.integrations.gguf import _repacker, _tencent_q1_0, _tensor_mapping + from mobius.integrations.gguf._builder import _load_quantized_state_dict + + fused_values = [] + + def _repack(values, **_kwargs): + fused_values.append(values.copy()) + return SimpleNamespace( + weight=np.zeros((18, 1, 16), dtype=np.uint8), + scales=np.ones((18, 1), dtype=np.float32), + zero_points=np.zeros((18, 1), dtype=np.uint8), + ) + + names = { + "k": "model.layers.0.self_attn.kv_b_proj.k_proj.weight", + "v": "model.layers.0.self_attn.kv_b_proj.v_proj.weight", + } + monkeypatch.setattr(_tencent_q1_0, "is_tencent_q1_0_layout", lambda _model: False) + monkeypatch.setattr( + _tensor_mapping, + "map_gguf_to_hf_names", + lambda name, _arch: names[name], + ) + monkeypatch.setattr(_repacker, "repack_dequantized_tensor", _repack) + + k_proj = np.arange(2 * 3 * 4, dtype=np.float32).reshape(2, 3, 4) + v_proj = np.arange(2 * 5 * 3, dtype=np.float32).reshape(2, 5, 3) + tensors = [ + ("k", k_proj, GGMLQuantizationType.F32, k_proj.shape), + ("v", v_proj, GGMLQuantizationType.F32, v_proj.shape), + ] + model = SimpleNamespace( + _tensor_index={"k": object(), "v": object()}, + tensor_items_raw=lambda: tensors, + dequantize_raw_tensor=lambda raw, _qtype, _shape: raw, + ) + module = SimpleNamespace(named_modules=list) + config = SimpleNamespace( + quantization=SimpleNamespace(bits=4, group_size=32, sym=False), + num_attention_heads=1, + num_key_value_heads=1, + model_type="glm_moe_dsa", + ) + + result = _load_quantized_state_dict(model, "glm-dsa", module, config) + + expected = np.concatenate((k_proj.transpose(0, 2, 1), v_proj), axis=1).reshape(18, 3) + np.testing.assert_array_equal(fused_values[0], expected) + assert "model.layers.0.self_attn.kv_b_proj.weight" in result + assert "model.layers.0.self_attn.kv_b_proj.scales" in result + assert "model.layers.0.self_attn.kv_b_proj.zero_points" in result diff --git a/src/mobius/integrations/gguf/_config_mapping.py b/src/mobius/integrations/gguf/_config_mapping.py index 82187b18..86502252 100644 --- a/src/mobius/integrations/gguf/_config_mapping.py +++ b/src/mobius/integrations/gguf/_config_mapping.py @@ -44,6 +44,7 @@ "qwen3_moe": "qwen3_moe", "qwen35": "qwen3_5_text", "qwen35moe": "qwen3_5_moe", + "glm-dsa": "glm_moe_dsa", "gemma2": "gemma2", "gemma3": "gemma3_text", # Gemma 4 GGUF contains the text backbone only — no vision or audio encoder. @@ -115,7 +116,24 @@ "hyper_connection.sinkhorn_iterations": "hc_sinkhorn_iters", "hyper_connection.epsilon": "hc_eps", "hash_layer_count": "num_hash_layers", - } + }, + "glm-dsa": { + "leading_dense_block_count": "first_k_dense_replace", + "expert_shared_count": "n_shared_experts", + "expert_group_count": "n_group", + "expert_group_used_count": "topk_group", + "expert_weights_scale": "routed_scaling_factor", + "expert_weights_norm": "norm_topk_prob", + "expert_gating_func": "expert_gating_func", + "attention.q_lora_rank": "q_lora_rank", + "attention.kv_lora_rank": "kv_lora_rank", + "attention.key_length_mla": "qk_head_dim", + "attention.value_length_mla": "v_head_dim", + "attention.indexer.head_count": "index_n_heads", + "attention.indexer.key_length": "index_head_dim", + "attention.indexer.top_k": "index_topk", + "nextn_predict_layers": "num_nextn_predict_layers", + }, } @@ -189,6 +207,11 @@ def _extract_config_fields( len(hf_fields), ) + for gguf_suffix, hf_key in _ARCH_KEY_MAPS.get(gguf_arch, {}).items(): + full_key = f"{gguf_arch}.{gguf_suffix}" + if full_key in metadata: + hf_fields[hf_key] = metadata[full_key] + # Extract vocab_size from tokenizer token list if not in metadata if "vocab_size" not in hf_fields: tokens = metadata.get("tokenizer.ggml.tokens") @@ -400,19 +423,29 @@ def gguf_to_config( shared_expert_intermediate_size=hf_fields.get("shared_expert_intermediate_size"), n_shared_experts=hf_fields.get("n_shared_experts"), norm_topk_prob=hf_fields.get("norm_topk_prob", True), + n_group=hf_fields.get("n_group", 1), + topk_group=hf_fields.get("topk_group", 1), routed_scaling_factor=hf_fields.get("routed_scaling_factor", 1.0), scoring_func=( "sqrtsoftplus" if gguf_arch == "deepseek4" else hf_fields.get("scoring_func", "softmax") ), + topk_method=hf_fields.get("topk_method", "greedy"), + first_k_dense_replace=hf_fields.get("first_k_dense_replace", 0), + # Multi-head Latent Attention fields q_lora_rank=hf_fields.get("q_lora_rank"), + kv_lora_rank=hf_fields.get("kv_lora_rank"), + qk_nope_head_dim=hf_fields.get("qk_nope_head_dim"), qk_rope_head_dim=hf_fields.get("qk_rope_head_dim"), + v_head_dim=hf_fields.get("v_head_dim"), + # DSA / MTP metadata + index_topk=hf_fields.get("index_topk"), + index_head_dim=hf_fields.get("index_head_dim"), + index_n_heads=hf_fields.get("index_n_heads"), + num_nextn_predict_layers=hf_fields.get("num_nextn_predict_layers", 0), o_groups=hf_fields.get("o_groups", 1), o_lora_rank=hf_fields.get("o_lora_rank"), - index_n_heads=hf_fields.get("index_n_heads"), - index_head_dim=hf_fields.get("index_head_dim"), - index_topk=hf_fields.get("index_topk"), compress_ratios=hf_fields.get("compress_ratios"), compress_rope_theta=hf_fields.get("compress_rope_theta"), hc_mult=hf_fields.get("hc_mult", 1), @@ -719,12 +752,70 @@ def _gemma3_postprocess( return config +def _glm_moe_dsa_postprocess( + config: ArchitectureConfig, + metadata: dict[str, Any], +) -> ArchitectureConfig: + """Restore GLM-5.2 fields that GGUF represents with MLA-specific keys.""" + arch = "glm-dsa" + n_mtp = int(metadata.get(f"{arch}.nextn_predict_layers", 0)) + num_hidden_layers = config.num_hidden_layers - n_mtp + qk_head_dim = int(metadata[f"{arch}.attention.key_length_mla"]) + qk_rope_head_dim = int(metadata[f"{arch}.rope.dimension_count"]) + qk_nope_head_dim = qk_head_dim - qk_rope_head_dim + n_shared_experts = int(metadata.get(f"{arch}.expert_shared_count", 0)) + first_k_dense_replace = int(metadata.get(f"{arch}.leading_dense_block_count", 0)) + + # Official GLM-5.2 schedule: full indexers in layers 0, 1, 2, then + # 6, 10, ...; each later full layer shares its selection with the next + # three layers. + indexer_types = [ + "full" if i <= 2 or (i >= 6 and (i - 6) % 4 == 0) else "shared" + for i in range(num_hidden_layers) + ] + mlp_layer_types = [ + "dense" if i < first_k_dense_replace else "sparse" for i in range(num_hidden_layers) + ] + + return dataclasses.replace( + config, + num_hidden_layers=num_hidden_layers, + num_key_value_heads=config.num_attention_heads, + head_dim=qk_rope_head_dim, + partial_rotary_factor=1.0, + rope_interleave=True, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=int(metadata[f"{arch}.attention.value_length_mla"]), + q_lora_rank=int(metadata[f"{arch}.attention.q_lora_rank"]), + kv_lora_rank=int(metadata[f"{arch}.attention.kv_lora_rank"]), + n_shared_experts=n_shared_experts, + shared_expert_intermediate_size=( + config.moe_intermediate_size * n_shared_experts + if config.moe_intermediate_size is not None + else None + ), + first_k_dense_replace=first_k_dense_replace, + mlp_layer_types=mlp_layer_types, + scoring_func="sigmoid", + topk_method="noaux_tc", + use_dsa=True, + index_topk_freq=4, + index_skip_topk_offset=3, + indexer_rope_interleave=True, + indexer_types=indexer_types, + index_share_for_mtp_iteration=n_mtp > 0, + num_nextn_predict_layers=n_mtp, + ) + + # Architecture-specific config postprocessors. # Each takes a base ArchitectureConfig + raw metadata and returns # an architecture-specific config subclass. _CONFIG_POSTPROCESSORS: dict[str, Any] = { "gemma3_text": _gemma3_postprocess, "gemma4_text": _gemma4_postprocess, + "glm_moe_dsa": _glm_moe_dsa_postprocess, } diff --git a/src/mobius/integrations/gguf/_reader_test.py b/src/mobius/integrations/gguf/_reader_test.py index 0ec0e108..fde61c3c 100644 --- a/src/mobius/integrations/gguf/_reader_test.py +++ b/src/mobius/integrations/gguf/_reader_test.py @@ -10,6 +10,7 @@ from __future__ import annotations from pathlib import Path +from typing import ClassVar import numpy as np import pytest @@ -411,6 +412,77 @@ def test_arch_to_model_type_mapping(self): assert GGUF_ARCH_TO_MODEL_TYPE["qwen35"] == "qwen3_5_text" assert GGUF_ARCH_TO_MODEL_TYPE["phi3"] == "phi3" assert GGUF_ARCH_TO_MODEL_TYPE["deepseek4"] == "deepseek_v4" + assert GGUF_ARCH_TO_MODEL_TYPE["glm-dsa"] == "glm_moe_dsa" + + def test_glm_dsa_config_extraction(self): + class _FakeGlmDsa: + architecture = "glm-dsa" + tensor_names: ClassVar[list[str]] = ["token_embd.weight", "output.weight"] + metadata: ClassVar[dict] = { + "glm-dsa.embedding_length": 6144, + "glm-dsa.feed_forward_length": 12288, + "glm-dsa.block_count": 79, + "glm-dsa.attention.head_count": 64, + "glm-dsa.attention.head_count_kv": 1, + "glm-dsa.attention.layer_norm_rms_epsilon": 1e-5, + "glm-dsa.attention.q_lora_rank": 2048, + "glm-dsa.attention.kv_lora_rank": 512, + "glm-dsa.attention.key_length_mla": 256, + "glm-dsa.attention.value_length_mla": 256, + "glm-dsa.attention.indexer.head_count": 32, + "glm-dsa.attention.indexer.key_length": 128, + "glm-dsa.attention.indexer.top_k": 2048, + "glm-dsa.expert_count": 256, + "glm-dsa.expert_used_count": 8, + "glm-dsa.expert_feed_forward_length": 2048, + "glm-dsa.expert_shared_count": 1, + "glm-dsa.expert_group_count": 1, + "glm-dsa.expert_group_used_count": 1, + "glm-dsa.expert_weights_scale": 2.5, + "glm-dsa.expert_weights_norm": True, + "glm-dsa.expert_gating_func": 2, + "glm-dsa.leading_dense_block_count": 3, + "glm-dsa.nextn_predict_layers": 1, + "glm-dsa.rope.dimension_count": 64, + "glm-dsa.rope.freq_base": 8_000_000.0, + "glm-dsa.context_length": 1_048_576, + "glm-dsa.vocab_size": 154880, + } + + def get_metadata(self, key, default=None): + return self.metadata.get(key, default) + + config = gguf_to_config(_FakeGlmDsa()) + + assert config.model_type == "glm_moe_dsa" + assert config.num_hidden_layers == 78 + assert config.hidden_size == 6144 + assert config.num_attention_heads == 64 + assert config.num_key_value_heads == 64 + assert config.head_dim == 64 + assert config.q_lora_rank == 2048 + assert config.kv_lora_rank == 512 + assert config.qk_nope_head_dim == 192 + assert config.qk_rope_head_dim == 64 + assert config.v_head_dim == 256 + assert config.num_local_experts == 256 + assert config.num_experts_per_tok == 8 + assert config.moe_intermediate_size == 2048 + assert config.shared_expert_intermediate_size == 2048 + assert config.first_k_dense_replace == 3 + assert config.routed_scaling_factor == pytest.approx(2.5) + assert config.scoring_func == "sigmoid" + assert config.rope_interleave is True + assert config.max_position_embeddings == 1_048_576 + assert config.indexer_types[:7] == [ + "full", + "full", + "full", + "shared", + "shared", + "shared", + "full", + ] def test_deepseek4_config_extraction(self): metadata = { diff --git a/src/mobius/integrations/gguf/_repacker.py b/src/mobius/integrations/gguf/_repacker.py index c54df705..951f3eb1 100644 --- a/src/mobius/integrations/gguf/_repacker.py +++ b/src/mobius/integrations/gguf/_repacker.py @@ -518,10 +518,10 @@ def repack_dequantized_tensor( """ if values.ndim != 2: raise ValueError(f"Expected 2D values (N, K), got shape {values.shape}") - if bits != 4 or block_size != _BLOCK_SIZE: + if bits not in (4, 8) or block_size != _BLOCK_SIZE: raise ValueError( "Float requantization currently supports only " - f"bits=4, block_size={_BLOCK_SIZE}; got bits={bits}, " + f"bits=4/8, block_size={_BLOCK_SIZE}; got bits={bits}, " f"block_size={block_size}" ) @@ -532,6 +532,31 @@ def repack_dequantized_tensor( padded[:, :k_in] = values.astype(np.float32, copy=False) blocks = padded.reshape(n_out, n_blocks, block_size) + if bits == 8: + block_min = np.minimum(blocks.min(axis=-1), 0.0) + block_max = np.maximum(blocks.max(axis=-1), 0.0) + if symmetric: + scales = np.maximum(-block_min / 128.0, block_max / 127.0) + zero_points_arr = np.full_like(scales, 128, dtype=np.uint8) + else: + scales = (block_max - block_min) / 255.0 + safe_scales = np.where(scales != 0, scales, 1.0) + zero_points_arr = np.clip( + np.rint(-block_min / safe_scales), 0, 255 + ).astype(np.uint8) + safe_scales = np.where(scales != 0, scales, 1.0) + quants = np.rint(blocks / safe_scales[:, :, None]) + quants += zero_points_arr[:, :, None] + weight = np.clip(quants, 0, 255).astype(np.uint8) + weight = np.where(scales[:, :, None] != 0, weight, 0).astype(np.uint8) + return RepackedTensor( + weight=weight, + scales=scales.astype(np.float32), + zero_points=None if symmetric else zero_points_arr, + block_size=block_size, + bits=bits, + ) + if symmetric: # MatMulNBits' implicit uint4 zero-point is 8. Choose a scale that # covers the asymmetric signed code range [-8, 7]. diff --git a/src/mobius/integrations/gguf/_repacker_test.py b/src/mobius/integrations/gguf/_repacker_test.py index d2f6e059..94ece26f 100644 --- a/src/mobius/integrations/gguf/_repacker_test.py +++ b/src/mobius/integrations/gguf/_repacker_test.py @@ -712,6 +712,37 @@ def test_asymmetric_q4_round_trip_bound(self): assert result.weight.shape == (3, 3, 16) assert result.zero_points.shape == (3, 2) + def test_asymmetric_q8_round_trip_bound(self): + rng = np.random.default_rng(1) + values = rng.normal(size=(3, 70)).astype(np.float32) + result = repack_dequantized_tensor(values, bits=8, block_size=32) + + n_blocks = result.scales.shape[1] + dequantized = ( + result.weight.astype(np.float32) + - result.zero_points[:, :, None].astype(np.float32) + ) * result.scales[:, :, None] + dequantized = dequantized.reshape(3, n_blocks * 32)[:, : values.shape[1]] + + max_abs_diff = float(np.max(np.abs(dequantized - values))) + assert max_abs_diff <= float(result.scales.max()) * 0.51 + 1e-6 + assert result.weight.shape == (3, 3, 32) + assert result.zero_points.shape == (3, 3) + + def test_symmetric_q8_omits_zero_points(self): + values = np.array([[-128.0, 0.0, 127.0] + [0.0] * 29], dtype=np.float32) + result = repack_dequantized_tensor( + values, bits=8, block_size=32, symmetric=True + ) + + assert result.zero_points is None + dequantized = ( + result.weight.astype(np.float32) - 128.0 + ) * result.scales[:, :, None] + np.testing.assert_allclose( + dequantized.reshape(values.shape), values, atol=result.scales.max() * 0.51 + ) + # ---- Q1_0 tests ---- diff --git a/src/mobius/integrations/gguf/_tensor_mapping.py b/src/mobius/integrations/gguf/_tensor_mapping.py index 7d61e5f6..9e0fef90 100644 --- a/src/mobius/integrations/gguf/_tensor_mapping.py +++ b/src/mobius/integrations/gguf/_tensor_mapping.py @@ -193,7 +193,48 @@ "blk.{bid}.ffn_down_shexp": ("model.layers.{bid}.mlp.shared_expert.down_proj"), } -# Qwen3.5 hybrid extensions: DeltaNet (SSM) + full-attention. +# GLM-5.2 uses MLA with split K/V decompression weights in GGUF. +# The temporary ``kv_b_proj.{k,v}_proj`` names are fused by +# ``_normalize_gguf_weights`` before model-specific preprocessing. +_GLM_DSA_MAPPING: dict[str, str] = { + "token_embd": "model.embed_tokens", + "output": "lm_head", + "output_norm": "model.norm", + "blk.{bid}.attn_norm": "model.layers.{bid}.input_layernorm", + "blk.{bid}.attn_q_a": "model.layers.{bid}.self_attn.q_a_proj", + "blk.{bid}.attn_q_b": "model.layers.{bid}.self_attn.q_b_proj", + "blk.{bid}.attn_kv_a_mqa": "model.layers.{bid}.self_attn.kv_a_proj_with_mqa", + "blk.{bid}.attn_k_b": "model.layers.{bid}.self_attn.kv_b_proj.k_proj", + "blk.{bid}.attn_v_b": "model.layers.{bid}.self_attn.kv_b_proj.v_proj", + "blk.{bid}.attn_q_a_norm": "model.layers.{bid}.self_attn.q_a_layernorm", + "blk.{bid}.attn_kv_a_norm": "model.layers.{bid}.self_attn.kv_a_layernorm", + "blk.{bid}.attn_output": "model.layers.{bid}.self_attn.o_proj", + "blk.{bid}.ffn_norm": "model.layers.{bid}.post_attention_layernorm", + "blk.{bid}.ffn_gate": "model.layers.{bid}.mlp.gate_proj", + "blk.{bid}.ffn_up": "model.layers.{bid}.mlp.up_proj", + "blk.{bid}.ffn_down": "model.layers.{bid}.mlp.down_proj", + "blk.{bid}.ffn_gate_inp": "model.layers.{bid}.mlp.gate", + "blk.{bid}.exp_probs_b": "model.layers.{bid}.mlp.gate.e_score_correction_bias", + # GLM-5.2 IndexShare DSA. + "blk.{bid}.indexer.attn_k": "model.layers.{bid}.self_attn.indexer.wk", + "blk.{bid}.indexer.attn_q_b": "model.layers.{bid}.self_attn.indexer.wq_b", + "blk.{bid}.indexer.k_norm": "model.layers.{bid}.self_attn.indexer.k_norm", + "blk.{bid}.indexer.proj": "model.layers.{bid}.self_attn.indexer.weights_proj", + # GLM-5.2 improved MTP. The shared embedding/head tensors are tied to + # token_embd/output and therefore intentionally remain unmapped. + "blk.{bid}.nextn.eh_proj": "model.layers.{bid}.eh_proj", + "blk.{bid}.nextn.enorm": "model.layers.{bid}.enorm", + "blk.{bid}.nextn.hnorm": "model.layers.{bid}.hnorm", + "blk.{bid}.nextn.shared_head_norm": "model.layers.{bid}.shared_head.norm", + "blk.{bid}.ffn_gate_exps": "model.layers.{bid}.mlp.experts.gate_proj", + "blk.{bid}.ffn_up_exps": "model.layers.{bid}.mlp.experts.up_proj", + "blk.{bid}.ffn_down_exps": "model.layers.{bid}.mlp.experts.down_proj", + "blk.{bid}.ffn_gate_shexp": "model.layers.{bid}.mlp.shared_experts.gate_proj", + "blk.{bid}.ffn_up_shexp": "model.layers.{bid}.mlp.shared_experts.up_proj", + "blk.{bid}.ffn_down_shexp": "model.layers.{bid}.mlp.shared_experts.down_proj", +} + +# Qwen3.5-MoE hybrid extensions: DeltaNet (SSM) + full-attention + MoE. # DeltaNet layers use linear_attn.* naming; full-attention layers add # q_norm/k_norm under self_attn; both use post_attention_layernorm. _QWEN35_HYBRID_EXTRAS: dict[str, str] = { @@ -366,13 +407,15 @@ def _build_mapping( result.update(_QWEN35_HYBRID_EXTRAS) elif arch == "deepseek4": result = dict(_DEEPSEEK4_MAPPING) + elif arch == "glm-dsa": + result = dict(_GLM_DSA_MAPPING) else: supported = sorted( _LLAMA_FAMILY | _GEMMA_FAMILY | _MOE_FAMILY | _HUNYUAN_FAMILY - | {"deepseek4", "gemma4", "phi3", "falcon", "gpt2", "mamba"} + | {"deepseek4", "gemma4", "glm-dsa", "phi3", "falcon", "gpt2", "mamba"} ) raise ValueError( f"Unsupported GGUF architecture: {architecture!r}. " diff --git a/src/mobius/integrations/gguf/_tensor_mapping_test.py b/src/mobius/integrations/gguf/_tensor_mapping_test.py index ac5c34ed..39563923 100644 --- a/src/mobius/integrations/gguf/_tensor_mapping_test.py +++ b/src/mobius/integrations/gguf/_tensor_mapping_test.py @@ -306,6 +306,33 @@ def test_moe_extras(self) -> None: "model.layers.0.mlp.experts.gate_proj.weight" ) + def test_glm_dsa_mapping(self) -> None: + assert map_gguf_to_hf_names("blk.3.attn_q_a.weight", "glm-dsa") == ( + "model.layers.3.self_attn.q_a_proj.weight" + ) + assert map_gguf_to_hf_names("blk.3.attn_k_b.weight", "glm-dsa") == ( + "model.layers.3.self_attn.kv_b_proj.k_proj.weight" + ) + assert map_gguf_to_hf_names("blk.3.ffn_gate_exps.weight", "glm-dsa") == ( + "model.layers.3.mlp.experts.gate_proj.weight" + ) + assert map_gguf_to_hf_names("blk.3.exp_probs_b.bias", "glm-dsa") == ( + "model.layers.3.mlp.gate.e_score_correction_bias.bias" + ) + assert map_gguf_to_hf_names("blk.3.indexer.attn_k.weight", "glm-dsa") == ( + "model.layers.3.self_attn.indexer.wk.weight" + ) + assert map_gguf_to_hf_names("blk.3.indexer.k_norm.bias", "glm-dsa") == ( + "model.layers.3.self_attn.indexer.k_norm.bias" + ) + assert map_gguf_to_hf_names("blk.78.nextn.eh_proj.weight", "glm-dsa") == ( + "model.layers.78.eh_proj.weight" + ) + assert ( + map_gguf_to_hf_names("blk.78.nextn.shared_head_norm.weight", "glm-dsa") + == "model.layers.78.shared_head.norm.weight" + ) + def test_deepseek4_mapping(self) -> None: expected = { "blk.2.attn_q_a.weight": "model.layers.2.self_attn.q_a_proj.weight", @@ -329,6 +356,7 @@ def test_deepseek4_mapping(self) -> None: for source, target in expected.items(): assert map_gguf_to_hf_names(source, "deepseek4") == target + # ---- Unsupported architecture ---- def test_unsupported_raises(self) -> None: diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 9c2a0c3c..0694adf0 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -12,6 +12,7 @@ from __future__ import annotations +import json import logging import os from typing import Any @@ -38,6 +39,41 @@ _DENOISER_KEYS = ("denoiser", "transformer", "unet") +def _write_mtp_config(output_dir: str, config: Any) -> str: + path = os.path.join(output_dir, "mtp_config.json") + with open(path, "w", encoding="utf-8") as handle: + json.dump( + { + "model": {"filename": "mtp/model.onnx"}, + "inputs": [ + "inputs_embeds", + "hidden_states", + "attention_mask", + "position_ids", + "past_key_values.0.key", + "past_key_values.0.value", + ], + "outputs": [ + "mtp_hidden", + "present.0.key", + "present.0.value", + "topk_indices", + ], + "num_nextn_predict_layers": getattr(config, "num_nextn_predict_layers", 0), + "index_share_for_mtp_iteration": getattr( + config, "index_share_for_mtp_iteration", False + ), + "shared_embedding": "model.embed_tokens", + "shared_lm_head": "lm_head", + "runtime_orchestration": "external", + }, + handle, + indent=2, + ) + handle.write("\n") + return path + + def _add_explicit_io_to_file(path: str, pkg: Any, config: Any) -> None: """Augment an emitted sidecar with roles derived from the actual ONNX ports.""" try: @@ -522,14 +558,28 @@ def write_onnx_genai_config( artifacts["tokenizer"] = tokenizer_path return artifacts - # Fallback: a single-component decoder language model. A multi-component - # package that matched none of the composite shapes above would be silently - # mis-emitted as a bare decoder — fail loudly instead so an unsupported shape - # is obvious rather than producing wrong metadata. try: component_names = sorted(pkg.keys()) except (AttributeError, TypeError): component_names = [] + + if component_names == ["model", "mtp"]: + path = write_decoder_metadata( + output_dir, config=resolved_config, kv_native_dtype=kv_native_dtype + ) + artifacts = { + "inference_metadata": path, + "mtp_config": _write_mtp_config(output_dir, resolved_config), + } + tokenizer_path = _write_hf_tokenizer(output_dir, source) + if tokenizer_path is not None: + artifacts["tokenizer"] = tokenizer_path + return artifacts + + # Fallback: a single-component decoder language model. A multi-component + # package that matched none of the composite shapes above would be silently + # mis-emitted as a bare decoder — fail loudly instead so an unsupported shape + # is obvious rather than producing wrong metadata. if len(component_names) > 1: raise ValueError( "onnx-genai config emission does not recognize this multi-component " diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 45c85fd9..fec0eb12 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -6,6 +6,7 @@ from __future__ import annotations import dataclasses +import json import onnx_ir as ir import pytest @@ -24,6 +25,8 @@ class _Cfg: max_position_embeddings: int = 8192 sliding_window: int | None = None model_type: str = "qwen" + num_nextn_predict_layers: int = 1 + index_share_for_mtp_iteration: bool = True @dataclasses.dataclass @@ -451,6 +454,20 @@ def test_unrecognized_multi_component_package_fails_loudly(tmp_path): write_onnx_genai_config(pkg, str(tmp_path)) +def test_mtp_package_emits_decoder_metadata_and_sidecar(tmp_path): + pkg = _MultimodalPkg({"model": object(), "mtp": object()}) + + artifacts = write_onnx_genai_config(pkg, str(tmp_path)) + + assert set(artifacts) == {"inference_metadata", "mtp_config"} + with open(artifacts["mtp_config"]) as handle: + mtp = json.load(handle) + assert mtp["model"]["filename"] == "mtp/model.onnx" + assert mtp["num_nextn_predict_layers"] == 1 + assert mtp["index_share_for_mtp_iteration"] is True + assert mtp["runtime_orchestration"] == "external" + + def test_decoder_emits_tokenizer_from_source(tmp_path): # A text-producing package emits tokenizer.json from its HF source so the # onnx-genai package is self-contained. diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 588906f3..a4467628 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -87,6 +87,7 @@ # ORT GenAI (see onnxruntime-genai/src/models/model_type.h LLM list). "hunyuan_v1_dense": "decoder", "deepseek_v4": "decoder", + "glm_moe_dsa": "decoder", # Qwen VL models all use the same GenAI pipeline as qwen2_5_vl "qwen2_vl": "qwen2_5_vl", "qwen3_vl": "qwen2_5_vl", @@ -1086,22 +1087,39 @@ def write_ort_genai_config( if "mtp" in pkg: mtp_model = pkg["mtp"] + mtp_inputs = [ + value.name for value in mtp_model.graph.inputs if value.name is not None + ] + if not mtp_inputs: + mtp_inputs = [ + "inputs_embeds", + "hidden_states", + "attention_mask", + "position_ids", + "past_key_values.0.key", + "past_key_values.0.value", + ] + mtp_outputs = [ + value.name for value in mtp_model.graph.outputs if value.name is not None + ] + if not mtp_outputs: + mtp_outputs = [ + "mtp_hidden", + "present.0.key", + "present.0.value", + "topk_indices", + ] mtp_path = os.path.join(directory, "mtp_config.json") with open(mtp_path, "w") as f: json.dump( { "model": {"filename": "mtp/model.onnx"}, - "inputs": [ - value.name - for value in mtp_model.graph.inputs - if value.name is not None - ], - "outputs": [ - value.name - for value in mtp_model.graph.outputs - if value.name is not None - ], + "inputs": mtp_inputs, + "outputs": mtp_outputs, "num_nextn_predict_layers": getattr(config, "num_nextn_predict_layers", 0), + "index_share_for_mtp_iteration": getattr( + config, "index_share_for_mtp_iteration", False + ), "shared_embedding": "model.embed_tokens", "shared_lm_head": "lm_head", "runtime_orchestration": "external", diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index f143090f..386dafa6 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -75,6 +75,9 @@ def test_hunyuan_v1_dense_maps_to_decoder(self): # decoder-only causal LM not in its built-in registry. assert _resolve_ort_genai_model_type("hunyuan_v1_dense") == "decoder" + def test_glm_moe_dsa_maps_to_decoder(self): + assert _resolve_ort_genai_model_type("glm_moe_dsa") == "decoder" + def test_unknown_model_type_passthrough(self): assert _resolve_ort_genai_model_type("my_custom") == "my_custom" @@ -622,6 +625,22 @@ def test_genai_config_json_is_written(self, tmp_path): assert "model" in data assert data["model"]["type"] == "qwen2" + def test_mtp_component_writes_sidecar_and_nested_decoder_path(self, tmp_path): + pkg = self._make_pkg() + pkg["mtp"] = mock.MagicMock() + + result = write_ort_genai_config(pkg, str(tmp_path)) + + assert "mtp_config" in result + with open(result["genai_config"]) as f: + genai = json.load(f) + assert genai["model"]["decoder"]["filename"] == "model/model.onnx" + with open(result["mtp_config"]) as f: + mtp = json.load(f) + assert mtp["model"]["filename"] == "mtp/model.onnx" + assert mtp["outputs"][-1] == "topk_indices" + assert mtp["runtime_orchestration"] == "external" + def test_processor_config_written_with_vision(self, tmp_path): """image_processor.json is written when pkg.config.vision is set.""" import dataclasses diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index ec45276f..d34c7cc5 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -56,6 +56,7 @@ "Glm4CausalLMModel", "Glm4MoECausalLMModel", "GlmCausalLMModel", + "GlmMoeDsaCausalLMModel", "GraniteCausalLMModel", "GraniteMoECausalLMModel", "GraniteMoeHybridCausalLMModel", @@ -190,6 +191,7 @@ ) from mobius.models.gemma4_assistant import Gemma4AssistantCausalLMModel from mobius.models.glm import Glm4CausalLMModel, GlmCausalLMModel +from mobius.models.glm_moe_dsa import GlmMoeDsaCausalLMModel from mobius.models.gpt2 import GPT2CausalLMModel from mobius.models.gpt_neox import GPTNeoXCausalLMModel, GPTNeoXJapaneseCausalLMModel from mobius.models.gptj_codegen import CodeGenCausalLMModel, GPTJCausalLMModel diff --git a/src/mobius/models/deepseek.py b/src/mobius/models/deepseek.py index c631845f..1236ecd5 100644 --- a/src/mobius/models/deepseek.py +++ b/src/mobius/models/deepseek.py @@ -20,14 +20,17 @@ MLP, Attention, Embedding, + FusedQuantizedMoE, Linear, MoELayer, + QuantizedEmbedding, RMSNorm, create_attention_bias, initialize_rope, + make_quantized_linear_factory, ) from mobius.components._deepseek_mla import DeepSeekMLA -from mobius.components._quantized_linear import make_quantized_linear_factory +from mobius.components._moe import pack_fused_quantized_moe_weights from mobius.models.base import CausalLMModel @@ -44,6 +47,21 @@ def _linear_factory(config: ArchitectureConfig): ) +def _linear_class(config: ArchitectureConfig) -> type: + qc = config.quantization + if qc is None or qc.quant_method == "none": + return Linear + import onnx_ir as ir + + zero_point_dtype = config.dtype if qc.float_zero_point else ir.DataType.UINT8 + return make_quantized_linear_factory( + bits=qc.bits, + block_size=qc.group_size, + has_zero_point=not qc.sym, + zero_point_dtype=zero_point_dtype, + ) + + class DeepSeekMoEGate(nn.Module): """Expert routing gate for DeepSeek-V2/V3 MoE. @@ -87,25 +105,15 @@ def __init__(self, config: ArchitectureConfig): def forward(self, op: OpBuilder, hidden_states: ir.Value): scores, scores_for_choice = self._routing_scores(op, hidden_states) - - # Select top-k experts - k_val = op.Constant(value_ints=[self.top_k]) - _, selected_experts = op.TopK(scores_for_choice, k_val, axis=-1, _outputs=2) - - # Gather original scores (without bias) for selected experts - routing_weights = op.GatherElements(scores, selected_experts, axis=-1) - - # Normalize weights (V3 with norm_topk_prob=True) - if self.norm_topk_prob: - weight_sum = op.ReduceSum(routing_weights, [-1], keepdims=True) - eps = 1e-20 - routing_weights = op.Div(routing_weights, op.Add(weight_sum, eps)) - - # Apply routing scale - routing_weights = op.Mul(routing_weights, float(self.routed_scaling_factor)) - + routing_weights, selected_experts = self._topk_weights(op, scores, scores_for_choice) return routing_weights, selected_experts + def route_for_qmoe(self, op: OpBuilder, hidden_states: ir.Value): + """Routing outputs shaped for the fused ``com.microsoft::QMoE`` op.""" + scores, scores_for_choice = self._routing_scores(op, hidden_states) + routing_weights, selected_experts = self._topk_weights(op, scores, scores_for_choice) + return scores_for_choice, routing_weights, selected_experts + def qmoe_routing(self, op: OpBuilder, hidden_states: ir.Value): """Return distinct expert-selection and aggregation scores for QMoE.""" scores, scores_for_choice = self._routing_scores(op, hidden_states) @@ -116,6 +124,17 @@ def qmoe_routing(self, op: OpBuilder, hidden_states: ir.Value): float(self.routed_scaling_factor), ) + def _topk_weights(self, op: OpBuilder, scores: ir.Value, scores_for_choice: ir.Value): + k_val = op.Constant(value_ints=[self.top_k]) + _, selected_experts = op.TopK(scores_for_choice, k_val, axis=-1, _outputs=2) + routing_weights = op.GatherElements(scores, selected_experts, axis=-1) + if self.norm_topk_prob: + weight_sum = op.ReduceSum(routing_weights, [-1], keepdims=True) + eps = 1e-20 + routing_weights = op.Div(routing_weights, op.Add(weight_sum, eps)) + routing_weights = op.Mul(routing_weights, float(self.routed_scaling_factor)) + return routing_weights, selected_experts + def _routing_scores(self, op: OpBuilder, hidden_states: ir.Value): weight_t = op.Transpose(self.weight, perm=[1, 0]) router_logits = op.MatMul( @@ -123,17 +142,13 @@ def _routing_scores(self, op: OpBuilder, hidden_states: ir.Value): op.Cast(weight_t, to=1), ) - # Score computation depends on scoring function if self.scoring_func == "sigmoid": - scores = op.Sigmoid(router_logits) # (B*S, num_experts) - # Add correction bias for expert selection (V3) + scores = op.Sigmoid(router_logits) scores_for_choice = op.Add(scores, self.e_score_correction_bias) else: - # Softmax scoring (V2) scores = op.Softmax(router_logits, axis=-1) scores_for_choice = scores - # Expert selection: group-based or simple TopK if self.n_group > 1 and self.topk_method != "greedy": scores_for_choice = self._group_topk_selection(op, scores_for_choice) return scores, scores_for_choice @@ -193,9 +208,15 @@ class DeepSeekMLADecoderLayer(nn.Module): Forward signature matches DecoderLayer for compatibility with TextModel. """ - def __init__(self, config: ArchitectureConfig, is_moe: bool = False): + def __init__( + self, + config: ArchitectureConfig, + is_moe: bool = False, + linear_class: type | None = None, + ): super().__init__() - linear_class = _linear_factory(config) + if linear_class is None: + linear_class = _linear_factory(config) self.self_attn = DeepSeekMLA(config, linear_class=linear_class) if is_moe: gate = DeepSeekMoEGate(config) @@ -241,9 +262,15 @@ class _DeepSeekStandardDecoderLayer(nn.Module): Forward signature matches DeepSeekMLADecoderLayer for compatibility. """ - def __init__(self, config: ArchitectureConfig, is_moe: bool = False): + def __init__( + self, + config: ArchitectureConfig, + is_moe: bool = False, + linear_class: type | None = None, + ): super().__init__() - linear_class = _linear_factory(config) + if linear_class is None: + linear_class = _linear_factory(config) self.self_attn = Attention(config, linear_class=linear_class) if is_moe: gate = DeepSeekMoEGate(config) @@ -290,17 +317,29 @@ class _DeepSeekMoEFFN(nn.Module): """ def __init__( - self, config: ArchitectureConfig, gate: nn.Module, linear_class: type | None = None + self, + config: ArchitectureConfig, + gate: nn.Module, + linear_class: type | None = None, ): super().__init__() assert config.num_local_experts is not None assert config.moe_intermediate_size is not None - self.moe = MoELayer(config, gate=gate) + qc = config.quantization + use_fused_qmoe = ( + config.fused_quantized_moe and qc is not None and qc.quant_method != "none" + ) + if use_fused_qmoe: + self.moe = FusedQuantizedMoE(config, gate=gate) + else: + self.moe = MoELayer(config, gate=gate, linear_class=linear_class) # Shared expert uses moe_intermediate_size * n_shared_experts n_shared = config.n_shared_experts or 1 shared_intermediate = config.moe_intermediate_size * n_shared self.shared_experts = _SharedExpertMLP( - config, shared_intermediate, linear_class=linear_class + config, + shared_intermediate, + linear_class=linear_class, ) def forward(self, op: OpBuilder, hidden_states: ir.Value): @@ -346,12 +385,26 @@ class DeepSeekV3TextModel(nn.Module): def __init__(self, config: ArchitectureConfig): super().__init__() - self.embed_tokens = Embedding(config.vocab_size, config.hidden_size) + linear_class = _linear_class(config) + qc = config.quantization + if qc is not None and qc.quantize_embeddings: + self.embed_tokens = QuantizedEmbedding( + config.vocab_size, + config.hidden_size, + bits=qc.bits, + block_size=qc.group_size, + has_zero_point=not qc.sym, + padding_idx=config.pad_token_id, + ) + else: + self.embed_tokens = Embedding(config.vocab_size, config.hidden_size) self._dtype = config.dtype # Detect MLA vs standard attention use_mla = config.qk_nope_head_dim is not None and config.qk_nope_head_dim > 0 - LayerClass = DeepSeekMLADecoderLayer if use_mla else _DeepSeekStandardDecoderLayer # noqa: N806 + LayerClass = ( # noqa: N806 + DeepSeekMLADecoderLayer if use_mla else _DeepSeekStandardDecoderLayer + ) # Build layers: dense for first k, MoE for rest first_k = config.first_k_dense_replace @@ -360,7 +413,7 @@ def __init__(self, config: ArchitectureConfig): first_k = config.num_hidden_layers self.layers = nn.ModuleList( [ - LayerClass(config, is_moe=(i >= first_k)) + LayerClass(config, is_moe=(i >= first_k), linear_class=linear_class) for i in range(config.num_hidden_layers) ] ) @@ -425,7 +478,11 @@ def __init__(self, config: ArchitectureConfig): nn.Module.__init__(self) self.config = config self.model = DeepSeekV3TextModel(config) - self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) + qc = config.quantization + lm_head_class = ( + _linear_class(config) if qc is not None and qc.quantize_lm_head else Linear + ) + self.lm_head = lm_head_class(config.hidden_size, config.vocab_size, bias=False) def preprocess_weights( self, state_dict: dict[str, torch.Tensor] @@ -439,23 +496,35 @@ def preprocess_weights( - MoE expert weights: HF stores all experts in a single fused tensor experts.gate_up_proj: (n_experts, 2*intermediate, hidden) experts.down_proj: (n_experts, hidden, intermediate) - These are split into per-expert ONNX weights: + Static MoE exports split these into per-expert ONNX weights: moe.experts.{i}.gate_proj.weight: (intermediate, hidden) moe.experts.{i}.up_proj.weight: (intermediate, hidden) moe.experts.{i}.down_proj.weight: (hidden, intermediate) + Fused quantized exports instead pack expert-major FC1/FC2 tensors for QMoE. """ renamed = {} + routed_experts = {} + qc = self.config.quantization + use_fused_qmoe = ( + self.config.fused_quantized_moe + and qc is not None + and qc.quant_method != "none" + ) use_qmoe = ( - self.config.quantization is not None - and self.config.quantization.bits == 4 - and self.config.quantization.quant_method in {"gptq", "awq"} - and not self.config.quantization.float_zero_point + not use_fused_qmoe + and qc is not None + and qc.bits == 4 + and qc.quant_method in {"gptq", "awq"} + and not qc.float_zero_point ) for key, value in state_dict.items(): new_key = key # Remap MoE layer names: mlp.gate.* → mlp.moe.gate.* new_key = new_key.replace(".mlp.gate.", ".mlp.moe.gate.") + if use_fused_qmoe and ".mlp.experts." in new_key: + routed_experts[new_key] = value + continue # HF stores all routed experts in fused tensors: # layers.N.mlp.experts.gate_up_proj (n_experts, 2*mid, hidden) @@ -476,9 +545,10 @@ def preprocess_weights( renamed[new_key] = value - # Handle weight tying and GPTQ/AWQ conversion before flattening the - # expert-major MatMulNBits blobs into the QMoE ABI. + # Handle weight tying and pack routed expert tensors for the selected QMoE path. processed = super().preprocess_weights(renamed) - if use_qmoe: + if use_fused_qmoe: + processed.update(pack_fused_quantized_moe_weights(routed_experts, self.config)) + elif use_qmoe: processed = pack_qmoe_expert_weights(processed) return processed diff --git a/src/mobius/models/glm_moe_dsa.py b/src/mobius/models/glm_moe_dsa.py new file mode 100644 index 00000000..451d4337 --- /dev/null +++ b/src/mobius/models/glm_moe_dsa.py @@ -0,0 +1,620 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""GLM-5.2 MLA+MoE with IndexShare deep sparse attention and MTP export.""" + +from __future__ import annotations + +import re +from typing import TYPE_CHECKING + +import onnx_ir as ir +from onnxscript import OpBuilder, nn + +from mobius._configs import ArchitectureConfig +from mobius.components import ( + MLP, + DeepSeekMLA, + Embedding, + LayerNorm, + Linear, + RMSNorm, + create_attention_bias, + initialize_rope, +) +from mobius.components._rotary_embedding import apply_rotary_pos_emb +from mobius.models.deepseek import ( + DeepSeekMoEGate, + DeepSeekV3CausalLMModel, + DeepSeekV3TextModel, + _DeepSeekMoEFFN, + _linear_class, +) + +if TYPE_CHECKING: + import torch + + +_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$") + + +def _indexer_types(config: ArchitectureConfig) -> list[str]: + """Return the authoritative full/shared IndexShare schedule.""" + if config.indexer_types is not None: + indexer_types = list(config.indexer_types) + if len(indexer_types) != config.num_hidden_layers: + raise ValueError( + "indexer_types must contain exactly one entry per hidden layer " + f"(expected {config.num_hidden_layers}, got {len(indexer_types)})" + ) + return indexer_types + full_layers = {0, 1, 2} + full_layers.update( + range( + config.index_skip_topk_offset * 2, + config.num_hidden_layers, + config.index_topk_freq, + ) + ) + return ["full" if i in full_layers else "shared" for i in range(config.num_hidden_layers)] + + +class GlmMoeDsaIndexer(nn.Module): + """GLM-5.2 fp32 token indexer used by full IndexShare layers.""" + + def __init__(self, config: ArchitectureConfig, linear_class: type): + super().__init__() + assert config.q_lora_rank is not None + assert config.index_head_dim is not None + assert config.index_n_heads is not None + self.index_head_dim = config.index_head_dim + self.index_n_heads = config.index_n_heads + self.index_topk = int(config.index_topk or 2048) + self.rope_interleave = config.indexer_rope_interleave + self.wq_b = linear_class( + config.q_lora_rank, + self.index_n_heads * self.index_head_dim, + bias=False, + ) + self.wk = linear_class(config.hidden_size, self.index_head_dim, bias=False) + self.k_norm = LayerNorm(self.index_head_dim, eps=1e-6) + self.weights_proj = linear_class(config.hidden_size, self.index_n_heads, bias=False) + + def project_key( + self, + op: OpBuilder, + hidden_states: ir.Value, + position_embeddings: tuple, + ) -> ir.Value: + key = self.k_norm(op, self.wk(op, hidden_states)) + # The indexer key shares the model's rotary_emb cos/sin cache (sized + # qk_rope_head_dim / 2). RoPE rotates the *entire* index_head_dim, so + # rotary_embedding_dim must be 0 (full rotation) — matching the main + # MLA q_rope/k_rope calls. Passing index_head_dim // 2 here makes the + # opset-24 RotaryEmbedding op expect a cos cache of that_value / 2, + # which mismatches the shared cache and fails shape validation. + return apply_rotary_pos_emb( + op, + key, + position_embeddings, + num_heads=1, + rotary_embedding_dim=0, + interleaved=self.rope_interleave, + ) + + def select( + self, + op: OpBuilder, + hidden_states: ir.Value, + q_resid: ir.Value, + all_index_keys: ir.Value, + attention_bias: ir.Value, + ) -> ir.Value: + query = self.wq_b(op, q_resid) + query = op.Reshape(query, [0, 0, self.index_n_heads, self.index_head_dim]) + query = op.Cast(query, to=ir.DataType.FLOAT) + keys = op.Cast(all_index_keys, to=ir.DataType.FLOAT) + scores = op.MatMul(query, op.Transpose(keys, perm=[0, 2, 1])) + scores = op.Relu(scores) + + weights = self.weights_proj(op, hidden_states) + weights = op.Cast(weights, to=ir.DataType.FLOAT) + weights = op.Mul(weights, self.index_n_heads**-0.5) + scores = op.ReduceSum(op.Mul(scores, op.Unsqueeze(weights, [3])), [2], keepdims=False) + # The additive causal mask is sized to the attention key axis, which the + # native decoder exposes at fixed KV capacity while the indexer key cache + # grows with the logical prefix. Slice the mask down to the score key + # length so the add stays broadcast-compatible under both the logical + # (ORT) and fixed-capacity (native single-token decode) exposures. + key_length = op.Shape(scores, start=2, end=3) + causal = op.Squeeze(attention_bias, [1]) + causal = op.Slice( + causal, op.Constant(value_ints=[0]), key_length, op.Constant(value_ints=[2]) + ) + scores = op.Add(scores, op.Cast(causal, to=ir.DataType.FLOAT)) + + k = op.Min(key_length, op.Constant(value_ints=[self.index_topk])) + _, indices = op.TopK(scores, k, axis=-1, largest=1, sorted=0, _outputs=2) + # ``pkg.nxrt::IndexShare`` requires the selected key positions to be + # strictly increasing per row. TopK returns them in score order, so sort + # the chosen positions ascending. The runtime TopK kernel is f32-only, so + # round-trip through float before restoring the Int64 index dtype. + indices = op.Cast(indices, to=ir.DataType.FLOAT) + indices, _ = op.TopK(indices, k, axis=-1, largest=0, sorted=1, _outputs=2) + indices = op.Cast(indices, to=ir.DataType.INT64) + return indices + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + q_resid: ir.Value, + position_embeddings: tuple, + past_index_key: ir.Value | None, + attention_bias: ir.Value, + ): + current_index_key = self.project_key(op, hidden_states, position_embeddings) + all_index_keys = ( + current_index_key + if past_index_key is None + else op.Concat(past_index_key, current_index_key, axis=1) + ) + indices = self.select( + op, + hidden_states, + q_resid, + all_index_keys, + attention_bias, + ) + return all_index_keys, indices + + +class GlmMoeDsaAttention(DeepSeekMLA): + """MLA attention with a packed indexer-key cache and sparse additive mask.""" + + def __init__( + self, + config: ArchitectureConfig, + indexer_type: str, + linear_class: type, + ): + super().__init__(config, linear_class=linear_class) + self.indexer_type = indexer_type + self.dtype = config.dtype + self.main_key_dim = self.num_heads * self.qk_head_dim + self.main_value_dim = self.num_heads * self.v_head_dim + self.index_head_dim = int(config.index_head_dim or 0) + if indexer_type == "full": + self.indexer = GlmMoeDsaIndexer(config, linear_class) + + def _unpack_past(self, op: OpBuilder, past_key_value: tuple): + packed_key, packed_value = past_key_value + key_tokens = op.Squeeze(packed_key, [1]) + main_key = op.Slice(key_tokens, [0], [self.main_key_dim], [2]) + main_key = op.Transpose( + op.Reshape(main_key, [0, 0, self.num_heads, self.qk_head_dim]), + perm=[0, 2, 1, 3], + ) + value_tokens = op.Squeeze(packed_value, [1]) + main_value = op.Transpose( + op.Reshape(value_tokens, [0, 0, self.num_heads, self.v_head_dim]), + perm=[0, 2, 1, 3], + ) + index_key = None + if self.indexer_type == "full": + index_key = op.Slice( + key_tokens, + [self.main_key_dim], + [self.main_key_dim + self.index_head_dim], + [2], + ) + return (main_key, main_value), index_key + + def _pack_present( + self, + op: OpBuilder, + present: tuple, + index_keys: ir.Value | None, + ) -> tuple: + key = op.Reshape( + op.Transpose(present[0], perm=[0, 2, 1, 3]), [0, 0, self.main_key_dim] + ) + if index_keys is not None: + key = op.Concat(key, index_keys, axis=-1) + value = op.Reshape( + op.Transpose(present[1], perm=[0, 2, 1, 3]), + [0, 0, self.main_value_dim], + ) + return op.Unsqueeze(key, [1]), op.Unsqueeze(value, [1]) + + def _index_share_attention( + self, + op: OpBuilder, + q: ir.Value, + key: ir.Value, + value: ir.Value, + main_past: tuple | None, + topk_indices: ir.Value, + ) -> tuple[ir.Value, ir.Value, ir.Value]: + """Device-resident selected-token attention via ``pkg.nxrt::IndexShare``. + + Replaces the dense ``ScatterElements`` sparse-mask + ``Attention`` lowering + with the frozen ``IndexShare`` op, which gathers only the indexer's + selected key positions. Causality is carried entirely by the (already + causal) ``topk_indices``, so no dense additive mask island — and thus no + logical-length-vs-fixed-capacity broadcast — is emitted here. + + ``IndexShare`` requires a homogeneous head size across query/key/value. + MLA's value head (``v_head_dim``) is narrower than the query/key head + (``qk_head_dim``), so the value is zero-padded up to ``qk_head_dim`` for + the kernel and the padding columns are sliced back off the output. The + returned present key/value stay in their native (unpadded) BNSH layout so + the packed KV cache is unchanged. + """ + q_bnsh = op.Transpose( + op.Reshape(q, [0, 0, self.num_heads, self.qk_head_dim]), perm=[0, 2, 1, 3] + ) + cur_key = op.Transpose( + op.Reshape(key, [0, 0, self.num_heads, self.qk_head_dim]), perm=[0, 2, 1, 3] + ) + cur_value = op.Transpose( + op.Reshape(value, [0, 0, self.num_heads, self.v_head_dim]), perm=[0, 2, 1, 3] + ) + present_key = cur_key + present_value = cur_value + if main_past is not None: + present_key = op.Concat(main_past[0], cur_key, axis=2) + present_value = op.Concat(main_past[1], cur_value, axis=2) + + value_pad = self.qk_head_dim - self.v_head_dim + padded_value = present_value + if value_pad > 0: + # Zero-pad the value head up to ``qk_head_dim`` without ``Pad`` (which + # the CUDA EP has no handler for). Build a device-resident zero block + # of the value's own dtype and concatenate it onto the head axis. + pad_shape = op.Concat( + op.Shape(present_value, start=0, end=3), + op.Constant(value_ints=[value_pad]), + axis=0, + ) + zeros = op.Expand(op.Cast(op.Constant(value_float=0.0), to=self.dtype), pad_shape) + padded_value = op.Concat(present_value, zeros, axis=3) + + selected_indices = op.Unsqueeze(topk_indices, [1]) + attn_output = op.IndexShare( + q_bnsh, + present_key, + padded_value, + None, + None, + selected_indices, + num_heads=self.num_heads, + kv_num_heads=self.num_heads, + scale=self.scaling, + _domain="pkg.nxrt", + _outputs=1, + ) + if value_pad > 0: + attn_output = op.Slice(attn_output, [0], [self.v_head_dim], [3]) + attn_output = op.Reshape( + op.Transpose(attn_output, perm=[0, 2, 1, 3]), + [0, 0, self.num_heads * self.v_head_dim], + ) + return attn_output, present_key, present_value + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + attention_bias: ir.Value, + position_embeddings: tuple, + past_key_value: tuple | None = None, + shared_topk_indices: ir.Value | None = None, + ): + q_resid = self.q_a_layernorm(op, self.q_a_proj(op, hidden_states)) + q = self.q_b_proj(op, q_resid) + q = op.Reshape(q, [0, 0, self.num_heads, self.qk_head_dim]) + q_nope, q_rope = op.Split( + q, + [self.qk_nope_head_dim, self.qk_rope_head_dim], + axis=-1, + _outputs=2, + ) + q_rope = op.Reshape(q_rope, [0, 0, -1]) + q_rope = apply_rotary_pos_emb( + op, + q_rope, + position_embeddings, + num_heads=self.num_heads, + interleaved=self._rope_interleave, + ) + q = op.Concat( + op.Reshape(q_nope, [0, 0, self.num_heads, self.qk_nope_head_dim]), + op.Reshape(q_rope, [0, 0, self.num_heads, self.qk_rope_head_dim]), + axis=-1, + ) + q = op.Reshape(q, [0, 0, self.num_heads * self.qk_head_dim]) + + compressed_kv = self.kv_a_proj_with_mqa(op, hidden_states) + kv_latent, k_rope = op.Split( + compressed_kv, + [self.kv_lora_rank, self.qk_rope_head_dim], + axis=-1, + _outputs=2, + ) + kv_latent = self.kv_a_layernorm(op, kv_latent) + kv = self.kv_b_proj(op, kv_latent) + kv = op.Reshape( + kv, + [0, 0, self.num_heads, self.qk_nope_head_dim + self.v_head_dim], + ) + k_nope, value = op.Split( + kv, + [self.qk_nope_head_dim, self.v_head_dim], + axis=-1, + _outputs=2, + ) + value = op.Reshape(value, [0, 0, -1]) + k_rope = apply_rotary_pos_emb( + op, + k_rope, + position_embeddings, + num_heads=1, + interleaved=self._rope_interleave, + ) + k_rope = op.Tile(k_rope, [1, 1, self.num_heads]) + key = op.Concat( + op.Reshape(k_nope, [0, 0, self.num_heads, self.qk_nope_head_dim]), + op.Reshape(k_rope, [0, 0, self.num_heads, self.qk_rope_head_dim]), + axis=-1, + ) + key = op.Reshape(key, [0, 0, self.num_heads * self.qk_head_dim]) + + main_past = None + past_index_key = None + if past_key_value is not None: + main_past, past_index_key = self._unpack_past(op, past_key_value) + + all_index_keys = None + topk_indices = shared_topk_indices + if self.indexer_type == "full": + all_index_keys, topk_indices = self.indexer( + op, + hidden_states, + q_resid, + position_embeddings, + past_index_key, + attention_bias, + ) + if topk_indices is None: + raise ValueError( + "Shared GLM DSA layers require top-k indices from a preceding full layer" + ) + + attn_output, present_key, present_value = self._index_share_attention( + op, + q, + key, + value, + main_past, + topk_indices, + ) + attn_output = self.o_proj(op, attn_output) + present = self._pack_present( + op, + (present_key, present_value), + all_index_keys, + ) + return attn_output, present, topk_indices + + +class GlmMoeDsaDecoderLayer(nn.Module): + def __init__( + self, + config: ArchitectureConfig, + indexer_type: str, + is_moe: bool, + linear_class: type, + ): + super().__init__() + self.self_attn = GlmMoeDsaAttention(config, indexer_type, linear_class) + if is_moe: + self.mlp = _DeepSeekMoEFFN( + config, + DeepSeekMoEGate(config), + linear_class=linear_class, + ) + else: + self.mlp = MLP(config, linear_class=linear_class) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + attention_bias: ir.Value, + position_embeddings: tuple, + past_key_value: tuple | None, + shared_topk_indices: ir.Value | None, + ): + residual = hidden_states + hidden_states = self.input_layernorm(op, hidden_states) + hidden_states, present, topk_indices = self.self_attn( + op, + hidden_states, + attention_bias, + position_embeddings, + past_key_value, + shared_topk_indices, + ) + hidden_states = op.Add(residual, hidden_states) + residual = hidden_states + hidden_states = self.post_attention_layernorm(op, hidden_states) + hidden_states = self.mlp(op, hidden_states) + return op.Add(residual, hidden_states), present, topk_indices + + +class GlmMoeDsaTextModel(nn.Module): + def __init__(self, config: ArchitectureConfig): + super().__init__() + self.config = config + self.embed_tokens = Embedding(config.vocab_size, config.hidden_size) + self.rotary_emb = initialize_rope(config) + linear_class = _linear_class(config) + types = _indexer_types(config) + self.layers = nn.ModuleList( + [ + GlmMoeDsaDecoderLayer( + config, + types[i], + is_moe=i >= config.first_k_dense_replace, + linear_class=linear_class, + ) + for i in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + op: OpBuilder, + input_ids: ir.Value, + attention_mask: ir.Value, + position_ids: ir.Value, + past_key_values: list | None, + ): + hidden_states = self.embed_tokens(op, input_ids) + position_embeddings = self.rotary_emb(op, position_ids) + attention_bias = create_attention_bias( + op, + input_ids, + attention_mask, + dtype=self.config.dtype, + ) + presents = [] + shared_topk_indices = None + for i, layer in enumerate(self.layers): + past = past_key_values[i] if past_key_values is not None else None + hidden_states, present, shared_topk_indices = layer( + op, + hidden_states, + attention_bias, + position_embeddings, + past, + shared_topk_indices, + ) + presents.append(present) + return self.norm(op, hidden_states), presents + + +class _SharedHead(nn.Module): + def __init__(self, config: ArchitectureConfig): + super().__init__() + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + +class GlmMoeDsaMtp(nn.Module): + """Improved GLM-5.2 MTP head: embed/hidden fusion plus a full DSA MoE block.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + linear_class = _linear_class(config) + self.config = config + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.eh_proj = linear_class(config.hidden_size * 2, config.hidden_size, bias=False) + self.layer = GlmMoeDsaDecoderLayer( + config, + "full", + is_moe=True, + linear_class=linear_class, + ) + self.shared_head = _SharedHead(config) + self.rotary_emb = initialize_rope(config) + + def forward( + self, + op: OpBuilder, + inputs_embeds: ir.Value, + hidden_states: ir.Value, + attention_mask: ir.Value, + position_ids: ir.Value, + past_key_value: tuple, + ): + fused = self.eh_proj( + op, + op.Concat(self.enorm(op, inputs_embeds), self.hnorm(op, hidden_states), axis=-1), + ) + position_embeddings = self.rotary_emb(op, position_ids) + attention_bias = create_attention_bias( + op, + position_ids, + attention_mask, + dtype=self.config.dtype, + ) + output, present, topk_indices = self.layer( + op, + fused, + attention_bias, + position_embeddings, + past_key_value, + None, + ) + return self.shared_head.norm(op, output), present, topk_indices + + +class GlmMoeDsaCausalLMModel(DeepSeekV3CausalLMModel): + """GLM-5.2 target model with portable IndexShare and an exported MTP graph.""" + + default_task = "glm-moe-dsa" + + def __init__(self, config: ArchitectureConfig): + if config.indexer_types is None: + config.indexer_types = _indexer_types(config) + nn.Module.__init__(self) + self.config = config + self.model = ( + GlmMoeDsaTextModel(config) if config.use_dsa else DeepSeekV3TextModel(config) + ) + qc = config.quantization + lm_head_class = ( + _linear_class(config) if qc is not None and qc.quantize_lm_head else Linear + ) + self.lm_head = lm_head_class(config.hidden_size, config.vocab_size, bias=False) + if config.num_nextn_predict_layers > 0: + if config.num_nextn_predict_layers != 1: + raise ValueError("GLM MTP export currently supports exactly one NextN layer") + if not config.use_dsa: + raise ValueError( + "GLM-5.2 MTP export currently requires DSA; set num_nextn_predict_layers=0 " + "when using the full-attention fallback" + ) + self.mtp = GlmMoeDsaMtp(config) + else: + self.mtp = None + + def preprocess_weights( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + state_dict = DeepSeekV3CausalLMModel.preprocess_weights(self, state_dict) + result: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + match = _LAYER_RE.match(key) + if match is not None and int(match.group(1)) >= self.config.num_hidden_layers: + layer_idx = int(match.group(1)) + if ( + self.mtp is None + or layer_idx + >= self.config.num_hidden_layers + self.config.num_nextn_predict_layers + ): + continue + suffix = match.group(2) + if suffix.startswith(("enorm.", "hnorm.", "eh_proj.", "shared_head.")): + key = f"mtp.{suffix}" + else: + key = f"mtp.layer.{suffix}" + elif not self.config.use_dsa and ".indexer." in key: + continue + result[key.replace(".mlp.experts.", ".mlp.moe.experts.")] = value + return result diff --git a/src/mobius/models/glm_moe_dsa_test.py b/src/mobius/models/glm_moe_dsa_test.py new file mode 100644 index 00000000..72ec95e2 --- /dev/null +++ b/src/mobius/models/glm_moe_dsa_test.py @@ -0,0 +1,247 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import pytest +import torch +from transformers.models.glm_moe_dsa.configuration_glm_moe_dsa import GlmMoeDsaConfig + +from mobius._builder import build_from_module +from mobius._config_resolver import _default_task_for_model +from mobius._configs import ArchitectureConfig, QuantizationConfig +from mobius._registry import registry +from mobius._testing import count_op_type, make_config +from mobius.models.glm_moe_dsa import GlmMoeDsaCausalLMModel, _indexer_types + + +def _tiny_glm_moe_dsa_config(**overrides): + defaults = dict( + model_type="glm_moe_dsa", + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=4, + q_lora_rank=16, + kv_lora_rank=8, + qk_nope_head_dim=12, + qk_rope_head_dim=4, + v_head_dim=8, + head_dim=4, + rope_interleave=True, + num_local_experts=2, + num_experts_per_tok=1, + moe_intermediate_size=32, + shared_expert_intermediate_size=32, + n_shared_experts=1, + first_k_dense_replace=1, + norm_topk_prob=True, + routed_scaling_factor=2.5, + scoring_func="sigmoid", + topk_method="noaux_tc", + index_topk=8, + index_head_dim=8, + index_n_heads=2, + indexer_types=["full", "shared", "shared", "shared"], + num_nextn_predict_layers=1, + ) + defaults.update(overrides) + return make_config(**defaults) + + +class TestGlmMoeDsaExport: + def test_hf_config_extraction(self): + hf_config = GlmMoeDsaConfig( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=4, + q_lora_rank=16, + kv_lora_rank=8, + qk_nope_head_dim=12, + qk_rope_head_dim=4, + v_head_dim=8, + n_routed_experts=2, + num_experts_per_tok=1, + moe_intermediate_size=32, + n_shared_experts=1, + first_k_dense_replace=1, + index_topk=8, + index_head_dim=8, + index_n_heads=2, + index_topk_freq=4, + index_skip_topk_offset=3, + index_share_for_mtp_iteration=True, + ) + + config = ArchitectureConfig.from_transformers(hf_config) + + assert config.model_type == "glm_moe_dsa" + assert config.num_local_experts == 2 + assert config.shared_expert_intermediate_size == 32 + assert config.qk_nope_head_dim == 12 + assert config.qk_rope_head_dim == 4 + assert config.index_topk == 8 + assert config.index_topk_freq == 4 + assert config.index_skip_topk_offset == 3 + assert config.index_share_for_mtp_iteration is True + assert config.rope_interleave is True + + def test_registry(self): + assert registry.get("glm_moe_dsa") is GlmMoeDsaCausalLMModel + assert _default_task_for_model("glm_moe_dsa") == "glm-moe-dsa" + + def test_indexer_types_length_must_match_layers(self): + config = _tiny_glm_moe_dsa_config(indexer_types=["full"]) + + with pytest.raises(ValueError, match="expected 4, got 1"): + _indexer_types(config) + + def test_builds_index_share_mla_moe_graph(self): + config = _tiny_glm_moe_dsa_config() + package = build_from_module( + GlmMoeDsaCausalLMModel(config), + config, + task="glm-moe-dsa", + ) + graph = package["model"].graph + + assert count_op_type(graph, "IndexShare") == config.num_hidden_layers + assert count_op_type(graph, "Attention") == 0 + assert count_op_type(graph, "ScatterElements") == 0 + assert count_op_type(graph, "Sigmoid") >= config.num_hidden_layers - 1 + assert any( + "layers.0.self_attn.indexer.wk.weight" in name for name in graph.initializers + ) + assert not any("layers.1.self_attn.indexer" in name for name in graph.initializers) + assert {value.name for value in graph.outputs} >= { + "logits", + "present.0.key", + "present.0.value", + } + assert set(package) == {"model", "mtp"} + + mtp_graph = package["mtp"].graph + assert count_op_type(mtp_graph, "IndexShare") == 1 + assert count_op_type(mtp_graph, "Attention") == 0 + assert count_op_type(mtp_graph, "ScatterElements") == 0 + assert {value.name for value in mtp_graph.outputs} == { + "mtp_hidden", + "present.0.key", + "present.0.value", + "topk_indices", + } + + def test_indexer_rotary_uses_full_rotation(self): + """Regression: indexer RoPE must rotate the full index_head_dim. + + The indexer key shares the model's rotary_emb cos/sin cache (sized + qk_rope_head_dim / 2). The opset-24 RotaryEmbedding op validates that + the cos cache last dim equals head_size / 2 or rotary_embedding_dim / 2. + Passing ``rotary_embedding_dim = index_head_dim // 2`` made the op + expect a cache of that_value / 2, which mismatched the shared cache and + raised "cos_cache dimension 2 should be same as head_size / 2 or + rotary_embedding_dim / 2" at runtime. Full rotation (0) matches the + shared cache like the main MLA q_rope/k_rope calls do. + """ + config = _tiny_glm_moe_dsa_config() + package = build_from_module( + GlmMoeDsaCausalLMModel(config), + config, + task="glm-moe-dsa", + ) + graph = package["model"].graph + + indexer_rope_nodes = [ + node + for node in graph + if node.op_type == "RotaryEmbedding" and "indexer" in (node.name or "") + ] + assert indexer_rope_nodes, "expected at least one indexer RotaryEmbedding node" + for node in indexer_rope_nodes: + attr = node.attributes.get("rotary_embedding_dim") + dim = attr.value if attr is not None else 0 + assert dim == 0, ( + f"indexer RotaryEmbedding {node.name!r} must use full rotation " + f"(rotary_embedding_dim=0) to match the shared cos cache, got {dim}" + ) + + def test_quantized_graph_uses_matmul_nbits(self): + config = _tiny_glm_moe_dsa_config( + quantization=QuantizationConfig( + bits=4, + group_size=32, + quant_method="gguf", + sym=False, + ) + ) + graph = build_from_module( + GlmMoeDsaCausalLMModel(config), + config, + task="glm-moe-dsa", + )["model"].graph + assert count_op_type(graph, "MatMulNBits") > 0 + + def test_fused_quantized_moe_emits_qmoe(self): + config = _tiny_glm_moe_dsa_config( + quantization=QuantizationConfig( + bits=4, + group_size=32, + quant_method="gguf", + sym=True, + ), + fused_quantized_moe=True, + ) + graph = build_from_module( + GlmMoeDsaCausalLMModel(config), + config, + task="glm-moe-dsa", + )["model"].graph + # Routed experts collapse into a single fused QMoE node per MoE layer. + n_moe_layers = config.num_hidden_layers - config.first_k_dense_replace + assert count_op_type(graph, "QMoE") == n_moe_layers + # Routed experts no longer emit per-expert MatMulNBits: the only + # MatMulNBits left come from attention, the dense layer, shared experts, + # and lm_head — none named under the routed-expert ``moe.experts`` scope. + routed = [ + node.name + for node in graph + if node.op_type == "MatMulNBits" and "moe.experts" in (node.name or "") + ] + assert routed == [], f"routed experts still emit MatMulNBits: {routed}" + config = _tiny_glm_moe_dsa_config() + model = GlmMoeDsaCausalLMModel(config) + state_dict = { + "model.layers.1.mlp.gate.weight": torch.zeros(2, 64), + "model.layers.1.mlp.experts.0.gate_proj.weight": torch.zeros(32, 64), + "model.layers.0.self_attn.indexer.wk.weight": torch.zeros(8, 64), + "model.layers.4.self_attn.q_a_proj.weight": torch.zeros(16, 64), + "model.layers.4.eh_proj.weight": torch.zeros(64, 128), + "model.layers.4.shared_head.norm.weight": torch.zeros(64), + } + + result = model.preprocess_weights(state_dict) + + assert "model.layers.1.mlp.moe.gate.weight" in result + assert "model.layers.1.mlp.moe.experts.0.gate_proj.weight" in result + assert "model.layers.0.self_attn.indexer.wk.weight" in result + assert "mtp.layer.self_attn.q_a_proj.weight" in result + assert "mtp.eh_proj.weight" in result + assert "mtp.shared_head.norm.weight" in result + + def test_full_attention_fallback(self): + config = _tiny_glm_moe_dsa_config( + use_dsa=False, + num_nextn_predict_layers=0, + ) + package = build_from_module( + GlmMoeDsaCausalLMModel(config), + config, + task="glm-moe-dsa", + ) + graph = package["model"].graph + + assert set(package) == {"model"} + assert count_op_type(graph, "Attention") == config.num_hidden_layers + assert count_op_type(graph, "ScatterElements") == 0 + assert not any("indexer" in name for name in graph.initializers) diff --git a/src/mobius/models/llada_test.py b/src/mobius/models/llada_test.py index 9e990a78..2af1f747 100644 --- a/src/mobius/models/llada_test.py +++ b/src/mobius/models/llada_test.py @@ -168,11 +168,13 @@ def _build_onnx_session(config: ArchitectureConfig, state: dict[str, torch.Tenso model = MaskedDiffusionTask().build(module, config)["model"] apply_weights(model, module.preprocess_weights(state)) - # Windows keeps the ORT model file mapped; ignore cleanup errors so the dir can be removed. with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir: model_path = Path(temp_dir) / "model.onnx" onnx_ir.save(model, model_path) - return ort.InferenceSession(model_path) + # onnxruntime reads the model fully at construction, so the temp file + # can be removed as soon as the session exists. + return ort.InferenceSession(str(model_path)) + def test_llada_matches_torch_reference(): diff --git a/src/mobius/models/unet_parity_test.py b/src/mobius/models/unet_parity_test.py index 95d9ab26..833588da 100644 --- a/src/mobius/models/unet_parity_test.py +++ b/src/mobius/models/unet_parity_test.py @@ -64,11 +64,11 @@ def test_cross_attention_block_matches_diffusers(): model = _make_model(graph) apply_weights(model, _remap_transformer(hf.state_dict())) - # Windows keeps the ORT model file mapped; ignore cleanup errors so the dir can be removed. with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir: model_path = Path(temp_dir) / "model.onnx" onnx_ir.save(model, model_path) - session = ort.InferenceSession(model_path) + session = ort.InferenceSession(str(model_path)) + actual = session.run(None, {"hidden": hidden.numpy(), "context": context.numpy()})[0] assert np.abs(actual - expected).max() < 1e-4 @@ -118,11 +118,11 @@ def test_unet_matches_diffusers(): model = DenoisingTask().build(module, config)["model"] apply_weights(model, module.preprocess_weights(dict(hf.state_dict()))) - # Windows keeps the ORT model file mapped; ignore cleanup errors so the dir can be removed. with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir: model_path = Path(temp_dir) / "model.onnx" onnx_ir.save(model, model_path) - session = ort.InferenceSession(model_path) + session = ort.InferenceSession(str(model_path)) + actual = session.run( None, { @@ -190,11 +190,11 @@ def test_unet_sd1x_mixed_block_types_matches_diffusers(): model = DenoisingTask().build(module, config)["model"] apply_weights(model, module.preprocess_weights(dict(hf.state_dict()))) - # Windows keeps the ORT model file mapped; ignore cleanup errors so the dir can be removed. with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir: model_path = Path(temp_dir) / "model.onnx" onnx_ir.save(model, model_path) - session = ort.InferenceSession(model_path) + session = ort.InferenceSession(str(model_path)) + actual = session.run( None, { @@ -278,11 +278,11 @@ def test_unet_lora_gate_parity(): weights.update(remap_diffusers_unet_lora(lora_state, "test")) apply_weights(model, weights) - # Windows keeps the ORT model file mapped; ignore cleanup errors so the dir can be removed. with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir: model_path = Path(temp_dir) / "model.onnx" onnx_ir.save(model, model_path) - session = ort.InferenceSession(model_path) + session = ort.InferenceSession(str(model_path)) + feed = { "sample": sample.numpy(), "timestep": timestep.numpy().astype(np.int64), diff --git a/src/mobius/rewrite_rules/_decompose_attention_test.py b/src/mobius/rewrite_rules/_decompose_attention_test.py index 9db4dac6..e5a99fbd 100644 --- a/src/mobius/rewrite_rules/_decompose_attention_test.py +++ b/src/mobius/rewrite_rules/_decompose_attention_test.py @@ -167,9 +167,9 @@ def _feeds( (4, 0, 8, 8, 0.0, 1, True, False, 0.0, "prefill_mha"), (4, 0, 8, 2, 0.0, 1, True, False, 0.0, "prefill_gqa"), (1, 3, 8, 1, 0.0, 1, True, True, 0.0, "decode_gqa_past"), - (4, 0, 8, 2, 30.0, 1, True, False, 1e-5, "softcap"), + (4, 0, 8, 2, 30.0, 1, True, False, 5e-3, "softcap"), (4, 0, 8, 2, 0.0, 0, True, False, 0.0, "mask_only_noncausal"), - (1, 5, 8, 1, 30.0, 1, True, True, 1e-5, "decode_softcap_past"), + (1, 5, 8, 1, 30.0, 1, True, True, 5e-3, "decode_softcap_past"), ] diff --git a/src/mobius/rewrite_rules/_group_query_attention_test.py b/src/mobius/rewrite_rules/_group_query_attention_test.py index 6d95812d..2bcd4c1a 100644 --- a/src/mobius/rewrite_rules/_group_query_attention_test.py +++ b/src/mobius/rewrite_rules/_group_query_attention_test.py @@ -180,6 +180,31 @@ ) +def _build_tiny_mla_graph(v_head_dim: int): + """Build a CPU-only MLA graph with configurable K/V head dimensions.""" + config = ArchitectureConfig( + model_type="deepseek_v2", + hidden_size=64, + intermediate_size=128, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=16, + num_hidden_layers=1, + vocab_size=256, + max_position_embeddings=128, + hidden_act="silu", + rms_norm_eps=1e-6, + rope_type="default", + rope_theta=10000.0, + q_lora_rank=16, + kv_lora_rank=8, + qk_nope_head_dim=12, + qk_rope_head_dim=4, + v_head_dim=v_head_dim, + ) + return build_from_module(registry.get("deepseek_v2")(config), config)["model"] + + class TestGroupQueryAttentionRules: def test_rules_returns_rule_set(self): rules = group_query_attention_rules() @@ -350,6 +375,7 @@ def test_deepseek_v2_lite_int4_uses_one_qmoe_per_moe_layer(self): ] assert len(shared_matmuls) == 3 + def test_fallback_attention_to_gqa_no_rope(self): """AttentionToGQA fallback fires when applied in isolation (do_rotary=0). diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index 3abfc09d..02509276 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -38,6 +38,7 @@ "Gemma4AssistantTask", "Gemma4Task", "Gemma4UnifiedTask", + "GlmMoeDsaTask", "Gemma4TextCausalLMTask", "HybridCausalLMTask", "HybridQwenVLTask", @@ -99,6 +100,7 @@ Gemma4UnifiedTask, ) from mobius.tasks._gemma4_assistant import Gemma4AssistantTask +from mobius.tasks._glm_moe_dsa import GlmMoeDsaTask from mobius.tasks._hunyuan_vl_mot import HunYuanVLMoTTask from mobius.tasks._image_classification import ImageClassificationTask from mobius.tasks._masked_diffusion import MaskedDiffusionTask @@ -162,6 +164,7 @@ "gemma4-text-generation": Gemma4TextCausalLMTask, "gemma4-unified": Gemma4UnifiedTask, "gemma4-assistant": Gemma4AssistantTask, + "glm-moe-dsa": GlmMoeDsaTask, "hunyuan-vl-mot": HunYuanVLMoTTask, "multimodal": MultiModalTask, "phi4mm-multimodal": Phi4MMMultiModalTask, diff --git a/src/mobius/tasks/_base.py b/src/mobius/tasks/_base.py index b09276fc..38b5e938 100644 --- a/src/mobius/tasks/_base.py +++ b/src/mobius/tasks/_base.py @@ -117,7 +117,7 @@ def _make_graph( [], nodes=[], name=name, - opset_imports={"": OPSET_VERSION, "com.microsoft": 1}, + opset_imports={"": OPSET_VERSION, "com.microsoft": 1, "pkg.nxrt": 1}, ) return graph, GraphBuilder(graph) diff --git a/src/mobius/tasks/_glm_moe_dsa.py b/src/mobius/tasks/_glm_moe_dsa.py new file mode 100644 index 00000000..484acd20 --- /dev/null +++ b/src/mobius/tasks/_glm_moe_dsa.py @@ -0,0 +1,140 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Graph tasks for GLM-5.2 IndexShare DSA and its MTP head.""" + +from __future__ import annotations + +from typing import ClassVar + +import onnx_ir as ir + +from mobius._configs import ArchitectureConfig +from mobius._model_package import ModelPackage +from mobius.tasks._base import ModelTask, _make_graph, _make_model + + +class GlmMoeDsaTask(ModelTask): + """Build the target decoder and optional GLM-5.2 MTP component.""" + + model_roles: ClassVar[dict[str, str]] = {"model": "decoder", "mtp": "decoder"} + + @staticmethod + def _cache_dims(config: ArchitectureConfig, layer_idx: int) -> tuple[int, int]: + qk_dim = config.num_attention_heads * ( + int(config.qk_nope_head_dim or 0) + int(config.qk_rope_head_dim or 0) + ) + value_dim = config.num_attention_heads * int(config.v_head_dim or 0) + if config.use_dsa and config.indexer_types: + if config.indexer_types[layer_idx] == "full": + qk_dim += int(config.index_head_dim or 0) + return qk_dim, value_dim + + def _build_target(self, module, config: ArchitectureConfig): + graph, builder = _make_graph("glm_moe_dsa") + input_ids = builder.input( + "input_ids", ir.DataType.INT64, ["batch_size", "sequence_length"] + ) + attention_mask = builder.input( + "attention_mask", ir.DataType.INT64, ["batch_size", "total_sequence_length"] + ) + position_ids = builder.input( + "position_ids", ir.DataType.INT64, ["batch_size", "sequence_length"] + ) + + past_key_values = [] + for i in range(config.num_hidden_layers): + key_dim, value_dim = self._cache_dims(config, i) + if config.use_dsa: + key_shape = ["batch_size", 1, "past_sequence_length", key_dim] + value_shape = ["batch_size", 1, "past_sequence_length", value_dim] + else: + key_shape = [ + "batch_size", + config.num_attention_heads, + "past_sequence_length", + int(config.qk_nope_head_dim or 0) + int(config.qk_rope_head_dim or 0), + ] + value_shape = [ + "batch_size", + config.num_attention_heads, + "past_sequence_length", + int(config.v_head_dim or 0), + ] + key = builder.input( + f"past_key_values.{i}.key", + config.dtype, + key_shape, + ) + value = builder.input( + f"past_key_values.{i}.value", + config.dtype, + value_shape, + ) + past_key_values.append((key, value)) + + logits, presents = module( + builder.op, + input_ids, + attention_mask, + position_ids, + past_key_values, + ) + builder.add_output(logits, "logits") + for i, (key, value) in enumerate(presents): + builder.add_output(key, f"present.{i}.key") + builder.add_output(value, f"present.{i}.value") + return _make_model(graph) + + def _build_mtp(self, module, config: ArchitectureConfig): + graph, builder = _make_graph("glm_moe_dsa_mtp") + inputs_embeds = builder.input( + "inputs_embeds", + config.dtype, + ["batch_size", "sequence_length", config.hidden_size], + ) + hidden_states = builder.input( + "hidden_states", + config.dtype, + ["batch_size", "sequence_length", config.hidden_size], + ) + attention_mask = builder.input( + "attention_mask", ir.DataType.INT64, ["batch_size", "total_sequence_length"] + ) + position_ids = builder.input( + "position_ids", ir.DataType.INT64, ["batch_size", "sequence_length"] + ) + key_dim = config.num_attention_heads * ( + int(config.qk_nope_head_dim or 0) + int(config.qk_rope_head_dim or 0) + ) + int(config.index_head_dim or 0) + value_dim = config.num_attention_heads * int(config.v_head_dim or 0) + past_key = builder.input( + "past_key_values.0.key", + config.dtype, + ["batch_size", 1, "past_sequence_length", key_dim], + ) + past_value = builder.input( + "past_key_values.0.value", + config.dtype, + ["batch_size", 1, "past_sequence_length", value_dim], + ) + + mtp_hidden, present, topk_indices = module.mtp( + builder.op, + inputs_embeds, + hidden_states, + attention_mask, + position_ids, + (past_key, past_value), + ) + builder.add_output(mtp_hidden, "mtp_hidden") + builder.add_output(present[0], "present.0.key") + builder.add_output(present[1], "present.0.value") + builder.add_output(topk_indices, "topk_indices") + return _make_model(graph) + + def build(self, module, config: ArchitectureConfig) -> ModelPackage: + models = {"model": self._build_target(module, config)} + if getattr(module, "mtp", None) is not None: + models["mtp"] = self._build_mtp(module, config) + return ModelPackage(models, config=config) diff --git a/tests/_test_configs.py b/tests/_test_configs.py index c7ebb020..ac09a789 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -587,6 +587,34 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: }, True, ), + ( + "glm_moe_dsa", + { + "num_key_value_heads": TINY_HEADS, + "q_lora_rank": 32, + "kv_lora_rank": 16, + "qk_nope_head_dim": 16, + "qk_rope_head_dim": 8, + "v_head_dim": 16, + "num_local_experts": 4, + "num_experts_per_tok": 2, + "moe_intermediate_size": 32, + "n_group": 1, + "topk_group": 1, + "routed_scaling_factor": 2.5, + "scoring_func": "sigmoid", + "topk_method": "noaux_tc", + "first_k_dense_replace": 1, + "n_shared_experts": 1, + "rope_interleave": True, + "index_topk": 8, + "index_head_dim": 8, + "index_n_heads": 2, + "indexer_types": ["full", "full"], + "num_nextn_predict_layers": 0, + }, + True, + ), ( "deepseek_v4", { diff --git a/tests/cli_test.py b/tests/cli_test.py index d1475a28..8c9dd294 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -125,6 +125,30 @@ def test_text_only_skips_diffusers_autodetect(self): mock_build.assert_called_once() assert mock_build.call_args.kwargs.get("text_only") is True + def test_glm_full_attention_passes_config_overrides(self): + """--glm-full-attention disables DSA and MTP for registry builds.""" + with ( + tempfile.TemporaryDirectory() as tmpdir, + mock.patch("mobius.__main__.build", return_value=mock.MagicMock()) as mock_build, + mock.patch("mobius.__main__._save_package"), + ): + main( + [ + "build", + "--model", + "THUDM/GLM-5.2-Air", + tmpdir, + "--no-weights", + "--glm-full-attention", + ] + ) + + mock_build.assert_called_once() + assert mock_build.call_args.kwargs.get("config_overrides") == { + "use_dsa": False, + "num_nextn_predict_layers": 0, + } + def test_build_static_cache(self): with tempfile.TemporaryDirectory() as tmpdir: main( @@ -254,6 +278,39 @@ def test_runtime_ort_genai_calls_write_ort_genai_config(self): call_kwargs = mock_export.call_args assert call_kwargs.kwargs.get("hf_model_id") == "Qwen/Qwen2.5-0.5B" + def test_runtime_onnx_genai_calls_write_onnx_genai_config(self, capsys): + """--runtime onnx-genai calls the unified config writer.""" + with ( + tempfile.TemporaryDirectory() as tmpdir, + mock.patch( + "mobius.integrations.onnx_genai.write_onnx_genai_config", + return_value={ + "inference_metadata": "inference_metadata.yaml", + "mtp_config": "mtp_config.json", + }, + ) as mock_export, + mock.patch( + "mobius.integrations.ort_genai.write_ort_genai_config", + return_value={}, + ) as mock_ort, + mock.patch("mobius._model_package.ModelPackage.save"), + ): + main( + [ + "build", + "--model", + "Qwen/Qwen2.5-0.5B", + tmpdir, + "--no-weights", + "--runtime", + "onnx-genai", + ] + ) + + mock_export.assert_called_once() + mock_ort.assert_not_called() + assert "mtp_config: mtp_config.json" in capsys.readouterr().out + def test_runtime_onnx_genai_uses_native_vlm_emitter(self): pkg = mock.MagicMock() pkg.items.return_value = [] @@ -330,6 +387,7 @@ def test_runtime_onnx_genai_does_not_fallback_for_unsupported_vlm(self): native_writer.assert_called_once() generic_writer.assert_not_called() + def test_no_runtime_does_not_call_write_ort_genai_config(self): """Omitting --runtime does NOT call write_ort_genai_config().""" with ( diff --git a/tests/model_coverage_test.py b/tests/model_coverage_test.py index 80a9da92..cf5d4163 100644 --- a/tests/model_coverage_test.py +++ b/tests/model_coverage_test.py @@ -209,6 +209,7 @@ def _all_registered_with_test_id() -> dict[str, str]: "dbrx": "Large MoE (132B) — no small public checkpoint", "deepseek_v3": "Very large MoE (671B) — no small public checkpoint", "deepseek_v4": "Very large MoE (284B) — no small public checkpoint", + "glm_moe_dsa": "Very large MoE — no small public checkpoint or golden data", "llama4_text": "Very large MoE (109B) — no small public checkpoint", "qwen3_5_moe": "Large MoE (22B) — no small public checkpoint", # --- Models without test_model_id ---