#!/usr/bin/env python3 """LoRA DPO for Qwen2.5-VL on fine-grained video caption negatives. Preference: chosen = correct dense caption, rejected = same caption with ONE fine-grained attribute flipped (from build_fg_negatives.py + Stage D hard filter). Targets the discrimination bottleneck found in the probes. Memory-safe by construction: - LoRA on the LLM decoder layers only; vision tower frozen (no grad through it). - reference logprobs from the SAME model with the adapter disabled (no 2nd copy). - the video is processed ONCE per sample; chosen/rejected reuse its pixel_values and only append different response tokens. - batch size 1 + grad accumulation + gradient checkpointing, bf16. Run a memory check first: python dpo_train.py ... --max_steps 4 --dry_run """ from __future__ import annotations import argparse, json, os, time from pathlib import Path os.environ.setdefault("HF_HOME", "/mnt/local-fast/opd_zt/hf_cache") os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") import numpy as np, torch import torch.nn.functional as F ROOT = Path("/mnt/local-fast/opd_zt") M7B = str(ROOT / "hf_cache/hub/models--Qwen--Qwen2.5-VL-7B-Instruct/snapshots/" "cc594898137f460bfe9f0759e9844b3ce807cfb5") NF = 32 MAXP = 128 * 28 * 28 MINP = 16 * 28 * 28 def decode_frames(path): from decord import VideoReader, cpu try: vr = VideoReader(path, ctx=cpu(0), num_threads=1) except Exception: return None total = len(vr) if total < 1: return None idx = np.linspace(0, total - 1, NF).round().astype(int).clip(0, total - 1) return vr.get_batch(idx.tolist()).asnumpy() def main(): ap = argparse.ArgumentParser() ap.add_argument("--data", default=str(ROOT / "data/fg_dpo/train.jsonl")) ap.add_argument("--out", default=str(ROOT / "ckpts/fg_dpo_lora")) ap.add_argument("--beta", type=float, default=0.1) ap.add_argument("--lr", type=float, default=5e-6) ap.add_argument("--epochs", type=int, default=1) ap.add_argument("--grad_accum", type=int, default=8) ap.add_argument("--lora_r", type=int, default=16) ap.add_argument("--lora_alpha", type=int, default=32) ap.add_argument("--max_resp_tokens", type=int, default=320) ap.add_argument("--max_steps", type=int, default=0, help="cap optimizer steps (dry run)") ap.add_argument("--dry_run", action="store_true", help="report peak memory, no checkpoint save") ap.add_argument("--device", default="cuda:0") args = ap.parse_args() from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration from peft import LoraConfig, get_peft_model proc = AutoProcessor.from_pretrained(M7B, trust_remote_code=True, max_pixels=MAXP, min_pixels=MINP) tok = proc.tokenizer model = Qwen2_5_VLForConditionalGeneration.from_pretrained( M7B, torch_dtype=torch.bfloat16, attn_implementation="sdpa", trust_remote_code=True ).to(args.device) # freeze vision tower (no grad, no LoRA there) -> big activation-memory saving n_vis = 0 for name, p in model.named_parameters(): if "visual" in name: p.requires_grad_(False); n_vis += 1 # LoRA on LLM decoder layers only (regex avoids the vision blocks) lcfg = LoraConfig( r=args.lora_r, lora_alpha=args.lora_alpha, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM", target_modules=r".*\.layers\.\d+\.(self_attn|mlp)\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)$", ) model = get_peft_model(model, lcfg) model.print_trainable_parameters() model.gradient_checkpointing_enable() model.enable_input_require_grads() model.config.use_cache = False model.train() rows = [json.loads(l) for l in open(args.data)] print(f"[data] {len(rows)} preference pairs") eos = tok.eos_token or "<|im_end|>" def build_sample(r): frames = decode_frames(r["video_abs"]) if frames is None: return None from PIL import Image pil = [Image.fromarray(f) for f in frames] prompt_text = proc.apply_chat_template( [{"role": "user", "content": [{"type": "video"}, {"type": "text", "text": r["prompt"]}]}], tokenize=False, add_generation_prompt=True) enc = proc(text=[prompt_text], videos=[pil], return_tensors="pt") pid = enc["input_ids"][0] Lp = pid.shape[0] def cat_resp(resp): rid = tok(resp.strip() + eos, add_special_tokens=False, return_tensors="pt")["input_ids"][0] rid = rid[: args.max_resp_tokens] ids = torch.cat([pid, rid]) labels = torch.full((ids.shape[0],), -100, dtype=torch.long) labels[Lp:] = ids[Lp:] return ids, labels cw, lw = cat_resp(r["chosen"]) cl, ll = cat_resp(r["rejected"]) vid_kw = {"pixel_values_videos": enc["pixel_values_videos"], "video_grid_thw": enc["video_grid_thw"]} if "second_per_grid_ts" in enc: vid_kw["second_per_grid_ts"] = enc["second_per_grid_ts"] return (cw, lw), (cl, ll), vid_kw def seq_logp(ids, labels, vid_kw): ids = ids.unsqueeze(0).to(args.device) attn = torch.ones_like(ids) kw = {k: (v.to(args.device) if torch.is_tensor(v) else v) for k, v in vid_kw.items()} out = model(input_ids=ids, attention_mask=attn, **kw) logits = out.logits[:, :-1] lab = labels.unsqueeze(0).to(args.device)[:, 1:] mask = lab != -100 lp = F.log_softmax(logits.float(), -1).gather(-1, lab.clamp(min=0).unsqueeze(-1)).squeeze(-1) return (lp * mask).sum() opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=args.lr) n_steps = 0 t0 = time.time() accum = 0 opt.zero_grad() for ep in range(args.epochs): for i, r in enumerate(rows): s = build_sample(r) if s is None: continue (cw, lw), (cl, ll), vid = s # policy logprobs (with grad, adapter on) pol_w = seq_logp(cw, lw, vid) pol_l = seq_logp(cl, ll, vid) # reference logprobs (adapter off, no grad) with torch.no_grad(), model.disable_adapter(): ref_w = seq_logp(cw, lw, vid) ref_l = seq_logp(cl, ll, vid) logits = args.beta * ((pol_w - ref_w) - (pol_l - ref_l)) loss = -F.logsigmoid(logits) / args.grad_accum loss.backward() accum += 1 if accum % args.grad_accum == 0: torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], 1.0) opt.step(); opt.zero_grad(); n_steps += 1 acc = (logits > 0).float().item() mem = torch.cuda.max_memory_allocated(args.device) / 1e9 print(f"[step {n_steps}] loss={loss.item()*args.grad_accum:.4f} " f"margin={logits.item():.3f} chosen>rej={acc:.0f} " f"peak_mem={mem:.1f}GB ({time.time()-t0:.0f}s, {i+1}/{len(rows)})", flush=True) if args.max_steps and n_steps >= args.max_steps: print(f"[done] reached max_steps={args.max_steps}, peak_mem={mem:.1f}GB") if not args.dry_run: model.save_pretrained(args.out) return if not args.dry_run: Path(args.out).mkdir(parents=True, exist_ok=True) model.save_pretrained(args.out) print(f"[save] LoRA adapter -> {args.out}") print(f"[done] peak_mem={torch.cuda.max_memory_allocated(args.device)/1e9:.1f}GB") if __name__ == "__main__": main()