train.py

"""
Training script for PaliGemma with support for fine-tuning and QLoRA.
"""
import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.optim import AdamW
try:
    from transformers import get_linear_schedule_with_warmup
except ImportError:
    # Fallback implementation if transformers not available
    from torch.optim.lr_scheduler import LambdaLR
    def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):
        """
        Create a schedule with a learning rate that decreases linearly from the initial lr
        set in the optimizer to 0, after a warmup period during which it increases
        linearly from 0 to the initial lr set in the optimizer.
        """
        def lr_lambda(current_step: int):
            if current_step < num_warmup_steps:
                return float(current_step) / float(max(1, num_warmup_steps))
            return max(0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps)))
        
        return LambdaLR(optimizer, lr_lambda, last_epoch)
import logging
from tqdm import tqdm
from typing import Optional, Dict, Any
import json
from pathlib import Path

from modeling_gemma import PaliGemmaForConditionalGeneration
from processing_paligemma import PaliGemmaProcessor
from utils import load_hf_model
from data_utils import PaliGemmaDataset, create_dataloader
from config import Config, TrainingConfig
from utils_optimization import apply_gradient_checkpointing, count_parameters
import wandb

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class QLoRALinear(nn.Module):
    """
    QLoRA (Quantized Low-Rank Adaptation) linear layer.
    Efficient fine-tuning with 4-bit quantization and LoRA adapters.
    """
    def __init__(
        self,
        base_layer: nn.Linear,
        r: int = 64,
        alpha: int = 16,
        dropout: float = 0.05,
    ):
        super().__init__()
        self.base_layer = base_layer
        self.r = r
        self.alpha = alpha
        self.scaling = alpha / r
        
        # LoRA adapters
        self.lora_A = nn.Parameter(torch.randn(r, base_layer.in_features) * 0.02)
        self.lora_B = nn.Parameter(torch.zeros(base_layer.out_features, r))
        self.lora_dropout = nn.Dropout(dropout)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Base layer output (quantized if applicable)
        base_output = self.base_layer(x)
        
        # LoRA adaptation
        lora_output = self.lora_dropout(x) @ self.lora_A.T @ self.lora_B.T
        lora_output = lora_output * self.scaling
        
        return base_output + lora_output


def apply_qlora_to_model(
    model: nn.Module,
    target_modules: Optional[list] = None,
    r: int = 64,
    alpha: int = 16,
    dropout: float = 0.05,
) -> nn.Module:
    """
    Apply QLoRA to specified modules in the model.
    
    Args:
        model: Model to apply QLoRA to
        target_modules: List of module names to target (e.g., ['q_proj', 'v_proj'])
        r: LoRA rank
        alpha: LoRA alpha scaling factor
        dropout: LoRA dropout rate
    
    Returns:
        Model with QLoRA applied
    """
    if target_modules is None:
        target_modules = ['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj']
    
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            # Check if this module should be replaced
            module_name = name.split('.')[-1]
            if module_name in target_modules:
                # Replace with QLoRA version
                parent_name = '.'.join(name.split('.')[:-1])
                parent_module = model
                for part in parent_name.split('.'):
                    if part:
                        parent_module = getattr(parent_module, part)
                
                qlora_layer = QLoRALinear(module, r=r, alpha=alpha, dropout=dropout)
                setattr(parent_module, module_name, qlora_layer)
                logger.info(f"Applied QLoRA to {name}")
    
    return model


def compute_loss(
    model: nn.Module,
    batch: Dict[str, torch.Tensor],
    device: str,
) -> torch.Tensor:
    """Compute training loss"""
    input_ids = batch["input_ids"].to(device)
    attention_mask = batch["attention_mask"].to(device)
    pixel_values = batch["pixel_values"].to(device)
    labels = batch["labels"].to(device)
    
    # Forward pass
    outputs = model(
        input_ids=input_ids,
        pixel_values=pixel_values,
        attention_mask=attention_mask,
    )
    
    logits = outputs["logits"]
    
    # Shift labels for language modeling
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()
    
    # Flatten for loss computation
    loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
    loss = loss_fct(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1)
    )
    
    return loss


