| """ |
| Training loop with: |
| - Mixed precision (AMP) for memory savings |
| - Gradient checkpointing for OOM prevention |
| - Cosine annealing with warmup |
| - Early stopping |
| - Gradient accumulation for larger effective batch |
| - TensorBoard logging |
| """ |
| import logging |
| import math |
| import time |
| from pathlib import Path |
| from typing import Optional |
|
|
| import torch |
| import torch.nn as nn |
| from torch.utils.data import DataLoader |
| from torch.utils.tensorboard import SummaryWriter |
|
|
| from src.s01_config import TrainConfig, PathConfig, get_device |
| from src.s04_model import MusicTransformer |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| class CosineWarmupScheduler: |
| """Cosine annealing LR with linear warmup.""" |
|
|
| def __init__(self, optimizer, warmup_steps: int, total_steps: int, min_lr: float = 1e-6): |
| self.optimizer = optimizer |
| self.warmup_steps = warmup_steps |
| self.total_steps = total_steps |
| self.min_lr = min_lr |
| self.base_lrs = [pg["lr"] for pg in optimizer.param_groups] |
| self.step_count = 0 |
|
|
| def step(self): |
| self.step_count += 1 |
| for pg, base_lr in zip(self.optimizer.param_groups, self.base_lrs): |
| if self.step_count < self.warmup_steps: |
| lr = base_lr * self.step_count / max(1, self.warmup_steps) |
| else: |
| progress = (self.step_count - self.warmup_steps) / max( |
| 1, self.total_steps - self.warmup_steps |
| ) |
| lr = self.min_lr + (base_lr - self.min_lr) * 0.5 * (1 + math.cos(math.pi * progress)) |
| pg["lr"] = lr |
|
|
| def get_lr(self) -> float: |
| return self.optimizer.param_groups[0]["lr"] |
|
|
|
|
| class Trainer: |
| """Handles the full training pipeline with memory-efficient techniques.""" |
|
|
| def __init__( |
| self, |
| model: MusicTransformer, |
| train_loader: DataLoader, |
| val_loader: DataLoader, |
| train_config: TrainConfig, |
| path_config: PathConfig, |
| ): |
| self.model = model |
| self.train_loader = train_loader |
| self.val_loader = val_loader |
| self.config = train_config |
| self.paths = path_config |
| self.device = get_device() |
|
|
| |
| if train_config.grad_checkpoint: |
| self.model.grad_checkpoint = True |
| logger.info("Gradient checkpointing ENABLED") |
|
|
| self.model.to(self.device) |
|
|
| |
| self.optimizer = torch.optim.AdamW( |
| self.model.parameters(), |
| lr=train_config.learning_rate, |
| weight_decay=train_config.weight_decay, |
| betas=(0.9, 0.95), |
| fused=torch.cuda.is_available(), |
| ) |
|
|
| |
| total_steps = len(train_loader) * train_config.max_epochs // train_config.grad_accum_steps |
| self.scheduler = CosineWarmupScheduler( |
| self.optimizer, train_config.warmup_steps, total_steps |
| ) |
|
|
| |
| use_amp = train_config.use_amp and torch.cuda.is_available() |
| self.scaler = torch.amp.GradScaler("cuda", enabled=use_amp) |
| if use_amp: |
| self.amp_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 |
| else: |
| self.amp_dtype = torch.float32 |
| self.use_amp = use_amp |
|
|
| |
| self.writer = SummaryWriter(log_dir=str(path_config.log_dir)) |
|
|
| |
| self.global_step = 0 |
| self.best_val_loss = float("inf") |
| self.patience_counter = 0 |
|
|
| def train(self): |
| """Main training loop.""" |
| logger.info(f"Starting training on {self.device}") |
| logger.info(f"Model params: {self.model.count_parameters():,}") |
| logger.info(f"Train batches: {len(self.train_loader)}, Val batches: {len(self.val_loader)}") |
|
|
| for epoch in range(1, self.config.max_epochs + 1): |
| t0 = time.time() |
| train_loss = self._train_epoch(epoch) |
| val_loss = self._validate() |
| elapsed = time.time() - t0 |
|
|
| logger.info( |
| f"Epoch {epoch}/{self.config.max_epochs} | " |
| f"Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f} | " |
| f"LR: {self.scheduler.get_lr():.2e} | Time: {elapsed:.1f}s" |
| ) |
|
|
| self.writer.add_scalars("loss", {"train": train_loss, "val": val_loss}, epoch) |
| self.writer.add_scalar("lr", self.scheduler.get_lr(), epoch) |
|
|
| |
| if val_loss < self.best_val_loss - self.config.min_delta: |
| self.best_val_loss = val_loss |
| self.patience_counter = 0 |
| self._save_checkpoint("best.pt", epoch, val_loss) |
| logger.info(f" New best model saved (val_loss={val_loss:.4f})") |
| else: |
| self.patience_counter += 1 |
| if self.patience_counter >= self.config.patience: |
| logger.info(f"Early stopping at epoch {epoch} (patience={self.config.patience})") |
| break |
|
|
| |
| if epoch % 5 == 0: |
| self._save_checkpoint(f"epoch_{epoch}.pt", epoch, val_loss) |
|
|
| self.writer.close() |
| logger.info("Training complete!") |
|
|
| def _train_epoch(self, epoch: int) -> float: |
| self.model.train() |
| total_loss = 0.0 |
| n_batches = 0 |
| self.optimizer.zero_grad(set_to_none=True) |
|
|
| for batch_idx, (input_ids, targets) in enumerate(self.train_loader): |
| input_ids = input_ids.to(self.device, non_blocking=True) |
| targets = targets.to(self.device, non_blocking=True) |
|
|
| |
| with torch.amp.autocast( |
| device_type=self.device.type, |
| dtype=self.amp_dtype, |
| enabled=self.use_amp, |
| ): |
| _, loss = self.model(input_ids, targets) |
| loss = loss / self.config.grad_accum_steps |
|
|
| |
| self.scaler.scale(loss).backward() |
|
|
| if (batch_idx + 1) % self.config.grad_accum_steps == 0: |
| self.scaler.unscale_(self.optimizer) |
| nn.utils.clip_grad_norm_(self.model.parameters(), self.config.max_grad_norm) |
| self.scaler.step(self.optimizer) |
| self.scaler.update() |
| self.optimizer.zero_grad(set_to_none=True) |
| self.scheduler.step() |
| self.global_step += 1 |
|
|
| total_loss += loss.item() * self.config.grad_accum_steps |
| n_batches += 1 |
|
|
| if (batch_idx + 1) % self.config.log_interval == 0: |
| avg = total_loss / n_batches |
| logger.info( |
| f" Epoch {epoch} [{batch_idx+1}/{len(self.train_loader)}] " |
| f"loss={avg:.4f} lr={self.scheduler.get_lr():.2e}" |
| ) |
|
|
| return total_loss / max(1, n_batches) |
|
|
| @torch.no_grad() |
| def _validate(self) -> float: |
| self.model.eval() |
| total_loss = 0.0 |
| n_batches = 0 |
|
|
| for input_ids, targets in self.val_loader: |
| input_ids = input_ids.to(self.device, non_blocking=True) |
| targets = targets.to(self.device, non_blocking=True) |
|
|
| with torch.amp.autocast( |
| device_type=self.device.type, |
| dtype=self.amp_dtype, |
| enabled=self.use_amp, |
| ): |
| _, loss = self.model(input_ids, targets) |
|
|
| total_loss += loss.item() |
| n_batches += 1 |
|
|
| return total_loss / max(1, n_batches) |
|
|
| def _save_checkpoint(self, name: str, epoch: int, val_loss: float): |
| path = self.paths.checkpoint_dir / name |
| torch.save( |
| { |
| "epoch": epoch, |
| "model_state_dict": self.model.state_dict(), |
| "optimizer_state_dict": self.optimizer.state_dict(), |
| "val_loss": val_loss, |
| "global_step": self.global_step, |
| "config": self.model.config, |
| }, |
| path, |
| ) |
|
|
| def load_checkpoint(self, path: Path): |
| ckpt = torch.load(path, map_location=self.device, weights_only=False) |
| self.model.load_state_dict(ckpt["model_state_dict"]) |
| self.optimizer.load_state_dict(ckpt["optimizer_state_dict"]) |
| self.global_step = ckpt.get("global_step", 0) |
| logger.info(f"Loaded checkpoint: {path} (epoch={ckpt['epoch']}, val_loss={ckpt['val_loss']:.4f})") |
|
|
