train.py

"""
ArcisVLM Training — Two-stage training for VL-JEPA + MoE.

Stage 1: JEPA Pretraining — InfoNCE contrastive learning
Stage 2: Supervised Finetuning — MoE decoder on VQA data

Supports Apple MPS (M1/M2) and CUDA.
"""

import os
import time
import yaml
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm

from model.vlm import VLJEPAModel
from model.tokenizer import BPETokenizer
from data.dataset import CaptionDataset, VQADataset


def get_device() -> torch.device:
    """Get best available device."""
    if torch.cuda.is_available():
        return torch.device("cuda")
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def train_stage1(model: VLJEPAModel, dataloader: DataLoader, config: dict, device: torch.device):
    """
    Stage 1: JEPA Pretraining with InfoNCE loss.

    Trains X-Encoder + Predictor + Y-Encoder to align visual+query embeddings
    with text embeddings in a shared 1536-D space.
    """
    cfg = config["train_stage1"]
    model.train()
    model.to(device)

    # Y-Encoder gets slower learning rate
    y_encoder_params = list(model.y_encoder.parameters())
    other_params = [p for n, p in model.named_parameters()
                    if not n.startswith("y_encoder") and not n.startswith("decoder") and p.requires_grad]

    optimizer = torch.optim.AdamW([
        {"params": other_params, "lr": cfg["learning_rate"]},
        {"params": y_encoder_params, "lr": cfg["learning_rate"] * config["y_encoder"]["lr_multiplier"]},
    ], weight_decay=0.01)

    print(f"\n{'='*60}")
    print(f"Stage 1: JEPA Pretraining")
    print(f"Device: {device}")
    print(f"Batch size: {cfg['batch_size']}")
    print(f"Learning rate: {cfg['learning_rate']}")
    print(f"{'='*60}\n")

    global_step = 0
    for epoch in range(cfg["max_epochs"]):
        epoch_loss = 0.0
        num_batches = 0

        pbar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{cfg['max_epochs']}")
        for batch in pbar:
            images = batch["image"].to(device)
            caption_ids = batch["caption_ids"].to(device)
            caption_mask = batch["caption_mask"].to(device)

            # Forward pass (no query for captioning, just image → caption embedding)
            output = model.forward_stage1(
                images=images,
                query_ids=None,
                query_padding_mask=None,
                answer_ids=caption_ids,
                answer_padding_mask=caption_mask,
            )

            loss = output["loss"]

            # Backward pass
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg["gradient_clip"])
            optimizer.step()

            epoch_loss += loss.item()
            num_batches += 1
            global_step += 1

            pbar.set_postfix({"loss": f"{loss.item():.4f}"})

            if global_step % 100 == 0:
                avg_loss = epoch_loss / num_batches
                print(f"  Step {global_step}: avg_loss={avg_loss:.4f}")

        avg_loss = epoch_loss / max(num_batches, 1)
        print(f"Epoch {epoch+1} complete: avg_loss={avg_loss:.4f}")

        # Save checkpoint
        if (epoch + 1) % 5 == 0:
            ckpt_path = f"checkpoints/stage1_epoch{epoch+1}.pt"
            os.makedirs("checkpoints", exist_ok=True)
            torch.save({
                "epoch": epoch + 1,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "loss": avg_loss,
            }, ckpt_path)
            print(f"  Saved checkpoint: {ckpt_path}")


