training/trainer.py

"""
Training loop abstraction.

Single Responsibility: Handles the training loop logic.
Open/Closed: Base class can be extended for different training stages.
Liskov Substitution: Stage1Trainer and Stage2Trainer are interchangeable where Trainer is expected.

Optimizations implemented:
- OOM handling with proper recovery (PyTorch FAQ best practices)
- Gradient NaN/Inf detection and skipping
- Gradient norm monitoring for stability
- Dynamic sequence length based on text ratio
- CUDA memory fragmentation reduction
"""

import os
import gc
import math
import torch
import torch.nn.functional as F
from torch.optim.lr_scheduler import CosineAnnealingLR
from accelerate import Accelerator
from accelerate.utils import set_seed
from tqdm import tqdm
from typing import Optional, Dict, Any, Tuple
from abc import ABC, abstractmethod

# Reduce CUDA memory fragmentation (PyTorch recommendation)
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

from .config import TrainingConfig, GPUConfig
from .data import load_datasets, create_dataloader
from .models import SpeechAdapter, ModelFactory
from .interleaving import (
    get_text_ratio,
    calculate_dynamic_decay_steps,
    apply_interleaving,
)
from .checkpoint import CheckpointManager, TrainingState
from .utils import (
    log,
    setup_cuda_optimizations,
    setup_hf_login,
    load_tokenizer,
    should_enable_gradient_checkpointing,
    write_step,
    get_device_info,
    limit_ram_usage,
    get_ram_info,
)


