GLM-4 GPTQ import: encoder remap + fused-QKV split

src/mobius/_builder.py

@@ -91,13 +91,14 @@ def _cast_module_dtype(module: nn.Module, dtype: ir.DataType) -> None:
     caches), the underlying data is also cast.
 
     Only recasts parameters that are currently FLOAT (float32). Integer
-    parameters and non-float types are left unchanged.
+    parameters, non-float types, and parameters marked ``_keep_float32`` are
+    left unchanged.
     """
     if dtype == ir.DataType.FLOAT:
         return
     torch_dtype = tensor_adapters.to_torch_dtype(dtype)
     for param in module.parameters():
-        if param.dtype != ir.DataType.FLOAT:
+        if param.dtype != ir.DataType.FLOAT or getattr(param, "_keep_float32", False):
             continue
         param.type = ir.TensorType(dtype)
         if param.const_value is not None:

src/mobius/_configs/_base.py

@@ -51,7 +51,7 @@ def _resolve_hidden_act(config, model_type: str) -> str | None:
       dense_act_fn          — some BERT variants
       activation            — generic fallback
       afn                   — older BERT configs
-      "silu"  (qwen)        — Qwen v1 hardcodes silu; no activation attr
+      "silu"  (qwen and chatglm) — Qwen v1 and ChatGLM hardcode silu; no activation attr
       "gelu"  (XLM)         — gelu_activation=True is a boolean flag
       "relu"  (ctrl)        — CTRL hardcodes relu; no hidden_act attr
     """
@@ -65,7 +65,7 @@ def _resolve_hidden_act(config, model_type: str) -> str | None:
         or getattr(config, "afn", None)
         # LLaDA/OLMo expose the activation as ``activation_type`` (e.g. "silu").
         or getattr(config, "activation_type", None)
-        or ("silu" if model_type in ("qwen",) else None)
+        or ("silu" if model_type in ("qwen", "chatglm") else None)
         # gelu_activation is a boolean (XLM) — must be after all string
         # attrs so it cannot override an explicit hidden_act.
         or ("gelu" if getattr(config, "gelu_activation", False) else None)
@@ -295,7 +295,7 @@ def _extract_vision_config(config, parent_config, model_type: str) -> dict:
     hooks live under :mod:`mobius._configs.per_model` and are
     registered with :mod:`mobius._configs._extractors` at import time.
     """
-    from mobius._configs import per_model  # noqa: F401  - side-effect import
+    from mobius._configs import per_model  # noqa: F401 - imported for registration side effect
     from mobius._configs._extractors import extract_vision_config as _dispatch
 
     return _dispatch(config, parent_config, model_type)
@@ -308,7 +308,7 @@ def _extract_audio_config(config, parent_config, model_type: str) -> dict:
     hooks live under :mod:`mobius._configs.per_model` and are
     registered with :mod:`mobius._configs._extractors` at import time.
     """
-    from mobius._configs import per_model  # noqa: F401  - side-effect import
+    from mobius._configs import per_model  # noqa: F401 - imported for registration side effect
     from mobius._configs._extractors import extract_audio_config as _dispatch
 
     return _dispatch(config, parent_config, model_type)
@@ -549,6 +549,7 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig:
             rope_config = RoPEConfig(
                 rope_type="default",
                 rope_theta=_IMPLICIT_ROPE_DEFAULTS[model_type],
+                partial_rotary_factor=0.5 if model_type == "chatglm" else 1.0,
             )
 
         # Some hierarchical models (Segformer, Swin) use plural list attrs
@@ -568,6 +569,7 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig:
         num_hidden_layers = (
             getattr(config, "num_hidden_layers", None)
             or getattr(config, "n_layers", None)
+            or getattr(config, "num_layers", None)
             or getattr(config, "num_encoder_blocks", None)
             or 0
         )
@@ -590,12 +592,18 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig:
                 config.head_dim
                 if (hasattr(config, "head_dim") and config.head_dim is not None)
                 else getattr(config, "d_kv", None)
