| """deg-MoE trainer: the W (writer) experts of a task-routed MoE (results/deg/moe-v1). |
| |
| One model, one context prefill: every decoder layer of a full fine-tune deg checkpoint |
| (deg2b-1) gets a second FFN expert ``mlp_w`` initialised from the base model's MLP, plus |
| an optional writer LoRA on the attention / linear-attention projections |
| (:func:`deg.model.add_moe`). Routing is by token role: the context and every |
| decide/extract/count/anchor fork run the reader (R = the checkpoint as it is); the KNOWN |
| VALUES block, the generate header and the generated tokens run W and read the R trunk's |
| K/V and recurrent state (:class:`deg.model.WriteRoute`). |
| |
| Training freezes everything R had (embeddings, attention, norms, LM head, R experts, |
| heads), so decide/extract outputs are unchanged by construction, and trains only the W |
| experts (+ writer LoRA) with the x-generate LM loss in the base layout |
| (:func:`deg.infer.base_generate_header`; the same items as :mod:`deg.gentrain`: |
| gen-v1 judge-kept fields, open-weight writer and judge only, suite guard on) plus |
| ``--kl-base W`` * KL(base || MoE) on the same target positions, where the base Qwen |
| re-reads the whole sequence (the base copy exists in training only). |
| |
| Each micro-batch reads every distinct context ONCE through R without autograd (R is |
| frozen and no W parameter touches context tokens); the W rows (KNOWN VALUES + header + |
| target) continue that trunk state, gradient-checkpointed. The fp32 master weights and |
| the AdamW moments live on ``--optim-device`` (default cpu: 12 bytes per W parameter |
| would not fit beside the bodies on one 24 GB L4); bf16 gradients are copied to pinned |
| host buffers once per optimizer step. The best ``val`` NLL checkpoint is saved as one |
| deg-MoE checkpoint directory (backbone.safetensors = R + shared, moe.safetensors = W |
| experts + writer LoRA, heads, deg.json ``moe``), loadable by |
| :func:`deg.model.load_model`, the server and the suite (default ``generate`` = moe). |
| |
| CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=<GPU UUID, never GPU 0> \\ |
| python -m deg.moetrain --ckpt /data/deg/runs/deg2b-1/step-0004000 \\ |
| --out /data/deg/runs/deg2b-moe-1 --kl-base 0.5 |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| import os |
| import random |
| import time |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| from deg.gentrain import GEN_RECORDS, GenItem, gen_items, lr_at |
| from deg.model import ( |
| DegModel, |
| add_moe, |
| attach_base, |
| load_base_lm, |
| load_model, |
| moe_params, |
| ) |
| from deg.train import read_jsonl |
|
|
| USER_GPU_UUID = "GPU-019676ed-23c9-ad9a-20cb-7bdc13ac61ca" |
|
|
|
|
| def check_device(dev: torch.device) -> None: |
| """CUDA runs must name their GPU (CUDA_VISIBLE_DEVICES = one UUID), never GPU 0.""" |
| if dev.type != "cuda": |
| return |
| cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip() |
| if not cvd or cvd == "0" or "," in cvd or USER_GPU_UUID.startswith(cvd): |
| raise SystemExit("set CUDA_VISIBLE_DEVICES=<one GPU UUID> (GPU 0 is the user's; " |
| f"got {cvd!r})") |
|
|
|
|
| def micro_batches(items: list[GenItem], batch_tokens: int, rng: random.Random | None |
| ) -> list[list[GenItem]]: |
| """Groups of whole records (a record's fields share its context read), length-sorted, |
| at most ``batch_tokens`` padded tokens counted as contexts x longest context + rows x |
| longest W part (a record over budget is a group alone); shuffled order with ``rng``.""" |
| by: dict[str, list[GenItem]] = {} |
| for it in items: |
| by.setdefault(it.uid, []).append(it) |
| recs = sorted(by.values(), key=lambda g: (g[0].ctx_len, max(len(x.ids) for x in g))) |
| groups: list[list[GenItem]] = [] |
| cur: list[list[GenItem]] = [] |
|
|
| def cost(rs: list[list[GenItem]]) -> int: |
| ctx = max(g[0].ctx_len for g in rs) |
| rows = [x for g in rs for x in g] |
| return len(rs) * ctx + len(rows) * max(len(x.ids) - x.ctx_len for x in rows) |
|
|
| for g in recs: |
| if cur and cost(cur + [g]) > batch_tokens: |
| groups.append([x for r in cur for x in r]) |
| cur = [] |
| cur.append(g) |
| if cur: |
| groups.append([x for r in cur for x in r]) |
| if rng is not None: |
| rng.shuffle(groups) |
| return groups |
|
|
|
|
| def _pack(group: list[GenItem], pad: int, dev: torch.device): |
| """(ctx ids [n, Lc], ctx lengths, W ids [rows, Lw], W lengths, row -> ctx row, |
| target row / column / token): column c of row r predicts W token c + 1.""" |
| uids: list[str] = [] |
| for x in group: |
| if x.uid not in uids: |
| uids.append(x.uid) |
| first = {x.uid: x for x in reversed(group)} |
| ctxs = [first[u].ids[:first[u].ctx_len] for u in uids] |
| for x in group: |
| assert x.ids[:x.ctx_len] == ctxs[uids.index(x.uid)], f"{x.uid}: contexts differ" |
| Lc = max(len(c) for c in ctxs) |
| cids = torch.full((len(ctxs), Lc), pad, dtype=torch.long) |
| for i, c in enumerate(ctxs): |
| cids[i, :len(c)] = torch.tensor(c) |
| ws = [x.ids[x.ctx_len:] for x in group] |
| Lw = max(len(w) for w in ws) |
| wids = torch.full((len(ws), Lw), pad, dtype=torch.long) |
| rows, cols, tgts = [], [], [] |
| for r, (x, w) in enumerate(zip(group, ws)): |
| wids[r, :len(w)] = torch.tensor(w) |
| for p in range(x.prompt_len - x.ctx_len, len(w)): |
| rows.append(r) |
| cols.append(p - 1) |
| tgts.append(w[p]) |
| t = lambda v: torch.tensor(v, dtype=torch.long, device=dev) |
| return (cids.to(dev), t([len(c) for c in ctxs]), wids.to(dev), t([len(w) for w in ws]), |
| t([uids.index(x.uid) for x in group]), t(rows), t(cols), t(tgts)) |
|
|
|
|
| def _target_logits(bb, cids, clens, wids, wlens, row_of, ri, ci, sel): |
| """Logits at the target positions: the contexts read once without autograd (R; the |
| base copy for the KL anchor), then the W rows continue them with route ``sel``.""" |
| with torch.no_grad(): |
| _, state = bb(cids, clens) |
| h, _ = bb(wids, wlens, state, rows=row_of, sel=sel) |
| del state |
| return bb.lm_logits(h[ri, ci]).float() |
|
|
|
|
| def batch_losses(model: DegModel, group: list[GenItem], kl_base: float = 0.0 |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, int]: |
| """(summed target NLL, target logits, summed KL(base || MoE) or None, target tokens) |
| for one micro-batch of W rows.""" |
| dev = model.device |
| cids, clens, wids, wlens, row_of, ri, ci, tgt = _pack(group, model.renderer.pad_id, dev) |
| logits = _target_logits(model.backbone, cids, clens, wids, wlens, row_of, ri, ci, |
| model.gen_route(len(group))) |
| nll = F.cross_entropy(logits, tgt, reduction="sum") |
| kl = None |
| if kl_base > 0: |
| if model.base is None: |
| raise RuntimeError("--kl-base needs the base copy: attach_base(model) first") |
| with torch.no_grad(): |
| base = _target_logits(model.base, cids, clens, wids, wlens, row_of, ri, ci, |
| None).log_softmax(-1) |
| kl = F.kl_div(logits.log_softmax(-1), base, log_target=True, reduction="sum") |
| return nll, logits, kl, int(tgt.shape[0]) |
|
|
|
|
| @torch.no_grad() |
| def evaluate(model: DegModel, items: list[GenItem], batch_tokens: int) -> float: |
| tot, n = 0.0, 0 |
| for g in micro_batches(items, batch_tokens, None): |
| nll, _, _, k = batch_losses(model, g) |
| tot, n = tot + float(nll), n + k |
| return tot / max(n, 1) |
|
|
|
|
| class OffloadAdamW: |
| """AdamW over fp32 master copies of ``params`` on ``device`` (cpu: pinned gradient |
| buffers, fused CPU AdamW); bf16 compute weights stay where they are. Two groups: |
| W experts (``lr``) and writer LoRA (``writer_lr``).""" |
|
|
| def __init__(self, params: list[torch.nn.Parameter], lora: list[bool], lr: float, |
| writer_lr: float, device: torch.device) -> None: |
| self.params = params |
| pin = device.type == "cpu" and torch.cuda.is_available() and params[0].is_cuda |
| self.master = [p.detach().to(device, torch.float32, copy=True).requires_grad_(True) |
| for p in params] |
| self.grads = [torch.zeros(m.shape, dtype=torch.float32, device=device, |
| pin_memory=pin) for m in self.master] |
| groups = [{"params": [m for m, lo in zip(self.master, lora) if not lo], "lr": lr, |
| "base_lr": lr}, |
| {"params": [m for m, lo in zip(self.master, lora) if lo], "lr": writer_lr, |
| "base_lr": writer_lr}] |
| groups = [g for g in groups if g["params"]] |
| try: |
| self.opt = torch.optim.AdamW(groups, weight_decay=0.0, fused=True) |
| except (RuntimeError, ValueError): |
| self.opt = torch.optim.AdamW(groups, weight_decay=0.0) |
|
|
| def set_lr(self, factor: float) -> None: |
| for g in self.opt.param_groups: |
| g["lr"] = g["base_lr"] * factor |
|
|
| def step(self, clip: float = 1.0) -> float: |
| for p, m, g in zip(self.params, self.master, self.grads): |
| if p.grad is None: |
| g.zero_() |
| else: |
| g.copy_(p.grad, non_blocking=True) |
| p.grad = None |
| m.grad = g |
| if self.params[0].is_cuda: |
| torch.cuda.synchronize() |
| norm = float(torch.nn.utils.clip_grad_norm_(self.master, clip)) |
| self.opt.step() |
| with torch.no_grad(): |
| for p, m in zip(self.params, self.master): |
| p.copy_(m, non_blocking=True) |
| return norm |
|
|
|
|
| def build(ckpt: Path | str, device: torch.device, writer_rank: int, |
| writer_alpha: float | None, kl_base: float, seed: int = 0, |
| base_lm: torch.nn.Module | None = None, lm: torch.nn.Module | None = None, |
| tokenizer: Any = None) -> tuple[DegModel, list[torch.nn.Parameter], list[bool]]: |
| """deg checkpoint -> deg-MoE with W = the base MLPs, R frozen, W trainable; the base |
| copy attached when ``kl_base`` > 0. ``base_lm``/``lm``/``tokenizer``: tests.""" |
| model = load_model(ckpt, device=device, generate="self", lm=lm, tokenizer=tokenizer) |
| if base_lm is None: |
| base_lm = load_base_lm(model.meta.get("backbone"), torch.bfloat16) |
| model.backbone.lm.requires_grad_(False) |
| model.heads.requires_grad_(False) |
| params = add_moe(model, base_lm, writer_rank, writer_alpha, seed) |
| for p in params: |
| p.requires_grad_(True) |
| lora = [p.dim() == 2 and any(p is getattr(m, "lora_A", None) or p is getattr(m, "lora_B", None) |
| for m in model.backbone.lm.modules()) for p in params] |
| if kl_base > 0: |
| attach_base(model, base_lm) |
| model.backbone.grad_ckpt = True |
| return model, params, lora |
|
|
|
|
| def main(argv: list[str] | None = None) -> None: |
| ap = argparse.ArgumentParser(description="train the W experts of a deg-MoE") |
| ap.add_argument("--ckpt", type=Path, required=True, help="full fine-tune deg checkpoint (R)") |
| ap.add_argument("--out", type=Path, required=True, help="deg-MoE checkpoint dir (best val)") |
| ap.add_argument("--records", type=Path, default=GEN_RECORDS) |
| ap.add_argument("--device", default="cuda") |
| ap.add_argument("--optim-device", default="cpu", help="fp32 master + AdamW moments") |
| ap.add_argument("--writer-rank", type=int, default=32, help="0 = no writer attention LoRA") |
| ap.add_argument("--writer-alpha", type=float, default=None, help="default 2 * rank") |
| ap.add_argument("--lr", type=float, default=1e-5, help="W experts") |
| ap.add_argument("--writer-lr", type=float, default=2e-4, help="writer LoRA") |
| ap.add_argument("--epochs", type=float, default=2.0) |
| ap.add_argument("--max-steps", type=int, default=None) |
| ap.add_argument("--batch-tokens", type=int, default=4096, help="padded tokens per micro") |
| ap.add_argument("--accum", type=int, default=4, help="micro-batches per optimizer step") |
| ap.add_argument("--warmup", type=int, default=20) |
| ap.add_argument("--eval-every", type=int, default=50) |
| ap.add_argument("--patience", type=int, default=4, help="evals without val improvement") |
| ap.add_argument("--kl-base", type=float, default=0.0) |
| ap.add_argument("--seed", type=int, default=0) |
| ap.add_argument("--limit", type=int, default=None, help="train records (smoke)") |
| ap.add_argument("--val-limit", type=int, default=None) |
| ap.add_argument("--no-save", action="store_true", help="smoke: never write --out") |
| args = ap.parse_args(argv) |
|
|
| torch.manual_seed(args.seed) |
| dev = torch.device(args.device) |
| check_device(dev) |
| if dev.type == "cuda" and dev.index is None: |
| dev = torch.device("cuda", torch.cuda.current_device()) |
| model, params, lora = build(args.ckpt, dev, args.writer_rank, args.writer_alpha, |
| args.kl_base, args.seed) |
| opt = OffloadAdamW(params, lora, args.lr, args.writer_lr, torch.device(args.optim_device)) |
| counts = moe_params(model) |
| def mem() -> float: |
| return (round(torch.cuda.max_memory_allocated(dev) / 2**30, 2) |
| if dev.type == "cuda" else 0.0) |
| print(json.dumps({"event": "model", **counts, "trainable": sum(p.numel() for p in params), |
| "base_copy": model.base is not None, "mem_gb": mem()}), flush=True) |
|
|
| stats = {"judge_drop": 0, "budget_drop": 0} |
| train = gen_items(model, read_jsonl(args.records, "train", args.limit), stats) |
| vstats = {"judge_drop": 0, "budget_drop": 0} |
| val = gen_items(model, read_jsonl(args.records, "val", args.val_limit), vstats) |
| print(json.dumps({"event": "data", "train": len(train), "val": len(val), |
| "train_stats": stats, "val_stats": vstats, |
| "train_target_tokens": sum(x.target_len for x in train)}), flush=True) |
|
|
| rng = random.Random(args.seed) |
| per_epoch = math.ceil(len(micro_batches(train, args.batch_tokens, None)) / args.accum) |
| steps = max(1, int(per_epoch * args.epochs)) |
| if args.max_steps: |
| steps = min(steps, args.max_steps) |
| t0 = time.time() |
| best = evaluate(model, val, args.batch_tokens) |
| print(json.dumps({"event": "val", "step": 0, "nll": round(best, 4), |
| "s": round(time.time() - t0, 1), "mem_gb": mem()}), flush=True) |
| meta = {"ckpt": str(args.ckpt), "records": str(args.records), "train_items": len(train), |
| "val_items": len(val), "lr": args.lr, "writer_lr": args.writer_lr, |
| "kl_base": args.kl_base, "seed": args.seed, "batch_tokens": args.batch_tokens, |
| "accum": args.accum, "layout": "base header (deg.infer.base_generate_header)", |
| "val_nll_init": round(best, 4), "params": counts} |
|
|
| def save(step: int, v: float) -> None: |
| if not args.no_save: |
| model.save(args.out, {"moe_train": {**meta, "step": step, "val_nll": round(v, 4)}}) |
|
|
| best_step, bad, step, stop, tok = 0, 0, 0, False, 0 |
| save(0, best) |
| while not stop: |
| groups = micro_batches(train, args.batch_tokens, rng) |
| for i in range(0, len(groups), args.accum): |
| chunk = groups[i:i + args.accum] |
| ntok = sum(x.target_len for g in chunk for x in g) |
| opt.set_lr(lr_at(step, steps, args.warmup)) |
| tot_nll = tot_kl = 0.0 |
| for g in chunk: |
| nll, _, kl, _ = batch_losses(model, g, args.kl_base) |
| loss = nll + (args.kl_base * kl if kl is not None else 0.0) |
| (loss / ntok).backward() |
| tot_nll += float(nll) |
| tot_kl += float(kl) if kl is not None else 0.0 |
| tok += sum(len(x.ids) for x in g) |
| gnorm = opt.step() |
| step += 1 |
| if step % 10 == 0 or step == 1: |
| s = time.time() - t0 |
| print(json.dumps({"event": "step", "step": step, "of": steps, |
| "nll": round(tot_nll / ntok, 4), "kl": round(tot_kl / ntok, 5), |
| "gnorm": round(gnorm, 3), |
| "lr": opt.opt.param_groups[0]["lr"], "s": round(s, 1), |
| "tok_s": round(tok / max(s, 1e-9), 1), "mem_gb": mem()}), |
| flush=True) |
| if step % args.eval_every == 0 or step >= steps: |
| v = evaluate(model, val, args.batch_tokens) |
| print(json.dumps({"event": "val", "step": step, "nll": round(v, 4)}), flush=True) |
| if v < best: |
| best, bad, best_step = v, 0, step |
| save(step, v) |
| else: |
| bad += 1 |
| if bad >= args.patience or step >= steps: |
| stop = True |
| break |
| print(json.dumps({"event": "done", "best_step": best_step, "best_val_nll": round(best, 4), |
| "out": str(args.out), "s": round(time.time() - t0, 1), "mem_gb": mem()}), |
| flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
