| """End-to-end norm tuning on the surrogate: every RMSNorm weight (F32, free in the GGUF) trained |
| against teacher trajectories with the eval's metric (KL + hidden-state distance), the same loss |
| the eval uses. Writes a new GGUF. |
| |
| .venv/bin/python qat_norms.py --data /root/train_samples.json --out /root/models/chal-norms |
| """ |
| 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.types import TYPES |
| from model_dedup.gguf_checkpoint import GgufCheckpoint |
|
|
| 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-norms") |
| ap.add_argument("--steps", type=int, default=15) |
| ap.add_argument("--lr", type=float, default=3e-3) |
| ap.add_argument("--alpha", type=float, default=1.0) |
| ap.add_argument("--max-tokens", type=int, default=2200) |
| args = ap.parse_args() |
|
|
| ck = GgufCheckpoint(args.king, "cpu") |
| g = ck.g |
| records = json.loads(Path(args.data).read_text()) |
| print("samples:", len(records)) |
|
|
| model = AutoModelForCausalLM.from_pretrained(args.teacher, torch_dtype=torch.bfloat16, device_map="cuda") |
| model.eval() |
| model.gradient_checkpointing_enable() |
| sd_king = torch.load("/dev/shm/king_hf.pt", map_location="cpu", weights_only=True) |
| model.load_state_dict(sd_king, strict=True) |
|
|
| from eval_service import metrics |
| from eval_service.config import LAYERS |
|
|
| heads = metrics.heads({"teacher": Path(args.teacher), "king": Path(args.king)}, ["cuda:0"])["cuda:0"] |
| head_t, head_k = heads["teacher"], heads["king"] |
|
|
| params = [] |
| for name, p in model.named_parameters(): |
| p.requires_grad_(False) |
| if any(k in name for k in ("layernorm.weight", "norm.weight")): |
| p.requires_grad_(True) |
| params.append(p) |
| print("trainable norms:", len(params)) |
|
|
| opt = torch.optim.Adam(params, lr=args.lr) |
| for step in range(args.steps): |
| torch.cuda.empty_cache() |
| total, n = 0.0, 0 |
| for rec in records: |
| tokens = rec["prompt"] + rec["answer"] |
| if len(tokens) > args.max_tokens: |
| continue |
| first = len(rec["prompt"]) - 1 |
| ids = torch.tensor([tokens], device="cuda") |
| out = model(input_ids=ids, output_hidden_states=True, use_cache=False) |
| h_k = torch.stack([out.hidden_states[l][0, first:].float() for l in LAYERS], 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).mean() |
| a, b = st[1:], h_k[1:] |
| dist = ((b - a).norm(dim=-1) / a.norm(dim=-1)).mean() |
| loss = kl + args.alpha * dist |
| opt.zero_grad() |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(params, 1.0) |
| opt.step() |
| total += loss.item() |
| n += 1 |
| del out, h_k, st, lp_t, lp_k, ids, loss |
| print(f"step {step}: loss {total / n:.5f} over {n} samples", flush=True) |
|
|
| |
| out_dir = Path(args.out) |
| out_dir.mkdir(parents=True, exist_ok=True) |
| dst = out_dir / "chal-norms.gguf" |
| shutil.copy(Path(args.king) / "Ingot-IX-27B-Q2.gguf", dst) |
|
|
| def write_tensor(tname, values): |
| t = g.tensors[tname] |
| size = int(np.prod(t.shape)) * 4 |
| raw = np.memmap(dst, np.uint8, "r+", g.data_start + t.offset, (size,)) |
| raw.view("<f4")[:] = values.astype(np.float32).ravel() |
|
|
| n = int(g.kv["qwen35.block_count"].value) |
| with torch.no_grad(): |
| for il in range(n): |
| for gguf_kind, hf_kind, plus1 in ( |
| ("attn_norm", "input_layernorm", True), |
| ("post_attention_norm", "post_attention_layernorm", True), |
| ("ssm_norm", "linear_attn.norm", False), |
| ): |
| if f"blk.{il}.{gguf_kind}.weight" not in g.tensors: |
| continue |
| hf = f"model.layers.{il}.{hf_kind}.weight" |
| mod = model |
| for part in hf.split(".")[:-1]: |
| mod = getattr(mod, part) |
| w = mod.weight.detach().cpu().float() |
| write_tensor(f"blk.{il}.{gguf_kind}.weight", w.numpy() + (1.0 if plus1 else 0.0)) |
| mod = model.model |
| w = mod.norm.weight.detach().cpu().float() |
| write_tensor("output_norm.weight", w.numpy() + 1.0) |
| print("saved", dst) |
|
|
