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