trainer.py

"""
Transformer模型训练器
支持分布式训练、混合精度、梯度累积等特性
训练日志会保存在TensorBoard中，使用命令查看: tensorboard --logdir=<log_dir>
"""
import os
import shutil
import time
import torch
from datetime import datetime
from typing import Optional
from accelerate import Accelerator, DistributedDataParallelKwargs
from transformers import get_cosine_schedule_with_warmup
from tqdm.auto import tqdm

class TransformerTrainer:
    """Transformer模型训练器

    封装了模型训练的完整流程，包括：
    - 优化器和学习率调度器初始化
    - 分布式训练支持
    - 训练和评估循环
    - 模型检查点保存
    - TensorBoard日志记录
    """

    def __init__(self, config, model, train_dataloader, test_dataloader=None, resume_path=None):
        self.config = config
        self.model = model
        self.train_dataloader = train_dataloader
        self.test_dataloader = test_dataloader
        self.global_step = 0
        self.best_eval_loss = float("inf")
        self.completed_epochs = 0

        self.optimizer = self._setup_optimizer()
        self.lr_scheduler = self._setup_lr_scheduler()
        self.accelerator = self._initialize_accelerator()
        self._prepare_training_components()
        self.accelerator.init_trackers(self.config.tensorboard_log_name)
        self.test_interval_steps = self._calculate_test_interval()

        # 断点续训
        if resume_path:
            self._load_training_state(resume_path)

    def _setup_optimizer(self) -> torch.optim.Optimizer:
        """配置优化器（bias / LayerNorm 不施加 weight decay）"""
        no_decay_keywords = {"bias", "layernorm", "layer_norm", "norm"}
        decay_params, no_decay_params = [], []
        for name, param in self.model.named_parameters():
            if not param.requires_grad:
                continue
            if any(kw in name.lower() for kw in no_decay_keywords):
                no_decay_params.append(param)
            else:
                decay_params.append(param)
        return torch.optim.AdamW(
            [
                {"params": decay_params, "weight_decay": 0.1},
                {"params": no_decay_params, "weight_decay": 0.0},
            ],
            lr=self.config.learning_rate,
            betas=(0.9, 0.99),
        )

    def _setup_lr_scheduler(self):
        """配置学习率调度器"""
        num_training_steps = int(
            (len(self.train_dataloader) * self.config.num_epochs * 1.5)
            / self.config.gradient_accumulation_steps
        )
        num_warmup_steps = int(num_training_steps * self.config.lr_warmup_ratio)
        return get_cosine_schedule_with_warmup(
            optimizer=self.optimizer,
            num_warmup_steps=num_warmup_steps,
            num_training_steps=num_training_steps,
        )

    def _initialize_accelerator(self) -> Accelerator:
        """初始化Accelerator用于分布式训练和混合精度"""
        ddp_kwargs = DistributedDataParallelKwargs()
        accelerator = Accelerator(
            mixed_precision=self.config.mixed_precision,
            gradient_accumulation_steps=self.config.gradient_accumulation_steps,
            log_with="tensorboard",
            project_dir=self.config.tensorboard_log_dir,
            device_placement=True,
            kwargs_handlers=[ddp_kwargs]
        )

        # 创建日志目录
        if accelerator.is_main_process and self.config.tensorboard_log_dir:
            os.makedirs(self.config.tensorboard_log_dir, exist_ok=True)

        return accelerator

    def _prepare_training_components(self):
        """准备训练组件（处理分布式、混合精度等）"""
        components = [self.model, self.optimizer, self.train_dataloader, self.lr_scheduler]

        if self.test_dataloader is not None:
            components.insert(3, self.test_dataloader)
            prepared = self.accelerator.prepare(*components)
            self.model, self.optimizer, self.train_dataloader, self.test_dataloader, self.lr_scheduler = prepared
        else:
            prepared = self.accelerator.prepare(*components)
            self.model, self.optimizer, self.train_dataloader, self.lr_scheduler = prepared

    def _calculate_test_interval(self) -> Optional[int]:
        """计算测试评估的间隔步数"""
        if self.config.use_test_set and self.test_dataloader is not None:
            interval = int(len(self.train_dataloader) * self.config.test_frequency)
            print(f"测试间隔: 每 {interval} 步测试一次 (约{self.config.test_frequency}个epoch)")
            return interval
        return None

    def train(self):
        """执行完整的训练流程"""
        for epoch in range(self.completed_epochs, self.config.num_epochs):
            self._train_one_epoch(epoch)
            self._save_checkpoint_if_needed(epoch)
            # 每个 epoch 结束保存可恢复的训练状态
            self.completed_epochs = epoch + 1
            self._save_training_state()

    def _train_one_epoch(self, epoch: int):
        self.model.train()
        progress_bar = tqdm(total=len(self.train_dataloader), disable=not self.config.log)
        progress_bar.set_description(f"Epoch {epoch}")
        epoch_start_step = self.global_step

        # 吞吐量计数器
        tokens_since_log = 0
        time_since_log = time.time()

        for batch in self.train_dataloader:
            loss = self._training_step(batch)
            tokens_since_log += batch["input_ids"].numel()

            # 记录日志（含吞吐量）
            if self.global_step % 10 == 0:
                elapsed = time.time() - time_since_log
                tokens_per_sec = tokens_since_log / max(elapsed, 1e-6)
                self._log_training_metrics(loss, tokens_per_sec)
                tokens_since_log = 0
                time_since_log = time.time()

            if self._should_evaluate(epoch_start_step):
                self._evaluate_test()

            if self._should_save_checkpoint():
                self._save_checkpoint(f"steps_{self.global_step}")

            if self.config.log:
                progress_bar.update(1)
                progress_bar.set_postfix(
                    loss=loss.detach().item(),
                    lr=self.lr_scheduler.get_last_lr()[0]
                )
            self.global_step += 1

    def _training_step(self, batch) -> torch.Tensor:
        """执行一个训练步骤

        Args:
            batch: 输入的批次数据

        Returns:
            当前批次的损失值
        """
        with self.accelerator.accumulate(self.model):
            outputs = self.model(
                input_ids=batch["input_ids"],
                labels=batch["labels"],
                attention_mask=batch["attention_mask"],
                position_ids=batch.get("position_ids"),
            )
            loss = outputs.loss

            # 反向传播
            self.accelerator.backward(loss)

            # 梯度裁剪
            if self.accelerator.sync_gradients:
                self.accelerator.clip_grad_norm_(self.model.parameters(), 1.0)

            # 优化器步进
            self.optimizer.step()
            self.lr_scheduler.step()
            self.optimizer.zero_grad()

        return loss

    def _log_training_metrics(self, loss: torch.Tensor, tokens_per_sec: float = 0):
        """记录训练指标到TensorBoard"""
        metrics = {
            "train/loss": loss.detach().item(),
            "train/learning_rate": self.lr_scheduler.get_last_lr()[0],
            "train/tokens_per_sec": tokens_per_sec,
            "train/step": self.global_step,
        }
        self.accelerator.log(metrics, step=self.global_step)

    def _should_evaluate(self, epoch_start_step: int) -> bool:
        """判断是否需要进行测试评估"""
        if self.test_interval_steps is None:
            return False

        return (
            self.global_step % self.test_interval_steps == 0
            and self.global_step > epoch_start_step
        )

    def _evaluate_test(self):
        """在测试集上评估模型"""
        if self.test_dataloader is None:
            return

        self.accelerator.wait_for_everyone()
        self.model.eval()
        total_loss = 0.0
        num_batches = 0

        if self.accelerator.is_main_process:
            print(f"\n开始测试集评估 (step {self.global_step})...")

        with torch.no_grad():
            eval_bar = tqdm(
                self.test_dataloader,
                desc="测试中",
                disable=not self.accelerator.is_main_process
            )
            for batch in eval_bar:
                outputs = self.model(
                    input_ids=batch["input_ids"],
                    labels=batch["labels"],
                    attention_mask=batch["attention_mask"],
                    position_ids=batch.get("position_ids"),
                )
                gathered_losses = self.accelerator.gather_for_metrics(
                    outputs.loss.detach().reshape(1)
                )
                total_loss += gathered_losses.sum().item()
                num_batches += gathered_losses.numel()

        # 计算平均损失和困惑度
        avg_loss = total_loss / max(num_batches, 1)
        perplexity = torch.exp(torch.tensor(avg_loss)).item()

        # 记录测试指标
        if self.config.test_save_results and self.accelerator.is_main_process:
            test_metrics = {
                "test/loss": avg_loss,
                "test/perplexity": perplexity,
                "test/step": self.global_step,
            }
            self.accelerator.log(test_metrics, step=self.global_step)

        if self.accelerator.is_main_process:
            print(f"测试集 - 损失: {avg_loss:.4f}, 困惑度: {perplexity:.4f}")

            if avg_loss < self.best_eval_loss:
                self.best_eval_loss = avg_loss
                self._save_checkpoint("best_model")
                print(f"best model 已更新: loss={self.best_eval_loss:.4f}")

        self.model.train()
        self.accelerator.wait_for_everyone()

    def _should_save_checkpoint(self) -> bool:
        """判断是否需要保存检查点（按步数）"""
        if self.config.save_steps is None or self.config.save_steps <= 0:
            return False

        return (
            self.global_step % self.config.save_steps == 0
            and self.accelerator.is_main_process
        )

    def _save_checkpoint_if_needed(self, epoch: int):
        """根据epoch判断是否需要保存检查点"""
        if not self.accelerator.is_main_process:
            return

        should_save = (
            (epoch + 1) % self.config.save_model_epochs == 0
            or epoch == self.config.num_epochs - 1
        )

        if should_save:
            self._save_checkpoint(f"epoch_{epoch}")

    def _save_checkpoint(self, prefix: str):
        """保存模型权重（用于推理）"""
        with torch.no_grad():
            unwrapped_model = self.accelerator.unwrap_model(self.model)
            # 兼容 torch.compile: 取出原始模型
            if hasattr(unwrapped_model, '_orig_mod'):
                unwrapped_model = unwrapped_model._orig_mod
            if prefix == "best_model":
                save_path = os.path.join(self.config.output_dir, prefix)
                if os.path.isdir(save_path):
                    shutil.rmtree(save_path)
            else:
                timestamp = datetime.now().strftime("%m%d_%H%M")
                save_path = f"{self.config.output_dir}/{prefix}_{timestamp}"
            unwrapped_model.save_pretrained(save_path)
            print(f"模型已保存至: {save_path}")

    def _save_training_state(self):
        """保存完整训练状态（可断点续训）"""
        save_path = os.path.join(self.config.output_dir, "latest")
        # Accelerate 保存 model / optimizer / scheduler / scaler / rng
        self.accelerator.save_state(save_path)
        if self.accelerator.is_main_process:
            torch.save({
                'global_step': self.global_step,
                'best_eval_loss': self.best_eval_loss,
                'completed_epochs': self.completed_epochs,
            }, os.path.join(save_path, 'training_meta.pt'))
            print(f"训练状态已保存至: {save_path}")

    def _load_training_state(self, load_path):
        """从检查点恢复完整训练状态"""
        self.accelerator.load_state(load_path)
        meta_path = os.path.join(load_path, 'training_meta.pt')
        if os.path.exists(meta_path):
            meta = torch.load(meta_path, map_location='cpu')
            self.global_step = meta['global_step']
            self.best_eval_loss = meta['best_eval_loss']
            self.completed_epochs = meta['completed_epochs']
        print(f"已恢复训练状态: epoch={self.completed_epochs}, step={self.global_step}, best_loss={self.best_eval_loss:.4f}")