def train_epoch(
    model: nn.Module,
    dataloader: DataLoader,
    optimizer: torch.optim.Optimizer,
    scheduler: Optional[Any],
    device: str,
    gradient_accumulation_steps: int = 1,
    max_grad_norm: float = 1.0,
    use_wandb: bool = False,
) -> Dict[str, float]:
    """Train for one epoch"""
    model.train()
    total_loss = 0.0
    num_steps = 0
    
    progress_bar = tqdm(dataloader, desc="Training")
    
    for step, batch in enumerate(progress_bar):
        loss = compute_loss(model, batch, device)
        loss = loss / gradient_accumulation_steps
        loss.backward()
        
        total_loss += loss.item() * gradient_accumulation_steps
        num_steps += 1
        
        if (step + 1) % gradient_accumulation_steps == 0:
            # Gradient clipping
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
            
            optimizer.step()
            if scheduler:
                scheduler.step()
            optimizer.zero_grad()
            
            if use_wandb:
                wandb.log({
                    "train/loss": loss.item() * gradient_accumulation_steps,
                    "train/learning_rate": scheduler.get_last_lr()[0] if scheduler else optimizer.param_groups[0]['lr'],
                })
        
        progress_bar.set_postfix({"loss": f"{loss.item() * gradient_accumulation_steps:.4f}"})
    
    avg_loss = total_loss / num_steps if num_steps > 0 else 0.0
    return {"loss": avg_loss}


def evaluate(
    model: nn.Module,
    dataloader: DataLoader,
    device: str,
) -> Dict[str, float]:
    """Evaluate model on validation set"""
    model.eval()
    total_loss = 0.0
    num_steps = 0
    
    with torch.no_grad():
        for batch in tqdm(dataloader, desc="Evaluating"):
            loss = compute_loss(model, batch, device)
            total_loss += loss.item()
            num_steps += 1
    
    avg_loss = total_loss / num_steps if num_steps > 0 else 0.0
    return {"eval_loss": avg_loss}


def save_checkpoint(
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    scheduler: Optional[Any],
    epoch: int,
    step: int,
    loss: float,
    output_dir: str,
    is_best: bool = False,
):
    """Save model checkpoint"""
    os.makedirs(output_dir, exist_ok=True)
    
    checkpoint = {
        "epoch": epoch,
        "step": step,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "loss": loss,
    }
    
    if scheduler:
        checkpoint["scheduler_state_dict"] = scheduler.state_dict()
    
    # Save regular checkpoint
    checkpoint_path = os.path.join(output_dir, f"checkpoint-{step}.pt")
    torch.save(checkpoint, checkpoint_path)
    
    # Save best model
    if is_best:
        best_path = os.path.join(output_dir, "best_model.pt")
        torch.save(checkpoint, best_path)
    
    logger.info(f"Saved checkpoint to {checkpoint_path}")


