diffusion_llm/training/trainer.py

"""
AAM Diffusion LLM — Trainer

Training loop for the AAM Diffusion Model.

Handles:
    - Training loop with gradient accumulation
    - Learning rate scheduling with warmup
    - Mixed precision training (AMP)
    - EMA model updates
    - Checkpoint saving/loading
    - Logging to console and Weights & Biases
    - Evaluation on validation set

Analogi: Seperti latihan fisik Jin Soun — berulang-ulang,
bertahap meningkat intensitas, dengan instruktur yang
mengawasi dan memberi koreksi.
"""

from __future__ import annotations

import json
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 diffusion_llm.config.model_config import AamDiffusionConfig
from diffusion_llm.model.aam_diffusion_model import AamDiffusionModel
from diffusion_llm.training.dataset import GraphNarrativeDataset, collate_fn
from diffusion_llm.tokenizer.aam_tokenizer import AamTokenizer
from diffusion_llm.training.losses import DiffusionLoss

logger = logging.getLogger(__name__)


class AamTrainer:
    """Trainer for the AAM Diffusion Model.

    Args:
        config: AamDiffusionConfig with training settings.
        model: AamDiffusionModel instance.
        tokenizer: AamTokenizer instance.
        train_dataset: Training dataset.
        val_dataset: Optional validation dataset.
    """

    def __init__(
        self,
        config: AamDiffusionConfig,
        model: AamDiffusionModel,
        tokenizer: AamTokenizer,
        train_dataset: GraphNarrativeDataset,
        val_dataset: Optional[GraphNarrativeDataset] = None,
    ):
        self.config = config
        self.model = model
        self.tokenizer = tokenizer
        self.train_dataset = train_dataset
        self.val_dataset = val_dataset

        # Device
        self.device = torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )
        self.model.to(self.device)
        logger.info("Training on device: %s", self.device)

        # Optimizer
        self.optimizer = torch.optim.AdamW(
            self.model.parameters(),
            lr=config.training.learning_rate,
            weight_decay=config.training.weight_decay,
            betas=(config.training.adam_beta1, config.training.adam_beta2),
            eps=config.training.adam_eps,
        )

        # Loss function
        self.loss_fn = DiffusionLoss(config.diffusion)

        # Data loaders
        self.train_loader = DataLoader(
            train_dataset,
            batch_size=config.training.batch_size,
            shuffle=True,
            num_workers=config.training.num_workers,
            collate_fn=collate_fn,
            pin_memory=True,
        )

        if val_dataset:
            self.val_loader = DataLoader(
                val_dataset,
                batch_size=config.training.batch_size,
                shuffle=False,
                num_workers=config.training.num_workers,
                collate_fn=collate_fn,
                pin_memory=True,
            )
        else:
            self.val_loader = None

        # LR scheduler
        self.scheduler = self._create_lr_scheduler()

        # AMP
        self.scaler = None
        if config.training.use_amp:
            dtype = torch.bfloat16 if config.training.amp_dtype == "bf16" else torch.float16
            self.scaler = torch.amp.GradScaler("cuda", enabled=(dtype == torch.float16))

        # EMA
        self.ema_model = None
        if config.training.use_ema:
            self.ema_model = self._create_ema_model()

        # State tracking
        self.global_step = 0
        self.best_val_loss = float("inf")
        self.train_losses: list[float] = []

        # Output directory
        self.output_dir = Path(config.output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        # Seed
        torch.manual_seed(config.seed)

    def _create_lr_scheduler(self):
        """Create learning rate scheduler with warmup."""
        total_steps = self.config.training.max_steps
        warmup_steps = self.config.training.warmup_steps

        def lr_lambda(step: int) -> float:
            if step < warmup_steps:
                return step / max(warmup_steps, 1)
            if self.config.training.lr_schedule == "cosine":
                progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1)
                return 0.5 * (1.0 + math.cos(math.pi * progress))
            elif self.config.training.lr_schedule == "linear":
                progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1)
                return 1.0 - progress
            else:
                return 1.0

        return torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda)

    def _create_ema_model(self) -> AamDiffusionModel:
        """Create EMA copy of the model."""
        import copy
        ema = copy.deepcopy(self.model)
        for param in ema.parameters():
            param.requires_grad = False
        return ema

    @torch.no_grad()
    def _update_ema(self) -> None:
        """Update EMA model weights."""
        if self.ema_model is None:
            return
        decay = self.config.training.ema_decay
        for ema_param, model_param in zip(
            self.ema_model.parameters(), self.model.parameters()
        ):
            ema_param.data.mul_(decay).add_(model_param.data, alpha=1 - decay)

    def train(self) -> None:
        """Main training loop.

        Runs for max_steps or max_epochs, whichever comes first.
        Saves checkpoints and runs evaluation periodically.
        """
        logger.info("Starting training...")
        logger.info("  Max steps: %d", self.config.training.max_steps)
        logger.info("  Batch size: %d", self.config.training.batch_size)
        logger.info("  Gradient accumulation: %d", self.config.training.gradient_accumulation_steps)
        logger.info("  Effective batch size: %d",
                     self.config.training.batch_size * self.config.training.gradient_accumulation_steps)

        start_time = time.time()
        epoch = 0

        while self.global_step < self.config.training.max_steps:
            epoch += 1
            if epoch > self.config.training.max_epochs:
                break

            logger.info("=== Epoch %d ===", epoch)
            epoch_loss = 0.0
            n_batches = 0

            for batch_idx, batch in enumerate(self.train_loader):
                loss = self._train_step(batch)
                epoch_loss += loss
                n_batches += 1

                # Logging
                if self.global_step % self.config.training.log_every_steps == 0:
                    avg_loss = epoch_loss / max(n_batches, 1)
                    lr = self.optimizer.param_groups[0]["lr"]
                    elapsed = time.time() - start_time
                    steps_per_sec = self.global_step / max(elapsed, 1)

                    logger.info(
                        "Step %d | Loss: %.4f | LR: %.2e | Speed: %.1f steps/s",
                        self.global_step, loss, lr, steps_per_sec,
                    )

                # Evaluation
                if (self.global_step % self.config.training.eval_every_steps == 0
                        and self.val_loader is not None):
                    val_loss = self.evaluate()
                    logger.info("Validation loss: %.4f", val_loss)
                    if val_loss < self.best_val_loss:
                        self.best_val_loss = val_loss
                        self._save_checkpoint("best.pt")

                # Checkpoint
                if self.global_step % self.config.training.save_every_steps == 0:
                    self._save_checkpoint(f"step_{self.global_step}.pt")

                # Stop condition
                if self.global_step >= self.config.training.max_steps:
                    break

            avg_epoch_loss = epoch_loss / max(n_batches, 1)
            logger.info("Epoch %d complete. Average loss: %.4f", epoch, avg_epoch_loss)

        # Final save
        self._save_checkpoint("final.pt")
        elapsed = time.time() - start_time
        logger.info(
            "Training complete! %d steps in %.1f hours",
            self.global_step, elapsed / 3600,
        )

    def _train_step(self, batch: dict[str, torch.Tensor]) -> float:
        """Single training step.

        Args:
            batch: Batch of training data.

        Returns:
            Loss value for this step.
        """
        self.model.train()

        # Move batch to device
        batch = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v
                 for k, v in batch.items()}

        # Sample random timesteps
        batch_size = batch["token_ids"].shape[0]
        t = torch.randint(
            0, self.config.diffusion.n_timesteps,
            (batch_size,), device=self.device,
        )

        # Forward pass
        if self.scaler is not None:
            with torch.amp.autocast("cuda", enabled=True):
                predicted, target = self.model(
                    token_ids=batch["token_ids"],
                    timestep=t,
                    evidence_ids=batch.get("evidence_ids"),
                    evidence_confidence=batch.get("evidence_confidence"),
                    anomaly_ids=batch.get("anomaly_ids"),
                    anomaly_confidence=batch.get("anomaly_confidence"),
                    reasoning_ids=batch.get("reasoning_ids"),
                    reasoning_confidence=batch.get("reasoning_confidence"),
                    source_trust=batch.get("source_trust"),
                )
                loss = self.model.compute_loss(predicted, target, t)
                loss = loss / self.config.training.gradient_accumulation_steps
        else:
            predicted, target = self.model(
                token_ids=batch["token_ids"],
                timestep=t,
                evidence_ids=batch.get("evidence_ids"),
                evidence_confidence=batch.get("evidence_confidence"),
                anomaly_ids=batch.get("anomaly_ids"),
                anomaly_confidence=batch.get("anomaly_confidence"),
                reasoning_ids=batch.get("reasoning_ids"),
                reasoning_confidence=batch.get("reasoning_confidence"),
                source_trust=batch.get("source_trust"),
            )
            loss = self.model.compute_loss(predicted, target, t)
            loss = loss / self.config.training.gradient_accumulation_steps

        # Backward pass
        if self.scaler is not None:
            self.scaler.scale(loss).backward()
        else:
            loss.backward()

        # Gradient accumulation
        if (self.global_step + 1) % self.config.training.gradient_accumulation_steps == 0:
            # Gradient clipping
            if self.scaler is not None:
                self.scaler.unscale_(self.optimizer)
            torch.nn.utils.clip_grad_norm_(
                self.model.parameters(),
                self.config.training.grad_clip_norm,
            )

            # Optimizer step
            if self.scaler is not None:
                self.scaler.step(self.optimizer)
                self.scaler.update()
            else:
                self.optimizer.step()

            # LR schedule
            self.scheduler.step()

            # Zero gradients
            self.optimizer.zero_grad()

            # EMA update
            self._update_ema()

        self.global_step += 1
        self.train_losses.append(loss.item())

        return loss.item()

    @torch.no_grad()
    def evaluate(self) -> float:
        """Evaluate on validation set.

        Returns:
            Average validation loss.
        """
        if self.val_loader is None:
            return float("inf")

        self.model.eval()
        total_loss = 0.0
        n_batches = 0

        for batch in self.val_loader:
            batch = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v
                     for k, v in batch.items()}

            batch_size = batch["token_ids"].shape[0]
            t = torch.randint(
                0, self.config.diffusion.n_timesteps,
                (batch_size,), device=self.device,
            )

            predicted, target = self.model(
                token_ids=batch["token_ids"],
                timestep=t,
                evidence_ids=batch.get("evidence_ids"),
                evidence_confidence=batch.get("evidence_confidence"),
                anomaly_ids=batch.get("anomaly_ids"),
                anomaly_confidence=batch.get("anomaly_confidence"),
                reasoning_ids=batch.get("reasoning_ids"),
                reasoning_confidence=batch.get("reasoning_confidence"),
                source_trust=batch.get("source_trust"),
            )
            loss = self.model.compute_loss(predicted, target, t)
            total_loss += loss.item()
            n_batches += 1

        avg_loss = total_loss / max(n_batches, 1)
        self.model.train()
        return avg_loss

    def _save_checkpoint(self, filename: str) -> None:
        """Save training checkpoint.

        Args:
            filename: Checkpoint filename.
        """
        path = self.output_dir / filename
        checkpoint = {
            "model_state_dict": self.model.state_dict(),
            "optimizer_state_dict": self.optimizer.state_dict(),
            "scheduler_state_dict": self.scheduler.state_dict(),
            "global_step": self.global_step,
            "best_val_loss": self.best_val_loss,
            "config": self.config.to_dict(),
        }
        if self.ema_model is not None:
            checkpoint["ema_state_dict"] = self.ema_model.state_dict()

        torch.save(checkpoint, path)
        logger.info("Checkpoint saved: %s", path)

        # Clean up old checkpoints
        self._cleanup_checkpoints()

    def _cleanup_checkpoints(self) -> None:
        """Remove old checkpoints, keeping only the last N."""
        keep_n = self.config.training.keep_last_n_checkpoints
        checkpoints = sorted(self.output_dir.glob("step_*.pt"))
        while len(checkpoints) > keep_n:
            oldest = checkpoints.pop(0)
            oldest.unlink()
            logger.info("Removed old checkpoint: %s", oldest)

    def load_checkpoint(self, path: str) -> None:
        """Load from checkpoint.

        Args:
            path: Checkpoint file path.
        """
        checkpoint = torch.load(path, map_location=self.device)
        self.model.load_state_dict(checkpoint["model_state_dict"])
        self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        self.scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
        self.global_step = checkpoint["global_step"]
        self.best_val_loss = checkpoint.get("best_val_loss", float("inf"))
        logger.info("Loaded checkpoint from step %d", self.global_step)