"""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()