+                or getattr(config, "kv_channels", None)
                 or _as_int(hidden_size) // _as_int(num_attention_heads)
             ),
             num_attention_heads=_as_int(num_attention_heads),
             num_key_value_heads=_as_int(
                 getattr(config, "num_key_value_heads", None)
                 or getattr(config, "n_kv_heads", None)
+                or (
+                    getattr(config, "multi_query_group_num", None)
+                    if getattr(config, "multi_query_attention", False)
+                    else None
+                )
                 or num_attention_heads
             ),
             num_hidden_layers=_as_int(num_hidden_layers),
@@ -643,7 +651,9 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig:
                 or 1e-6
             ),
             attn_qkv_bias=(
-                getattr(
+                getattr(config, "add_qkv_bias", False)
+                or getattr(config, "add_bias_linear", False)
+                or getattr(
                     config,
                     "attention_bias",
                     getattr(
@@ -677,7 +687,8 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig:
                 )
             ),
             attn_o_bias=(
-                getattr(
+                getattr(config, "add_bias_linear", False)
+                or getattr(
                     config,
                     "attention_bias",
                     getattr(
@@ -722,7 +733,8 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig:
             ),
             attn_qk_norm_full=(model_type in ("flex_olmo", "olmoe", "olmo2", "olmo3")),
             mlp_bias=(
-                getattr(
+                getattr(config, "add_bias_linear", False)
+                or getattr(
                     config,
                     "use_mlp_bias",
                     getattr(
@@ -765,6 +777,7 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig:
             max_position_embeddings=(
                 getattr(config, "max_position_embeddings", None)
                 or getattr(config, "max_sequence_length", None)
+                or getattr(config, "seq_length", None)
                 or 0
             ),
             tie_word_embeddings=(

src/mobius/_configs_test.py

@@ -188,6 +188,39 @@ class FakeQwen3Config:
         config = ArchitectureConfig.from_transformers(FakeQwen3Config())
         assert config.attn_qk_norm is True
 
+    def test_from_transformers_chatglm4_legacy_fields(self):
+        class FakeChatGLMConfig:
+            model_type = "chatglm"
+            num_attention_heads = 32
+            multi_query_attention = True
+            multi_query_group_num = 2
+            num_layers = 40
+            hidden_size = 4096
+            ffn_hidden_size = 13696
+            kv_channels = 128
+            padded_vocab_size = 151552
+            vocab_size = 151552
+            seq_length = 131072
+            layernorm_epsilon = 1.5625e-7
+            add_bias_linear = False
+            add_qkv_bias = True
+            pad_token_id = 151329
+            tie_word_embeddings = False
+
+        config = ArchitectureConfig.from_transformers(FakeChatGLMConfig())
+
+        assert config.num_hidden_layers == 40
+        assert config.num_key_value_heads == 2
+        assert config.head_dim == 128
+        assert config.intermediate_size == 13696
+        assert config.hidden_act == "silu"
+        assert config.max_position_embeddings == 131072
+        assert config.partial_rotary_factor == pytest.approx(0.5)
+        assert config.rope_interleave is True
+        assert config.attn_qkv_bias is True
+        assert config.attn_o_bias is False
+        assert config.mlp_bias is False
+
     def test_from_transformers_rope_scaling(self):
         class FakeConfig:
             model_type = "llama"

src/mobius/_weight_utils.py

@@ -557,19 +557,22 @@ def vlm_vision_weights(
 def _reshape_packed_qweight(value: torch.Tensor, blob_size: int) -> torch.Tensor:
     """Transpose and reshape a packed qweight tensor for MatMulNBits.
 
-    Converts ``[K_packed, N]`` int32 → ``[N, n_blocks, blob_size]`` uint8.
+    Converts ``[..., K_packed, N]`` int32 to
+    ``[..., N, n_blocks, blob_size]`` uint8.
     """
-    transposed = value.t().contiguous()
-    n = transposed.shape[0]
+    transposed = value.transpose(-1, -2).contiguous()
+    prefix = transposed.shape[:-2]
+    n = transposed.shape[-2]
     packed = transposed.view(torch.uint8)
-    n_blocks = packed.shape[1] // blob_size
-    return packed.reshape(n, n_blocks, blob_size)
+    n_blocks = packed.shape[-1] // blob_size
+    return packed.reshape(*prefix, n, n_blocks, blob_size)
 
 
 def _reshape_packed_qzeros(value: torch.Tensor, bits: int, n_blocks: int) -> torch.Tensor:
     """Transpose and unpack packed qzeros for MatMulNBits.
 
-    Converts ``[n_groups_packed, N]`` int32 → ``[N, zp_cols]`` uint8
+    Converts ``[..., n_groups_packed, N]`` int32 to
+    ``[..., N, zp_cols]`` uint8
     where ``zp_cols = ceil(n_blocks * bits / 8)``.  For 4-bit this
     packs two zero-point values per byte, matching ORT's expectation.
 
@@ -578,11 +581,56 @@ def _reshape_packed_qzeros(value: torch.Tensor, bits: int, n_blocks: int) -> tor
         bits: Quantization bit-width (4 or 8).
         n_blocks: Actual number of quantization blocks (``ceil(K / block_size)``).
     """
-    transposed = value.t().contiguous()
-    n = transposed.shape[0]
-    flat_uint8 = transposed.flatten().view(torch.uint8).reshape(n, -1)
+    transposed = value.transpose(-1, -2).contiguous()
+    prefix = transposed.shape[:-2]
+    n = transposed.shape[-2]
+    flat_uint8 = transposed.reshape(-1).view(torch.uint8).reshape(*prefix, n, -1)
     zp_cols = math.ceil(n_blocks * bits / 8)
-    return flat_uint8[:, :zp_cols]
+    return flat_uint8[..., :zp_cols]
+
+
+def pack_qmoe_expert_weights(
+    state_dict: dict[str, torch.Tensor],
+    *,
+    target_moe_path: str = ".mlp.moe",
+) -> dict[str, torch.Tensor]:
+    """Map fused expert-major quantized tensors to native QMoE parameters."""
+    packed: dict[str, torch.Tensor] = {}
+    projections = {
+        ".mlp.experts.gate_up_proj.weight": (
+            f"{target_moe_path}.fc1_experts_weights",
+            True,
+        ),
+        ".mlp.experts.gate_up_proj.scales": (
+            f"{target_moe_path}.fc1_scales",
+            False,
+        ),
+        ".mlp.experts.gate_up_proj.zero_points": (
+            f"{target_moe_path}.fc1_experts_zero_points",
+            False,
+        ),
+        ".mlp.experts.down_proj.weight": (
+            f"{target_moe_path}.fc2_experts_weights",
+            True,
+        ),
+        ".mlp.experts.down_proj.scales": (
+            f"{target_moe_path}.fc2_scales",
+            False,
+        ),
+        ".mlp.experts.down_proj.zero_points": (
+            f"{target_moe_path}.fc2_experts_zero_points",
+            False,
+        ),
+    }
+    for key, value in state_dict.items():
+        for source, (target, flatten_blocks) in projections.items():
+            if source in key:
+                key = key.replace(source, target)
+                if flatten_blocks:
+                    value = value.flatten(-2)
+                break
+        packed[key] = value
+    return packed
 
 
 def preprocess_gptq_weights(
@@ -645,12 +693,12 @@ def preprocess_gptq_weights(
                 raise ValueError(
                     f"Missing {qw_key} — qweight must be present alongside qzeros for {key}"
                 )
-            k = state_dict[qw_key].shape[0] * 32 // bits
+            k = state_dict[qw_key].shape[-2] * 32 // bits
             n_blocks = math.ceil(k / group_size)
             result[new_key] = _reshape_packed_qzeros(value, bits, n_blocks)
 
         elif key.endswith(".scales"):
-            result[key] = value.t().contiguous()
+            result[key] = value.transpose(-1, -2).contiguous()
 
         else:
             result[key] = value
@@ -792,7 +840,7 @@ def preprocess_awq_weights(
                 raise ValueError(
                     f"Missing {qw_key} — qweight must be present alongside qzeros for {key}"
                 )
-            k = state_dict[qw_key].shape[0] * 32 // bits
+            k = state_dict[qw_key].shape[-2] * 32 // bits
             n_blocks = math.ceil(k / group_size)
             # AWQ zero points have an implicit +1 offset; subtract it
             # before unpacking so MatMulNBits sees the raw value.
@@ -810,7 +858,7 @@ def preprocess_awq_weights(
                 result[new_key] = zp_int16.clamp(min=0).to(torch.uint8)
 
         elif key.endswith(".scales"):
-            result[key] = value.t().contiguous()
+            result[key] = value.transpose(-1, -2).contiguous()
 
         else:
             result[key] = value

src/mobius/_weight_utils_test.py

@@ -510,6 +510,42 @@ def test_scales_transposed(self):
         result = preprocess_gptq_weights(sd, bits=self.BITS, group_size=self.GROUP_SIZE)
         assert result["q_proj.scales"].shape == (self.N, self.N_GROUPS)
 
+    def test_expert_major_tensors_preserve_expert_axis(self):
+        num_experts = 64
+        sd = {
+            "experts.qweight": torch.randint(
+                0,
+                255,
+                (num_experts, self.K_PACKED, self.N),
+                dtype=torch.int32,
+            ),
+            "experts.qzeros": torch.randint(
+                0,
+                255,
+                (num_experts, max(1, self.N_GROUPS_PACKED), self.N),
+                dtype=torch.int32,
+            ),
+            "experts.scales": torch.randn(num_experts, self.N_GROUPS, self.N),
+        }
+        result = preprocess_gptq_weights(sd, bits=self.BITS, group_size=self.GROUP_SIZE)
+
+        assert result["experts.weight"].shape == (
+            num_experts,
+            self.N,
+            self.N_GROUPS,
+            self.BLOB_SIZE,
+        )
+        assert result["experts.zero_points"].shape == (
+            num_experts,
+            self.N,
+            self.N_GROUPS * self.BITS // 8,
+        )
+        assert result["experts.scales"].shape == (
+            num_experts,
+            self.N,
+            self.N_GROUPS,
+        )
+
     def test_non_gptq_keys_pass_through(self):
         t = torch.randn(4, 8)
         sd = {"model.embed_tokens.weight": t, "lm_head.weight": t.clone()}

src/mobius/components/_deepseek_mla.py

@@ -43,8 +43,11 @@ def __init__(
         self,
         config: ArchitectureConfig,
         scale: float | None = None,
+        linear_class: type | None = None,
     ):
         super().__init__()
+        if linear_class is None:
+            linear_class = Linear
         self.num_heads = config.num_attention_heads
         self.q_lora_rank = config.q_lora_rank
         self.kv_lora_rank = config.kv_lora_rank
@@ -55,37 +58,37 @@ def __init__(
 
         # Q path: LoRA compression → norm → decompression
         if self.q_lora_rank is not None and self.q_lora_rank > 0:
-            self.q_a_proj = Linear(config.hidden_size, self.q_lora_rank, bias=False)
+            self.q_a_proj = linear_class(config.hidden_size, self.q_lora_rank, bias=False)
             self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
-            self.q_b_proj = Linear(
+            self.q_b_proj = linear_class(
                 self.q_lora_rank,
                 self.num_heads * self.qk_head_dim,
                 bias=False,
             )
         else:
             # No LoRA — direct projection
-            self.q_proj = Linear(
+            self.q_proj = linear_class(
                 config.hidden_size,
                 self.num_heads * self.qk_head_dim,
                 bias=False,
             )
 
         # KV path: joint projection for latent KV + RoPE key
         # Output: (kv_lora_rank) for latent KV + (qk_rope_head_dim) for k_rope
-        self.kv_a_proj_with_mqa = Linear(
+        self.kv_a_proj_with_mqa = linear_class(
             config.hidden_size,
             self.kv_lora_rank + self.qk_rope_head_dim,
             bias=False,
         )
         self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
         # Decompresses latent KV into per-head k_nope + v
-        self.kv_b_proj = Linear(
+        self.kv_b_proj = linear_class(
             self.kv_lora_rank,
             self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
             bias=False,
         )
 
-        self.o_proj = Linear(
+        self.o_proj = linear_class(
             self.num_heads * self.v_head_dim,
             config.hidden_size,
             bias=False,