| """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: |
| 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) |
| 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") |
|
|
| |
| from transformers import AutoModelForCausalLM |
|
|
| 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]) |
|
|
| |
| 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") |
|
|
