DS metadata

slimshard.py

@@ -2711,8 +2711,9 @@ _LAYER_SUFFIX_TO_GGUF = {
     "self_attn.o_proj.weight": "attn_output.weight",
     "attn.o_proj.weight": "attn_output.weight",
     "attention.wo.weight": "attn_output.weight",
-    "attn.wo_b.weight": "attn_output.weight",
-    "attn.wo_a.weight": "attn_o_a.weight",
+    "attn.wo_b.weight": "attn_output_b.weight",
+    "attn.wo_a.weight": "attn_output_a.weight",
+    "attn.attn_sink": "attn_sinks.weight",
     "self_attn.q_a_proj.weight": "attn_q_a.weight",
     "attn.wq_a.weight": "attn_q_a.weight",
     "self_attn.q_a_layernorm.weight": "attn_q_a_norm.weight",
@@ -2720,11 +2721,17 @@ _LAYER_SUFFIX_TO_GGUF = {
     "self_attn.q_b_proj.weight": "attn_q_b.weight",
     "attn.wq_b.weight": "attn_q_b.weight",
     "self_attn.kv_a_proj_with_mqa.weight": "attn_kv_a_mqa.weight",
-    "attn.wkv.weight": "attn_kv_a_mqa.weight",
+    "attn.wkv.weight": "attn_kv.weight",
     "self_attn.kv_a_layernorm.weight": "attn_kv_a_norm.weight",
     "attn.kv_norm.weight": "attn_kv_a_norm.weight",
     "self_attn.kv_b_proj.weight": "attn_kv_b.weight",
     "self_attn.g_proj.weight": "attn_gate.weight",
+    "hc_attn_fn": "hc_attn_fn.weight",
+    "hc_attn_base": "hc_attn_base.weight",
+    "hc_attn_scale": "hc_attn_scale.weight",
+    "hc_ffn_fn": "hc_ffn_fn.weight",
+    "hc_ffn_base": "hc_ffn_base.weight",
+    "hc_ffn_scale": "hc_ffn_scale.weight",
     "mlp.gate_proj.weight": "ffn_gate.weight",
     "feed_forward.w1.weight": "ffn_gate.weight",
     "mlp.up_proj.weight": "ffn_up.weight",
@@ -2777,7 +2784,9 @@ def _build_gguf_kv_pairs(
     model_name: str,
 ) -> list[tuple[str, int, bytes]]:
     raw_mtype = str(arch.get("model_type", "llama")).lower()
-    if "deepseek" in raw_mtype:
+    if "deepseek_v4" in raw_mtype or "deepseek4" in raw_mtype:
+        g_arch = "deepseek4"
+    elif "deepseek" in raw_mtype:
         g_arch = "deepseek2"
     elif "qwen" in raw_mtype:
         g_arch = "qwen2moe" if int(arch.get("num_experts") or 0) > 0 else "qwen2"
@@ -2792,15 +2801,19 @@ def _build_gguf_kv_pairs(
         ("general.quantization_version", 4, struct.pack("<I", 2)),
     ]
 
-    def add_u32(k: str, v: int) -> None:
-        if v is not None and int(v) > 0:
+    def add_u32(k: str, v: int, allow_zero: bool = False) -> None:
+        if v is not None and (allow_zero or int(v) > 0):
             kvs.append((k, 4, struct.pack("<I", int(v))))
 
     def add_f32(k: str, v: float) -> None:
         if v is not None:
             kvs.append((k, 6, struct.pack("<f", float(v))))
 
-    add_u32(f"{g_arch}.block_count", int(arch.get("num_hidden_layers", 0)))
+    def add_bool(k: str, v: bool) -> None:
+        kvs.append((k, 7, struct.pack("<B", 1 if v else 0)))
+
+    n_layers = int(arch.get("num_hidden_layers", 0))
+    add_u32(f"{g_arch}.block_count", n_layers)
     add_u32(f"{g_arch}.embedding_length", int(arch.get("hidden_size", 0)))
     add_u32(
         f"{g_arch}.feed_forward_length",
@@ -2821,11 +2834,45 @@ def _build_gguf_kv_pairs(
         f"{g_arch}.expert_used_count",
         int(arch.get("num_experts_per_token") or arch.get("num_experts_per_tok") or 0),
     )
-    add_u32(f"{g_arch}.context_length", int(arch.get("max_position_embeddings", 32768)))
+    add_u32(f"{g_arch}.context_length", min(int(arch.get("max_position_embeddings", 32768)), 65536))
     add_u32(f"{g_arch}.vocab_size", int(arch.get("vocab_size", 0)))
     add_u32(f"{g_arch}.attention.q_lora_rank", int(arch.get("q_lora_rank") or 0))
     add_u32(f"{g_arch}.attention.kv_lora_rank", int(arch.get("kv_lora_rank") or 0))
 
+    if g_arch == "deepseek4":
+        head_dim = int(arch.get("head_dim", 512))
+        rope_dim = int(arch.get("qk_rope_head_dim", 64))
+        add_u32("deepseek4.attention.key_length", head_dim)
+        add_u32("deepseek4.attention.value_length", head_dim)
+        add_u32("deepseek4.rope.dimension_count", rope_dim)
+        add_f32("deepseek4.rope.freq_base", float(arch.get("rope_theta", 10000.0)))
+        add_u32("deepseek4.nextn_predict_layers", 0, allow_zero=True)
+        add_u32("deepseek4.attention.sliding_window", int(arch.get("sliding_window", 128)))
+        add_u32("deepseek4.expert_shared_count", int(arch.get("n_shared_experts", 1)))
+        add_f32("deepseek4.expert_weights_scale", float(arch.get("routed_scaling_factor", 1.5)))
+        add_bool("deepseek4.expert_weights_norm", bool(arch.get("norm_topk_prob", True)))
+        add_f32("deepseek4.swiglu_clamp_exp", float(arch.get("swiglu_limit", 10.0)))
+        add_f32("deepseek4.swiglu_clamp_shexp", float(arch.get("swiglu_limit", 10.0)))
+        add_u32("deepseek4.attention.indexer.head_count", int(arch.get("index_n_heads", 32)))
+        add_u32("deepseek4.attention.indexer.key_length", int(arch.get("index_head_dim", 128)))
+        add_u32("deepseek4.attention.indexer.top_k", int(arch.get("index_topk", 512)))
+        add_u32("deepseek4.attention.output_group_count", int(arch.get("o_groups", 8)))
+        add_u32("deepseek4.attention.output_lora_rank", int(arch.get("o_lora_rank", 1024)))
+        add_f32("deepseek4.attention.compress_rope_freq_base", float(arch.get("compress_rope_theta", 160000.0)))
+        add_u32("deepseek4.hyper_connection.count", int(arch.get("hc_mult", 4)))
+        add_u32("deepseek4.hyper_connection.sinkhorn_iterations", int(arch.get("hc_sinkhorn_iters", 20)))
+        add_f32("deepseek4.hyper_connection.epsilon", float(arch.get("hc_eps", 1.0e-6)))
+        add_u32("deepseek4.hash_layer_count", 0, allow_zero=True)
+        c_ratios = np.zeros(max(n_layers, 1), dtype="<u4")
+        kvs.append(
+            (
+                "deepseek4.attention.compress_ratios",
+                9,
+                struct.pack("<IQ", 4, len(c_ratios)) + c_ratios.tobytes(),
+            )
+        )
+        add_u32("deepseek4.expert_gating_func", 4)
+
     vocab_sz = int(arch.get("vocab_size", 0))
     if tokenizer_path is not None and tokenizer_path.exists():
         try:
@@ -2963,11 +3010,18 @@ def export_safetensors_to_gguf(
         for gk in expert_groups:
             expert_groups[gk].sort(key=lambda x: x[0])
 
+        is_dsv4 = "deepseek_v4" in str(arch.get("model_type", "")).lower()
+        o_groups = int(arch.get("o_groups", 8))
+        o_lora_rank = int(arch.get("o_lora_rank", 1024))
+        hc_mult = int(arch.get("hc_mult", 4))
+
         # Build work items: ("single", src_key, scale_key, gguf_name) or ("experts", (lidx, proj), items, gguf_name)
         work_items: list[tuple] = []
         for k in keys:
             if k in consumed_keys:
                 continue
+            if is_dsv4 and (".compressor." in k or ".indexer." in k):
+                continue
             sc_k = None
             if k.endswith(".weight"):
                 for sfx in (".scale", ".weight_scale_inv", ".scales", ".weight_scale"):
@@ -3003,6 +3057,33 @@ def export_safetensors_to_gguf(
         data_offset = 0
         payload_cache: dict[tuple, tuple[int, tuple[int, ...], int, int]] = {}
         stats = {"tensors": 0, "q2_k": 0, "q4_0": 0, "bf16": 0, "f32": 0}
+
+        if is_dsv4:
+            last_l = max(int(arch.get("num_hidden_layers", 40)) - 1, 0)
+            fn_k = f"layers.{last_l}.hc_ffn_fn"
+            bs_k = f"layers.{last_l}.hc_ffn_base"
+            sc_k = f"layers.{last_l}.hc_ffn_scale"
+            if fn_k in keys_set and bs_k in keys_set and sc_k in keys_set:
+                fn_t = f.get_tensor(fn_k).float()[:hc_mult].contiguous()
+                bs_t = f.get_tensor(bs_k).float()[:hc_mult].contiguous()
+                sc_t = f.get_tensor(sc_k).float()[:1].contiguous()
+                for g_hc_name, hc_t in (
+                    ("output_hc_fn.weight", fn_t),
+                    ("output_hc_base.weight", bs_t),
+                    ("output_hc_scale.weight", sc_t),
+                ):
+                    pad = _gguf_align(data_offset) - data_offset
+                    if pad > 0:
+                        tmp_data.write(b"\x00" * pad)
+                        data_offset += pad
+                    raw_hc = hc_t.cpu().numpy().astype("<f4").tobytes()
+                    gguf_tensors.append((g_hc_name, tuple(reversed(hc_t.shape)), _GGML_F32, data_offset))
+                    tmp_data.write(raw_hc)
+                    data_offset += len(raw_hc)
+                    stats["tensors"] += 1
+                    stats["f32"] += 1
+                tmp_data.flush()
+
         total_items = len(work_items)
 
         def _dequant_single_to_f32_torch(w_key: str, s_key: str | None) -> "torch.Tensor":
@@ -3075,6 +3156,8 @@ def export_safetensors_to_gguf(
                 _, w_key, s_key, gname = item
                 sl = f.get_slice(w_key)
                 t_shape = tuple(sl.get_shape())
+                if is_dsv4 and gname.endswith(".attn_output_a.weight") and len(t_shape) == 2 and t_shape[0] == o_groups * o_lora_rank:
+                    t_shape = (o_groups, o_lora_rank, t_shape[1])
                 g_dims = tuple(reversed(t_shape)) if t_shape else (1,)
                 tgt = _choose_gguf_target(gname, t_shape or (1,))