def train_stage2(model: VLJEPAModel, dataloader: DataLoader, config: dict, device: torch.device):
    """
    Stage 2: Supervised Finetuning with MoE Decoder.

    Freezes X-Encoder, trains Predictor + MoE Decoder on VQA data.
    """
    cfg = config["train_stage2"]
    model.freeze_x_encoder()
    model.train()
    model.to(device)

    # Only train predictor + decoder parameters
    trainable_params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.AdamW(trainable_params, lr=cfg["learning_rate"], weight_decay=0.01)

    # Cosine annealing scheduler
    total_steps = cfg["max_epochs"] * len(dataloader)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_steps)

    print(f"\n{'='*60}")
    print(f"Stage 2: Supervised Finetuning (MoE Decoder)")
    print(f"Device: {device}")
    print(f"Batch size: {cfg['batch_size']}")
    print(f"Learning rate: {cfg['learning_rate']}")
    print(f"X-Encoder: FROZEN")
    params = model.count_parameters()
    print(f"Trainable params: {params['trainable']:,}")
    print(f"{'='*60}\n")

    global_step = 0
    for epoch in range(cfg["max_epochs"]):
        epoch_loss = 0.0
        epoch_decode_loss = 0.0
        epoch_lb_loss = 0.0
        num_batches = 0

        pbar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{cfg['max_epochs']}")
        for batch in pbar:
            images = batch["image"].to(device)
            q_ids = batch["question_ids"].to(device)
            q_mask = batch["question_mask"].to(device)
            a_ids = batch["answer_ids"].to(device)

            output = model.forward_stage2(
                images=images,
                query_ids=q_ids,
                query_padding_mask=q_mask,
                answer_ids=a_ids,
                load_balance_weight=cfg["load_balance_weight"],
            )

            loss = output["loss"]

            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg["gradient_clip"])
            optimizer.step()
            scheduler.step()

            epoch_loss += loss.item()
            epoch_decode_loss += output["decode_loss"].item()
            epoch_lb_loss += output["load_balance_loss"].item()
            num_batches += 1
            global_step += 1

            pbar.set_postfix({
                "loss": f"{loss.item():.4f}",
                "decode": f"{output['decode_loss'].item():.4f}",
                "lb": f"{output['load_balance_loss'].item():.4f}",
            })

        avg_loss = epoch_loss / max(num_batches, 1)
        avg_decode = epoch_decode_loss / max(num_batches, 1)
        avg_lb = epoch_lb_loss / max(num_batches, 1)
        print(f"Epoch {epoch+1}: loss={avg_loss:.4f}, decode={avg_decode:.4f}, lb={avg_lb:.4f}")

        # Save checkpoint
        if (epoch + 1) % 5 == 0:
            ckpt_path = f"checkpoints/stage2_epoch{epoch+1}.pt"
            os.makedirs("checkpoints", exist_ok=True)
            torch.save({
                "epoch": epoch + 1,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "loss": avg_loss,
            }, ckpt_path)
            print(f"  Saved checkpoint: {ckpt_path}")


def main():
    # Load config
    with open("configs/default.yaml") as f:
        config = yaml.safe_load(f)

    device = get_device()
    print(f"Using device: {device}")

    # Initialize tokenizer
    tokenizer = BPETokenizer(vocab_size=config["decoder"]["vocab_size"])

    # Check if trained tokenizer exists
    tokenizer_path = "checkpoints/tokenizer.json"
    if os.path.exists(tokenizer_path):
        tokenizer.load(tokenizer_path)
        print(f"Loaded tokenizer from {tokenizer_path} ({len(tokenizer)} tokens)")
    else:
        print("WARNING: No trained tokenizer found. Train one first or provide training texts.")
        print("Using untrained tokenizer with special tokens only.")

    # Initialize model
    model = VLJEPAModel(config)
    params = model.count_parameters()
    print(f"\nModel parameters:")
    for name, count in params.items():
        print(f"  {name}: {count:,}")

    # Stage 1: JEPA Pretraining
    stage1_data_dir = config["data"]["flickr8k_dir"]
    if os.path.exists(stage1_data_dir):
        print(f"\nLoading Stage 1 dataset from {stage1_data_dir}...")
        caption_dataset = CaptionDataset(
            image_dir=os.path.join(stage1_data_dir, "Images"),
            captions_file=os.path.join(stage1_data_dir, "captions.txt"),
            tokenizer=tokenizer,
            img_size=config["vision"]["img_size"],
        )
        caption_loader = DataLoader(
            caption_dataset,
            batch_size=config["train_stage1"]["batch_size"],
            shuffle=True,
            num_workers=config["data"]["num_workers"],
            pin_memory=True,
        )
        print(f"Stage 1 dataset: {len(caption_dataset)} samples")
        train_stage1(model, caption_loader, config, device)
    else:
        print(f"\nSkipping Stage 1: {stage1_data_dir} not found")
        print("Download Flickr8k dataset first.")

    # Stage 2: Supervised Finetuning
    stage2_data_dir = config["data"]["vqav2_dir"]
    if os.path.exists(stage2_data_dir):
        print(f"\nLoading Stage 2 dataset from {stage2_data_dir}...")
        vqa_dataset = VQADataset(
            image_dir=os.path.join(stage2_data_dir, "images"),
            questions_file=os.path.join(stage2_data_dir, "questions.json"),
            annotations_file=os.path.join(stage2_data_dir, "annotations.json"),
            tokenizer=tokenizer,
            img_size=config["vision"]["img_size"],
        )
        vqa_loader = DataLoader(
            vqa_dataset,
            batch_size=config["train_stage2"]["batch_size"],
            shuffle=True,
            num_workers=config["data"]["num_workers"],
            pin_memory=True,
        )
        print(f"Stage 2 dataset: {len(vqa_dataset)} samples")
        train_stage2(model, vqa_loader, config, device)
    else:
        print(f"\nSkipping Stage 2: {stage2_data_dir} not found")
        print("Download VQAv2 dataset first.")


if __name__ == "__main__":
    main()