| """Measure train step time, inference latency and peak memory for several max_pixels settings."""
|
| import os; os.environ.setdefault("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "1")
|
| import argparse, json, statistics, time
|
| import torch
|
| from glint.model import build_student, load_processor, 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 is_box_question
|
| from glint.joint_schema_model import collate_records
|
| from glint.schema import encode_glint_record as encode_record
|
|
|
| def set_pixels(proc, max_pixels):
|
| ip = proc.image_processor
|
| ip.size = {"shortest_edge": 65536, "longest_edge": max_pixels}
|
| if hasattr(ip, "max_pixels"):
|
| ip.max_pixels = max_pixels
|
|
|
| def main():
|
| ap = argparse.ArgumentParser()
|
| ap.add_argument("--data", required=True)
|
| ap.add_argument("--model", default="Qwen/Qwen3.5-2B")
|
| ap.add_argument("--pixels", default="430080,1048576")
|
| ap.add_argument("--train-steps", type=int, default=8)
|
| ap.add_argument("--infer-runs", type=int, default=8)
|
| ap.add_argument("--dtype", default="bf16")
|
| ap.add_argument("--grad-ckpt", action="store_true")
|
| args = ap.parse_args(); dev = default_device(); sync = torch.cuda.synchronize if dev == "cuda" else (lambda: None)
|
| data = load_jsonl(args.data); results = []
|
| has_box = any(is_box_question(q) for ex in data for q in ex["questions"].values())
|
| model, proc, meta = build_student(args.model, 430080, dev, args.dtype, grad_ckpt=args.grad_ckpt,
|
| box_head_cfg={} if has_box else None)
|
| opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-5)
|
| for mp in [int(x) for x in args.pixels.split(",")]:
|
| set_pixels(proc, mp)
|
| encs = [(ex, encode_record(proc.tokenizer, to_record(ex), processor=proc)) for ex in data[: max(args.train_steps, args.infer_runs) + 2]]
|
| grid = encs[0][1].media["image_grid_thw"][0].tolist()
|
| img_tokens = grid[0] * grid[1] * grid[2] // 4
|
|
|
| model.train(); times = []
|
| if dev == "cuda":
|
| torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats()
|
| for i in range(args.train_steps + 2):
|
| ex, enc = encs[i % len(encs)]
|
| sync(); t0 = time.perf_counter()
|
| batch = collate_records([enc], proc.tokenizer.pad_token_id, torch.device(dev))
|
| if has_box:
|
| logits, boxes = model.decide(batch)
|
| loss, _ = decision_loss(logits[0], enc, normalize_teacher(ex), derive_presence_gold(ex))
|
| bl, _, _ = record_box_loss(boxes[0], normalize_box_gold(ex))
|
| loss = loss + (bl if bl is not None else 0.0)
|
| else:
|
| loss, _ = decision_loss(model(batch)[0], enc, normalize_teacher(ex), ex.get("gold") or {})
|
| loss.backward(); opt.step(); opt.zero_grad(set_to_none=True)
|
| sync(); times.append(time.perf_counter() - t0)
|
| train_peak = torch.cuda.max_memory_allocated() / 2**30 if dev == "cuda" else None
|
|
|
| model.eval(); lat = []
|
| if dev == "cuda":
|
| torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats()
|
| with torch.inference_mode():
|
| for i in range(args.infer_runs + 2):
|
| ex = data[i % len(data)]
|
| sync(); t0 = time.perf_counter()
|
| enc = encode_record(proc.tokenizer, to_record(ex), processor=proc)
|
| model.decide(collate_records([enc], proc.tokenizer.pad_token_id, torch.device(dev)))
|
| sync(); lat.append(time.perf_counter() - t0)
|
| infer_peak = torch.cuda.max_memory_allocated() / 2**30 if dev == "cuda" else None
|
| r = {"max_pixels": mp, "image_grid_thw": grid, "image_tokens": img_tokens, "seq_tokens": len(encs[0][1].input_ids),
|
| "train_step_s_median": round(statistics.median(times[2:]), 3), "train_peak_mem_gb": train_peak and round(train_peak, 2),
|
| "infer_ms_median": round(statistics.median(lat[2:]) * 1000, 1), "infer_ms_min": round(min(lat[2:]) * 1000, 1),
|
| "infer_peak_mem_gb": infer_peak and round(infer_peak, 2)}
|
| results.append(r); print(json.dumps(r), flush=True)
|
| print(json.dumps({"model": args.model, "dtype": args.dtype, "device": torch.cuda.get_device_name(0) if dev == "cuda" else "cpu",
|
| "results": results}, indent=1))
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
