File size: 8,708 Bytes
e379257 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | """Train Glint student: Qwen3.5 + LoRA + JointSchemaHead, distilled from teacher option distributions (+ gold CE)."""
import os; os.environ.setdefault("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "1") # ROCm: faster SDPA (see RESULTS.md); set to 0 to disable
import argparse, json, math, random, time
from pathlib import Path
import torch
from glint.model import build_student, save_checkpoint, default_device
from glint.data import load_jsonl, to_record, normalize_teacher, normalize_box_gold, derive_presence_gold
from glint.losses import decision_loss, record_box_loss
from glint.box import BoxLossWeights, DEFAULT_BINS, is_box_question
from glint.schema import encode_glint_record
from glint.joint_schema_model import collate_records
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--data", required=True)
ap.add_argument("--eval-data", default=None)
ap.add_argument("--out", default="checkpoints/run")
ap.add_argument("--model", default="Qwen/Qwen3.5-2B")
ap.add_argument("--max-pixels", type=int, default=430080, help="430080 ~= 390 tokens / 1048576 ~= 990 tokens (iPhone portrait)")
ap.add_argument("--steps", type=int, default=200, help="optimizer steps")
ap.add_argument("--accum", type=int, default=1)
ap.add_argument("--lr-lora", type=float, default=2e-4)
ap.add_argument("--lr-head", type=float, default=5e-4)
ap.add_argument("--warmup", type=int, default=10)
ap.add_argument("--temperature", type=float, default=2.0)
ap.add_argument("--alpha", type=float, default=0.7, help="KD weight when both teacher and gold exist")
ap.add_argument("--beta", type=float, default=0.1, help="Brier weight")
ap.add_argument("--lora-r", type=int, default=32)
ap.add_argument("--dtype", default="bf16", choices=["bf16", "fp16", "fp32"])
ap.add_argument("--grad-ckpt", action="store_true")
ap.add_argument("--max-length", type=int, default=4096)
ap.add_argument("--log-every", type=int, default=10)
ap.add_argument("--seed", type=int, default=0)
# Jev-style box questions (glint/box.py). auto = add a box head iff the training data has box questions.
ap.add_argument("--box-head", default="auto", choices=["auto", "on", "off"])
ap.add_argument("--box-weight", type=float, default=1.0, help="weight of the box loss in the total loss")
ap.add_argument("--box-dfl", type=float, default=1.0, help="DFL (two-bin soft CE) weight inside the box loss")
ap.add_argument("--box-l1", type=float, default=5.0, help="L1 on decoded edges weight inside the box loss")
ap.add_argument("--box-giou", type=float, default=2.0, help="(1 - GIoU) weight inside the box loss")
ap.add_argument("--box-conf", type=float, default=1.0, help="BCE(confidence, detached IoU) weight inside the box loss")
ap.add_argument("--box-bins", type=int, default=DEFAULT_BINS)
ap.add_argument("--init-from", default=None, help="warm start: load LoRA adapter, head and box head from this checkpoint dir")
args = ap.parse_args()
random.seed(args.seed); torch.manual_seed(args.seed)
dev = default_device()
data = load_jsonl(args.data)
has_box = any(is_box_question(q) for ex in data for q in ex["questions"].values())
use_box = args.box_head == "on" or (args.box_head == "auto" and has_box)
if has_box and not use_box:
raise SystemExit("training data has box questions but --box-head off")
model, proc, meta = build_student(args.model, args.max_pixels, dev, args.dtype, lora_r=args.lora_r,
lora_alpha=2 * args.lora_r, grad_ckpt=args.grad_ckpt,
box_head_cfg={"bins": args.box_bins} if use_box else None)
if args.init_from:
from safetensors.torch import load_file
from peft import set_peft_model_state_dict
ck = Path(args.init_from)
sd = load_file(ck / "adapter" / "adapter_model.safetensors")
res = set_peft_model_state_dict(model.language_model, sd)
bad = list(getattr(res, "unexpected_keys", []) or [])
if bad: raise SystemExit(f"init-from: {len(bad)} unexpected LoRA keys, e.g. {bad[:3]}")
model.head.load_state_dict(load_file(ck / "head.safetensors"), strict=True)
if use_box:
model.box_head.load_state_dict(load_file(ck / "box_head.safetensors"), strict=True)
print(f"init-from {ck}: lora_tensors={len(sd)} head ok box_head={'ok' if use_box else 'n/a'}", flush=True)
box_w = BoxLossWeights(dfl=args.box_dfl, l1=args.box_l1, giou=args.box_giou, conf=args.box_conf)
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"device={dev} trainable_params={trainable/1e6:.1f}M head={sum(p.numel() for p in model.head.parameters())/1e6:.1f}M"
+ (f" box_head={sum(p.numel() for p in model.box_head.parameters())/1e6:.1f}M" if use_box else ""), flush=True)
opt = torch.optim.AdamW([
{"params": [p for p in model.language_model.parameters() if p.requires_grad], "lr": args.lr_lora},
{"params": list(model.head.parameters()) + (list(model.box_head.parameters()) if use_box else []), "lr": args.lr_head}],
weight_decay=0.01)
sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / max(1, args.warmup)) *
0.5 * (1 + math.cos(math.pi * min(1.0, s / max(1, args.steps)))))
if dev == "cuda":
torch.cuda.reset_peak_memory_stats()
model.train(); order = []; hist = []; t_start = time.time(); step = 0; micro = 0
win = {"loss": 0.0, "kd": 0.0, "ce": 0.0, "brier": 0.0, "n": 0, "t": 0.0, "tok": 0}
box_keys = ("box", "dfl", "l1", "giou", "conf", "iou")
if use_box:
win.update({k: 0.0 for k in box_keys}, n_box=0)
while step < args.steps:
if not order:
order = list(range(len(data))); random.shuffle(order)
ex = data[order.pop()]
t0 = time.time()
enc = encode_glint_record(proc.tokenizer, to_record(ex, augment=True), max_length=args.max_length, processor=proc)
batch = collate_records([enc], proc.tokenizer.pad_token_id, torch.device(dev))
if use_box:
all_logits, boxes = model.decide(batch); logits = all_logits[0]
gold = derive_presence_gold(ex)
else:
logits = model(batch)[0]; gold = ex.get("gold") or {}
loss, parts = decision_loss(logits, enc, normalize_teacher(ex), gold,
T=args.temperature, alpha=args.alpha, beta=args.beta)
if use_box:
bl, bparts, nb = record_box_loss(boxes[0], normalize_box_gold(ex), box_w)
if bl is not None:
loss = loss + args.box_weight * bl
win["box"] += float(bl.detach()); win["n_box"] += 1
for k in box_keys[1:]:
win[k] += bparts[k]
if not (torch.is_tensor(loss) and loss.requires_grad): # e.g. only a present=false box question without gold
loss = sum(p.sum() for p in model.box_head.parameters()) * 0.0
(loss / args.accum).backward(); micro += 1
if micro % args.accum == 0:
torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], 1.0)
opt.step(); sched.step(); opt.zero_grad(set_to_none=True); step += 1
if dev == "cuda":
torch.cuda.synchronize()
win["t"] += time.time() - t0; win["loss"] += float(loss.detach()); win["n"] += 1; win["tok"] += len(enc.input_ids)
for k in ("kd", "ce", "brier"):
win[k] += parts[k]
if micro % args.accum == 0 and step % args.log_every == 0:
n = win["n"]
rec = {"step": step, "loss": win["loss"] / n, "kd": win["kd"] / n, "ce": win["ce"] / n, "brier": win["brier"] / n,
"sec_per_micro": win["t"] / n, "tokens": win["tok"] // n,
"peak_mem_gb": torch.cuda.max_memory_allocated() / 2**30 if dev == "cuda" else None}
if use_box:
nb = max(win["n_box"], 1)
rec.update({f"box_{k}" if k != "box" else "box": win[k] / nb for k in box_keys}, n_box=win["n_box"])
hist.append(rec); print(json.dumps(rec), flush=True)
win = {k: 0 for k in win}
out = Path(args.out)
save_checkpoint(model, meta, out, {"train_args": vars(args), "history": hist,
"train_seconds": time.time() - t_start})
print(f"saved checkpoint -> {out.resolve()}", flush=True)
if __name__ == "__main__":
main()
|