diff --git a/src/mobius/integrations/gguf/_builder.py b/src/mobius/integrations/gguf/_builder.py index 0f3e9de3..313a0600 100644 --- a/src/mobius/integrations/gguf/_builder.py +++ b/src/mobius/integrations/gguf/_builder.py @@ -486,7 +486,7 @@ def build_from_gguf( # This converts GGUF tensor quirks (stacked experts, 1D gates, 2D # conv weights, suffix artifacts) into the shapes that HF models # produce, so preprocess_weights only needs to handle HF→ONNX. - state_dict = _normalize_gguf_weights(state_dict) + state_dict = _normalize_gguf_weights(state_dict, gguf_arch, config) # 8. Run model-specific preprocess_weights (HF → ONNX names) if hasattr(module, "preprocess_weights"): @@ -612,8 +612,26 @@ def _replace_native_block_linears(module, gguf_model, gguf_arch: str) -> None: ) +#: GGUF architectures whose transformer RMSNorms are zero-centered +#: (``output = norm(x) * (1 + weight)``, mobius :class:`OffsetRMSNorm`). Their +#: llama.cpp converter bakes the ``+1`` into every ``*norm.weight`` *except* the +#: Gated-DeltaNet internal ``linear_attn.norm`` (a plain gated RMSNorm), so the +#: GGUF path must undo it — see :func:`_normalize_gguf_weights`. +_OFFSET_NORM_GGUF_ARCHS: frozenset[str] = frozenset({"qwen35", "qwen35moe"}) + +#: GGUF architectures whose llama.cpp converter reorders Gated-DeltaNet V-heads +#: from HuggingFace *grouped* order (``head = group * v_per_k + j``) into ggml +#: *tiled* order (``head = j * num_k_heads + group``) whenever the linear layer +#: is grouped (``num_value_heads != num_key_heads``). mobius's ``GatedDeltaNet`` +#: forward consumes the HF grouped order, so the GGUF path must undo the tiling — +#: see :func:`_reorder_deltanet_v_heads`. +_V_HEAD_REORDER_GGUF_ARCHS: frozenset[str] = frozenset({"qwen35", "qwen35moe"}) + + def _normalize_gguf_weights( state_dict: dict, + gguf_arch: str | None = None, + config=None, ) -> dict: """Normalize GGUF-specific weight shapes to match HF conventions. @@ -635,9 +653,33 @@ def _normalize_gguf_weights( - **dt_bias suffix**: GGUF ``ssm_dt.bias`` maps to ``dt_bias.bias`` after suffix splitting, but the model parameter is just ``dt_bias`` (an ``nn.Parameter``, not a module bias). + - **DeltaNet A_log**: GGUF stores the SSM decay pre-transformed as + ``ssm_a = -exp(A_log)``; mobius's ``GatedDeltaNet`` re-derives + ``-exp(A_log)`` at runtime, so the raw log is recovered via + ``A_log = log(-ssm_a)`` (scoped to ``linear_attn.A_log``). + - **Zero-centered RMSNorm** (``gguf_arch`` in + :data:`_OFFSET_NORM_GGUF_ARCHS`): the converter bakes ``+1`` into every + ``*norm.weight`` except ``linear_attn.norm.weight``; mobius applies the + ``1 +`` at runtime via :class:`OffsetRMSNorm`, so subtract ``1`` back out + to avoid double-counting. + - **Gated-DeltaNet V-head tiling** (``gguf_arch`` in + :data:`_V_HEAD_REORDER_GGUF_ARCHS`, grouped linear attention): the + converter reorders every V-indexed ``linear_attn`` tensor from HF grouped + order into ggml tiled order; mobius consumes grouped order, so the tiling + is undone via :func:`_reorder_deltanet_v_heads`. + + Args: + state_dict: Dequantized GGUF weights keyed by HF tensor names. + gguf_arch: The source GGUF architecture string (e.g. ``"qwen35"``), + used to gate architecture-specific value transforms such as the + zero-centered RMSNorm offset. + config: The resolved :class:`ArchitectureConfig`; supplies the + Gated-DeltaNet head counts / dims used to undo the V-head tiling. """ import torch + offset_norms = gguf_arch in _OFFSET_NORM_GGUF_ARCHS + result: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): # Stacked expert weights [num_experts, out, in] → per-expert @@ -668,6 +710,36 @@ def _normalize_gguf_weights( result[key[: -len(".bias")]] = value continue + # DeltaNet A_log: undo the converter's pre-transform. GGUF's converter + # stores the SSM decay already transformed as ``ssm_a = -exp(A_log)`` + # (llama.cpp applies ``-torch.exp`` to every ``.A_log`` tensor and the + # reference then uses it *directly* as the decay coefficient ``a`` in + # ``a * softplus(dt)``). mobius's ``GatedDeltaNet`` parameter is the raw + # ``A_log`` and recomputes ``g = -exp(A_log) * softplus(...)`` at + # runtime, so feeding it the already-negated-exp value squashes every + # head's decay to ``-exp(-exp(A_log)) ≈ -1`` and the linear-attention + # recurrence emits garbage. Invert to recover the raw log parameter, + # ``A_log = log(-ssm_a)``, so mobius's ``-exp(A_log)`` reproduces the + # original ``ssm_a`` exactly. Scoped to the GatedDeltaNet ``linear_attn`` + # projection so Mamba/PLaMo SSM modules (which consume ``A = -exp(A_log)`` + # directly) are left untouched. + if key.endswith(".linear_attn.A_log"): + result[key] = torch.log(-value) + continue + + # Zero-centered RMSNorm: undo the converter's baked-in ``+1`` so + # mobius's OffsetRMSNorm (which adds it back at runtime) does not + # double-count. The DeltaNet internal ``linear_attn.norm`` is a plain + # gated RMSNorm (no offset) and is excluded — mirroring exactly which + # tensors the llama.cpp converter transforms. + if ( + offset_norms + and key.endswith("norm.weight") + and not key.endswith(".linear_attn.norm.weight") + ): + result[key] = value - 1.0 + 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 @@ -678,9 +750,173 @@ def _normalize_gguf_weights( result[key] = value + # DeltaNet V-head tiling: undo the converter's grouped→tiled permutation of + # every V-indexed linear_attn tensor so mobius's GatedDeltaNet (which expects + # HF grouped order) reads consistent heads. Runs last so it operates on the + # already-normalized keys/shapes (renamed dt_bias, unsqueezed conv1d, ...). + if gguf_arch in _V_HEAD_REORDER_GGUF_ARCHS: + result = _reorder_deltanet_v_heads(result, config) + return result +def _reorder_deltanet_v_heads(state_dict: dict, config) -> dict: + """Undo the GGUF converter's grouped→tiled V-head permutation. + + llama.cpp's ``_LinearAttentionVReorderBase`` reorders every V-indexed + Gated-DeltaNet tensor from HuggingFace *grouped* order (V-head + ``s = group * v_per_k + j``) into ggml *tiled* order + (``t = j * num_k_heads + group``) whenever the linear layer is grouped + (``num_value_heads != num_key_heads``). mobius's ``GatedDeltaNet`` forward + reshapes the value stream as ``num_value_heads`` contiguous ``head_v_dim`` + blocks in the original grouped order, so the GGUF tensors must be permuted + back: grouped position ``s`` is fetched from tiled slot + ``perm[s] = (s % v_per_k) * num_key_heads + (s // v_per_k)``. + + The permutation is applied (all derived from ``config`` — no hardcoded head + counts) to: + + - ``in_proj_qkv`` output rows — V rows only (after ``2 * key_dim``); + - ``in_proj_z`` output rows — all rows; + - ``in_proj_a`` / ``in_proj_b`` output rows — one row per V-head; + - ``A_log`` / ``dt_bias`` — one element per V-head; + - ``conv1d`` channels — V channels only (after ``2 * key_dim``); + - ``out_proj`` input columns — all columns (a block-granular permutation of + the quantized ``K`` axis). + + Quantized projections are stored as MatMulNBits triplets + (``weight`` ``[N, K/block, block/2]``, ``scales`` ``[N, K/block]``, + ``zero_points`` ``[N, K/block/2]``). Output-row permutations reindex axis 0 + of all three; the ``out_proj`` input permutation reindexes the block axis + (axis 1), valid because ``head_v_dim`` is a whole number of quant blocks. + """ + import torch + + num_k_heads = getattr(config, "linear_num_key_heads", None) + num_v_heads = getattr(config, "linear_num_value_heads", None) + head_k_dim = getattr(config, "linear_key_head_dim", None) + head_v_dim = getattr(config, "linear_value_head_dim", None) + # Nothing to do unless this is a grouped linear-attention model. + if not (num_k_heads and num_v_heads and head_k_dim and head_v_dim): + return state_dict + if num_v_heads == num_k_heads or num_v_heads % num_k_heads != 0: + return state_dict + + v_per_k = num_v_heads // num_k_heads + key_dim = head_k_dim * num_k_heads + v_offset = 2 * key_dim # in_proj_qkv / conv1d layout is [Q | K | V] + + # perm[s] = tiled slot holding grouped V-head s. + head_perm = torch.tensor( + [(s % v_per_k) * num_k_heads + (s // v_per_k) for s in range(num_v_heads)], + dtype=torch.long, + ) + + def _expand(perm: "torch.Tensor", stride: int) -> "torch.Tensor": + # Expand a per-head permutation into a per-row/-channel index. + base = (perm * stride).unsqueeze(1) + torch.arange(stride) + return base.reshape(-1) + + v_rows = _expand(head_perm, head_v_dim) # length value_dim + + def _index_dim0(t: "torch.Tensor", idx: "torch.Tensor") -> "torch.Tensor": + return t.index_select(0, idx) + + def _index_dim1(t: "torch.Tensor", idx: "torch.Tensor") -> "torch.Tensor": + return t.index_select(1, idx) + + def _apply_rows(stem: str, idx: "torch.Tensor") -> None: + # Permute axis 0 of a float weight or a quantized triplet in place. + for suffix in (".weight", ".scales", ".zero_points"): + key = stem + suffix + if key in state_dict: + state_dict[key] = _index_dim0(state_dict[key], idx) + + def _apply_bare(key: str, idx: "torch.Tensor") -> None: + if key in state_dict: + state_dict[key] = _index_dim0(state_dict[key], idx) + + layer_stems = { + k.rsplit(".", 1)[0] + for k in state_dict + if ".linear_attn." in k + } + for stem in layer_stems: + name = stem.rsplit(".", 1)[-1] + if name == "in_proj_z": + _apply_rows(stem, v_rows) + elif name in ("in_proj_a", "in_proj_b"): + _apply_rows(stem, head_perm) + elif name == "in_proj_qkv": + n_rows = state_dict[stem + ".weight"].shape[0] + full = torch.cat([torch.arange(v_offset), v_offset + v_rows]) + assert full.numel() == n_rows, (n_rows, full.numel()) + _apply_rows(stem, full) + elif name == "out_proj": + _reorder_out_proj_cols(state_dict, stem, head_perm, head_v_dim) + + # Bare (non-".weight") linear_attn parameters. + for k in list(state_dict): + if k.endswith(".linear_attn.A_log") or k.endswith(".linear_attn.dt_bias"): + _apply_bare(k, head_perm) + elif k.endswith(".linear_attn.conv1d.weight"): + conv = state_dict[k] + n_ch = conv.shape[0] + full = torch.cat([torch.arange(v_offset), v_offset + v_rows]) + assert full.numel() == n_ch, (n_ch, full.numel()) + state_dict[k] = _index_dim0(conv, full) + + return state_dict + + +def _reorder_out_proj_cols( + state_dict: dict, stem: str, head_perm, head_v_dim: int +) -> None: + """Permute the quantized ``out_proj`` input (K) axis by V-head. + + ``out_proj`` maps ``value_dim -> hidden``; its input columns are the V + stream, so they carry the same head tiling. In MatMulNBits form the K axis is + the block axis (axis 1) of ``weight``/``scales`` and the packed block axis of + ``zero_points`` (two 4-bit blocks per byte). ``head_v_dim`` spans a whole + number of blocks, so the permutation is block-granular and lossless. + """ + import torch + + weight = state_dict.get(stem + ".weight") + if weight is None or weight.dim() < 2: + return + n_blocks = weight.shape[1] + if n_blocks % head_perm.numel() != 0: + raise ValueError( + f"{stem}: cannot map {n_blocks} quant blocks onto " + f"{head_perm.numel()} V-heads for column reorder" + ) + blocks_per_head = n_blocks // head_perm.numel() + + def _expand(perm, stride): + base = (perm * stride).unsqueeze(1) + torch.arange(stride) + return base.reshape(-1) + + blk_idx = _expand(head_perm, blocks_per_head) + state_dict[stem + ".weight"] = weight.index_select(1, blk_idx) + + scales = state_dict.get(stem + ".scales") + if scales is not None: + state_dict[stem + ".scales"] = scales.index_select(1, blk_idx) + + zp = state_dict.get(stem + ".zero_points") + if zp is not None and zp.dim() >= 2: + # zero_points pack two 4-bit blocks per byte along the block axis. + if blocks_per_head % 2 != 0: + raise ValueError( + f"{stem}: {blocks_per_head} blocks/head is not byte-aligned for " + "packed zero_points reorder" + ) + zp_bytes_per_head = blocks_per_head // 2 + zp_idx = _expand(head_perm, zp_bytes_per_head) + state_dict[stem + ".zero_points"] = zp.index_select(1, zp_idx) + + def _has_quantized_weights(gguf_model, gguf_arch: str) -> bool: """Return whether a GGUF has mapped weights with a quantized tensor type.""" from gguf import GGMLQuantizationType diff --git a/src/mobius/integrations/gguf/_builder_test.py b/src/mobius/integrations/gguf/_builder_test.py index 643bdb45..e27bf643 100644 --- a/src/mobius/integrations/gguf/_builder_test.py +++ b/src/mobius/integrations/gguf/_builder_test.py @@ -1050,3 +1050,234 @@ def test_local_split_metadata_is_rejected(self): with pytest.raises(NotImplementedError, match="cannot assemble split tensor tables"): _raise_for_sharded_gguf(source="model-00001-of-00002.gguf", split_count=2) + + +class TestNormalizeGgufWeights: + """Tests for GGUF-specific weight shape/value normalization.""" + + def test_deltanet_a_log_is_inverted_from_neg_exp(self): + """GGUF ssm_a = -exp(A_log); normalize must recover raw A_log = log(-ssm_a). + + The GatedDeltaNet module re-applies ``-exp(A_log)`` at runtime, so the + round-trip ``-exp(normalize(ssm_a))`` must reproduce the original + ``ssm_a`` the converter stored. + """ + import torch + + from mobius.integrations.gguf._builder import _normalize_gguf_weights + + # A representative raw A_log, and the value the GGUF converter stores. + a_log_raw = torch.tensor([-3.4688, -1.0703, -5.0, -0.5], dtype=torch.float32) + ssm_a = -torch.exp(a_log_raw) # what llama.cpp writes to blk.N.ssm_a + assert bool((ssm_a < 0).all()) # sanity: pre-transformed value is negative + + key = "model.layers.0.linear_attn.A_log" + out = _normalize_gguf_weights({key: ssm_a}) + + # The stored parameter must be the raw A_log again ... + assert torch.allclose(out[key], a_log_raw, atol=1e-5) + # ... so that the module's runtime -exp(A_log) recovers ssm_a exactly. + assert torch.allclose(-torch.exp(out[key]), ssm_a, atol=1e-6) + + def test_non_deltanet_a_log_is_untouched(self): + """Mamba/PLaMo SSM ``A_log`` (consumed as ``A`` directly) must not be inverted.""" + import torch + + from mobius.integrations.gguf._builder import _normalize_gguf_weights + + ssm_a = torch.tensor([-0.04, -0.5], dtype=torch.float32) + key = "backbone.layers.0.mixer.A_log" + out = _normalize_gguf_weights({key: ssm_a}) + assert torch.allclose(out[key], ssm_a) + + def test_zero_centered_norm_offset_removed_for_qwen35(self): + """qwen35 GGUF bakes +1 into transformer norms; normalize must strip it. + + mobius applies the ``1 +`` at runtime via OffsetRMSNorm, so the stored + weight must be the raw zero-centered value again. + """ + import torch + + from mobius.integrations.gguf._builder import _normalize_gguf_weights + + sd = { + "model.layers.0.input_layernorm.weight": torch.tensor([1.5, 2.0]), + "model.layers.3.self_attn.q_norm.weight": torch.tensor([1.25]), + "model.layers.3.self_attn.k_norm.weight": torch.tensor([1.1]), + "model.norm.weight": torch.tensor([1.94]), + # DeltaNet internal gated norm — converter did NOT add +1. + "model.layers.0.linear_attn.norm.weight": torch.tensor([0.87]), + } + out = _normalize_gguf_weights(dict(sd), gguf_arch="qwen35") + + assert torch.allclose( + out["model.layers.0.input_layernorm.weight"], torch.tensor([0.5, 1.0]) + ) + assert torch.allclose( + out["model.layers.3.self_attn.q_norm.weight"], torch.tensor([0.25]) + ) + assert torch.allclose(out["model.norm.weight"], torch.tensor([0.94])) + # linear_attn.norm is a plain gated RMSNorm — must be left untouched. + assert torch.allclose( + out["model.layers.0.linear_attn.norm.weight"], torch.tensor([0.87]) + ) + + def test_norm_offset_not_applied_for_non_offset_arch(self): + """Standard-RMSNorm archs (e.g. llama/qwen2) must not have norms shifted.""" + import torch + + from mobius.integrations.gguf._builder import _normalize_gguf_weights + + sd = { + "model.layers.0.input_layernorm.weight": torch.tensor([1.0, 1.0]), + "model.norm.weight": torch.tensor([1.0]), + } + out = _normalize_gguf_weights(dict(sd), gguf_arch="qwen2") + assert torch.allclose( + out["model.layers.0.input_layernorm.weight"], torch.tensor([1.0, 1.0]) + ) + assert torch.allclose(out["model.norm.weight"], torch.tensor([1.0])) + + +class TestReorderDeltaNetVHeads: + """Undo of the GGUF converter's grouped→tiled Gated-DeltaNet V-head order. + + The llama.cpp converter reorders every V-indexed ``linear_attn`` tensor from + HuggingFace *grouped* order into ggml *tiled* order (see + ``_LinearAttentionVReorderBase._reorder_v_heads``). mobius consumes grouped + order, so ``_reorder_deltanet_v_heads`` must be the exact inverse. + """ + + # Small grouped linear-attention geometry: 2 K-heads, 6 V-heads (v_per_k=3). + CFG = SimpleNamespace( + linear_num_key_heads=2, + linear_num_value_heads=6, + linear_key_head_dim=4, + linear_value_head_dim=4, + ) + + @staticmethod + def _converter_reorder(tensor, dim, num_k_heads, num_v_per_k, head_dim): + """Reference grouped→tiled reorder copied from llama.cpp's converter.""" + import torch # noqa: F401 + + shape = list(tensor.shape) + if dim < 0: + dim += len(shape) + new_shape = shape[:dim] + [num_k_heads, num_v_per_k, head_dim] + shape[dim + 1 :] + tensor = tensor.reshape(*new_shape) + perm = list(range(len(new_shape))) + perm[dim], perm[dim + 1] = perm[dim + 1], perm[dim] + return tensor.permute(*perm).contiguous().reshape(*shape) + + def test_row_tensors_roundtrip(self): + """Grouped weights survive tile→untile for every V-row projection.""" + import torch + + from mobius.integrations.gguf._builder import _reorder_deltanet_v_heads + + cfg = self.CFG + n_k, n_v = cfg.linear_num_key_heads, cfg.linear_num_value_heads + v_per_k = n_v // n_k + hd_k, hd_v = cfg.linear_key_head_dim, cfg.linear_value_head_dim + key_dim, value_dim = hd_k * n_k, hd_v * n_v + hidden = 5 + torch.manual_seed(0) + + p = "model.layers.0.linear_attn." + grouped = { + f"{p}in_proj_z.weight": torch.randn(value_dim, hidden), + f"{p}in_proj_a.weight": torch.randn(n_v, hidden), + f"{p}in_proj_b.weight": torch.randn(n_v, hidden), + f"{p}A_log": torch.randn(n_v), + f"{p}dt_bias": torch.randn(n_v), + f"{p}conv1d.weight": torch.randn(2 * key_dim + value_dim, 1, 4), + } + # in_proj_qkv: only the V rows (after 2*key_dim) are reordered. + qkv = torch.randn(2 * key_dim + value_dim, hidden) + grouped[f"{p}in_proj_qkv.weight"] = qkv + + # Build the tiled (GGUF) state by applying the converter's reorder. + tiled = {k: v.clone() for k, v in grouped.items()} + tiled[f"{p}in_proj_z.weight"] = self._converter_reorder( + grouped[f"{p}in_proj_z.weight"], 0, n_k, v_per_k, hd_v + ) + for name in ("in_proj_a", "in_proj_b"): + tiled[f"{p}{name}.weight"] = self._converter_reorder( + grouped[f"{p}{name}.weight"], 0, n_k, v_per_k, 1 + ) + for name in ("A_log", "dt_bias"): + tiled[f"{p}{name}"] = self._converter_reorder( + grouped[f"{p}{name}"], 0, n_k, v_per_k, 1 + ) + # V portion of qkv / conv1d. + v0 = 2 * key_dim + qv = self._converter_reorder(qkv[v0:], 0, n_k, v_per_k, hd_v) + tiled[f"{p}in_proj_qkv.weight"] = torch.cat([qkv[:v0], qv], dim=0) + conv = grouped[f"{p}conv1d.weight"] + cv = self._converter_reorder(conv[v0:], 0, n_k, v_per_k, hd_v) + tiled[f"{p}conv1d.weight"] = torch.cat([conv[:v0], cv], dim=0) + + out = _reorder_deltanet_v_heads({k: v.clone() for k, v in tiled.items()}, cfg) + + for k in grouped: + assert torch.allclose(out[k], grouped[k]), k + + def test_quantized_out_proj_columns_roundtrip(self): + """out_proj's quantized K axis (blocks + packed zero-points) round-trips.""" + import torch + + from mobius.integrations.gguf._builder import _reorder_deltanet_v_heads + + cfg = self.CFG + n_k, n_v = cfg.linear_num_key_heads, cfg.linear_num_value_heads + v_per_k = n_v // n_k + hd_v = cfg.linear_value_head_dim # 4 + value_dim = hd_v * n_v # 24 + block = 2 # 2 elems/block -> head_v_dim(4) = 2 blocks (even -> byte aligned) + n_blocks = value_dim // block # 12 + hidden = 5 + torch.manual_seed(1) + + p = "model.layers.0.linear_attn." + # Grouped quantized out_proj triplet: [hidden, K/block, block/2], etc. + gw = torch.randint(0, 255, (hidden, n_blocks, block // 2 + 7), dtype=torch.uint8) + gs = torch.randn(hidden, n_blocks, dtype=torch.float16) + gz = torch.randint(0, 255, (hidden, n_blocks // 2), dtype=torch.uint8) + # Provide a grouped row tensor so the head geometry is exercised too. + grouped = { + f"{p}out_proj.weight": gw, + f"{p}out_proj.scales": gs, + f"{p}out_proj.zero_points": gz, + } + blocks_per_head = n_blocks // n_v # 2 + tiled = { + f"{p}out_proj.weight": self._converter_reorder(gw, 1, n_k, v_per_k, blocks_per_head), + f"{p}out_proj.scales": self._converter_reorder(gs, 1, n_k, v_per_k, blocks_per_head), + f"{p}out_proj.zero_points": self._converter_reorder( + gz, 1, n_k, v_per_k, blocks_per_head // 2 + ), + } + + out = _reorder_deltanet_v_heads({k: v.clone() for k, v in tiled.items()}, cfg) + + for k in grouped: + assert torch.equal(out[k], grouped[k]), k + + def test_no_reorder_when_heads_equal(self): + """Ungrouped linear attention (num_v == num_k) is left untouched.""" + import torch + + from mobius.integrations.gguf._builder import _reorder_deltanet_v_heads + + cfg = SimpleNamespace( + linear_num_key_heads=4, + linear_num_value_heads=4, + linear_key_head_dim=4, + linear_value_head_dim=4, + ) + p = "model.layers.0.linear_attn." + sd = {f"{p}in_proj_z.weight": torch.randn(16, 5)} + ref = sd[f"{p}in_proj_z.weight"].clone() + out = _reorder_deltanet_v_heads({k: v.clone() for k, v in sd.items()}, cfg) + assert torch.equal(out[f"{p}in_proj_z.weight"], ref) diff --git a/src/mobius/integrations/gguf/_config_mapping.py b/src/mobius/integrations/gguf/_config_mapping.py index f1752b50..bea1a303 100644 --- a/src/mobius/integrations/gguf/_config_mapping.py +++ b/src/mobius/integrations/gguf/_config_mapping.py @@ -252,6 +252,27 @@ def gguf_to_config( f"GGUF file missing required metadata for '{field}'. Architecture: {gguf_arch}" ) + # Exclude Multi-Token-Prediction (MTP / "nextn") blocks from the decoder + # layer count. GGUF's ``block_count`` counts the trailing MTP prediction + # block(s) alongside the regular decoder layers (e.g. Qwen3.5/3.8 store + # ``block_count = num_hidden_layers + nextn_predict_layers``), but the base + # decode model does not build them. Their weights (``blk..nextn.*`` and + # the accompanying attention/FFN tensors of the trailing block) are skipped + # during tensor mapping. Without this correction the builder would create + # an extra decoder layer whose linear-attention / GQA initializers have no + # backing GGUF weights and fail the ``_check_weights`` invariant on save. + nextn_layers = metadata.get(f"{gguf_arch}.nextn_predict_layers") + if nextn_layers is not None and int(nextn_layers) > 0: + mtp_count = int(nextn_layers) + decoder_layers = int(hf_fields["num_hidden_layers"]) - mtp_count + if decoder_layers <= 0: + raise ValueError( + f"GGUF metadata inconsistent: block_count " + f"({hf_fields['num_hidden_layers']}) <= nextn_predict_layers " + f"({mtp_count}) for architecture {gguf_arch}." + ) + hf_fields["num_hidden_layers"] = decoder_layers + # Derive head_dim if not explicitly provided. # Prefer attention.key_length (the actual head dimension) over # rope.dimension_count (which may be just the rotary embedding @@ -308,11 +329,21 @@ def gguf_to_config( else: partial_rotary_factor = 1.0 - # Derive rope_interleave from rope.dimension_sections metadata. - rope_sections = metadata.get(f"{gguf_arch}.rope.dimension_sections") - rope_interleave = gguf_arch == "deepseek4" or ( - rope_sections is not None and any(s > 0 for s in rope_sections) - ) + # Derive rope_interleave. + # + # ``rope.dimension_sections`` encodes M-RoPE *section* sizes (the Qwen-VL + # family splits the rotary dimension across the temporal/height/width + # position axes, e.g. ``[11, 11, 10, 0]``). It does NOT select the GPT-J + # style adjacent-pair rotation that the flat ``rope_interleave`` flag + # controls: Qwen3.5 (and every other section-carrying arch here) rotates + # with split-half (NEOX / ``rotate_half``) semantics. Deriving + # ``rope_interleave`` from section presence therefore corrupts RoPE — the + # exported GroupQueryAttention/RotaryEmbedding gets ``rotary_interleaved=1`` + # and the full-attention layers produce garbage tokens. Section interleave, + # when a model needs it, is a distinct ``mrope_interleaved`` signal handled + # via ``mrope_section``. Only architectures that genuinely use adjacent-pair + # rotation set the flat flag here. + rope_interleave = gguf_arch == "deepseek4" # Derive rope_type from rope.scaling.type. GGUF stores the scaling # variant under ``.rope.scaling.type`` (or omits the key for the @@ -485,6 +516,7 @@ def gguf_to_config( def _gemma2_postprocess( config: ArchitectureConfig, metadata: dict[str, Any], + model: Any = None, ) -> Gemma2Config: """Convert a base config to Gemma2Config with architecture-specific fields. diff --git a/src/mobius/integrations/gguf/_config_mapping_test.py b/src/mobius/integrations/gguf/_config_mapping_test.py index 3d2f41d2..212e620b 100644 --- a/src/mobius/integrations/gguf/_config_mapping_test.py +++ b/src/mobius/integrations/gguf/_config_mapping_test.py @@ -147,6 +147,132 @@ def test_silu_default(self, model_type: str) -> None: assert _default_activation(model_type) == "silu" +class TestQwen35MtpBlockExclusion: + """Qwen3.5/3.8 GGUF ``block_count`` includes trailing MTP (nextn) blocks. + + ``gguf_to_config`` must subtract ``nextn_predict_layers`` so the decoder + builds only the real transformer layers; otherwise it fabricates an extra + layer whose weights are missing from the GGUF (the ``blk..nextn.*`` + prediction head is skipped during tensor mapping). + """ + + def _fake_model(self, metadata: dict) -> object: + class _FakeGGUF: + architecture = "qwen35" + + def __init__(self, md: dict) -> None: + self.metadata = md + + def get_metadata(self, key, default=None): + return self.metadata.get(key, default) + + @property + def tensor_names(self) -> list[str]: + return ["output.weight", "blk.0.attn_q.weight"] + + return _FakeGGUF(metadata) + + def _base_metadata(self, block_count: int) -> dict: + return { + "qwen35.embedding_length": 5120, + "qwen35.block_count": block_count, + "qwen35.attention.head_count": 24, + "qwen35.attention.head_count_kv": 4, + "qwen35.attention.key_length": 256, + "qwen35.attention.value_length": 256, + "qwen35.feed_forward_length": 17408, + "qwen35.vocab_size": 248320, + "qwen35.full_attention_interval": 4, + "qwen35.rope.dimension_count": 64, + } + + def test_nextn_layers_excluded_from_decoder_count(self) -> None: + from mobius.integrations.gguf._config_mapping import gguf_to_config + + md = self._base_metadata(block_count=65) + md["qwen35.nextn_predict_layers"] = 1 + config = gguf_to_config(self._fake_model(md)) + + assert config.num_hidden_layers == 64 + assert config.layer_types is not None + assert len(config.layer_types) == 64 + # 3 linear + 1 full pattern (full at every 4th, 1-indexed). + assert config.layer_types[3] == "full_attention" + assert config.layer_types[0] == "linear_attention" + + def test_no_nextn_metadata_leaves_count_unchanged(self) -> None: + from mobius.integrations.gguf._config_mapping import gguf_to_config + + md = self._base_metadata(block_count=64) + config = gguf_to_config(self._fake_model(md)) + + assert config.num_hidden_layers == 64 + + def test_nextn_not_greater_than_block_count(self) -> None: + from mobius.integrations.gguf._config_mapping import gguf_to_config + + md = self._base_metadata(block_count=1) + md["qwen35.nextn_predict_layers"] = 1 + with pytest.raises(ValueError, match="nextn_predict_layers"): + gguf_to_config(self._fake_model(md)) + + +class TestQwen35RopeInterleave: + """``rope.dimension_sections`` is M-RoPE section metadata, not a GPT-J + adjacent-pair rotation signal. + + Qwen3.5/3.8 rotate with split-half (NEOX) semantics. Deriving the flat + ``rope_interleave`` from section presence corrupts RoPE — the exported + GroupQueryAttention/RotaryEmbedding gets ``rotary_interleaved=1`` and the + full-attention layers emit garbage tokens. The mapping must keep + ``rope_interleave`` False for section-carrying non-interleaving arches. + """ + + def _fake_model(self, metadata: dict, architecture: str = "qwen35") -> object: + class _FakeGGUF: + def __init__(self, md: dict, arch: str) -> None: + self.metadata = md + self.architecture = arch + + def get_metadata(self, key, default=None): + return self.metadata.get(key, default) + + @property + def tensor_names(self) -> list[str]: + return ["output.weight", "blk.0.attn_q.weight"] + + return _FakeGGUF(metadata, architecture) + + def _base_metadata(self) -> dict: + return { + "qwen35.embedding_length": 5120, + "qwen35.block_count": 64, + "qwen35.attention.head_count": 24, + "qwen35.attention.head_count_kv": 4, + "qwen35.attention.key_length": 256, + "qwen35.attention.value_length": 256, + "qwen35.feed_forward_length": 17408, + "qwen35.vocab_size": 248320, + "qwen35.full_attention_interval": 4, + "qwen35.rope.dimension_count": 64, + "qwen35.rope.freq_base": 1e7, + } + + def test_dimension_sections_do_not_force_interleave(self) -> None: + from mobius.integrations.gguf._config_mapping import gguf_to_config + + md = self._base_metadata() + md["qwen35.rope.dimension_sections"] = [11, 11, 10, 0] + config = gguf_to_config(self._fake_model(md)) + + assert config.rope_interleave is False + + def test_no_sections_still_not_interleaved(self) -> None: + from mobius.integrations.gguf._config_mapping import gguf_to_config + + config = gguf_to_config(self._fake_model(self._base_metadata())) + + assert config.rope_interleave is False class TestMuseGlimmerPostprocess: """Muse Glimmer config postprocessing. diff --git a/src/mobius/rewrite_rules/_group_query_attention.py b/src/mobius/rewrite_rules/_group_query_attention.py index fc3c02f8..bf136536 100644 --- a/src/mobius/rewrite_rules/_group_query_attention.py +++ b/src/mobius/rewrite_rules/_group_query_attention.py @@ -347,6 +347,12 @@ def rewrite( # GLM4/ChatGLM use interleaved=1; most models use 0. Must not hardcode. q_rope_node = attn.inputs[0].producer() # RotaryEmbedding producing q_rot rotary_interleaved = q_rope_node.attributes.get_int("interleaved", 0) + # Preserve partial RoPE: the source RotaryEmbedding rotates only the + # first ``rotary_embedding_dim`` head elements (e.g. Qwen3.5 rotates 64 + # of 256). Omitting it makes the fused GQA default to the full head_dim + # and read past the partial cos/sin cache, corrupting every + # full-attention layer. Propagate it whenever the source op sets it. + rotary_embedding_dim = q_rope_node.attributes.get_int("rotary_embedding_dim", 0) # Trace cos/sin back through Gather to the cache table initializers if self._cos_cache is None: @@ -389,6 +395,8 @@ def rewrite( } if softcap: gqa_attrs["softcap"] = softcap + if rotary_embedding_dim: + gqa_attrs["rotary_embedding_dim"] = rotary_embedding_dim window = local_window_from_attention_bias(attention_bias).window if window is not None: gqa_attrs["local_window_size"] = window diff --git a/src/mobius/rewrite_rules/_group_query_attention_test.py b/src/mobius/rewrite_rules/_group_query_attention_test.py index a6261ffd..f6a3bbee 100644 --- a/src/mobius/rewrite_rules/_group_query_attention_test.py +++ b/src/mobius/rewrite_rules/_group_query_attention_test.py @@ -851,7 +851,55 @@ def test_softcap_absent_for_non_softcap_models(self): sc = node.attributes.get("softcap") assert sc is None, f"Unexpected softcap attribute on GQA node for Llama: {sc}" + def test_partial_rotary_embedding_dim_propagated(self): + """Partial-RoPE models must forward rotary_embedding_dim to the GQA. + Qwen3.5/3.8 rotate only the first ``head_dim * partial_rotary_factor`` + head elements (e.g. 64 of 256). If the fusion drops the dimension the + fused GQA defaults to the full head_dim and reads past the partial + cos/sin cache, corrupting every full-attention layer's output. + """ + cfg = dataclasses.replace(_QWEN3_CONFIG, partial_rotary_factor=0.5) + model = registry.get("qwen3")(cfg) + pkg = build_from_module(model, cfg) + m = pkg["model"] + + rewrite(m, pattern_rewrite_rules=group_query_attention_rules()) + + gqa_nodes = [n for n in m.graph if n.op_type == "GroupQueryAttention"] + assert len(gqa_nodes) > 0, "Expected GQA nodes after fusion" + + expected = int(cfg.head_dim * 0.5) + for node in gqa_nodes: + val = node.attributes.get("rotary_embedding_dim") + assert val is not None, ( + "rotary_embedding_dim missing on GQA node — partial RoPE would " + "silently rotate the full head_dim and corrupt attention" + ) + assert val.value == expected, ( + f"Expected rotary_embedding_dim={expected}, got {val.value}" + ) + + def test_full_rotary_omits_rotary_embedding_dim(self): + """Full-RoPE models (Qwen3 default) must not carry rotary_embedding_dim. + + Omitting it lets the GQA kernel default to the full head_dim, which is + correct; a spurious attribute could mislead consumers. + """ + model = registry.get("qwen3")(_QWEN3_CONFIG) + pkg = build_from_module(model, _QWEN3_CONFIG) + m = pkg["model"] + + rewrite(m, pattern_rewrite_rules=group_query_attention_rules()) + + gqa_nodes = [n for n in m.graph if n.op_type == "GroupQueryAttention"] + assert len(gqa_nodes) > 0 + + for node in gqa_nodes: + val = node.attributes.get("rotary_embedding_dim") + assert val is None, ( + f"Unexpected rotary_embedding_dim on full-RoPE GQA node: {val}" + ) class TestSlidingWindowSurvivesFusion: """A window baked into the attention bias must reach ``local_window_size``.