training/VLA/train_cot_belief_v2.py

#!/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} </|BELIEF|> <|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": "</|BELIEF|>",
                     "actions": ["<|ALERT|>","<|OBSERVE|>","<|SILENT|>"]}, f, indent=2)
    print(f"[done] final -> {final}")


if __name__ == "__main__":
    main()