add train.py

train.py

@@ -0,0 +1,258 @@
+import re
+import random
+import argparse
+from dataclasses import dataclass, field
+from typing import List
+
+import torch
+import wandb
+from tqdm import tqdm
+from PIL import Image
+from datasets import load_dataset
+from transformers import (
+    Qwen2_5_VLForConditionalGeneration,
+    AutoProcessor,
+    BitsAndBytesConfig,
+)
+from qwen_vl_utils import process_vision_info
+from peft import LoraConfig, get_peft_model
+from trl import SFTConfig, SFTTrainer
+
+
+def extract_question(raw_text: str) -> str:
+    pattern = r"<\|start_header_id\|>user<\|end_header_id\|>\s*(.*?)\s*<\|eot_id\|>"
+    m = re.search(pattern, raw_text, re.DOTALL)
+    return m.group(1).strip() if m else raw_text.strip()
+
+def format_data_spacethinker(sample):
+    system_message = {
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": (
+                    "You are VL-Thinking U+1F914, a helpful assistant with excellent reasoning ability.\n"
+                    "A user asks you a question, and you should try to solve it."
+                    "You should first think about the reasoning process in the mind and then provides the user with the answer.\n"
+                    "The reasoning process and answer are enclosed within <think></think> and <answer></answer> tags, respectively, i.e., <think> reasoning process here </think> <answer> answer here </answer>."
+                )
+            }
+        ]
+    }
+    formatted = [system_message]
+
+    user_msg = {"role": "user", "content": []}
+    question = extract_question(sample.get("input", ""))
+    if question:
+        user_msg["content"].append({"type": "text", "text": question})
+    images = sample.get("images") or []
+    if images:
+        user_msg["content"].append({"type": "image", "image": images[0]})
+    formatted.append(user_msg)
+
+    if sample.get("output"):
+        formatted.append({
+            "role": "assistant",
+            "content": [{"type": "text", "text": sample["output"]}]
+        })
+    return formatted
+
+
+def collate_fn(examples, processor):
+    # examples: list of formatted samples (list of message dicts)
+    texts = [processor.apply_chat_template(sample, tokenize=False) for sample in examples]
+    image_batches = [process_vision_info(sample)[0] for sample in examples]
+    batch = processor(text=texts, images=image_batches, return_tensors="pt", padding=True)
+    batch = {k: v.cpu() for k, v in batch.items()}
+
+    labels = batch["input_ids"].clone()
+    labels[labels == processor.tokenizer.pad_token_id] = -100
+
+    image_token_ids = (
+        [151652, 151653, 151655]
+        if hasattr(processor, "image_processor")
+        else [processor.tokenizer.convert_tokens_to_ids(processor.image_token)]
+    )
+    for tid in image_token_ids:
+        labels[labels == tid] = -100
+
+    batch["labels"] = labels
+    return batch
+
+
+@dataclass
+class TrainingConfig:
+    model_id: str = "UCSC-VLAA/VLAA-Thinker-Qwen2.5VL-3B"
+    lora_r: int = 128
+    lora_alpha: int = 256
+    lora_dropout: float = 0.05
+    target_modules: List[str] = field(default_factory=lambda: ["q_proj", "v_proj", "o_proj"])
+    num_train_epochs: int = 3
+    train_batch_size: int = 1
+    eval_batch_size: int = 1
+    gradient_accumulation_steps: int = 8
+    learning_rate: float = 2e-5
+    warmup_ratio: float = 0.03
+    output_dir: str = "spaceom"
+    wandb_project: str = "spaceom"
+    wandb_run_name: str = "spaceom"
+
+
+def parse_args() -> TrainingConfig:
+    default_cfg = TrainingConfig()
+    parser = argparse.ArgumentParser(description="Train a VL Spacethinker model with LoRA")
+    parser.add_argument("--model_id", default=default_cfg.model_id)
+    parser.add_argument("--lora_r", type=int, default=default_cfg.lora_r)
+    parser.add_argument("--lora_alpha", type=int, default=default_cfg.lora_alpha)
+    parser.add_argument("--lora_dropout", type=float, default=default_cfg.lora_dropout)
+    parser.add_argument(
+        "--target_modules",
+        default=','.join(default_cfg.target_modules),
+        help="Comma-separated list of target modules for LoRA"
+    )
+    parser.add_argument("--num_train_epochs", type=int, default=default_cfg.num_train_epochs)
+    parser.add_argument("--train_batch_size", type=int, default=default_cfg.train_batch_size)
+    parser.add_argument("--eval_batch_size", type=int, default=default_cfg.eval_batch_size)
+    parser.add_argument(
+        "--gradient_accumulation_steps", type=int, default=default_cfg.gradient_accumulation_steps
+    )
+    parser.add_argument("--learning_rate", type=float, default=default_cfg.learning_rate)
+    parser.add_argument("--warmup_ratio", type=float, default=default_cfg.warmup_ratio)
+    parser.add_argument("--output_dir", default=default_cfg.output_dir)
+    parser.add_argument("--wandb_project", default=default_cfg.wandb_project)
+    parser.add_argument("--wandb_run_name", default=default_cfg.wandb_run_name)
+
+    args = parser.parse_args()
+    return TrainingConfig(
+        model_id=args.model_id,
+        lora_r=args.lora_r,
+        lora_alpha=args.lora_alpha,
+        lora_dropout=args.lora_dropout,
+        target_modules=args.target_modules.split(","),
+        num_train_epochs=args.num_train_epochs,
+        train_batch_size=args.train_batch_size,
+        eval_batch_size=args.eval_batch_size,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        learning_rate=args.learning_rate,
+        warmup_ratio=args.warmup_ratio,
+        output_dir=args.output_dir,
+        wandb_project=args.wandb_project,
+        wandb_run_name=args.wandb_run_name,
+    )
+
+
+def prepare_datasets(cfg: TrainingConfig):
+    print(f"Loading dataset: SpaceThinker")
+    raw_train_spacethinker = load_dataset("remyxai/SpaceThinker", split="train")
+    raw_eval_spacethinker = load_dataset("remyxai/SpaceThinker", split="test")
+
+    print(f"Loading dataset: SpaceOm")
+    raw_train_spaceom = load_dataset("remyxai/SpaceOm", split="train")
+    raw_eval_spaceom = load_dataset("remyxai/SpaceOm", split="test")
+
+    print(f"Loading dataset: Robo2VLM")
+    raw_train_robo2vlm = load_dataset("remyxai/Robo2VLM-Reasoning", split="train")
+    raw_eval_robo2vlm = load_dataset("remyxai/Robo2VLM-Reasoning", split="test")
+
+    print("Formatting train samples…")
+    train_ds_spacethinker = [format_data_spacethinker(s) for s in tqdm(raw_train_spacethinker, desc="Train")]
+    train_ds_spaceom = [format_data_spacethinker(s) for s in tqdm(raw_train_spaceom, desc="Train")]
+    train_ds_robo2vlm = [format_data_spacethinker(s) for s in tqdm(raw_train_robo2vlm, desc="Train")]
+    print("Formatting eval samples…")
+    eval_ds_spacethinker = [format_data_spacethinker(s) for s in tqdm(raw_eval_spacethinker, desc="Eval")]
+    eval_ds_spaceom = [format_data_spacethinker(s) for s in tqdm(raw_eval_spaceom, desc="Eval")]
+    eval_ds_robo2vlm = [format_data_spacethinker(s) for s in tqdm(raw_eval_robo2vlm, desc="Eval")]
+
+    train_ds = train_ds_spacethinker + train_ds_spaceom + train_ds_robo2vlm
+    eval_ds = eval_ds_spacethinker + eval_ds_spaceom + eval_ds_robo2vlm
+    random.shuffle(train_ds)
+    random.shuffle(eval_ds)
+
+    return train_ds, eval_ds
+
+
+def prepare_model_and_optimizer(cfg: TrainingConfig):
+    bnb = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16
+    )
+    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+        cfg.model_id,
+        device_map="auto",
+        torch_dtype=torch.bfloat16,
+        quantization_config=bnb
+    )
+    processor = AutoProcessor.from_pretrained(cfg.model_id)
+
+    peft_cfg = LoraConfig(
+        r=cfg.lora_r,
+        lora_alpha=cfg.lora_alpha,
+        lora_dropout=cfg.lora_dropout,
+        bias="none",
+        target_modules=cfg.target_modules,
+        task_type="CAUSAL_LM",
+    )
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    peft_model = get_peft_model(model, peft_cfg).to(device)
+    peft_model.print_trainable_parameters()
+    return peft_model, processor, peft_cfg
+
+
+def main():
+    cfg = parse_args()
+    train_ds, eval_ds = prepare_datasets(cfg)
+    model, processor, peft_cfg = prepare_model_and_optimizer(cfg)
+
+    sft_args = SFTConfig(
+        output_dir=cfg.output_dir,
+        num_train_epochs=cfg.num_train_epochs,
+        per_device_train_batch_size=cfg.train_batch_size,
+        per_device_eval_batch_size=cfg.eval_batch_size,
+        gradient_accumulation_steps=cfg.gradient_accumulation_steps,
+        gradient_checkpointing=True,
+        optim="adamw_torch_fused",
+        learning_rate=cfg.learning_rate,
+        lr_scheduler_type="constant",
+        logging_steps=10,
+        eval_steps=10,
+        eval_strategy="steps",
+        save_strategy="steps",
+        save_steps=20,
+        metric_for_best_model="eval_loss",
+        greater_is_better=False,
+        load_best_model_at_end=True,
+        bf16=True,
+        tf32=True,
+        max_grad_norm=0.3,
+        warmup_ratio=cfg.warmup_ratio,
+        gradient_checkpointing_kwargs={"use_reentrant": False},
+        push_to_hub=True,
+        report_to="wandb",
+        dataset_kwargs={"skip_prepare_dataset": True},
+    )
+    sft_args.remove_unused_columns = False
+
+    wandb.init(
+        project=cfg.wandb_project,
+        name=cfg.wandb_run_name,
+        config=sft_args,
+    )
+
+    trainer = SFTTrainer(
+        model=model,
+        args=sft_args,
+        train_dataset=train_ds,
+        eval_dataset=eval_ds,
+        data_collator=lambda ex: collate_fn(ex, processor),
+        peft_config=peft_cfg,
+        tokenizer=processor.tokenizer,
+    )
+
+    trainer.train()
+    trainer.save_model(cfg.output_dir)
+
+
+if __name__ == "__main__":
+    main()