scripts/dpo_train.py

#!/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()