nexus/training/trainer.py

"""
Nexus Trainer - Training loop cho Nexus Coder
==============================================
Hỗ trợ:
- Training với kiến trúc MoE
- Auxiliary loss (load balancing)
- Checkpointing
- Logging
- Mixed precision (fp16/bf16)
"""
import os
import json
import time
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from typing import Optional, Dict, Callable
from tqdm.auto import tqdm

from ..model.nexus_coder import NexusCoderForCausalLM
from ..config import NexusConfig
from .dataset import NexusDataset, AUTHOR_TRAINING_DATA


def get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps: int,
    num_training_steps: int,
    num_cycles: float = 0.5,
    last_epoch: int = -1,
):
    """Cosine LR schedule với warmup."""
    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        progress = float(current_step - num_warmup_steps) / float(
            max(1, num_training_steps - num_warmup_steps)
        )
        return max(0.0, 0.5 * (1.0 + __import__("math").cos(__import__("math").pi * num_cycles * 2.0 * progress)))

    return LambdaLR(optimizer, lr_lambda, last_epoch)


class NexusTrainer:
    """Trainer cho Nexus Coder."""

    def __init__(
        self,
        model: NexusCoderForCausalLM,
        config: NexusConfig,
        train_dataset: NexusDataset,
        output_dir: str = "./checkpoints",
        learning_rate: float = 5e-4,
        weight_decay: float = 0.01,
        warmup_steps: int = 100,
        max_steps: int = 5000,
        per_device_batch_size: int = 4,
        gradient_accumulation_steps: int = 4,
        logging_steps: int = 10,
        save_steps: int = 500,
        use_amp: bool = False,
        amp_dtype: torch.dtype = torch.float16,
    ):
        self.model = model
        self.config = config
        self.train_dataset = train_dataset
        self.output_dir = output_dir
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay
        self.warmup_steps = warmup_steps
        self.max_steps = max_steps
        self.per_device_batch_size = per_device_batch_size
        self.gradient_accumulation_steps = gradient_accumulation_steps
        self.logging_steps = logging_steps
        self.save_steps = save_steps
        self.use_amp = use_amp
        self.amp_dtype = amp_dtype

        # Device
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = self.model.to(self.device)

        # DataLoader
        self.dataloader = DataLoader(
            train_dataset,
            batch_size=per_device_batch_size,
            shuffle=True,
            num_workers=0,
            pin_memory=True,
        )

        # Optimizer
        self.optimizer = AdamW(
            model.parameters(),
            lr=learning_rate,
            weight_decay=weight_decay,
            betas=(0.9, 0.95),
            eps=1e-8,
        )

        # Scheduler
        total_steps = max_steps
        self.scheduler = get_cosine_schedule_with_warmup(
            self.optimizer,
            num_warmup_steps=warmup_steps,
            num_training_steps=total_steps,
        )

        # AMP scaler
        self.scaler = torch.amp.GradScaler("cuda") if use_amp and torch.cuda.is_available() else None

        # Logging
        os.makedirs(output_dir, exist_ok=True)
        self.log_history = []

    def train(self, resume_from_checkpoint: Optional[str] = None) -> Dict:
        """Bắt đầu training."""
        print("=" * 60)
        print(f"  Nexus Coder Training")
        print(f"  Tác giả: {self.config.author}")
        print(f"  Device:  {self.device}")
        print(f"  Steps:   {self.max_steps}")
        print("=" * 60)

        global_step = 0
        if resume_from_checkpoint and os.path.exists(resume_from_checkpoint):
            global_step = self._load_checkpoint(resume_from_checkpoint)

        self.model.train()
        start_time = time.time()

        # Training loop
        dataloader_iter = iter(self.dataloader)
        accumulated_loss = 0.0

        progress_bar = tqdm(range(global_step, self.max_steps), desc="Training")
        for step in progress_bar:
            try:
                batch = next(dataloader_iter)
            except StopIteration:
                dataloader_iter = iter(self.dataloader)
                batch = next(dataloader_iter)

            input_ids = batch["input_ids"].to(self.device)
            labels = batch["labels"].to(self.device)
            attention_mask = batch["attention_mask"].to(self.device)

            # Forward
            if self.use_amp and torch.cuda.is_available():
                with torch.amp.autocast("cuda", dtype=self.amp_dtype):
                    outputs = self.model(
                        input_ids=input_ids,
                        attention_mask=attention_mask,
                        labels=labels,
                    )
                    loss = outputs["loss"] / self.gradient_accumulation_steps
                self.scaler.scale(loss).backward()
            else:
                outputs = self.model(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    labels=labels,
                )
                loss = outputs["loss"] / self.gradient_accumulation_steps
                loss.backward()

            accumulated_loss += loss.item()

            # Optimizer step
            if (step + 1) % self.gradient_accumulation_steps == 0:
                if self.use_amp and torch.cuda.is_available():
                    self.scaler.unscale_(self.optimizer)
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                    self.scaler.step(self.optimizer)
                    self.scaler.update()
                else:
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                    self.optimizer.step()
                self.optimizer.zero_grad()
                self.scheduler.step()
                global_step += 1

                # Logging
                if global_step % self.logging_steps == 0:
                    avg_loss = accumulated_loss / self.logging_steps
                    elapsed = time.time() - start_time
                    lr = self.scheduler.get_last_lr()[0]
                    log_entry = {
                        "step": global_step,
                        "loss": avg_loss,
                        "learning_rate": lr,
                        "elapsed_seconds": elapsed,
                    }
                    self.log_history.append(log_entry)
                    progress_bar.set_postfix({
                        "loss": f"{avg_loss:.4f}",
                        "lr": f"{lr:.2e}",
                    })
                    accumulated_loss = 0.0

                # Save checkpoint
                if global_step % self.save_steps == 0:
                    self._save_checkpoint(global_step)

            if global_step >= self.max_steps:
                break

        # Final save
        self._save_checkpoint(global_step, final=True)

        # Save log
        self._save_logs()

        elapsed = time.time() - start_time
        print(f"\n✅ Training hoàn thành trong {elapsed:.1f}s")
        return {"global_step": global_step, "elapsed": elapsed}

    def _save_checkpoint(self, step: int, final: bool = False) -> None:
        """Lưu checkpoint (v0.4: include AMP scaler state for safe resume)."""
        suffix = "final" if final else f"step-{step}"
        path = os.path.join(self.output_dir, f"nexus_coder-{suffix}.pt")
        # v0.4 fix: persist GradScaler state so AMP can resume safely without
        # scale-factor NaNs on first few steps.
        scaler_state = None
        scaler = getattr(self, "scaler", None)
        if scaler is not None and hasattr(scaler, "state_dict"):
            try:
                scaler_state = scaler.state_dict()
            except Exception:
                scaler_state = None
        torch.save({
            "model_state_dict": self.model.state_dict(),
            "optimizer_state_dict": self.optimizer.state_dict(),
            "scheduler_state_dict": self.scheduler.state_dict(),
            "scaler_state_dict": scaler_state,
            "step": step,
            "config": self.config.__dict__,
        }, path)
        print(f"  Checkpoint saved: {path}")

    def _load_checkpoint(self, path: str) -> int:
        """Load checkpoint (v0.4: also restore AMP scaler if present)."""
        checkpoint = torch.load(path, map_location=self.device, weights_only=False)
        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"])
        # v0.4 fix: restore scaler state if present
        scaler_state = checkpoint.get("scaler_state_dict")
        scaler = getattr(self, "scaler", None)
        if scaler_state is not None and scaler is not None and hasattr(scaler, "load_state_dict"):
            try:
                scaler.load_state_dict(scaler_state)
            except Exception:
                pass
        step = checkpoint.get("step", 0)
        print(f"  Resumed from checkpoint at step {step}")
        return step

    def _save_logs(self) -> None:
        """Lưu training logs."""
        log_path = os.path.join(self.output_dir, "training_log.json")
        with open(log_path, "w", encoding="utf-8") as f:
            json.dump(self.log_history, f, ensure_ascii=False, indent=2)
        print(f"  📝 Saved training log: {log_path}")