| """QAT on the validated surrogate: per layer, straight-through ternary codes + block scales + norm |
| gains, distilled toward teacher trajectories with the eval's own metric (KL + hidden-state |
| distance). Writes a new GGUF at --out. |
| |
| .venv/bin/python qat.py --teacher /root/models/teacher --king /root/models/king-ix \ |
| --data /root/train_samples.json --out /root/models/chal-v2 [--layers all|7] |
| """ |
| import argparse |
| import json |
| import shutil |
| import sys |
| from pathlib import Path |
|
|
| sys.path.insert(0, "/root/ingot/src") |
| import numpy as np |
| import torch |
| from transformers import AutoModelForCausalLM |
|
|
| from common.gguf import read_header |
| from common.gguf.dequant import hadamard, rotation |
| from common.gguf.types import PQ2_0, TYPES |
| from model_dedup.gguf_checkpoint import GgufCheckpoint, canonical |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--teacher", default="/root/models/teacher") |
| ap.add_argument("--king", default="/root/models/king-ix") |
| ap.add_argument("--data", default="/root/train_samples.json") |
| ap.add_argument("--out", default="/root/models/chal-v2") |
| ap.add_argument("--layers", default="7,8,15,16,23,24,31,32,39,40,47,48,55,56,63") |
| ap.add_argument("--steps", type=int, default=25) |
| ap.add_argument("--lr", type=float, default=2e-3) |
| ap.add_argument("--alpha", type=float, default=1.0) |
| ap.add_argument("--seed", type=int, default=0) |
| args = ap.parse_args() |
|
|
| torch.manual_seed(args.seed) |
| ck = GgufCheckpoint(args.king, "cpu") |
| g = ck.g |
| hd = ck._heads |
| block_size = int(g.kv["prism.hadamard.block_size"].value) |
| H = torch.from_numpy(hadamard(block_size)).float() |
|
|
| |
| records = json.loads(Path(args.data).read_text()) |
| print("training samples:", len(records)) |
|
|
| |
| model = AutoModelForCausalLM.from_pretrained(args.teacher, torch_dtype=torch.bfloat16, device_map="cuda") |
| model.eval() |
| sd_king = torch.load("/dev/shm/king_hf.pt", map_location="cpu", weights_only=True) |
|
|
| sys.path.insert(0, "/root/ingot/src") |
| from eval_service import metrics |
| from eval_service.config import LAYERS as CAP_LAYERS |
|
|
| heads = metrics.heads({"teacher": Path(args.teacher), "king": Path(args.king)}, ["cuda:0"])["cuda:0"] |
| head_t, head_k = heads["teacher"], heads["king"] |
|
|
| |
| src_gguf = Path(args.king) / "Ingot-IX-27B-Q2.gguf" |
| out_dir = Path(args.out) |
| out_dir.mkdir(parents=True, exist_ok=True) |
| dst = out_dir / "chal-v2.gguf" |
| shutil.copy(src_gguf, dst) |
|
|
|
|
| def gguf_bytes(tname): |
| t = g.tensors[tname] |
| n = int(np.prod(t.shape)) |
| size = n // TYPES[t.ggml_type].block_elems * TYPES[t.ggml_type].block_bytes |
| return np.memmap(dst, np.uint8, "r+", g.data_start + t.offset, (size,)) |
|
|
|
|
| def read_blocks(tname): |
| t = g.tensors[tname] |
| kind = TYPES[t.ggml_type] |
| n = int(np.prod(t.shape)) |
| size = n // kind.block_elems * kind.block_bytes |
| raw = np.memmap(dst, np.uint8, "r", g.data_start + t.offset, (size,)) |
| blocks = raw.reshape(-1, 34) |
| d = blocks[:, :2].copy().view("<f2").astype(np.float32) |
| qs = blocks[:, 2:] |
| codes = np.stack([(qs >> (2 * k)) & 3 for k in range(4)], axis=-1).reshape(-1, 128) |
| return codes, d |
|
|
|
|
| def unrotate(w, tname): |
| """get()'s transform: HF-layout tensor from the stored one. w: (rows, width).""" |
| folded = rotation(g, tname) |
| if folded is not None: |
| block, signs = folded |
| rows, width = w.shape |
| w = (w.reshape(rows, width // block, block) @ H).reshape(rows, width) |
| w = w * torch.from_numpy(signs).float() |
| return ck._hf_order(tname, w, folded is not None) |
|
|
|
|
| def hf_name(tname): |
| c = canonical(tname) |
| return c if c.startswith(("lm_head", "norm.")) else "model." + c |
|
|
|
|
| def hf_module(hf): |
| parts = hf.split(".") |
| m = model |
| for p in parts[:-1]: |
| m = getattr(m, p) |
| return m |
|
|
|
|
| def layer_tensors(il): |
| out = [] |
| for kind in ("attn_qkv", "attn_gate", "ssm_alpha", "ssm_beta", "ssm_out", "attn_q", "attn_k", "attn_v", "attn_output", "ffn_gate", "ffn_up", "ffn_down"): |
| tname = f"blk.{il}.{kind}.weight" |
| if tname in g.tensors and g.tensors[tname].ggml_type == PQ2_0: |
| out.append(tname) |
| return out |
|
|
|
|
| def train_layer(il): |
| torch.cuda.empty_cache() |
| print(f"--- layer {il}") |
| tensors = layer_tensors(il) |
| zs, ss = {}, {} |
| for tname in tensors: |
| codes, d = read_blocks(tname) |
| t = g.tensors[tname] |
| rows, width = int(np.prod(t.shape) // t.shape[0]), t.shape[0] |
| nblk = rows * width // 128 |
| zs[tname] = torch.tensor(codes, dtype=torch.float32, device="cuda", requires_grad=True).reshape(nblk, 128) |
| ss[tname] = torch.tensor(d, dtype=torch.float32, device="cuda", requires_grad=True).reshape(nblk) |
| hooks = [] |
| for tname in tensors: |
| hf = hf_name(tname) |
| module = hf_module(hf) |
| t = g.tensors[tname] |
| rows, width = int(np.prod(t.shape) // t.shape[0]), t.shape[0] |
|
|
| def make_pre(mod, tname=tname, rows=rows, width=width): |
| def pre(mod, inp): |
| z, s = zs[tname], ss[tname] |
| c = torch.clamp(torch.round(z), 0, 2) |
| W = (s * (c - 1.0)).reshape(rows, width) |
| mod._orig_weight = mod.weight |
| mod.weight = unrotate(W, tname).bfloat16() |
| return pre |
|
|
| def make_post(mod): |
| def post(mod, inp, out): |
| mod.weight = mod._orig_weight |
| return post |
|
|
| module.register_forward_pre_hook(make_pre(module)) |
| module.register_forward_hook(make_post(module)) |
| hooks.append(module) |
|
|
| opt = torch.optim.Adam(list(zs.values()) + list(ss.values()), lr=args.lr) |
| for step in range(args.steps): |
| torch.cuda.empty_cache() |
| total = torch.zeros((), device="cuda") |
| n = 0 |
| for rec in records: |
| tokens = rec["prompt"] + rec["answer"] |
| if len(tokens) > 2600: |
| continue |
| first = len(rec["prompt"]) - 1 |
| ids = torch.tensor([tokens], device="cuda") |
| out = model(input_ids=ids, output_hidden_states=True, use_cache=False) |
| hs = [out.hidden_states[l][0, first:].float() for l in CAP_LAYERS] |
| h_k = torch.stack(hs, 1) |
| st = metrics.load_states(rec["states"], tokens[:-1], first).float().to("cuda") |
| m = min(h_k.shape[0], st.shape[0]) |
| h_k, st = h_k[:m], st[:m] |
| lp_t = head_t.logprobs(st[:, -1]) |
| lp_k = head_k.logprobs(h_k[:, -1]) |
| kl = (lp_t.exp() * (lp_t - lp_k)).sum(-1) |
| a = st[1:] |
| b = h_k[1:] |
| dist = ((b - a).norm(dim=-1) / a.norm(dim=-1)).mean(1) |
| loss = kl.mean() + args.alpha * dist.mean() |
| opt.zero_grad() |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(list(zs.values()) + list(ss.values()), 1.0) |
| opt.step() |
| total += loss.detach() |
| n += 1 |
| del out, hs, h_k, st, lp_t, lp_k |
| print(f" step {step}: loss {total.item() / n:.5f} over {n} samples", flush=True) |
|
|
| for tname in tensors: |
| z, s = zs[tname].detach().cpu().numpy(), ss[tname].detach().cpu().numpy() |
| c = np.clip(np.round(z), 0, 2).astype(np.uint8) |
| rows, width = int(np.prod(g.tensors[tname].shape) // g.tensors[tname].shape[0]), g.tensors[tname].shape[0] |
| nblk = rows * width // 128 |
| raw = gguf_bytes(tname).reshape(-1, 34) |
| raw[:, :2] = s.astype(np.float16).reshape(-1, 1).astype("<f2").view(np.uint8).reshape(-1, 2) |
| qs = np.zeros((nblk, 32), np.uint8) |
| for j in range(128): |
| qs[:, j // 4] |= (c[:, j] << (2 * (j % 4))).astype(np.uint8) |
| raw[:, 2:] = qs |
| print(f" wrote {tname}", flush=True) |
| for mod in hooks: |
| mod._forward_pre_hooks.clear() |
| mod._forward_hooks.clear() |
| del zs, ss |
| torch.cuda.empty_cache() |
|
|
|
|
| for il in (int(x) for x in args.layers.split(",") if x): |
| train_layer(il) |
|
|
| print("saved", dst) |
|
|