def main(
    model_path: str,
    train_data_path: str,
    val_data_path: Optional[str] = None,
    output_dir: str = "./checkpoints",
    num_epochs: int = 3,
    batch_size: int = 4,
    gradient_accumulation_steps: int = 4,
    learning_rate: float = 2e-4,
    warmup_steps: int = 100,
    save_steps: int = 500,
    eval_steps: int = 250,
    use_qlora: bool = True,
    qlora_r: int = 64,
    qlora_alpha: int = 16,
    use_wandb: bool = True,
    wandb_project: str = "paligemma",
    resume_from_checkpoint: Optional[str] = None,
    only_cpu: bool = False,
):
    """Main training function"""
    
    # Setup device
    device = "cpu"
    if not only_cpu:
        if torch.cuda.is_available():
            device = "cuda"
        elif torch.backends.mps.is_available():
            device = "mps"
    
    logger.info(f"Using device: {device}")
    
    # Load model
    logger.info("Loading model...")
    model, tokenizer = load_hf_model(model_path, device)
    
    # Apply QLoRA if requested
    if use_qlora:
        logger.info("Applying QLoRA...")
        model = apply_qlora_to_model(model, r=qlora_r, alpha=qlora_alpha)
    
    # Enable gradient checkpointing for memory efficiency
    apply_gradient_checkpointing(model, enable=True)
    
    # Print parameter counts
    param_counts = count_parameters(model)
    logger.info(f"Trainable parameters: {param_counts['trainable_millions']:.2f}M")
    logger.info(f"Total parameters: {param_counts['total_millions']:.2f}M")
    
    # Setup processor
    num_image_tokens = model.config.vision_config.num_image_tokens
    image_size = model.config.vision_config.image_size
    processor = PaliGemmaProcessor(tokenizer, num_image_tokens, image_size)
    
    # Create datasets
    train_dataset = PaliGemmaDataset(
        train_data_path,
        processor,
        image_augmentation=True,
    )
    train_dataloader = create_dataloader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=4,
    )
    
    val_dataloader = None
    if val_data_path:
        val_dataset = PaliGemmaDataset(
            val_data_path,
            processor,
            image_augmentation=False,
        )
        val_dataloader = create_dataloader(
            val_dataset,
            batch_size=batch_size,
            shuffle=False,
            num_workers=4,
        )
    
    # Setup optimizer
    optimizer = AdamW(
        model.parameters(),
        lr=learning_rate,
        weight_decay=0.01,
        betas=(0.9, 0.999),
    )
    
    # Setup scheduler
    total_steps = len(train_dataloader) * num_epochs // gradient_accumulation_steps
    scheduler = get_linear_schedule_with_warmup(
        optimizer,
        num_warmup_steps=warmup_steps,
        num_training_steps=total_steps,
    )
    
    # Initialize wandb
    if use_wandb:
        wandb.init(
            project=wandb_project,
            config={
                "model_path": model_path,
                "num_epochs": num_epochs,
                "batch_size": batch_size,
                "learning_rate": learning_rate,
                "use_qlora": use_qlora,
            }
        )
    
    # Resume from checkpoint if provided
    start_epoch = 0
    global_step = 0
    best_eval_loss = float('inf')
    
    if resume_from_checkpoint:
        logger.info(f"Resuming from checkpoint: {resume_from_checkpoint}")
        checkpoint = torch.load(resume_from_checkpoint, map_location=device)
        model.load_state_dict(checkpoint["model_state_dict"])
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        if scheduler and "scheduler_state_dict" in checkpoint:
            scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
        start_epoch = checkpoint.get("epoch", 0)
        global_step = checkpoint.get("step", 0)
    
    # Training loop
    model.to(device)
    for epoch in range(start_epoch, num_epochs):
        logger.info(f"Epoch {epoch + 1}/{num_epochs}")
        
        # Train
        train_metrics = train_epoch(
            model,
            train_dataloader,
            optimizer,
            scheduler,
            device,
            gradient_accumulation_steps=gradient_accumulation_steps,
            use_wandb=use_wandb,
        )
        
        logger.info(f"Train loss: {train_metrics['loss']:.4f}")
        
        # Evaluate
        if val_dataloader and (global_step % eval_steps == 0):
            eval_metrics = evaluate(model, val_dataloader, device)
            logger.info(f"Eval loss: {eval_metrics['eval_loss']:.4f}")
            
            if use_wandb:
                wandb.log(eval_metrics)
            
            # Save best model
            is_best = eval_metrics['eval_loss'] < best_eval_loss
            if is_best:
                best_eval_loss = eval_metrics['eval_loss']
        
        # Save checkpoint
        if global_step % save_steps == 0:
            save_checkpoint(
                model,
                optimizer,
                scheduler,
                epoch,
                global_step,
                train_metrics['loss'],
                output_dir,
                is_best=False,
            )
        
        global_step += len(train_dataloader)
    
    # Final save
    save_checkpoint(
        model,
        optimizer,
        scheduler,
        num_epochs - 1,
        global_step,
        train_metrics['loss'],
        output_dir,
        is_best=False,
    )
    
    logger.info("Training completed!")
    if use_wandb:
        wandb.finish()


if __name__ == "__main__":
    import fire
    fire.Fire(main)