Upload moe/kd_moe_train_v4.py with huggingface_hub

moe/kd_moe_train_v4.py

@@ -0,0 +1,348 @@
+#!/usr/bin/env python3
+"""Teach the 14.7B what the parent believes, not just what the answer was.
+
+Note on the targets: the teacher pass ran the parent at nf4, because bf16
+weights left too little of an 80GB card for the gated-delta-rule activations.
+That perturbs the deep tail of the recorded distribution, which is why the
+head is kept at top-128 rather than pushed further: past that point we would
+be recording quantisation noise rather than the parent's beliefs.
+
+Cross-entropy on hard labels gives one bit per token, and this corpus is spent:
+the student already sits near 0.14 there. The parent's distribution says how much
+probability belongs on every plausible continuation, which is the signal that decides
+whether a model obeys "answer in one word" or wanders. That is the failure we measured.
+
+Loss = KL(teacher || student) over the teacher's recorded top-k, plus a small
+cross-entropy anchor so the argmax stays pinned to the real answer. Both are taken
+on assistant turns only, and only where the turn's terminator was visible, so nothing
+here teaches the model to run on.
+
+  python kd_student_train.py --teacher-logits teacher_top128.npz --minutes 180
+"""
+import argparse, json, os, time
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--base", default="logic65/Qwen3.8-Whittle-tri-14.7B")
+ap.add_argument("--subfolder", default="bf16")
+ap.add_argument("--teacher-logits", required=True)
+ap.add_argument("--data-repo", default="logic65/Qwen3.8-Whittle-dev")
+ap.add_argument("--data-file", default="data/heal60_mix.jsonl")
+ap.add_argument("--out", default="tri-kd-lora")
+ap.add_argument("--seq", type=int, default=4096)
+ap.add_argument("--tok-budget", type=int, default=8192)
+ap.add_argument("--accum", type=int, default=4)
+ap.add_argument("--lr", type=float, default=1e-4)
+ap.add_argument("--kd-weight", type=float, default=0.9, help="rest goes to the CE anchor")
+ap.add_argument("--temperature", type=float, default=1.0)
+ap.add_argument("--use-topk", type=int, default=0, help="trim teacher head to this k")
+ap.add_argument("--minutes", type=float, default=180.0)
+ap.add_argument("--rank", type=int, default=32)
+ap.add_argument("--binary-weight", type=float, default=0.5,
+                help="weight on the set-mass (binary) KL term")
+ap.add_argument("--quant", default="bf16", choices=("bf16", "nf4"))
+ap.add_argument("--push", default="")
+A = ap.parse_args()
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from huggingface_hub import hf_hub_download, snapshot_download
+from transformers import (AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig,
+                          Trainer, TrainingArguments, TrainerCallback)
+from peft import LoraConfig, get_peft_model
+
+t0 = time.time()
+# accept a local checkpoint directory (the merged SFT model) as well as a hub repo
+if os.path.isdir(A.base):
+    base = os.path.join(A.base, A.subfolder) if os.path.isdir(os.path.join(A.base, A.subfolder)) else A.base
+else:
+    local = snapshot_download(A.base, allow_patterns=[f"{A.subfolder}/*", "*.json"])
+    base = os.path.join(local, A.subfolder)
+mix = (A.data_file if os.path.exists(A.data_file)
+       else hf_hub_download(A.data_repo, A.data_file, repo_type="model"))
+tok = AutoTokenizer.from_pretrained(base)
+IM_START = tok.convert_tokens_to_ids("<|im_start|>")
+IM_END = tok.convert_tokens_to_ids("<|im_end|>")
+ASSIST = tok("assistant", add_special_tokens=False)["input_ids"]
+
+Z = np.load(A.teacher_logits)
+TOPK = int(Z["topk"][0])
+USE_K = min(A.use_topk, TOPK) if A.use_topk else TOPK
+row_ids, lengths = Z["row_ids"], Z["lengths"]
+offs = np.concatenate([[0], np.cumsum(lengths)])
+T_IDX, T_VAL = Z["idx"], Z["val"]
+print(f"teacher targets: {len(row_ids)} rows, {lengths.sum()/1e6:.2f}M tokens, top-{TOPK}",
+      flush=True)
+
+raw = [json.loads(l)["input_ids"][:A.seq] for l in open(mix)]
+pos = {int(r): k for k, r in enumerate(row_ids)}
+
+
+def mask_for(ids):
+    """Assistant turns only, and only those whose <|im_end|> is inside the window."""
+    lab = [-100]*len(ids)
+    if IM_START not in ids:
+        return list(ids)
+    i, n = 0, len(ids)
+    while i < n:
+        if ids[i] == IM_START and ids[i+1:i+1+len(ASSIST)] == ASSIST:
+            j = i + 1 + len(ASSIST)
+            while j < n and ids[j] != IM_END:
+                j += 1
+            if j < n and ids[j] == IM_END:
+                for k in range(i+1+len(ASSIST), j+1):
+                    lab[k] = ids[k]
+            i = j
+        i += 1
+    return lab if any(v != -100 for v in lab) else list(ids)
+
+
+data = []
+for r_i, ids in enumerate(raw):
+    if r_i not in pos:
+        continue                      # teacher never reached this row
+    k = pos[r_i]
+    n = min(len(ids), int(lengths[k]))
+    data.append({"input_ids": ids[:n], "labels": mask_for(ids)[:n],
+                 "t_idx": T_IDX[offs[k]:offs[k]+n, :USE_K],
+                 "t_val": T_VAL[offs[k]:offs[k]+n, :USE_K]})
+data.sort(key=lambda d: len(d["input_ids"]))
+sup = sum(sum(1 for v in d["labels"] if v != -100) for d in data)
+print(f"{len(data)} conversations | {sum(len(d['input_ids']) for d in data)/1e6:.2f}M tokens "
+      f"| {sup/1e6:.2f}M supervised", flush=True)
+
+GROUPS, cur = [], []
+for d in data:
+    L = len(d["input_ids"])
+    if cur and max(L, max(len(x["input_ids"]) for x in cur)) * (len(cur)+1) > A.tok_budget:
+        GROUPS.append(cur); cur = []
+    cur.append(d)
+if cur:
+    GROUPS.append(cur)
+print(f"{len(GROUPS)} batches, budget {A.tok_budget} tok", flush=True)
+
+
+class Chats(torch.utils.data.Dataset):
+    def __len__(self):
+        return len(GROUPS)
+
+    def __getitem__(self, i):
+        return GROUPS[i]
+
+
+PAD = tok.pad_token_id if tok.pad_token_id is not None else IM_END
+
+
+def collate(batch):
+    b = batch[0] if (len(batch) == 1 and isinstance(batch[0], list)) else batch
+    m = max(len(x["input_ids"]) for x in b)
+    ids = torch.full((len(b), m), PAD, dtype=torch.long)
+    lab = torch.full((len(b), m), -100, dtype=torch.long)
+    att = torch.zeros((len(b), m), dtype=torch.long)
+    ti = torch.zeros((len(b), m, USE_K), dtype=torch.long)
+    tv = torch.full((len(b), m, USE_K), -1e4, dtype=torch.float)
+    for k, x in enumerate(b):
+        n = len(x["input_ids"])
+        ids[k, :n] = torch.tensor(x["input_ids"])
+        lab[k, :n] = torch.tensor(x["labels"])
+        att[k, :n] = 1
+        ti[k, :n] = torch.tensor(x["t_idx"].astype(np.int64))
+        tv[k, :n] = torch.tensor(x["t_val"].astype(np.float32))
+    return {"input_ids": ids, "labels": lab, "attention_mask": att,
+            "t_idx": ti, "t_val": tv}
+
+
+kw = dict(dtype=torch.bfloat16, attn_implementation="sdpa", device_map={"": 0})
+if A.quant == "nf4":
+    # Keep the routers OUT of nf4. They are the parameters this run exists to train,
+    # and a 4-bit tensor cannot carry gradients. Naming them in full matters: a bare
+    # "gate" would substring-match every expert's gate_proj and drag 12B of weights
+    # back into bf16. All 64 routers together are ~21M parameters.
+    import json as _json
+    _cfg = _json.load(open(os.path.join(base, "config.json")))
+    _nl = _cfg.get("num_hidden_layers") or _cfg["text_config"]["num_hidden_layers"]
+    _skip = ([f"model.layers.{i}.mlp.gate" for i in range(_nl)]
+             + [f"model.layers.{i}.mlp.shared_expert_gate" for i in range(_nl)]
+             + [f"model.layers.{i}.mlp.shared_expert"
+                for i in range(int(os.environ.get("SHARED_FROM", "40")), _nl)]
+             + ["lm_head"])
+    kw["quantization_config"] = BitsAndBytesConfig(
+        load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True,
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        llm_int8_skip_modules=_skip)
+model = AutoModelForCausalLM.from_pretrained(base, **kw)
+model.gradient_checkpointing_enable()
+model.enable_input_require_grads()
+model.config.use_cache = False
+# MoE targeting. The routed experts are exact slices of the parent FFN: the knowledge
+# in them is already correct, so they are left frozen. What did NOT exist in the parent
+# is the router — nothing yet knows which slice holds what, and misrouting is
+# indistinguishable from knowledge loss at the output. So the routers are trained FULLY
+# (all 64 gates together are only ~21M parameters; a rank-128 adapter on a 5120x64
+# matrix would be a straitjacket, not a compression) while attention and the shared
+# expert get LoRA to smooth the seams. Routers move WITH the experts' consumers, never
+# alone: the routers-only run doubled Unknown facts precisely because gates drifted
+# away from everything else.
+#
+# LoRA on the 12288 routed-expert projections was the alternative and it is a trap:
+# each is only 5120x192, so r=128 adapters would add ~8.4B trainable parameters.
+import re
+TARGET_RE = re.compile(r"(linear_attn\.(in_proj_qkv|in_proj_z|in_proj_a|in_proj_b|out_proj)"
+                       r"|self_attn\.(q_proj|k_proj|v_proj|o_proj)"
+                       r")$")   # shared experts train FULL this round, not LoRA
+targets = sorted({n for n, _ in model.named_modules() if TARGET_RE.search(n)})
+# All 64 shared experts full-rank is 5B trainable = 10GB weights + 10GB grads +
+# 10GB optimiser states on top of the nf4 base: OOM at 79GB. The planner showed
+# shared-expert activation mass is 80-90% in the DEEP layers and ~36% mid-stack,
+# and stopping behaviour is late-layer, so train the deep half only (~1.9B).
+SHARED_FROM = int(os.environ.get("SHARED_FROM", "40"))
+import re as _re
+def _deep_shared(n):
+    m = _re.search(r"layers\.(\d+)\.mlp\.shared_expert$", n)
+    return m is not None and int(m.group(1)) >= SHARED_FROM
+gates = sorted({n for n, _ in model.named_modules()
+                if n.endswith("mlp.gate") or n.endswith("mlp.shared_expert_gate")
+                or _deep_shared(n)})
+print(f"full-rank modules: {len(gates)} (routers + shared experts from L{SHARED_FROM})", flush=True)
+assert targets and gates, (len(targets), len(gates))
+assert not any(".experts." in t for t in targets), "routed experts must stay frozen"
+print(f"LoRA on {len(targets)} modules | fully-trained gates: {len(gates)}", flush=True)
+model = get_peft_model(model, LoraConfig(
+    r=A.rank, lora_alpha=2*A.rank, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM",
+    target_modules=targets,
+    modules_to_save=gates + ["input_layernorm", "post_attention_layernorm", "norm"]))
+model.print_trainable_parameters()
+
+
+class KD(Trainer):
+    """KL against the teacher's recorded top-k, with a CE anchor. Assistant tokens only."""
+    t0 = time.time()
+    seen = 0
+
+    def compute_loss(self, model, inputs, return_outputs=False, **kw):
+        t_idx = inputs.pop("t_idx")
+        t_val = inputs.pop("t_val")
+        labels = inputs.pop("labels")
+        KD.seen += int(inputs["input_ids"].numel())
+        out = model(**inputs)
+        logits = out.logits[:, :-1]                      # predict token t+1
+        tgt_i = t_idx[:, :-1]
+        tgt_v = t_val[:, :-1]
+        keep = labels[:, 1:] != -100                     # supervise assistant turns only
+        if keep.any():
+            # only the teacher's top-k student log-probs are needed, so never build the
+            # full 248320-wide fp32 log_softmax: log_softmax(x)_i = x_i - logsumexp(x).
+            # The naive version costs ~8GB per batch to keep 128 numbers per token.
+            kept = logits[keep]                      # bf16 [N, V]
+            ti, tv, lb = tgt_i[keep], tgt_v[keep], labels[:, 1:][keep]
+
+            # At 8k positions the fp32 ops on [N, 248320] cost ~16GB stored for
+            # backward. Chunk over positions under torch.utils.checkpoint: fp32
+            # intermediates are recomputed per-chunk in backward, never stored.
+            import torch.utils.checkpoint as ckpt
+
+            def _loss_chunk(kc, tic, tvc, lbc):
+                lse = torch.logsumexp(kc.float(), dim=-1, keepdim=True)
+                s_top = torch.gather(kc, -1, tic).float() - lse
+                t_in = torch.logsumexp(tvc.float(), -1, keepdim=True).clamp(max=-1e-6)
+                s_in = torch.logsumexp(s_top, -1, keepdim=True).clamp(max=-1e-6)
+                t_out = torch.log1p(-torch.exp(t_in).clamp(max=1 - 1e-6))
+                s_out = torch.log1p(-torch.exp(s_in).clamp(max=1 - 1e-6))
+                kdb = (torch.exp(t_in)*(t_in-s_in) + torch.exp(t_out)*(t_out-s_out)).sum(-1).sum()
+                t_c = torch.softmax(tvc.float()/A.temperature, -1)
+                s_c = s_top/A.temperature
+                s_c = s_c - torch.logsumexp(s_c, -1, keepdim=True)
+                kdc = (t_c*(torch.log(t_c+1e-9)-s_c)).sum(-1).sum()
+                cec = F.cross_entropy(kc.float(), lbc, reduction="sum")
+                return kdc, kdb, cec
+
+            CH, N = 1024, kept.shape[0]
+            kdc_s = kdb_s = ce_s = 0.0
+            for c0 in range(0, N, CH):
+                kdc_c, kdb_c, ce_c = ckpt.checkpoint(
+                    _loss_chunk, kept[c0:c0+CH], ti[c0:c0+CH], tv[c0:c0+CH],
+                    lb[c0:c0+CH], use_reentrant=False)
+                kdc_s = kdc_s + kdc_c; kdb_s = kdb_s + kdb_c; ce_s = ce_s + ce_c
+            kd_cond, kd_bin, ce = kdc_s/N, kdb_s/N, ce_s/N
+            kd = kd_cond + A.binary_weight * kd_bin
+        else:
+            kd = ce = logits.sum() * 0.0
+        loss = A.kd_weight * kd + (1.0 - A.kd_weight) * ce
+        if self.state.global_step % 5 == 0:
+            self.log({"kd": float(kd), "kdb": float(kd_bin), "ce": float(ce)})
+        return (loss, out) if return_outputs else loss
+
+    def training_step(self, *a, **k):
+        if time.time() - self.t0 > A.minutes * 60:
+            self.control.should_training_stop = True
+        return super().training_step(*a, **k)
+
+
+class DriveLog(TrainerCallback):
+    def __init__(self, path):
+        self.path = path
+        with open(path, "w") as f:
+            f.write("step,epoch,loss,kd,ce,lr\n")
+        self.last = {}
+
+    def on_log(self, args, state, control, logs=None, **kw):
+        if not logs:
+            return
+        self.last.update(logs)
+        if "loss" in logs:
+            with open(self.path, "a") as f:
+                f.write("%d,%.4f,%.5f,%.5f,%.5f,%.3e\n" % (
+                    state.global_step, logs.get("epoch", 0), logs["loss"],
+                    self.last.get("kd", 0), self.last.get("ce", 0),
+                    logs.get("learning_rate", 0)))
+
+
+import inspect
+_want = dict(output_dir=A.out, per_device_train_batch_size=1,
+             gradient_accumulation_steps=A.accum, num_train_epochs=4,
+             learning_rate=A.lr, lr_scheduler_type="cosine", warmup_steps=20,
+             bf16=True, logging_steps=5, save_steps=25, save_total_limit=2,
+             save_strategy="steps", optim="adamw_torch", report_to="none",
+             gradient_checkpointing=True, dataloader_num_workers=2,
+             remove_unused_columns=False, label_names=["labels"])
+_ok = set(inspect.signature(TrainingArguments.__init__).parameters)
+_drop = [k for k in _want if k not in _ok]
+if _drop:
+    print("dropping unsupported TrainingArguments:", _drop, flush=True)
+args = TrainingArguments(**{k: v for k, v in _want.items() if k in _ok})
+
+CSV = os.path.join(os.path.dirname(A.out) or ".",
+                   "kd_loss_%s.csv" % os.path.basename(A.out.rstrip("/")))
+trainer = KD(model=model, args=args, train_dataset=Chats(), data_collator=collate,
+             callbacks=[DriveLog(CSV)])
+trainer.train()
+el = time.time() - KD.t0
+print(f"\ntrained {KD.seen/1e6:.2f}M tokens in {el/60:.1f} min ({KD.seen/el:.0f} tok/s)",
+      flush=True)
+
+try:
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+    h = [x for x in trainer.state.log_history if "loss" in x]
+    if h:
+        fig, ax = plt.subplots(figsize=(9, 4.5))
+        ax.plot([x["step"] for x in h], [x["loss"] for x in h], lw=1.2, color="#3b6ea5")
+        ax.set_xlabel("step"); ax.set_ylabel("KD loss"); ax.grid(alpha=.25)
+        ax.set_title("Whittle 14.7B: logit distillation from the FP8 parent")
+        fig.tight_layout()
+        fig.savefig(os.path.join(os.path.dirname(A.out) or ".", "kd_loss_curve.png"), dpi=120)
+        print("curve saved | first %.3f last %.3f" % (h[0]["loss"], h[-1]["loss"]), flush=True)
+except Exception as e:
+    print("plot skipped:", type(e).__name__, flush=True)
+
+model.save_pretrained(A.out)
+tok.save_pretrained(A.out)
+if A.push:
+    from huggingface_hub import HfApi
+    api = HfApi()
+    api.create_repo(A.push, repo_type="model", private=False, exist_ok=True)
+    api.upload_folder(folder_path=A.out, repo_id=A.push,
+                      ignore_patterns=["checkpoint-*", "*.log"])
+    print("pushed ->", A.push, flush=True)
+print(f"total {(time.time()-t0)/60:.1f} min", flush=True)