training/pretrain/trainer_v2.py

#!/usr/bin/env python3
"""
改进的多任务训练器（修正版）
修复/增强点：
1) Scheduler/ warmup 步数按 optimizer update steps（考虑 gradient_accumulation_steps）
2) eval_steps / save_steps 按 global_step(update step) 触发，而不是按 epoch
3) 支持最后不足 accumulation 的尾 batch 也执行一次 optimizer.step()
4) 支持 AMP（bf16 / fp16），fp16 使用 GradScaler
5) optimizer 仅包含 requires_grad=True 参数（适配 LoRA）
6) 训练/验证指标写入 jsonl，便于可视化与复现实验
7) 尝试启用 use_reentrant=False 的 gradient checkpointing（更推荐的实现）
"""

import os
import math
import json
from dataclasses import asdict
from datetime import datetime
from pathlib import Path
from contextlib import nullcontext

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from transformers import get_linear_schedule_with_warmup
from tqdm import tqdm
import numpy as np

try:
    import wandb
except Exception:
    wandb = None

from model_loader import (
    load_model_and_processor,
    prepare_model_inputs
)
from config import PretrainConfig


class FocalLoss(nn.Module):
    """
    Focal Loss for accident detection (保留，当前训练主损失仍使用 outputs.loss)
    """
    def __init__(self, alpha=0.25, gamma=2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(self, inputs, targets):
        ce_loss = F.cross_entropy(inputs, targets, reduction='none')
        p_t = torch.exp(-ce_loss)
        focal_loss = self.alpha * (1 - p_t) ** self.gamma * ce_loss
        return focal_loss.mean()


class MultiTaskTrainer:
    """
    多任务训练器
    - 支持单帧任务 + 序列任务（取决于 batch 是否包含 "single_frame"/"sequence"）
    - 使用 LM loss (outputs.loss) 作为主要训练信号
    """

    def __init__(
        self,
        config: PretrainConfig,
        train_loader: DataLoader,
        val_loader: DataLoader,
        pretrained_lora_path: str | None = None,
    ):
        self.config = config
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.pretrained_lora_path = pretrained_lora_path

        # 设备
        self.device = torch.device(config.training.device)
        print(f"使用设备: {self.device}")

        # 训练统计
        self.global_step = 0  # optimizer update steps
        self.best_val_loss = float("inf")

        # 输出目录
        self.output_dir = Path(self.config.training.output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        # 日志落盘（jsonl）
        self.train_log_path = self.output_dir / "train_metrics.jsonl"
        self.val_log_path = self.output_dir / "val_metrics.jsonl"

        # Loss functions（保留）
        self.ce_loss = nn.CrossEntropyLoss()
        self.focal_loss = FocalLoss(alpha=0.25, gamma=2.0)

        # Task weights（你的事故任务 + 兼容 bdd_* 默认 1.0）
        self.task_weights = {
            "scene_understanding": self.config.data.task1_weight,
            "binary_detection": self.config.data.task2_weight,
            "accident_description": self.config.data.task3_weight,
            "sequence_prediction": self.config.data.task3_weight,
            # bdd_attributes / bdd_detection / bdd_drivable / bdd_risk 等，缺省走 1.0
        }

        # Curriculum state（保留：注意你目前并未实际重建 dataloader）
        self.current_stage = 0
        self.stage_epochs = [1, 2, 2]

        # AMP 配置
        self.use_bf16 = bool(getattr(self.config.training, "bf16", False))
        self.use_fp16 = bool(getattr(self.config.training, "fp16", False)) and (not self.use_bf16)
        self.autocast_dtype = torch.bfloat16 if self.use_bf16 else (torch.float16 if self.use_fp16 else None)
        self.scaler = torch.cuda.amp.GradScaler(enabled=self.use_fp16)

        # 初始化
        self._init_model()
        self._init_optimizer()

        # wandb
        if self.config.training.use_wandb:
            if wandb is None:
                print("⚠️  wandb 未安装或不可用，但 use_wandb=True；将跳过 wandb 记录。")
                self.config.training.use_wandb = False
            else:
                wandb.init(
                    project=self.config.training.wandb_project,
                    name=self.config.training.wandb_run_name or f"{self.config.model.model_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
                    config=asdict(self.config),
                )

    # ------------------------- utils -------------------------
    def _autocast_ctx(self):
        if self.autocast_dtype is None:
            return nullcontext()
        return torch.cuda.amp.autocast(dtype=self.autocast_dtype)

    def _write_jsonl(self, path: Path, record: dict):
        record = dict(record)
        record.setdefault("time", datetime.now().isoformat(timespec="seconds"))
        with open(path, "a", encoding="utf-8") as f:
            f.write(json.dumps(record, ensure_ascii=False) + "\n")

    def _move_batch_to_device(self, inputs: dict):
        """
        将 processor 输出的 BatchFeature/dict 移动到 device，并对浮点张量对齐到模型 dtype。
        """
        moved = {}
        for k, v in inputs.items():
            if torch.is_tensor(v):
                if v.is_floating_point():
                    moved[k] = v.to(self.device, dtype=self.model.dtype)
                else:
                    moved[k] = v.to(self.device)
            else:
                moved[k] = v
        return moved

    # ------------------------- init model/optim -------------------------
    def _init_model(self):
        print("=" * 60)
        print("加载模型...")

        self.model, self.processor = load_model_and_processor(self.config.model)

        # 可选：Stage B 加载 Stage A 的 LoRA adapter 继续训练
        if self.pretrained_lora_path:
            try:
                from peft import PeftModel
                self.model = PeftModel.from_pretrained(
                    self.model,
                    self.pretrained_lora_path,
                    is_trainable=True
                )
                print(f"✓ 已加载LoRA权重: {self.pretrained_lora_path}")
            except Exception as e:
                print(f"⚠️  加载LoRA权重失败（将继续使用当前模型的LoRA状态）: {e}")

        self.model.to(self.device)

        # tokenizer pad_token
        if self.processor.tokenizer.pad_token is None:
            self.processor.tokenizer.pad_token = self.processor.tokenizer.eos_token
            self.processor.tokenizer.pad_token_id = self.processor.tokenizer.eos_token_id

        # 尽量启用更推荐的 checkpoint 实现（避免 use_reentrant 默认变化/限制）
        if hasattr(self.model, "gradient_checkpointing_enable"):
            try:
                self.model.gradient_checkpointing_enable(
                    gradient_checkpointing_kwargs={"use_reentrant": False}
                )
            except TypeError:
                # 旧版本 transformers/pytorch 不支持 kwargs
                try:
                    self.model.gradient_checkpointing_enable()
                except Exception:
                    pass

        try:
            if hasattr(self.model, "config"):
                self.model.config.use_cache = False
        except Exception:
            pass

        print(f"✓ 模型加载完成: {self.config.model.model_name}")
        print("=" * 60)

    def _init_optimizer(self):
        # 只优化 requires_grad=True 的参数（LoRA 正常情况下只有 adapter 参数是 True）
        trainable_params = [p for p in self.model.parameters() if p.requires_grad]
        if len(trainable_params) == 0:
            raise RuntimeError("没有可训练参数（requires_grad=True 为 0）。请检查 LoRA/冻结策略。")

        self.optimizer = torch.optim.AdamW(
            trainable_params,
            lr=self.config.training.learning_rate,
            weight_decay=self.config.training.weight_decay
        )

        # Scheduler：关键修复 —— total steps 必须按 optimizer update steps 计算
        grad_acc = max(1, int(self.config.training.gradient_accumulation_steps))
        if len(self.train_loader) == 0:
            raise RuntimeError("train_loader 为空（len=0），请检查数据是否正确生成。")

        updates_per_epoch = math.ceil(len(self.train_loader) / grad_acc)
        num_training_steps = updates_per_epoch * int(self.config.training.num_epochs)
        num_warmup_steps = int(num_training_steps * float(self.config.training.warmup_ratio))

        self.scheduler = get_linear_schedule_with_warmup(
            self.optimizer,
            num_warmup_steps=num_warmup_steps,
            num_training_steps=num_training_steps
        )

        print("✓ 优化器初始化完成")
        print(f"  batches/epoch: {len(self.train_loader)}")
        print(f"  grad_accum: {grad_acc}")
        print(f"  updates/epoch: {updates_per_epoch}")
        print(f"  total update steps: {num_training_steps}")
        print(f"  warmup steps: {num_warmup_steps}")

    # ------------------------- label construction helpers -------------------------
    def _concat_answers_to_prompt_inputs(self, prompt_inputs, labels_text):
        """
        Fallback：在无法用 processor 重新编码 full_texts 时，把答案 tokens 拼接到 prompt input_ids 后面并构造 labels。
        """
        tokenizer = self.processor.tokenizer
        pad_id = tokenizer.pad_token_id
        eos_id = tokenizer.eos_token_id

        input_ids = prompt_inputs["input_ids"]
        attention_mask = prompt_inputs["attention_mask"]

        if input_ids.dim() != 2 or attention_mask.dim() != 2:
            raise ValueError("prompt_inputs 必须包含二维的 input_ids 和 attention_mask")

        B, L = input_ids.shape
        prompt_lens = attention_mask.sum(dim=1).tolist()

        answer_ids_list = [
            tokenizer.encode(ans, add_special_tokens=False) + ([eos_id] if eos_id is not None else [])
            for ans in labels_text
        ]

        max_full_len = max(int(pl) + len(ans_ids) for pl, ans_ids in zip(prompt_lens, answer_ids_list))

        new_input_ids = input_ids.new_full((B, max_full_len), pad_id)
        new_attention_mask = attention_mask.new_zeros((B, max_full_len))
        new_labels = input_ids.new_full((B, max_full_len), -100)

        for i, (pl, ans_ids) in enumerate(zip(prompt_lens, answer_ids_list)):
            pl = int(pl)
            ans_tensor = torch.tensor(ans_ids, device=input_ids.device, dtype=input_ids.dtype)
            seq = torch.cat([input_ids[i, :pl], ans_tensor], dim=0)
            seq_len = seq.size(0)

            new_input_ids[i, :seq_len] = seq
            new_attention_mask[i, :seq_len] = 1

            new_labels[i, :seq_len] = seq
            new_labels[i, :pl] = -100

        out = {}
        for k, v in prompt_inputs.items():
            if k in ("input_ids", "attention_mask", "labels"):
                continue
            if torch.is_tensor(v) and v.dim() == 2 and v.shape[0] == B and v.shape[1] == L:
                continue
            out[k] = v

        out["input_ids"] = new_input_ids
        out["attention_mask"] = new_attention_mask
        out["labels"] = new_labels
        return out

    def prepare_inputs_and_labels(self, batch_data):
        """
        单帧任务：labels 与模型真实 input_ids 对齐（包含视觉 token）
        """
        images = batch_data["images"]
        user_prompts = batch_data["user_prompts"]
        labels_text = batch_data["labels"]
        task_types = batch_data["task"]

        prompt_inputs = prepare_model_inputs(
            self.processor,
            self.config.model.model_type,
            images,
            user_prompts,
            self.device
        )

        prompt_texts = prompt_inputs.pop("__prompt_texts__", None)
        if prompt_texts is None:
            raise RuntimeError("prepare_model_inputs 未返回 __prompt_texts__，无法构造对齐 labels")

        full_texts = [
            prompt + answer + self.processor.tokenizer.eos_token
            for prompt, answer in zip(prompt_texts, labels_text)
        ]

        try:
            prompt_encodings = self.processor(
                text=prompt_texts,
                images=images,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=512
            )
            full_inputs = self.processor(
                text=full_texts,
                images=images,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=512
            )

            labels = full_inputs["input_ids"].clone()
            for i in range(labels.size(0)):
                prompt_len = int(prompt_encodings["attention_mask"][i].sum().item())
                labels[i, :prompt_len] = -100
            labels[full_inputs["attention_mask"] == 0] = -100

            full_inputs["labels"] = labels
            full_inputs = self._move_batch_to_device(full_inputs)
            return full_inputs, labels_text, task_types

        except Exception:
            fallback_inputs = self._concat_answers_to_prompt_inputs(prompt_inputs, labels_text)
            fallback_inputs = self._move_batch_to_device(fallback_inputs)
            return fallback_inputs, labels_text, task_types

    def prepare_sequence_inputs_and_labels(self, batch_data):
        """
        序列任务：images 是 List[List[PIL]] 或等价格式
        """
        images_list = batch_data["image_sequences"]
        user_prompts = batch_data["user_prompts"]
        labels_text = batch_data["labels"]
        task_types = batch_data["task"]

        prompt_inputs = prepare_model_inputs(
            self.processor,
            self.config.model.model_type,
            images_list,
            user_prompts,
            self.device
        )

        prompt_texts = prompt_inputs.pop("__prompt_texts__", None)
        if prompt_texts is None:
            raise RuntimeError("prepare_model_inputs 未返回 __prompt_texts__，无法构造对齐 labels")

        full_texts = [
            prompt + answer + self.processor.tokenizer.eos_token
            for prompt, answer in zip(prompt_texts, labels_text)
        ]

        try:
            prompt_encodings = self.processor(
                text=prompt_texts,
                images=images_list,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=512
            )
            full_inputs = self.processor(
                text=full_texts,
                images=images_list,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=512
            )

            labels = full_inputs["input_ids"].clone()
            for i in range(labels.size(0)):
                prompt_len = int(prompt_encodings["attention_mask"][i].sum().item())
                labels[i, :prompt_len] = -100
            labels[full_inputs["attention_mask"] == 0] = -100

            full_inputs["labels"] = labels
            full_inputs = self._move_batch_to_device(full_inputs)
            return full_inputs, labels_text, task_types

        except Exception:
            fallback_inputs = self._concat_answers_to_prompt_inputs(prompt_inputs, labels_text)
            fallback_inputs = self._move_batch_to_device(fallback_inputs)
            return fallback_inputs, labels_text, task_types

    # ------------------------- loss / forward -------------------------
    def compute_loss(self, batch):
        """
        计算 batch 总 loss（支持单帧 + 序列）
        """
        total_loss = 0.0
        task_losses = {}
        n_tasks = 0

        with self._autocast_ctx():
            if "single_frame" in batch:
                sf_data = batch["single_frame"]
                inputs, _, task_types = self.prepare_inputs_and_labels(sf_data)
                outputs = self.model(**inputs)
                loss = outputs.loss

                task_type = task_types[0]
                weight = float(self.task_weights.get(task_type, 1.0))
                total_loss = total_loss + (loss * weight)

                task_losses[task_type] = float(loss.detach().item())
                n_tasks += 1

            if "sequence" in batch:
                seq_data = batch["sequence"]
                inputs, _, task_types = self.prepare_sequence_inputs_and_labels(seq_data)
                outputs = self.model(**inputs)
                loss = outputs.loss

                task_type = task_types[0]
                weight = float(self.task_weights.get(task_type, 1.0))
                total_loss = total_loss + (loss * weight)

                task_losses[task_type] = float(loss.detach().item())
                n_tasks += 1

        if n_tasks > 0:
            total_loss = total_loss / n_tasks

        return total_loss, task_losses

    # ------------------------- eval / save -------------------------
    @torch.no_grad()
    def evaluate(self, epoch: int):
        self.model.eval()
        val_loss_sum = 0.0
        val_task_losses = {}

        for batch in tqdm(self.val_loader, desc="Validation"):
            loss, task_losses = self.compute_loss(batch)
            val_loss_sum += float(loss.detach().item())
            for t, v in task_losses.items():
                val_task_losses.setdefault(t, []).append(v)

        val_loss = val_loss_sum / max(1, len(self.val_loader))
        avg_task_losses = {t: float(np.mean(vs)) for t, vs in val_task_losses.items()}

        record = {"step": self.global_step, "epoch": epoch, "val/loss": float(val_loss)}
        for t, v in avg_task_losses.items():
            record[f"val/{t}"] = float(v)

        self._write_jsonl(self.val_log_path, record)
        if self.config.training.use_wandb and wandb is not None:
            wandb.log(record, step=self.global_step)

        return val_loss, avg_task_losses

    def _rotate_checkpoints(self):
        limit = int(getattr(self.config.training, "save_total_limit", 0) or 0)
        if limit <= 0:
            return

        ckpts = sorted(
            [p for p in self.output_dir.glob("checkpoint-*") if p.is_dir()],
            key=lambda p: int(p.name.split("-")[-1]) if p.name.split("-")[-1].isdigit() else 0
        )

        if len(ckpts) <= limit:
            return

        for p in ckpts[:-limit]:
            try:
                for sub in sorted(p.rglob("*"), reverse=True):
                    if sub.is_file():
                        sub.unlink()
                    elif sub.is_dir():
                        sub.rmdir()
                p.rmdir()
            except Exception:
                pass

    def save_checkpoint(self, tag: str, is_best: bool = False):
        if is_best:
            save_dir = self.output_dir / "best_model"
        else:
            save_dir = self.output_dir / tag

        save_dir.mkdir(parents=True, exist_ok=True)

        # 保存模型（LoRA/PEFT 会保存 adapter 权重）
        if hasattr(self.model, "save_pretrained"):
            self.model.save_pretrained(save_dir)
        else:
            torch.save(self.model.state_dict(), save_dir / "pytorch_model.bin")

        # 保存 processor
        self.processor.save_pretrained(save_dir)

        # 保存 trainer state
        torch.save({
            "epoch_tag": tag,
            "global_step": self.global_step,
            "optimizer_state_dict": self.optimizer.state_dict(),
            "scheduler_state_dict": self.scheduler.state_dict(),
            "best_val_loss": self.best_val_loss,
        }, save_dir / "trainer_state.pt")

        print(f"✓ 保存checkpoint: {save_dir}")

        if not is_best:
            self._rotate_checkpoints()

    # ------------------------- train loop -------------------------
    def train(self):
        print("\n" + "=" * 60)
        print("开始训练")
        print("=" * 60)

        grad_acc = max(1, int(self.config.training.gradient_accumulation_steps))
        logging_steps = max(1, int(self.config.training.logging_steps))

        # 注意：这里的 eval_steps / save_steps 现在是按 global_step(update step) 生效
        eval_steps = int(self.config.training.eval_steps) if self.config.training.eval_steps else 0
        save_steps = int(self.config.training.save_steps) if self.config.training.save_steps else 0

        total_epochs = int(self.config.training.num_epochs)

        # logging window stats
        window_task_sum = {}
        window_task_cnt = {}
        window_loss_sum = 0.0
        window_updates = 0

        for epoch in range(total_epochs):
            self.model.train()

            print(f"\n{'='*60}")
            print(f"Epoch {epoch+1}/{total_epochs}")
            print(f"Curriculum Stage: {self.current_stage} ({['easy','medium','hard','all'][min(self.current_stage,3)]})")
            print("=" * 60)

            self.optimizer.zero_grad(set_to_none=True)
            pbar = tqdm(self.train_loader, desc=f"Epoch {epoch+1}/{total_epochs}")

            for step, batch in enumerate(pbar):
                loss, task_losses = self.compute_loss(batch)

                # 梯度累积：缩放 loss
                loss_for_backward = loss / grad_acc

                if self.scaler.is_enabled():
                    self.scaler.scale(loss_for_backward).backward()
                else:
                    loss_for_backward.backward()

                # accumulate task stats（记录未缩放 loss）
                for t, v in task_losses.items():
                    window_task_sum[t] = window_task_sum.get(t, 0.0) + float(v)
                    window_task_cnt[t] = window_task_cnt.get(t, 0) + 1

                # 是否该更新（包含尾 batch flush）
                do_update = ((step + 1) % grad_acc == 0) or ((step + 1) == len(self.train_loader))
                if not do_update:
                    continue

                # clip grad
                if self.scaler.is_enabled():
                    self.scaler.unscale_(self.optimizer)

                torch.nn.utils.clip_grad_norm_(
                    self.model.parameters(),
                    float(self.config.training.max_grad_norm)
                )

                # optimizer.step + scheduler.step（顺序很重要：先 optimizer 再 scheduler）
                if self.scaler.is_enabled():
                    self.scaler.step(self.optimizer)
                    self.scaler.update()
                else:
                    self.optimizer.step()

                self.scheduler.step()
                self.optimizer.zero_grad(set_to_none=True)

                # global step（update step）
                self.global_step += 1

                # window logging
                window_updates += 1
                window_loss_sum += float(loss.detach().item())

                if self.global_step % logging_steps == 0:
                    avg_loss = window_loss_sum / max(1, window_updates)
                    log_dict = {
                        "step": self.global_step,
                        "epoch": epoch,
                        "train/loss": float(avg_loss),
                        "train/lr": float(self.scheduler.get_last_lr()[0]),
                    }
                    for t in window_task_sum:
                        log_dict[f"train/{t}"] = float(window_task_sum[t] / max(1, window_task_cnt.get(t, 1)))

                    # GPU memory（可选但很实用）
                    if torch.cuda.is_available():
                        log_dict["train/gpu_mem_alloc_mb"] = float(torch.cuda.memory_allocated() / 1024 / 1024)
                        log_dict["train/gpu_mem_reserved_mb"] = float(torch.cuda.memory_reserved() / 1024 / 1024)

                    self._write_jsonl(self.train_log_path, log_dict)
                    if self.config.training.use_wandb and wandb is not None:
                        wandb.log(log_dict, step=self.global_step)

                    pbar.set_postfix({
                        "loss": f"{avg_loss:.4f}",
                        "lr": f"{log_dict['train/lr']:.2e}"
                    })

                    # reset window
                    window_task_sum, window_task_cnt = {}, {}
                    window_loss_sum, window_updates = 0.0, 0

                # save by steps
                if save_steps > 0 and (self.global_step % save_steps == 0):
                    self.save_checkpoint(tag=f"checkpoint-{self.global_step}", is_best=False)

                # eval by steps（Stage A 验证集很大，eval_steps 不要设太小）
                if eval_steps > 0 and (self.global_step % eval_steps == 0):
                    val_loss, val_task_losses = self.evaluate(epoch)
                    print("\nValidation Results:")
                    print(f"  Overall Loss: {val_loss:.4f}")
                    for t, v in val_task_losses.items():
                        print(f"  {t}: {v:.4f}")

                    if val_loss < self.best_val_loss:
                        self.best_val_loss = val_loss
                        self.save_checkpoint(tag="best_model", is_best=True)
                        print(f"✓ 新的最佳模型! Val Loss: {val_loss:.4f}")

            # 每个 epoch 结束强制 eval（更稳妥）
            val_loss, val_task_losses = self.evaluate(epoch)
            print("\n[Epoch End] Validation Results:")
            print(f"  Overall Loss: {val_loss:.4f}")
            for t, v in val_task_losses.items():
                print(f"  {t}: {v:.4f}")

            if val_loss < self.best_val_loss:
                self.best_val_loss = val_loss
                self.save_checkpoint(tag="best_model", is_best=True)
                print(f"✓ 新的最佳模型! Val Loss: {val_loss:.4f}")

            # Curriculum stage 更新（仍保留提示；如要真正生效应重建 dataloader）
            if self.current_stage < 3:
                if epoch + 1 == sum(self.stage_epochs[: self.current_stage + 1]):
                    self.current_stage += 1
                    print(f"\n>>> Curriculum升级到 Stage {self.current_stage} <<<\n")

        # 最终保存
        self.save_checkpoint(tag=f"checkpoint-{self.global_step}", is_best=False)

        print("\n" + "=" * 60)
        print("训练完成!")
        print(f"最佳验证Loss: {self.best_val_loss:.4f}")
        print(f"模型保存在: {self.output_dir}")
        print("=" * 60)

        if self.config.training.use_wandb and wandb is not None:
            wandb.finish()