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