feat: add Ling 3.0 LongHaul support

conversion/__init__.py

@@ -27,6 +27,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "BaichuanForCausalLM": "baichuan",
     "BailingMoeForCausalLM": "bailingmoe",
     "BailingMoeV2ForCausalLM": "bailingmoe",
+    "BailingMoeV3ForCausalLM": "bailingmoe",
     "BambaForCausalLM": "granite",
     "BertForMaskedLM": "bert",
     "BertForSequenceClassification": "bert",

conversion/bailingmoe.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import math
+
 from typing import Callable, Iterable, TYPE_CHECKING
 
 import torch
@@ -188,6 +190,184 @@ class BailingMoeV2Model(TextModel):
                 raise ValueError(f"Unprocessed experts: {experts}")
 
 
+@ModelBase.register("BailingMoeV3ForCausalLM")
+class BailingMoeV3Model(TextModel):
+    """Ling 3.0 - hybrid KDA (linear) + gated MLA (full) attention with a bailing MoE"""
+    model_arch = gguf.MODEL_ARCH.BAILINGMOE3
+
+    _experts: list[dict[str, Tensor]] | None = None
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # the MTP block sits right after the last decoder layer and is not converted
+        self.block_count = self.hparams["num_hidden_layers"]
+        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
+
+    def is_mla_layer(self, il: int) -> bool:
+        # every layer_group_size-th layer is a full attention (MLA) layer, the rest are KDA
+        return (il + 1) % self.hparams["layer_group_size"] == 0
+
+    def set_vocab(self):
+        # identical tokenizer to Ling 2.0, the bailingmoe2 pre-tokenizer hash matches
+        self._set_vocab_gpt2()
+
+    def set_gguf_parameters(self):
+        hparams = self.hparams
+
+        # note: to enable the MLA KV cache, attention is converted into MQA (ie: GQA with 1 group)
+        hparams["num_key_value_heads"] = 1
+
+        super().set_gguf_parameters()
+
+        self.gguf_writer.add_vocab_size(hparams["vocab_size"])
+
+        # n_head_kv == 0 marks a KDA (recurrent) layer, > 0 marks an MLA (attention) layer
+        self.gguf_writer.add_head_count_kv([
+            1 if self.is_mla_layer(il) else 0 for il in range(self.block_count)
+        ])
+
+        # KDA
+        self.gguf_writer.add_ssm_conv_kernel(hparams["short_conv_kernel_size"])
+        self.gguf_writer.add_kda_head_dim(hparams["head_dim"])
+        # safe gate: g = lower_bound * sigmoid(exp(A_log) * (f_proj(x) + dt_bias))
+        assert hparams.get("kda_safe_gate", False), "only the safe gate form is implemented"
+        self.gguf_writer.add_kda_gate_lower_bound(hparams["kda_lower_bound"])
+
+        # MLA - converted into MQA with larger heads, then decompressed to MHA
+        kv_lora_rank = hparams["kv_lora_rank"]
+        qk_rope_head_dim = hparams["qk_rope_head_dim"]
+        self.gguf_writer.add_kv_lora_rank(kv_lora_rank)
+        self.gguf_writer.add_key_length(kv_lora_rank + qk_rope_head_dim)
+        self.gguf_writer.add_value_length(kv_lora_rank)
+        self.gguf_writer.add_key_length_mla(hparams["qk_nope_head_dim"] + qk_rope_head_dim)
+        self.gguf_writer.add_value_length_mla(hparams["v_head_dim"])
+        self.gguf_writer.add_rope_dimension_count(qk_rope_head_dim)
+
+        # MoE
+        self.gguf_writer.add_leading_dense_block_count(hparams["first_k_dense_replace"])
+        self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"])
+        self.gguf_writer.add_expert_shared_feed_forward_length(
+            hparams["moe_shared_expert_intermediate_size"] * hparams["num_shared_experts"])
+        self.gguf_writer.add_expert_shared_count(hparams["num_shared_experts"])
+        self.gguf_writer.add_expert_weights_scale(hparams["routed_scaling_factor"])
+        self.gguf_writer.add_expert_weights_norm(hparams["norm_topk_prob"])
+        self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID)
+        self.gguf_writer.add_expert_group_count(hparams["n_group"])
+        self.gguf_writer.add_expert_group_used_count(hparams["topk_group"])
+
+        # Optional per-layer SwiGLU clamps (vLLM SwigluStepAndMul):
+        #   out = silu(gate).clamp(max=limit) * up.clamp(-limit, limit)
+        # 0.0 (or a missing/null entry) means no clamping for that layer.
+        def _clamp_limits(key: str) -> list[float] | None:
+            if (limits := hparams.get(key)) is None:
+                return None
+            limits = [0.0 if v is None else float(v) for v in limits[:self.block_count]]
+            return limits + [0.0] * (self.block_count - len(limits))
+
+        if (limits := _clamp_limits("expert_swiglu_limit_list")) is not None:
+            self.gguf_writer.add_swiglu_clamp_exp(limits)
+        if (limits := _clamp_limits("share_expert_swiglu_limit_list")) is not None:
+            self.gguf_writer.add_swiglu_clamp_shexp(limits)
+
+        if (nextn_layers := hparams.get("num_nextn_predict_layers")) is not None:
+            self.gguf_writer.add_nextn_predict_layers(nextn_layers)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+
+        if name.endswith(".expert_bias"):
+            name = name.replace(".expert_bias", ".expert_bias.bias")
+
+        return super().filter_tensors((name, gen))
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # drop the MTP block
+        if bid is not None and bid >= self.block_count:
+            return
+
+        n_head = self.hparams["num_attention_heads"]
+        head_dim = self.hparams["head_dim"]
+
+        if name.endswith(".A_log"):
+            # the safe gate uses -exp(A_log) only through exp(A_log), see the graph:
+            #   g = kda_lower_bound * sigmoid(exp(A_log) * (f_proj(x) + dt_bias))
+            # {n_head} -> ggml ne = [1, n_head]
+            data_torch = torch.exp(data_torch.float())
+            data_torch = data_torch.reshape(-1, 1)
+        elif name.endswith(".dt_bias"):
+            name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias"
+        elif name.endswith((".q_conv1d.weight", ".k_conv1d.weight", ".v_conv1d.weight")):
+            # HF {d_inner, [1,] d_conv} -> numpy (1, d_inner, 1, d_conv) -> ggml ne = [d_conv, 1, d_inner, 1]
+            d_conv = data_torch.shape[-1]
+            d_inner = math.prod(data_torch.shape[:-1])
+            data_torch = data_torch.reshape(1, d_inner, 1, d_conv)
+        elif name.endswith("attention.g_proj.weight"):
+            assert bid is not None
+            # the very same HF name is the KDA output gate on linear layers and the
+            # attention output gate on MLA layers
+            tensor = gguf.MODEL_TENSOR.ATTN_GATE if self.is_mla_layer(bid) else gguf.MODEL_TENSOR.SSM_G
+            yield from super().modify_tensors(data_torch, self.format_tensor_name(tensor, bid), bid)
+            return
+        elif name.endswith("attention.kv_b_proj.weight"):
+            assert bid is not None
+            # MLA with the absorption optimization needs these two split and k_b transposed
+            v_head_dim = self.hparams["v_head_dim"]
+            qk_nope_head_dim = self.hparams["qk_nope_head_dim"]
+            assert data_torch.shape[0] == n_head * (v_head_dim + qk_nope_head_dim)
+
+            kv_b = data_torch.view(n_head, qk_nope_head_dim + v_head_dim, data_torch.shape[-1])
+            k_b, v_b = torch.split(kv_b, [qk_nope_head_dim, v_head_dim], dim=1)
+            k_b = k_b.transpose(1, 2)
+
+            yield from super().modify_tensors(k_b, self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K_B, bid), bid)
+            yield from super().modify_tensors(v_b, self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V_B, bid), bid)
+            return
+        elif name.endswith("attention.dense.weight"):
+            assert bid is not None
+            # MLA output projection (KDA layers call it o_proj)
+            yield from super().modify_tensors(data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_OUT, bid), bid)
+            return
+        elif "mlp.experts" in name:
+            n_experts = self.hparams["num_experts"]
+            assert bid is not None
+
+            if self._experts is None:
+                self._experts = [{} for _ in range(self.block_count)]
+
+            self._experts[bid][name] = data_torch
+
+            if len(self._experts[bid]) >= n_experts * 3:
+                # merge the experts into a single 3d tensor
+                for w_name in ["down_proj", "gate_proj", "up_proj"]:
+                    datas: list[Tensor] = []
+
+                    for xid in range(n_experts):
+                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
+                        datas.append(self._experts[bid][ename])
+                        del self._experts[bid][ename]
+
+                    data_torch = torch.stack(datas, dim=0)
+
+                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
+
+                    yield from super().modify_tensors(data_torch, merged_name, bid)
+            return
+
+        del head_dim  # only used for the asserts above
+
+        yield from super().modify_tensors(data_torch, name, bid)
+
+    def prepare_tensors(self):
+        super().prepare_tensors()
+
+        if self._experts is not None:
+            experts = [k for d in self._experts for k in d.keys()]
+            if len(experts) > 0:
+                raise ValueError(f"Unprocessed experts: {experts}")
+
+
+
 @ModelBase.register("SarvamMoEForCausalLM", "modeling_sarvam_moe.SarvamMoEForCausalLM")
 class SarvamMoEModel(BailingMoeV2Model):
     model_arch = gguf.MODEL_ARCH.BAILINGMOE2