training/VLA/train_vla_cot.py

"""LoRA fine-tune Qwen2.5-VL-3B-Instruct on Nexar CoT JSON outputs.

Minimal trainer — single-GPU bf16 LoRA. Smoke-test friendly.

Run:
  python -m training.VLA.train_vla_cot \
      --cot_jsonl data/vla_cot/train_cot.jsonl \
      --video_dir nexar-collision-prediction/train \
      --out_dir checkpoints/VLA/qwen_cot_smoke \
      --lora_r 32 --lr 2e-4 --epochs 1 --batch_size 1 --grad_accum 4
"""
from __future__ import annotations

import argparse
import json
import math
import os
import sys
from functools import partial
from pathlib import Path

import torch
from torch.optim import AdamW
from torch.utils.data import DataLoader
from tqdm import tqdm

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from peft import LoraConfig, get_peft_model
from transformers import AutoProcessor, AutoModelForImageTextToText
from transformers.optimization import get_cosine_schedule_with_warmup

from training.VLA.cot_dataset import NexarCoTDataset, collate_fn


def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model_name", default="Qwen/Qwen2.5-VL-3B-Instruct")
    ap.add_argument("--cot_jsonl", required=True)
    ap.add_argument("--video_dir", required=True)
    ap.add_argument("--out_dir", required=True)
    ap.add_argument("--lora_r", type=int, default=32)
    ap.add_argument("--lora_alpha", type=int, default=16)
    ap.add_argument("--lora_dropout", type=float, default=0.05)
    ap.add_argument("--lr", type=float, default=2e-4)
    ap.add_argument("--epochs", type=int, default=1)
    ap.add_argument("--batch_size", type=int, default=1)
    ap.add_argument("--grad_accum", type=int, default=4)
    ap.add_argument("--warmup_ratio", type=float, default=0.03)
    ap.add_argument("--n_frames", type=int, default=8)
    ap.add_argument("--resize_short", type=int, default=336)
    ap.add_argument("--max_len", type=int, default=4096)
    ap.add_argument("--supervise", default="assistant", choices=["assistant", "verdict_only"],
                    help="'assistant' = supervise all CoT tokens (original); "
                         "'verdict_only' = supervise ONLY the yes/no token (concentrated gradient)")
    ap.add_argument("--log_every", type=int, default=1)
    ap.add_argument("--save_every_epoch", action="store_true")
    ap.add_argument("--seed", type=int, default=0)
    return ap.parse_args()


def main():
    args = parse_args()
    torch.manual_seed(args.seed)
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    print(f"[train] loading processor/model from {args.model_name}")
    processor = AutoProcessor.from_pretrained(args.model_name, trust_remote_code=True)
    model = AutoModelForImageTextToText.from_pretrained(
        args.model_name,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        attn_implementation="sdpa",
    )

    # Freeze the vision tower — LoRA only on the LLM.
    if hasattr(model, "visual"):
        for p in model.visual.parameters():
            p.requires_grad = False

    # CRITICAL for (frozen vision + LoRA on LLM + gradient_checkpointing):
    # force input embeddings to require grad so backward can flow through
    # the checkpointed LLM layers.
    if hasattr(model, "enable_input_require_grads"):
        model.enable_input_require_grads()
    else:
        # fallback: register a hook on the input embedding
        try:
            emb = model.get_input_embeddings()
            def _make_inputs_require_grad(module, inp, out):
                out.requires_grad_(True)
            emb.register_forward_hook(_make_inputs_require_grad)
        except Exception:
            pass

    model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})

    lora_cfg = LoraConfig(
        r=args.lora_r,
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        bias="none",
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, lora_cfg)
    model.print_trainable_parameters()
    model.to("cuda")

    # Keep inputs that don't require grad in bf16 to match the model.
    model.config.use_cache = False

    ds = NexarCoTDataset(
        jsonl_path=args.cot_jsonl,
        video_dir=args.video_dir,
        processor=processor,
        n_frames=args.n_frames,
        resize_short=args.resize_short,
        max_len=args.max_len,
        supervise=args.supervise,
    )
    print(f"[train] dataset size = {len(ds)}")
    if len(ds) == 0:
        raise SystemExit("empty dataset — check your CoT jsonl")

    pad_id = processor.tokenizer.pad_token_id or processor.tokenizer.eos_token_id
    dl = DataLoader(
        ds,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=0,  # Qwen processor is not fork-safe; keep single-process
        collate_fn=partial(collate_fn, pad_token_id=pad_id),
        pin_memory=True,
    )

    trainable = [p for p in model.parameters() if p.requires_grad]
    opt = AdamW(trainable, lr=args.lr, betas=(0.9, 0.95), weight_decay=0.0)
    total_steps = math.ceil(len(dl) * args.epochs / args.grad_accum)
    warmup_steps = max(1, int(total_steps * args.warmup_ratio))
    sched = get_cosine_schedule_with_warmup(opt, warmup_steps, total_steps)
    print(f"[train] total_updates={total_steps} warmup={warmup_steps} lr={args.lr}")

    global_step = 0
    model.train()
    for epoch in range(args.epochs):
        pbar = tqdm(enumerate(dl), total=len(dl), desc=f"ep{epoch}")
        running = 0.0
        running_n = 0
        for step, batch in pbar:
            input_ids = batch["input_ids"].to("cuda", non_blocking=True)
            attn = batch["attention_mask"].to("cuda", non_blocking=True)
            labels = batch["labels"].to("cuda", non_blocking=True)
            pix = batch["pixel_values"].to("cuda", dtype=torch.bfloat16, non_blocking=True)
            grid = batch["image_grid_thw"].to("cuda", non_blocking=True)

            out = model(
                input_ids=input_ids,
                attention_mask=attn,
                labels=labels,
                pixel_values=pix,
                image_grid_thw=grid,
            )
            loss = out.loss / args.grad_accum
            loss.backward()
            running += out.loss.detach().float().item()
            running_n += 1

            if (step + 1) % args.grad_accum == 0 or (step + 1) == len(dl):
                torch.nn.utils.clip_grad_norm_(trainable, 1.0)
                opt.step()
                sched.step()
                opt.zero_grad(set_to_none=True)
                global_step += 1
                if global_step % args.log_every == 0:
                    pbar.set_postfix(loss=running / max(1, running_n), lr=sched.get_last_lr()[0])
                    running, running_n = 0.0, 0

        if args.save_every_epoch or epoch == args.epochs - 1:
            ep_dir = out_dir / f"epoch_{epoch}"
            ep_dir.mkdir(parents=True, exist_ok=True)
            model.save_pretrained(ep_dir)
            processor.save_pretrained(ep_dir)
            with (ep_dir / "train_args.json").open("w") as f:
                json.dump(vars(args), f, indent=2)
            print(f"[train] saved -> {ep_dir}")

    # final save
    final = out_dir / "best"
    final.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(final)
    processor.save_pretrained(final)
    with (final / "train_args.json").open("w") as f:
        json.dump(vars(args), f, indent=2)
    print(f"[train] done. final -> {final}")


if __name__ == "__main__":
    main()