Upload slimshard.py

slimshard.py

@@ -941,6 +941,312 @@ def _run_generic_transformer_forward(
     return outputs
 
 
+_FP4_LUT_CPU = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0]
+
+
+def _ds4_dequant_fp8_32x32(w_fp8: "torch.Tensor", s_u8: "torch.Tensor") -> "torch.Tensor":
+    out_d, in_d = w_fp8.shape
+    s = torch.exp2(s_u8.view(torch.uint8).float() - 127.0)
+    b0, b1 = s.shape
+    g0, g1 = out_d // b0, in_d // b1
+    return (w_fp8.float().view(b0, g0, b1, g1) * s[:, None, :, None]).reshape(out_d, in_d)
+
+
+def _ds4_dequant_fp4_1x32(w_i8: "torch.Tensor", s_u8: "torch.Tensor", lut: "torch.Tensor") -> "torch.Tensor":
+    out_d, half_in = w_i8.shape
+    in_d = half_in * 2
+    u = w_i8.view(torch.uint8)
+    low = (u & 0x0F).long()
+    high = ((u >> 4) & 0x0F).long()
+    vals = torch.stack([lut[low], lut[high]], dim=-1).reshape(out_d, in_d // 32, 32)
+    s = torch.exp2(s_u8.view(torch.uint8).float() - 127.0).unsqueeze(-1)
+    return (vals * s).reshape(out_d, in_d)
+
+
+def _ds4_hc_mixes(
+    x: "torch.Tensor",
+    hc_fn: "torch.Tensor",
+    hc_scale: "torch.Tensor",
+    hc_base: "torch.Tensor",
+    hc_mult: int = 4,
+    sinkhorn_iters: int = 20,
+    norm_eps: float = 1.0e-20,
+    hc_eps: float = 1.0e-6,
+) -> tuple["torch.Tensor", "torch.Tensor", "torch.Tensor"]:
+    xf = x.flatten(2).float()
+    rsqrt = torch.rsqrt(xf.square().mean(-1, keepdim=True) + norm_eps)
+    mixes = F.linear(xf, hc_fn.float()) * rsqrt
+    s0, s1, s2 = hc_scale.float()[0], hc_scale.float()[1], hc_scale.float()[2]
+    b_f = hc_base.float()
+    pre = torch.sigmoid(mixes[..., :hc_mult] * s0 + b_f[:hc_mult]) + hc_eps
+    post = 2.0 * torch.sigmoid(mixes[..., hc_mult : 2 * hc_mult] * s1 + b_f[hc_mult : 2 * hc_mult])
+    comb = (mixes[..., 2 * hc_mult :] * s2 + b_f[2 * hc_mult :]).reshape(*mixes.shape[:-1], hc_mult, hc_mult)
+    comb = torch.softmax(comb, dim=-1) + hc_eps
+    comb = comb / (comb.sum(dim=-2, keepdim=True) + hc_eps)
+    for _ in range(sinkhorn_iters - 1):
+        comb = comb / (comb.sum(dim=-1, keepdim=True) + hc_eps)
+        comb = comb / (comb.sum(dim=-2, keepdim=True) + hc_eps)
+    return pre, post, comb
+
+
+def _ds4_apply_rope(x: "torch.Tensor", freqs_cis: "torch.Tensor", inverse: bool = False) -> "torch.Tensor":
+    rd = freqs_cis.shape[-1] * 2
+    tail = x[..., -rd:].float().unflatten(-1, (-1, 2))
+    c = torch.view_as_complex(tail)
+    fc = freqs_cis.conj() if inverse else freqs_cis
+    if x.ndim == 3:
+        fc = fc.view(1, x.size(1), -1)
+    else:
+        fc = fc.view(1, x.size(1), 1, -1)
+    rot = torch.view_as_real(c * fc).flatten(-2).to(x.dtype)
+    if rd == x.shape[-1]:
+        return rot
+    return torch.cat([x[..., :-rd], rot], dim=-1)
+
+
+def _run_deepseek_v41_forward(
+    read_fn,
+    has_key_fn,
+    token_seqs: list[list[int]],
+    coords: np.ndarray | None,
+    basis_in: "torch.Tensor | None",
+    scale_in: float,
+    arch: dict,
+    dev: "torch.device",
+) -> list["torch.Tensor"]:
+    """Execute DeepSeek-V4.1 causal forward pass (4-stream Sinkhorn HC + MLA + FP4/FP8 MoE)."""
+    num_layers = int(arch.get("num_hidden_layers", 40))
+    dim = int(arch.get("hidden_size", 5120))
+    eps = float(arch.get("rms_norm_eps", 1.0e-20))
+    hc_mult = int(arch.get("hc_mult", 4))
+    hc_iters = int(arch.get("hc_sinkhorn_iters", 20))
+    hc_eps = float(arch.get("hc_eps", 1.0e-6))
+    n_heads = int(arch.get("num_attention_heads", 64))
+    head_dim = int(arch.get("head_dim", 512))
+    rope_dim = int(arch.get("qk_rope_head_dim", 64))
+    o_groups = int(arch.get("o_groups", 8))
+    o_lora_rank = int(arch.get("o_lora_rank", 1024))
+    top_k = int(arch.get("num_experts_per_tok", arch.get("num_experts_per_token", 6)))
+    route_scale = float(arch.get("routed_scaling_factor", 1.5))
+    swiglu_limit = float(arch.get("swiglu_limit", 10.0))
+    rope_theta = float(arch.get("rope_theta", 10000.0))
+    softmax_scale = head_dim ** -0.5
+
+    lut = torch.tensor(_FP4_LUT_CPU, dtype=torch.float32, device=dev)
+    max_t = max(len(s) for s in token_seqs)
+    freqs = 1.0 / (rope_theta ** (torch.arange(0, rope_dim, 2, dtype=torch.float32, device=dev) / rope_dim))
+    freqs_cis_all = torch.polar(torch.ones((max_t, rope_dim // 2), device=dev), torch.outer(torch.arange(max_t, device=dev, dtype=torch.float32), freqs))
+
+    streams = []
+    pre_mixes = []
+    if coords is not None and basis_in is not None:
+        for seq in token_seqs:
+            c_seq = torch.from_numpy(coords[seq]).to(dev, dtype=torch.float32)
+            inp = (c_seq @ basis_in * scale_in).to(torch.bfloat16).unsqueeze(0)
+            h0 = inp.unsqueeze(2).repeat(1, 1, hc_mult, 1)
+            pm = torch.zeros(1, len(seq), hc_mult, dtype=torch.float32, device=dev)
+            pm[:, :, 0] = 1.0
+            streams.append(h0)
+            pre_mixes.append(pm)
+    else:
+        emb_w = read_fn("embed.weight")
+        for seq in token_seqs:
+            ids_t = torch.tensor([seq], dtype=torch.long, device=dev)
+            inp = F.embedding(ids_t, emb_w).to(torch.bfloat16)
+            h0 = inp.unsqueeze(2).repeat(1, 1, hc_mult, 1)
+            pm = torch.zeros(1, len(seq), hc_mult, dtype=torch.float32, device=dev)
+            pm[:, :, 0] = 1.0
+            streams.append(h0)
+            pre_mixes.append(pm)
+        del emb_w
+
+    # Cache layer weights if multiple layers are identical (e.g. single-shard anchor recovery)
+    cached_layer_ signature = None
+    cached_layer_data = None
+
+    for lidx in range(num_layers):
+        lp = f"layers.{lidx}"
+        if not has_key_fn(f"{lp}.attn_norm.weight"):
+            lp = f"model.layers.{lidx}"
+        w_anorm = read_fn(f"{lp}.attn_norm.weight")
+        w_fnorm = read_fn(f"{lp}.ffn_norm.weight")
+        hc_afn = read_fn(f"{lp}.hc_attn_fn")
+        hc_asc = read_fn(f"{lp}.hc_attn_scale")
+        hc_ab = read_fn(f"{lp}.hc_attn_base")
+        hc_ffn = read_fn(f"{lp}.hc_ffn_fn")
+        hc_fsc = read_fn(f"{lp}.hc_ffn_scale")
+        hc_fb = read_fn(f"{lp}.hc_ffn_base")
+
+        sig = float(hc_ab.float().sum().item())
+        if cached_layer_data is not None and cached_layer_signature == sig:
+            wq_a, wq_b, wkv, wo_a, wo_b, q_norm_w, kv_norm_w, attn_sink, sw1, sw2, sw3, gw, gb = cached_layer_data
+        else:
+            wq_a = _ds4_dequant_fp8_32x32(read_fn(f"{lp}.attn.wq_a.weight"), read_fn(f"{lp}.attn.wq_a.scale")).to(torch.bfloat16)
+            wq_b = _ds4_dequant_fp8_32x32(read_fn(f"{lp}.attn.wq_b.weight"), read_fn(f"{lp}.attn.wq_b.scale")).to(torch.bfloat16)
+            wkv = _ds4_dequant_fp8_32x32(read_fn(f"{lp}.attn.wkv.weight"), read_fn(f"{lp}.attn.wkv.scale")).to(torch.bfloat16)
+            wo_a = _ds4_dequant_fp8_32x32(read_fn(f"{lp}.attn.wo_a.weight"), read_fn(f"{lp}.attn.wo_a.scale")).to(torch.bfloat16)
+            wo_b = _ds4_dequant_fp8_32x32(read_fn(f"{lp}.attn.wo_b.weight"), read_fn(f"{lp}.attn.wo_b.scale")).to(torch.bfloat16)
+            q_norm_w = read_fn(f"{lp}.attn.q_norm.weight")
+            kv_norm_w = read_fn(f"{lp}.attn.kv_norm.weight")
+            attn_sink = read_fn(f"{lp}.attn.attn_sink").float()
+            sw1 = _ds4_dequant_fp8_32x32(read_fn(f"{lp}.ffn.shared_experts.w1.weight"), read_fn(f"{lp}.ffn.shared_experts.w1.scale")).to(torch.bfloat16)
+            sw2 = _ds4_dequant_fp8_32x32(read_fn(f"{lp}.ffn.shared_experts.w2.weight"), read_fn(f"{lp}.ffn.shared_experts.w2.scale")).to(torch.bfloat16)
+            sw3 = _ds4_dequant_fp8_32x32(read_fn(f"{lp}.ffn.shared_experts.w3.weight"), read_fn(f"{lp}.ffn.shared_experts.w3.scale")).to(torch.bfloat16)
+            gw = read_fn(f"{lp}.ffn.gate.weight").float()
+            gb = read_fn(f"{lp}.ffn.gate.bias").float()
+            cached_layer_signature = sig
+            cached_layer_data = (wq_a, wq_b, wkv, wo_a, wo_b, q_norm_w, kv_norm_w, attn_sink, sw1, sw2, sw3, gw, gb)
+
+        n_routed_eff = gw.shape[0]
+        top_k_eff = min(top_k, n_routed_eff)
+        wo_a_grouped = wo_a.view(o_groups, o_lora_rank, -1)
+
+        in_ffn_list = []
+        ffn_meta = []
+
+        for s_idx, h in enumerate(streams):
+            B, T, _, _ = h.shape
+            fc = freqs_cis_all[:T]
+            residual = h
+            attn_pre, attn_post, attn_comb = _ds4_hc_mixes(h, hc_afn, hc_asc, hc_ab, hc_mult, hc_iters, eps, hc_eps)
+            x = torch.sum(pre_mixes[s_idx].unsqueeze(-1) * h.float(), dim=2).to(torch.bfloat16)
+            x = _rms_norm_torch(x, w_anorm, eps)
+
+            qr = _rms_norm_torch(F.linear(x, wq_a), q_norm_w, eps)
+            q = F.linear(qr, wq_b).view(B, T, n_heads, head_dim)
+            q = _ds4_apply_rope(q, fc, inverse=False)
+            kv = _rms_norm_torch(F.linear(x, wkv), kv_norm_w, eps)
+            kv = _ds4_apply_rope(kv, fc, inverse=False)
+
+            # Causal attention with attn_sink
+            scores = torch.einsum("bthd,bsd->bhts", q.float(), kv.float()) * softmax_scale
+            if T > 1:
+                causal_mask = torch.triu(torch.ones((T, T), dtype=torch.bool, device=dev), diagonal=1)
+                scores = scores.masked_fill(causal_mask[None, None, :, :], -1.0e30)
+            s_max = torch.maximum(scores.amax(dim=-1, keepdim=True), attn_sink.view(1, n_heads, 1, 1))
+            exp_s = torch.exp(scores - s_max)
+            denom = exp_s.sum(dim=-1, keepdim=True) + torch.exp(attn_sink.view(1, n_heads, 1, 1) - s_max)
+            probs_a = (exp_s / denom).to(torch.bfloat16)
+            o = torch.einsum("bhts,bsd->bthd", probs_a, kv)
+            o = _ds4_apply_rope(o, fc, inverse=True)
+            o_grp = o.reshape(B, T, o_groups, -1)
+            o_proj = torch.einsum("bsgd,grd->bsgr", o_grp, wo_a_grouped).flatten(2)
+            attn_out = F.linear(o_proj, wo_b)
+
+            h = (attn_post.unsqueeze(-1) * attn_out.float().unsqueeze(-2) + torch.sum(attn_comb.unsqueeze(-1) * residual.float().unsqueeze(-2), dim=2)).to(torch.bfloat16)
+            residual2 = h
+            ffn_pre, ffn_post, ffn_comb = _ds4_hc_mixes(h, hc_ffn, hc_fsc, hc_fb, hc_mult, hc_iters, eps, hc_eps)
+            x_ffn = _rms_norm_torch(torch.sum(attn_pre.unsqueeze(-1) * h.float(), dim=2).to(torch.bfloat16), w_fnorm, eps)
+            in_ffn_list.append(x_ffn.reshape(B * T, dim))
+            ffn_meta.append((s_idx, B, T, residual2, ffn_pre, ffn_post, ffn_comb))
+
+        flat_all = torch.cat(in_ffn_list, dim=0)
+        g_scores = F.softplus(F.linear(flat_all.float(), gw)).sqrt()
+        _, top_idx = (g_scores + gb).topk(top_k_eff, dim=-1)
+        top_w = g_scores.gather(1, top_idx)
+        top_w = (top_w / (top_w.sum(dim=-1, keepdim=True) + 1.0e-20)) * route_scale
+
+        y_all = torch.zeros_like(flat_all, dtype=torch.float32)
+        hit = torch.bincount(top_idx.flatten(), minlength=n_routed_eff)
+        for eidx in hit.nonzero().flatten().tolist():
+            erows, eslot = torch.where(top_idx == eidx)
+            xe = flat_all.index_select(0, erows)
+            ep = f"{lp}.ffn.experts.{eidx}"
+            ew1 = _ds4_dequant_fp4_1x32(read_fn(f"{ep}.w1.weight"), read_fn(f"{ep}.w1.scale"), lut).to(torch.bfloat16)
+            ew2 = _ds4_dequant_fp4_1x32(read_fn(f"{ep}.w2.weight"), read_fn(f"{ep}.w2.scale"), lut).to(torch.bfloat16)
+            ew3 = _ds4_dequant_fp4_1x32(read_fn(f"{ep}.w3.weight"), read_fn(f"{ep}.w3.scale"), lut).to(torch.bfloat16)
+            egate = F.linear(xe, ew1).float().clamp(max=swiglu_limit)
+            eup = F.linear(xe, ew3).float().clamp(min=-swiglu_limit, max=swiglu_limit)
+            eact = (F.silu(egate) * eup) * top_w[erows, eslot, None]
+            y_all.index_add_(0, erows, F.linear(eact.to(torch.bfloat16), ew2).float())
+
+        sgate = F.linear(flat_all, sw1).float().clamp(max=swiglu_limit)
+        sup = F.linear(flat_all, sw3).float().clamp(min=-swiglu_limit, max=swiglu_limit)
+        y_all = y_all + F.linear((F.silu(sgate) * sup).to(torch.bfloat16), sw2).float()
+
+        off = 0
+        for s_idx, B, T, residual2, ffn_pre, ffn_post, ffn_comb in ffn_meta:
+            nt = B * T
+            fout = y_all[off : off + nt].reshape(B, T, dim)
+            off += nt
+            h_next = (ffn_post.unsqueeze(-1) * fout.unsqueeze(-2) + torch.sum(ffn_comb.unsqueeze(-1) * residual2.float().unsqueeze(-2), dim=2)).to(torch.bfloat16)
+            streams[s_idx] = h_next
+            pre_mixes[s_idx] = ffn_pre
+
+    norm_key = next((k for k in ("norm.weight", "model.norm.weight") if has_key_fn(k)), None)
+    w_final = read_fn(norm_key) if norm_key else torch.ones(dim, dtype=torch.bfloat16, device=dev)
+    if w_final.numel() != dim:
+        w_final = torch.ones(dim, dtype=torch.bfloat16, device=dev)
+    outputs = []
+    for s_idx, h in enumerate(streams):
+        h_col = torch.sum(pre_mixes[s_idx].unsqueeze(-1) * h.float(), dim=2).to(torch.bfloat16)
+        out_h = _rms_norm_torch(h_col, w_final, eps)
+        outputs.append(out_h[0].float())
+    return outputs
+
+
+def run_deepseek_v41_generate(
+    safetensors_path: str | Path,
+    config_path: str | Path,
+    tokenizer_dir: str | Path,
+    prompt: str = "hello",
+    max_new_tokens: int = 16,
+    use_chat_template: bool = False,
+    device: str = "auto",
+) -> tuple[str, list[int]]:
+    """Run autoregressive greedy generation on a recovered DeepSeek-V4.1 safetensors model."""
+    from transformers import AutoTokenizer
+
+    tok = AutoTokenizer.from_pretrained(str(tokenizer_dir), trust_remote_code=True, local_files_only=True)
+    raw_cfg = json.loads(Path(config_path).read_text())
+    arch = _normalize_arch_config(raw_cfg)
+    dev = torch.device("cuda:0" if (device in ("auto", "cuda") and torch.cuda.is_available()) else "cpu")
+
+    if use_chat_template and getattr(tok, "chat_template", None):
+        ids = tok.apply_chat_template(
+            [{"role": "user", "content": prompt}],
+            add_generation_prompt=True,
+            tokenize=True,
+        )
+        if isinstance(ids, dict):
+            ids = ids["input_ids"]
+    else:
+        ids = tok.encode(prompt, add_special_tokens=False)
+        if not ids:
+            ids = [0]
+
+    prompt_len = len(ids)
+    cur_ids = list(ids)
+    eos_id = tok.eos_token_id if tok.eos_token_id is not None else int(arch.get("eos_token_id", 1))
+    im_end_id = tok.convert_tokens_to_ids("<|im_end|>") if "<|im_end|>" in tok.get_vocab() else -1
+
+    with safe_open(str(safetensors_path), framework="pt", device="cpu") as sf:
+        keys_set = set(sf.keys())
+        head_key = next((k for k in _HEAD_NAMES if k in keys_set), "head.weight")
+        head_w = sf.get_tensor(head_key).to(dev, dtype=torch.bfloat16)
+
+        def read_fn(k: str) -> "torch.Tensor":
+            return sf.get_tensor(k).to(dev, non_blocking=True)
+
+        def has_key_fn(k: str) -> bool:
+            return k in keys_set
+
+        with torch.no_grad():
+            for _ in range(max_new_tokens):
+                h_outs = _run_deepseek_v41_forward(
+                    read_fn, has_key_fn, [cur_ids], None, None, 1.0, arch, dev
+                )
+                last_h = h_outs[0][-1].to(torch.bfloat16)
+                logits = F.linear(last_h, head_w).float()
+                next_id = int(torch.argmax(logits).item())
+                cur_ids.append(next_id)
+                if next_id in (eos_id, im_end_id):
+                    break
+
+    text = tok.decode(cur_ids[prompt_len:] if use_chat_template else cur_ids, skip_special_tokens=False)
+    return text, cur_ids
+
+
 def _synthesize_embed_tokens_checkfree(
     f_handle,
     present: set[str],