#!/usr/bin/env python3 """VLAlert-X v2 SFT on Qwen3-VL-4B-Instruct. Adapts training/VLA/train_cot_belief.py to use CoTBeliefDatasetV2 with the new prompt format where BELIEF tags wrap per-frame REASONING TEXT and action tokens sit AFTER the closing tag. Per-frame assistant string: <|BELIEF|> {reasoning text} <|ACTION_i|> (×8) CE loss is on all assistant tokens. Action token positions optionally get extra weight via --action_token_weight (default 2.0). Run: python -m training.VLA.train_cot_belief_v2 \ --train_jsonl data/cot_corpus_v2/vlalert_x_perframe_v2_train.jsonl \ --val_jsonl data/cot_corpus_v2/vlalert_x_perframe_v2_val.jsonl \ --out_dir checkpoints/sft_x_v2 \ --epochs 5 --batch_size 1 --grad_accum 4 \ --lora_r 128 --lora_alpha 32 --lr 1e-4 For two-stage LR ("broad + fine"): Run once with --lr 1e-4 --epochs 3, then re-run with --resume checkpoints/sft_x_v2/best --lr 2e-5 --epochs 2. """ from __future__ import annotations import sys from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parents[2])) # Conv3d→Linear PR patch (PR/qwen3vl_patch_embed_conv3d_slowdown.md). # Must run BEFORE any Qwen3VL import — patches the class-level forward so # every later .from_pretrained() call picks up the fast Linear path. import torch # noqa: F401 — keep early so patch can typecheck from tools import run_train_cot_belief_fast # noqa: F401 (side-effect: applies patch) import argparse import json import math from functools import partial from torch.optim import AdamW from torch.utils.data import DataLoader from tqdm import tqdm 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_belief_dataset_v2 import ( CoTBeliefDatasetV2, CollatorV2, ALL_SPECIAL, ) def parse_args(): ap = argparse.ArgumentParser() ap.add_argument("--model_name", default="PROJECT_ROOT/models/Qwen3-VL-4B-Instruct") ap.add_argument("--train_jsonl", default="data/cot_corpus_v2/vlalert_x_perframe_v2_train.jsonl") ap.add_argument("--val_jsonl", default="data/cot_corpus_v2/vlalert_x_perframe_v2_val.jsonl") ap.add_argument("--out_dir", required=True) ap.add_argument("--lora_r", type=int, default=128) ap.add_argument("--lora_alpha", type=int, default=32) ap.add_argument("--lora_dropout", type=float, default=0.05) ap.add_argument("--lr", type=float, default=1e-4) ap.add_argument("--epochs", type=int, default=5) 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("--action_token_weight", type=float, default=2.0, help="Extra CE weight on the 3 action token positions") ap.add_argument("--max_samples", type=int, default=0, help="Cap dataset size for smoke (0 = all)") ap.add_argument("--log_every", type=int, default=20) ap.add_argument("--save_every_epoch", action="store_true") ap.add_argument("--seed", type=int, default=0) ap.add_argument("--resume", type=str, default="", help="Warm-start LoRA from this adapter directory") return ap.parse_args() def add_special_tokens_and_resize(processor, model): tok = processor.tokenizer before = len(tok) added = tok.add_special_tokens({"additional_special_tokens": ALL_SPECIAL}) after = len(tok) print(f"[tokens] vocab {before} → {after} ({added} new)") if added == 0: return model.resize_token_embeddings(after) emb = model.get_input_embeddings() with torch.no_grad(): mean_vec = emb.weight[:before].mean(dim=0) for tok_str in ALL_SPECIAL: tid = tok.convert_tokens_to_ids(tok_str) emb.weight[tid] = mean_vec + 0.01 * torch.randn_like(mean_vec) out_emb = model.get_output_embeddings() if out_emb is not None and out_emb.weight.data_ptr() != emb.weight.data_ptr(): with torch.no_grad(): mean_out = out_emb.weight[:before].mean(dim=0) for tok_str in ALL_SPECIAL: tid = tok.convert_tokens_to_ids(tok_str) out_emb.weight[tid] = mean_out + 0.01 * torch.randn_like(mean_out) def weighted_ce_loss(logits, labels, action_mask, action_weight: float): """Causal-LM CE on labels with extra weight at action_mask=True positions. CRITICAL: applies the standard next-token shift — position t's logits predict position (t+1)'s label. Forgetting this shift collapses the objective to a trivial copy task (the answer is in the input via the residual stream), driving the train loss to near-zero while the model never learns next-token prediction. Args: logits: [B, T, V] labels: [B, T] (-100 at masked positions) action_mask: [B, T] (True at the position holding an action token) """ # Shift so that predicting at position t targets label at position t+1. shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() # Action mask aligns to the LABEL side (the action token at t+1). shift_amask = action_mask[..., 1:].contiguous() V = shift_logits.size(-1) flat_logits = shift_logits.view(-1, V) flat_labels = shift_labels.view(-1) flat_amask = shift_amask.view(-1) valid = flat_labels != -100 if not valid.any(): return flat_logits.sum() * 0.0 loss_per = torch.nn.functional.cross_entropy( flat_logits[valid], flat_labels[valid], reduction="none") w = torch.where(flat_amask[valid], torch.full_like(loss_per, action_weight, dtype=loss_per.dtype), torch.ones_like(loss_per)) return (loss_per * w).sum() / w.sum() 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", ) add_special_tokens_and_resize(processor, model) # Freeze vision tower for attr in ("visual", "vision_tower"): if hasattr(model, attr): for p in getattr(model, attr).parameters(): p.requires_grad = False if hasattr(model, "enable_input_require_grads"): model.enable_input_require_grads() model.gradient_checkpointing_enable( gradient_checkpointing_kwargs={"use_reentrant": False}) if args.resume: from peft import PeftModel print(f"[resume] loading PEFT adapter from {args.resume}") model = PeftModel.from_pretrained(model, args.resume, is_trainable=True) else: 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", modules_to_save=["embed_tokens", "lm_head"], ) model = get_peft_model(model, lora_cfg) model.print_trainable_parameters() model.to("cuda") model.config.use_cache = False ds = CoTBeliefDatasetV2( jsonl_path=args.train_jsonl, processor=processor, n_frames=args.n_frames, resize_short=args.resize_short, max_len=args.max_len, action_token_weight=args.action_token_weight, ) if args.max_samples > 0 and len(ds) > args.max_samples: from torch.utils.data import Subset ds = Subset(ds, list(range(args.max_samples))) print(f"[smoke] truncated to {len(ds)}") print(f"[train] dataset size = {len(ds)}") collator = CollatorV2(processor, n_frames=args.n_frames) dl = DataLoader(ds, batch_size=args.batch_size, shuffle=True, num_workers=0, collate_fn=collator, 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_updates = math.ceil(len(dl) * args.epochs / args.grad_accum) warmup = max(1, int(total_updates * args.warmup_ratio)) sched = get_cosine_schedule_with_warmup(opt, warmup, total_updates) print(f"[train] total_updates={total_updates} warmup={warmup} 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}", ncols=80, leave=True) running = 0.0; running_n = 0 for step, batch in pbar: input_ids = batch["input_ids"].to("cuda", non_blocking=True) labels = batch["labels"].to("cuda", non_blocking=True) amask = batch["action_token_mask"].to("cuda", non_blocking=True) attn = batch.get("attention_mask") if attn is not None: attn = attn.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) fwd_kwargs = dict(input_ids=input_ids, pixel_values=pix, image_grid_thw=grid) if attn is not None: fwd_kwargs["attention_mask"] = attn out = model(**fwd_kwargs) loss = weighted_ce_loss( out.logits, labels, amask, args.action_token_weight ) / args.grad_accum loss.backward() running += loss.detach().float().item() * args.grad_accum 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"[save] -> {ep_dir}") # Final "best" 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) with (final / "belief_tokens.json").open("w") as f: json.dump({"special_tokens": ALL_SPECIAL, "belief_open": "<|BELIEF|>", "belief_close": "", "actions": ["<|ALERT|>","<|OBSERVE|>","<|SILENT|>"]}, f, indent=2) print(f"[done] final -> {final}") if __name__ == "__main__": main()