"""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()