File size: 4,838 Bytes
98eacce 0f79c28 98eacce 0f79c28 98eacce 0f79c28 98eacce | 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 | """LoRA fine-tune a VLM judge on the train split.
No checkpoint selection on cal or test: the final adapter after a fixed budget
is the one that gets evaluated. Selecting on the calibration split would
contaminate the confidence layer; selecting on test would make the numbers
meaningless.
Example (Kaggle T4, 16 GB):
python scripts/train.py --manifests manifests/ --out adapters/qwen3b-lora \
--load-in-4bit --steps 1500 --grad-accum 8
"""
import argparse
import json
import random
import sys
import time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from judgecal.data import load_manifest # noqa: E402
from judgecal.judge import QwenVLJudge # noqa: E402
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--manifests", default="manifests")
ap.add_argument("--model-id", default="Qwen/Qwen2.5-VL-3B-Instruct")
ap.add_argument("--out", required=True)
ap.add_argument("--steps", type=int, default=1500, help="optimizer steps")
ap.add_argument("--batch-size", type=int, default=1)
ap.add_argument("--grad-accum", type=int, default=8)
ap.add_argument("--lr", type=float, default=1e-4)
ap.add_argument("--lora-r", type=int, default=16)
ap.add_argument("--load-in-4bit", action="store_true")
ap.add_argument("--max-pixels", type=int, default=256 * 28 * 28)
ap.add_argument("--seed", type=int, default=0)
ap.add_argument("--checkpoint-every", type=int, default=100,
help="save the adapter every N steps; 0 disables")
a = ap.parse_args()
import torch
import torch.nn.functional as F
from peft import LoraConfig, get_peft_model
torch.manual_seed(a.seed)
rng = random.Random(a.seed)
train = load_manifest(Path(a.manifests) / "train.jsonl")
pos = [e for e in train if e.label == 1]
neg = [e for e in train if e.label == 0]
print(f"train: {len(pos)} success, {len(neg)} failure")
judge = QwenVLJudge(a.model_id, max_pixels=a.max_pixels, load_in_4bit=a.load_in_4bit)
if a.load_in_4bit:
from peft import prepare_model_for_kbit_training
judge.model = prepare_model_for_kbit_training(judge.model, use_gradient_checkpointing=True)
else:
judge.model.gradient_checkpointing_enable()
judge.model.enable_input_require_grads()
# q/k/v/o_proj names match the language model only; the vision tower is left frozen.
cfg = LoraConfig(r=a.lora_r, lora_alpha=2 * a.lora_r, lora_dropout=0.05,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], task_type="CAUSAL_LM")
judge.model = get_peft_model(judge.model, cfg)
judge.model.print_trainable_parameters()
judge.model.train()
opt = torch.optim.AdamW([p for p in judge.model.parameters() if p.requires_grad], lr=a.lr, weight_decay=0.0)
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=a.lr, total_steps=a.steps, pct_start=0.05)
def sample_batch():
# Class-balanced sampling: failures are often the minority.
return [rng.choice(pos if rng.random() < 0.5 else neg) for _ in range(a.batch_size)]
t0, running = time.time(), None
for step in range(a.steps):
opt.zero_grad(set_to_none=True)
for _ in range(a.grad_accum):
batch = sample_batch()
z = judge.forward_logits(batch)
y = torch.tensor([e.label for e in batch], dtype=z.dtype, device=z.device)
loss = F.binary_cross_entropy_with_logits(z, y) / a.grad_accum
loss.backward()
running = loss.item() * a.grad_accum if running is None else 0.98 * running + 0.02 * loss.item() * a.grad_accum
torch.nn.utils.clip_grad_norm_(judge.model.parameters(), 1.0)
opt.step()
sched.step()
if step % 25 == 0 or step == a.steps - 1:
print(f"step {step:5d} loss(ema) {running:.4f} lr {sched.get_last_lr()[0]:.2e} {time.time() - t0:.0f}s",
flush=True)
# Save periodically. Saving only at the end means a session that is killed at its wall --
# which a long fine-tune on a fixed-length runtime can genuinely hit -- yields nothing at
# all, discarding every GPU-hour spent. This is NOT checkpoint selection: the adapter that
# gets evaluated is still whichever one the fixed budget ends on, and nothing here consults
# cal or test.
if a.checkpoint_every and (step + 1) % a.checkpoint_every == 0:
judge.model.save_pretrained(a.out)
Path(a.out).joinpath("PROGRESS.json").write_text(
json.dumps({"step": step + 1, "of": a.steps, "loss_ema": running}))
print(f" checkpointed at step {step + 1}", flush=True)
judge.model.save_pretrained(a.out)
print(f"saved adapter to {a.out}")
if __name__ == "__main__":
main()
|