import os import time import logging from pathlib import Path from typing import Dict, Optional import torch from torch.utils.data import Dataset, DataLoader from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR from torch.cuda.amp import autocast try: from torch.amp import GradScaler as AmpGradScaler _has_new_amp = True except ImportError: from torch.cuda.amp import GradScaler as AmpGradScaler _has_new_amp = False from .modeling import TinyDocVLMForConditionalGeneration from .processing import TinyDocVLMProcessor from .losses import CombinedLoss from .data import collate_fn logger = logging.getLogger(__name__) class TrainerConfig: def __init__( self, output_dir: str = "checkpoints", num_epochs: int = 3, batch_size: int = 8, gradient_accumulation_steps: int = 4, learning_rate: float = 1e-4, min_learning_rate: float = 1e-5, warmup_steps: int = 500, weight_decay: float = 0.01, max_grad_norm: float = 1.0, max_seq_length: int = 2048, stage: int = 1, use_fp16: bool = True, save_every_steps: int = 1000, eval_every_steps: int = 500, log_every_steps: int = 10, gradient_checkpointing: bool = True, num_workers: int = 4, ): self.output_dir = output_dir self.num_epochs = num_epochs self.batch_size = batch_size self.gradient_accumulation_steps = gradient_accumulation_steps self.learning_rate = learning_rate self.min_learning_rate = min_learning_rate self.warmup_steps = warmup_steps self.weight_decay = weight_decay self.max_grad_norm = max_grad_norm self.max_seq_length = max_seq_length self.stage = stage self.use_fp16 = use_fp16 self.save_every_steps = save_every_steps self.eval_every_steps = eval_every_steps self.log_every_steps = log_every_steps self.gradient_checkpointing = gradient_checkpointing self.num_workers = num_workers def to_dict(self) -> Dict: return {k: v for k, v in self.__dict__.items()} @classmethod def from_dict(cls, d: Dict) -> "TrainerConfig": return cls(**{k: v for k, v in d.items() if k in cls.__init__.__code__.co_varnames}) class TinyDocVLMTrainer: """ Trainer for TinyDoc-VLM across all 3 training stages. Supports FSDP, mixed precision, gradient checkpointing, and checkpointing. """ def __init__( self, model: TinyDocVLMForConditionalGeneration, processor: TinyDocVLMProcessor, train_dataset: Dataset, eval_dataset: Optional[Dataset] = None, config: Optional[TrainerConfig] = None, device: Optional[torch.device] = None, ): self.model = model self.processor = processor self.train_dataset = train_dataset self.eval_dataset = eval_dataset self.config = config or TrainerConfig() self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu") self.global_step = 0 self.epoch = 0 self.best_eval_loss = float("inf") os.makedirs(self.config.output_dir, exist_ok=True) if self.config.gradient_checkpointing: self.model.gradient_checkpointing_enable() self.model.to(self.device) no_decay = ["bias", "LayerNorm.weight", "layer_norm.weight"] optimizer_grouped_params = [ { "params": [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay)], "weight_decay": self.config.weight_decay, }, { "params": [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay)], "weight_decay": 0.0, }, ] self.optimizer = AdamW(optimizer_grouped_params, lr=self.config.learning_rate, betas=(0.9, 0.95), eps=1e-8) total_steps = len(self.train_dataset) // (self.config.batch_size * self.config.gradient_accumulation_steps) * self.config.num_epochs warmup_scheduler = LinearLR(self.optimizer, start_factor=0.05, end_factor=1.0, total_iters=self.config.warmup_steps) cosine_scheduler = CosineAnnealingLR(self.optimizer, T_max=max(1, total_steps - self.config.warmup_steps), eta_min=self.config.min_learning_rate) self.scheduler = SequentialLR(self.optimizer, schedulers=[warmup_scheduler, cosine_scheduler], milestones=[self.config.warmup_steps]) if _has_new_amp and torch.cuda.is_available(): self.scaler = AmpGradScaler('cuda', enabled=self.config.use_fp16) else: self.scaler = AmpGradScaler(enabled=self.config.use_fp16) self.loss_fn = CombinedLoss(stage=self.config.stage) self.train_loader = DataLoader( self.train_dataset, batch_size=self.config.batch_size, shuffle=True, collate_fn=lambda batch: collate_fn(batch, self.processor.tokenizer, self.processor.image_token_id, self.config.max_seq_length), num_workers=self.config.num_workers, pin_memory=True, ) if self.eval_dataset: self.eval_loader = DataLoader( self.eval_dataset, batch_size=self.config.batch_size, shuffle=False, collate_fn=lambda batch: collate_fn(batch, self.processor.tokenizer, self.processor.image_token_id, self.config.max_seq_length), num_workers=self.config.num_workers, pin_memory=True, ) def train_step(self, batch: Dict) -> Dict: input_ids = batch["input_ids"].to(self.device) attention_mask = batch["attention_mask"].to(self.device) pixel_values = batch["pixel_values"].to(self.device) labels = batch["labels"].to(self.device) task = batch.get("task", None) with autocast(enabled=self.config.use_fp16): if task and self.config.stage == 2: outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, attention_mask=attention_mask, labels=labels, task=task, ) head_outputs = outputs["head_outputs"] loss_dict = self.loss_fn( lm_logits=outputs["lm_outputs"].logits, lm_labels=labels, head_outputs=head_outputs, head_labels=batch.get("head_labels", None), ) loss = loss_dict["loss"] else: outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, attention_mask=attention_mask, labels=labels, ) loss = outputs.loss if hasattr(outputs, "loss") else outputs[0] if self.config.use_fp16: self.scaler.scale(loss).backward() else: loss.backward() return {"loss": loss.item()} def train_epoch(self) -> Dict: self.model.train() total_loss = 0.0 num_batches = 0 accumulation_loss = 0.0 start_time = time.time() for step, batch in enumerate(self.train_loader): step_loss = self.train_step(batch) accumulation_loss += step_loss["loss"] if (step + 1) % self.config.gradient_accumulation_steps == 0: if self.config.use_fp16: self.scaler.unscale_(self.optimizer) torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.max_grad_norm) if self.config.use_fp16: self.scaler.step(self.optimizer) self.scaler.update() else: self.optimizer.step() self.optimizer.zero_grad() self.scheduler.step() self.global_step += 1 avg_loss = accumulation_loss / self.config.gradient_accumulation_steps total_loss += avg_loss num_batches += 1 accumulation_loss = 0.0 if self.global_step % self.config.log_every_steps == 0: elapsed = time.time() - start_time lr = self.scheduler.get_last_lr()[0] logger.info( f"Epoch {self.epoch+1} | Step {self.global_step} | Loss: {avg_loss:.4f} | " f"LR: {lr:.2e} | Elapsed: {elapsed:.1f}s" ) if self.global_step % self.config.eval_every_steps == 0 and self.eval_loader: eval_metrics = self.evaluate() eval_loss = eval_metrics.get("eval_loss", float("inf")) logger.info(f"Eval loss: {eval_loss:.4f}") if isinstance(eval_loss, (int, float)) and eval_loss < self.best_eval_loss: self.best_eval_loss = eval_loss self.save_checkpoint("best") if self.global_step % self.config.save_every_steps == 0: self.save_checkpoint(f"step_{self.global_step}") avg_epoch_loss = total_loss / max(num_batches, 1) elapsed = time.time() - start_time logger.info(f"Epoch {self.epoch+1} complete. Avg loss: {avg_epoch_loss:.4f}. Elapsed: {elapsed:.1f}s") return {"loss": avg_epoch_loss, "epoch": self.epoch + 1, "steps": self.global_step} def evaluate(self, benchmark_name: Optional[str] = None) -> float: self.model.eval() total_loss = 0.0 num_batches = 0 all_metrics = {} if benchmark_name: try: from evaluation.evaluate import evaluate_model data_dir = Path("evaluation/data") if data_dir.exists(): results = evaluate_model(self.model, self.processor, [benchmark_name], data_dir) all_metrics = results.get(benchmark_name, {}) logger.info(f"Benchmark {benchmark_name}: {all_metrics}") else: logger.warning(f"Benchmark data not found at {data_dir}") except ImportError: logger.warning("Evaluation module not available") if self.eval_loader: with torch.no_grad(): for batch in self.eval_loader: input_ids = batch["input_ids"].to(self.device) attention_mask = batch["attention_mask"].to(self.device) pixel_values = batch["pixel_values"].to(self.device) labels = batch["labels"].to(self.device) with autocast(enabled=self.config.use_fp16): outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, attention_mask=attention_mask, labels=labels, ) loss = outputs.loss if hasattr(outputs, "loss") else outputs[0] total_loss += loss.item() num_batches += 1 self.model.train() eval_loss = total_loss / max(num_batches, 1) if num_batches > 0 else float("inf") all_metrics["eval_loss"] = eval_loss return all_metrics def train(self): logger.info(f"Starting training on {self.device}") logger.info(f"Train samples: {len(self.train_dataset)}") logger.info(f"Eval samples: {len(self.eval_dataset) if self.eval_dataset else 0}") logger.info(f"Config: {self.config.to_dict()}") self.save_checkpoint("init") for epoch in range(self.config.num_epochs): self.epoch = epoch self.train_epoch() self.save_checkpoint(f"epoch_{epoch+1}") logger.info("Training complete.") def save_checkpoint(self, tag: str): output_dir = Path(self.config.output_dir) / tag output_dir.mkdir(parents=True, exist_ok=True) self.model.save_pretrained(str(output_dir)) torch.save(self.model.state_dict(), output_dir / "model_state.pt") self.processor.tokenizer.save_pretrained(str(output_dir)) trainer_state = { "global_step": self.global_step, "epoch": self.epoch + 1, "best_eval_loss": self.best_eval_loss, "optimizer_state_dict": self.optimizer.state_dict(), "scheduler_state_dict": self.scheduler.state_dict(), "config": self.config.to_dict(), } torch.save(trainer_state, output_dir / "trainer_state.pt") saved_files = list(output_dir.iterdir()) logger.info(f"Checkpoint saved to {output_dir} ({len(saved_files)} files: {[f.name for f in saved_files]})") if not any(f.name.endswith('.bin') or f.name.endswith('.safetensors') or f.name == 'model_state.pt' for f in saved_files): logger.error(f"No weight file found in checkpoint! Files: {[f.name for f in saved_files]}") def load_checkpoint(self, checkpoint_dir: str): checkpoint_dir = Path(checkpoint_dir) model_state_path = checkpoint_dir / "model_state.pt" if model_state_path.exists(): model = TinyDocVLMForConditionalGeneration(self.model.config) model.load_state_dict(torch.load(str(model_state_path), map_location=self.device, weights_only=True)) self.model = model else: self.model = TinyDocVLMForConditionalGeneration.from_pretrained(str(checkpoint_dir)) self.model.to(self.device) trainer_state = torch.load(checkpoint_dir / "trainer_state.pt", map_location=self.device) self.global_step = trainer_state["global_step"] self.epoch = trainer_state.get("epoch", 1) - 1 self.best_eval_loss = trainer_state.get("best_eval_loss", float("inf")) self.optimizer.load_state_dict(trainer_state["optimizer_state_dict"]) self.scheduler.load_state_dict(trainer_state["scheduler_state_dict"]) logger.info(f"Checkpoint loaded from {checkpoint_dir}")