File size: 4,838 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 | """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") # ROCm: faster SDPA (see RESULTS.md); set to 0 to disable
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 # == Clef encode_record unless box questions exist
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()) # box head only when the data asks for boxes
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
# --- train
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
# --- inference (LoRA unmerged, same as training graph)
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()
|