File size: 2,634 Bytes
e7013d4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """Build the king's full HF-layout state dict (BF16) by dequantizing the GGUF with the repo's
own code, adding the ssm tensors the dedup ignores. Saves /dev/shm/king_hf.pt."""
import sys
sys.path.insert(0, "/root/ingot/src")
import torch
import numpy as np
from common.gguf import read_header
from common.gguf.dequant import tensor as deq_tensor
from model_dedup.gguf_checkpoint import GgufCheckpoint
GGUF_DIR = "/root/models/king-ix"
ck = GgufCheckpoint(GGUF_DIR, "cpu")
g = ck.g
def raw(key):
w = torch.from_numpy(deq_tensor(ck.file, g, key))
return w.float()
sd = {}
for name in ck.names():
sd[name] = ck.get(name)
n = int(g.kv["qwen35.block_count"].value)
hd = ck._heads
def untile(w: torch.Tensor, dim: int, head: int) -> torch.Tensor:
"""GgufCheckpoint._untile: stored tiled [rep, n_k, head] -> HF grouped [n_k, rep, head]."""
shape = list(w.shape)
x = w.reshape(*shape[:dim], hd.rep, hd.n_k, head, *shape[dim + 1 :])
return x.transpose(dim, dim + 1).contiguous().reshape(shape)
for il in range(n):
if f"blk.{il}.ssm_a" not in g.tensors: # full-attention layers have no ssm tensors
continue
p = f"layers.{il}.linear_attn."
a = raw(f"blk.{il}.ssm_a")
sd[p + "A_log"] = untile(torch.log(-a).unsqueeze(-1), 0, 1).squeeze(-1)
dt = raw(f"blk.{il}.ssm_dt.bias")
sd[p + "dt_bias"] = untile(dt.unsqueeze(-1), 0, 1).squeeze(-1)
conv = raw(f"blk.{il}.ssm_conv1d.weight").unsqueeze(1) # (channels, 1, 4)
qk_channels = 2 * hd.n_k * hd.k_dim
v = untile(conv[qk_channels:], 0, hd.v_dim)
sd[p + "conv1d.weight"] = torch.cat([conv[:qk_channels], v], dim=0)
sd[p + "norm.weight"] = raw(f"blk.{il}.ssm_norm.weight")
# check against the teacher's shapes
from transformers import AutoModelForCausalLM # noqa: E402
teacher = AutoModelForCausalLM.from_pretrained(
"/root/models/teacher", torch_dtype=torch.bfloat16, device_map="cpu"
)
td = teacher.state_dict()
bad = []
for k, v in sd.items():
hits = [x for x in td if x.endswith(k)]
if not hits:
bad.append((k, "no key"))
continue
if tuple(td[hits[0]].shape) != tuple(v.shape):
bad.append((k, f"shape {tuple(v.shape)} vs {tuple(td[hits[0]].shape)}"))
print("mismatches:", len(bad), bad[:10])
# overlay the king's tensors onto the teacher model (MTP/visual stay the teacher's)
out = {}
for k, v in td.items():
hit = next((x for x in sd if k.endswith(x)), None)
out[k] = (sd[hit].to(torch.bfloat16) if hit is not None else v.to(torch.bfloat16))
print("saving", len(out), "tensors")
torch.save(out, "/dev/shm/king_hf.pt")
print("saved /dev/shm/king_hf.pt")
|