class BaseTrainer(ABC):
    """
    Abstract base class for trainers.

    Implements Template Method pattern: defines training skeleton,
    subclasses implement specific steps.
    """

    def __init__(self, config: TrainingConfig):
        self.config = config
        self.accelerator: Optional[Accelerator] = None
        self.tokenizer = None
        self.adapter: Optional[SpeechAdapter] = None
        self.llm = None
        self.optimizer = None
        self.scheduler = None
        self.train_loader = None
        self.checkpoint_manager: Optional[CheckpointManager] = None

        # Training state
        self.global_step = 0
        self.start_epoch = 0
        self.best_loss = float("inf")
        self.current_text_ratio = config.initial_text_ratio

        # Stability monitoring
        self.nan_count = 0
        self.oom_count = 0
        self.last_grad_norm = 0.0
        self.max_grad_norm_seen = 0.0

    @property
    def is_main(self) -> bool:
        """Check if this is the main process."""
        return self.accelerator is None or self.accelerator.is_main_process

    @property
    def device(self):
        """Get current device."""
        return self.accelerator.device if self.accelerator else torch.device("cpu")

    def setup(self):
        """Setup training environment."""
        self._setup_accelerator()
        self._setup_memory()
        self._setup_tokenizer()
        self._setup_data()
        self._setup_models()
        self._setup_optimizer()
        self._setup_checkpoint_manager()
        self._resume_if_needed()
        self._prepare_for_training()

    def _setup_accelerator(self):
        """Initialize Accelerator."""
        mixed_precision = "bf16" if torch.cuda.is_available() else None

        self.accelerator = Accelerator(
            gradient_accumulation_steps=self.config.grad_accum,
            mixed_precision=mixed_precision,
        )

        set_seed(42)

        if self.is_main:
            self._log_setup_info()

    def _log_setup_info(self):
        """Log setup information."""
        device_info = get_device_info()
        gpu_config = GPUConfig.auto_detect()

        log("=" * 60)
        log(self._get_stage_name())
        log("=" * 60)
        log(f"Device: {device_info}")
        log(f"GPU: {gpu_config.name} ({gpu_config.vram_gb}GB)")
        log(f"Num processes: {self.accelerator.num_processes}")
        log(f"Batch: {self.config.batch_size}, Grad accum: {self.config.grad_accum}")
        log(f"LR: {self.config.learning_rate}, Epochs: {self.config.epochs}")
        if getattr(self.config, 'no_decay', False):
            log(f"Fixed text ratio: {self.config.initial_text_ratio} (no decay)")
        else:
            log(f"Initial text ratio: {self.config.initial_text_ratio}")
        log("[Optimizations] OOM recovery, NaN detection, grad monitoring, expandable_segments")

    def _setup_memory(self):
        """Configure memory settings."""
        if self.device.type == 'cuda':
            setup_cuda_optimizations(self.config.vram_fraction)

        ram_total, _ = get_ram_info()
        ram_limit = self.config.ram_limit_gb or (ram_total * 0.80)
        limit_ram_usage(ram_limit)

        if self.is_main:
            log(f"[Memory] VRAM: {self.config.vram_fraction*100:.0f}%, RAM: {ram_limit:.1f}GB")

    def _setup_tokenizer(self):
        """Load tokenizer."""
        setup_hf_login()
        self.tokenizer = load_tokenizer(self.config.model_path)

    def _setup_data(self):
        """Load datasets and create dataloader."""
        if self.is_main:
            log("\nLoading datasets...")

        dataset = load_datasets(
            self.config.data_paths,
            self.tokenizer,
            max_audio_len=self.config.max_audio_len,
            max_seq_len=self.config.max_seq_len,
            verbose=self.is_main,
        )

        # Apply demo/test mode
        if self.config.test_mode:
            dataset = torch.utils.data.Subset(dataset, range(min(5, len(dataset))))
            self.config.batch_size = min(self.config.batch_size, len(dataset))
            self.config.grad_accum = 1
        elif self.config.demo_mode:
            dataset = torch.utils.data.Subset(dataset, range(min(1000, len(dataset))))
            self.config.batch_size = min(4, self.config.batch_size)

        if self.is_main:
            log(f"Total samples: {len(dataset):,}")

        self.train_loader = create_dataloader(
            dataset,
            self.config.batch_size,
            shuffle=True,
            verbose=self.is_main,
        )

    @abstractmethod
    def _setup_models(self):
        """Setup models (implemented by subclasses)."""
        pass

    @abstractmethod
    def _setup_optimizer(self):
        """Setup optimizer (implemented by subclasses)."""
        pass

    @abstractmethod
    def _setup_checkpoint_manager(self):
        """Setup checkpoint manager (implemented by subclasses)."""
        pass

    def _resume_if_needed(self):
        """Resume from checkpoint if specified."""
        if not self.config.resume_from or not os.path.exists(self.config.resume_from):
            return

        if self.is_main:
            log(f"\nResuming from: {self.config.resume_from}")

        ckpt = self.checkpoint_manager.load(self.config.resume_from)
        self._load_checkpoint(ckpt)

    @abstractmethod
    def _load_checkpoint(self, ckpt: Dict[str, Any]):
        """Load checkpoint state (implemented by subclasses)."""
        pass

    def _prepare_for_training(self):
        """Prepare models and optimizer for training."""
        # Prepare with accelerator
        prepared = self.accelerator.prepare(
            self.adapter, self.llm, self.optimizer, self.train_loader
        )
        self.adapter, self.llm, self.optimizer, self.train_loader = prepared

        # Calculate steps
        self.steps_per_epoch = max(1, len(self.train_loader) // self.config.grad_accum)
        self.total_steps = max(1, self.steps_per_epoch * self.config.epochs)
        self.warmup_steps = int(self.total_steps * self.config.warmup_ratio)

        # Calculate decay steps
        if self.config.dynamic_decay:
            self.effective_decay_steps = calculate_dynamic_decay_steps(
                self.total_steps,
                steps_per_epoch=self.steps_per_epoch,  # Complete decay in epoch 1
                initial_ratio=self.config.initial_text_ratio,
            )
            if self.is_main:
                log(f"[Dynamic Decay] decay_steps={self.effective_decay_steps} (decay completes in epoch 1)")
        else:
            self.effective_decay_steps = self.config.decay_steps

        # Setup scheduler
        self.scheduler = CosineAnnealingLR(
            self.optimizer,
            T_max=max(1, self.total_steps - self.warmup_steps),
            eta_min=1e-6,
        )

        if self.is_main:
            log(f"Steps: {self.steps_per_epoch}/epoch, {self.total_steps} total, {self.warmup_steps} warmup")

    @abstractmethod
    def _get_stage_name(self) -> str:
        """Get stage name for logging."""
        pass

    def train(self):
        """Main training loop."""
        if self.is_main:
            log("\n" + "=" * 60)
            log(f"STARTING {self._get_stage_name()}")
            log("=" * 60)

        for epoch in range(self.start_epoch, self.config.epochs):
            self._train_epoch(epoch)

        self._finish_training()

    def _train_epoch(self, epoch: int):
        """Train one epoch."""
        self.adapter.train()
        epoch_loss = 0
        accum_loss = 0

        pbar = tqdm(
            self.train_loader,
            desc=f"Epoch {epoch+1}/{self.config.epochs}",
            disable=not self.is_main,
        )

        for batch_idx, raw_batch in enumerate(pbar):
            loss = self._train_step(raw_batch, batch_idx)
            accum_loss += loss

            if self.accelerator.sync_gradients:
                self._update_after_step(accum_loss, pbar)
                epoch_loss += accum_loss
                accum_loss = 0

        self._finish_epoch(epoch, epoch_loss)

    def _check_numerical_stability(
        self,
        loss: torch.Tensor,
        params,
    ) -> Tuple[bool, str]:
        """
        Check for NaN/Inf in loss and gradients.

        Returns:
            (is_stable, reason): Tuple of stability flag and reason if unstable
        """
        # Check loss
        if torch.isnan(loss) or torch.isinf(loss):
            return False, f"loss={loss.item()}"

        # Check gradients
        for name, param in params:
            if param.grad is not None:
                if torch.isnan(param.grad).any():
                    return False, f"NaN gradient in {name}"
                if torch.isinf(param.grad).any():
                    return False, f"Inf gradient in {name}"

        return True, ""

    def _compute_grad_norm(self, params) -> float:
        """Compute total gradient norm for monitoring."""
        total_norm = 0.0
        for _, param in params:
            if param.grad is not None:
                param_norm = param.grad.data.norm(2)
                total_norm += param_norm.item() ** 2
        return math.sqrt(total_norm)

    def _train_step(self, raw_batch: Dict[str, Any], batch_idx: int) -> float:
        """Single training step with OOM and NaN/Inf handling.

        Stability features:
        - OOM recovery outside except clause (PyTorch FAQ best practice)
        - NaN/Inf detection in loss and gradients
        - Gradient norm monitoring
        - gc.collect() before empty_cache() for better cleanup

        See: https://pytorch.org/docs/stable/notes/faq.html
        """
        # Update text ratio (skip if no_decay mode for Stage 1)
        if not getattr(self.config, 'no_decay', False):
            self.current_text_ratio = get_text_ratio(
                self.global_step,
                self.effective_decay_steps,
                self.config.initial_text_ratio,
            )
        # else: keep current_text_ratio fixed at initial value

        # Dynamic max_seq_len based on text ratio
        dynamic_max_seq = self._get_dynamic_max_seq()

        # Apply interleaving
        batch = apply_interleaving(
            raw_batch,
            self.current_text_ratio,
            tokenizer=self.tokenizer,
            max_seq_len=dynamic_max_seq,
        )

        # Get adapter dtype for proper casting
        adapter_dtype = next(self.adapter.parameters()).dtype
        whisper = batch["whisper"].to(self.device, dtype=adapter_dtype)
        interleaved = batch["interleaved"].to(self.device)

        # Clear cache periodically
        if batch_idx % 100 == 0 and self.device.type == 'cuda':
            torch.cuda.empty_cache()

        # Forward and backward with OOM and NaN handling
        # Use flag pattern to move recovery outside except (PyTorch FAQ recommendation)
        oom_error = False
        nan_error = False
        error_reason = ""
        seq_len_for_log = interleaved.shape[1]
        loss_value = 0.0

        try:
            with self.accelerator.accumulate(*self._get_accumulate_models()):
                loss = self._compute_loss(whisper, interleaved)

                # Check for NaN/Inf in loss before backward
                if torch.isnan(loss) or torch.isinf(loss):
                    nan_error = True
                    error_reason = f"loss={loss.item()}"
                else:
                    self.accelerator.backward(loss)

                    if self.accelerator.sync_gradients:
                        # Check gradient stability
                        trainable_params = list(self._get_trainable_params_named())
                        is_stable, reason = self._check_numerical_stability(loss, trainable_params)

                        if not is_stable:
                            nan_error = True
                            error_reason = reason
                        else:
                            # Track gradient norm for monitoring
                            self.last_grad_norm = self._compute_grad_norm(trainable_params)
                            self.max_grad_norm_seen = max(self.max_grad_norm_seen, self.last_grad_norm)

                            # Clip gradients
                            self.accelerator.clip_grad_norm_(
                                [p for _, p in trainable_params],
                                self.config.max_grad_norm,
                            )

                    if not nan_error:
                        self.optimizer.step()
                        self.optimizer.zero_grad(set_to_none=True)
                        loss_value = loss.item()

        except torch.cuda.OutOfMemoryError:
            oom_error = True

        # Recovery outside except clause (prevents memory leak from exception stack frame)
        if oom_error:
            self.oom_count += 1
            if self.is_main:
                log(f"[OOM #{self.oom_count}] Skipping batch {batch_idx} (seq_len={seq_len_for_log})")
            del whisper, interleaved, batch
            gc.collect()
            torch.cuda.empty_cache()
            self.optimizer.zero_grad(set_to_none=True)
            return 0.0

        if nan_error:
            self.nan_count += 1
            if self.is_main:
                log(f"[NaN #{self.nan_count}] Skipping batch {batch_idx}: {error_reason}")
            self.optimizer.zero_grad(set_to_none=True)
            return 0.0

        return loss_value

    def _get_trainable_params_named(self):
        """Get named trainable parameters for gradient checking."""
        # Default: adapter parameters
        for name, param in self.adapter.named_parameters():
            if param.requires_grad:
                yield f"adapter.{name}", param

    def _get_dynamic_max_seq(self) -> int:
        """Get dynamic max sequence length based on text ratio.

        More conservative limits to prevent CUDA OOM.
        RTX 4090 (24GB) at 80% VRAM can handle ~1280-1536 max seq with LLM.
        """
        # Base limit more conservative for memory safety
        base_limit = min(self.config.max_seq_len, 1536)

        if self.current_text_ratio >= 0.7:
            return base_limit
        elif self.current_text_ratio >= 0.5:
            return int(base_limit * 0.75)  # ~1152
        elif self.current_text_ratio >= 0.3:
            return int(base_limit * 0.6)   # ~922
        else:
            return int(base_limit * 0.5)   # ~768

    def _compute_loss(
        self,
        whisper: torch.Tensor,
        interleaved: torch.Tensor,
    ) -> torch.Tensor:
        """
        Compute training loss with numerical stability.

        Numerical stability measures:
        - Clamp logits to prevent extreme values
        - Use BF16 which has better dynamic range than FP16
        """
        unwrapped_llm = self.accelerator.unwrap_model(self.llm)

        # Forward through adapter
        audio_embeds = self.adapter(whisper)

        # Get token embeddings
        input_tokens = interleaved[:, :-1].clamp(min=0)
        with torch.no_grad():
            token_embeds = unwrapped_llm.model.embed_tokens(input_tokens)

        # Combine embeddings
        combined = torch.cat([audio_embeds, token_embeds], dim=1)

        # Forward through LLM
        outputs = self.llm(inputs_embeds=combined, use_cache=False)
        logits = outputs.logits

        # Compute loss with numerical stability
        audio_len = audio_embeds.shape[1]
        seq_len = interleaved.shape[1]
        seq_logits = logits[:, audio_len-1:audio_len-1+seq_len]

        # Clamp logits to prevent numerical issues in softmax
        # Large logits can cause overflow in exp() during cross_entropy
        seq_logits = seq_logits.clamp(min=-100, max=100)

        return F.cross_entropy(
            seq_logits.reshape(-1, logits.size(-1)),
            interleaved.reshape(-1),
            ignore_index=-100,
            label_smoothing=self.config.label_smoothing,
        )

    @abstractmethod
    def _get_trainable_params(self):
        """Get trainable parameters for gradient clipping."""
        pass

    def _get_accumulate_models(self):
        """Get models for gradient accumulation context.

        Override to include additional models (e.g., LLM for full fine-tuning).
        With FSDP, all trained models must be passed to accumulate() so that
        gradient sync is properly skipped during accumulation steps.
        """
        return (self.adapter,)

    def _update_after_step(self, accum_loss: float, pbar):
        """Update after gradient accumulation step."""
        # Learning rate schedule
        if self.global_step < self.warmup_steps:
            lr_scale = (self.global_step + 1) / self.warmup_steps
            for pg in self.optimizer.param_groups:
                pg["lr"] = self.config.learning_rate * lr_scale
        else:
            self.scheduler.step()

        self.global_step += 1
        if self.is_main:
            write_step(self.global_step)

        # Update progress bar with gradient norm for stability monitoring
        avg_loss = accum_loss / self.config.grad_accum
        pbar.set_postfix(
            loss=f"{avg_loss:.4f}",
            lr=f"{self.optimizer.param_groups[0]['lr']:.2e}",
            text_ratio=f"{self.current_text_ratio:.1f}",
            grad=f"{self.last_grad_norm:.1f}",
        )

        # Log warning if gradient norm is very high (potential instability)
        if self.last_grad_norm > 100 and self.global_step % 50 == 0 and self.is_main:
            log(f"[WARN] High gradient norm: {self.last_grad_norm:.1f} at step {self.global_step}")

        # Save checkpoint
        if self.global_step % self.config.save_steps == 0 and self.is_main:
            self._save_step_checkpoint(accum_loss)

        # Periodic status log (flushes to file, unlike tqdm \r)
        if self.global_step % 50 == 0 and self.is_main:
            avg = accum_loss / self.config.grad_accum if self.config.grad_accum > 0 else accum_loss
            log(f"[Step {self.global_step}/{self.total_steps}] loss={avg:.4f} lr={self.optimizer.param_groups[0]['lr']:.2e} grad={self.last_grad_norm:.1f} text_ratio={self.current_text_ratio:.1f}")

    @abstractmethod
    def _save_step_checkpoint(self, loss: float):
        """Save step checkpoint (implemented by subclasses)."""
        pass

    @abstractmethod
    def _finish_epoch(self, epoch: int, epoch_loss: float):
        """Finish epoch (implemented by subclasses)."""
        pass

    def _finish_training(self):
        """Finish training."""
        self.accelerator.wait_for_everyone()

        if self.is_main:
            self.checkpoint_manager.wait_for_saves()
            log("\n" + "=" * 60)
            log(f"{self._get_stage_name()} COMPLETE!")
            log(f"Best loss: {self.best_loss:.4f}")
            log(f"Final text_ratio: {self.current_text_ratio:.1f}")
            log(f"Max gradient norm seen: {self.max_grad_norm_seen:.2f}")
            if self.oom_count > 0:
                log(f"OOM errors recovered: {self.oom_count}")
            if self.nan_count > 0:
                log(f"NaN/Inf errors recovered: {self.nan_count}")
            log("=" * 60)