training/PRETRAIN/trainer.py

# """
# 多任务训练器
# 处理环境描述、事故检测、序列预测三个任务
# """

# import torch
# import torch.nn as nn
# from torch.utils.data import DataLoader
# from tqdm import tqdm
# from PIL import Image
# import os
# import json
# from typing import Dict, List
# from config import PretrainConfig
# from model_loader import load_model_and_processor, prepare_model_inputs


# class MultiTaskTrainer:
#     """多任务预训练器"""
    
#     def __init__(self, config: PretrainConfig, train_loader, val_loader):
#         self.config = config
#         self.train_loader = train_loader
#         self.val_loader = val_loader
        
#         # 加载模型
#         print("=" * 60)
#         print("初始化模型...")
#         self.model, self.processor = load_model_and_processor(config.model)
#         self.model.to(config.training.device)
        
#         # 优化器
#         self.optimizer = torch.optim.AdamW(
#             self.model.parameters(),
#             lr=config.training.learning_rate,
#             weight_decay=config.training.weight_decay
#         )
        
#         # 学习率调度器
#         total_steps = len(train_loader) * config.training.num_epochs // config.training.gradient_accumulation_steps
#         warmup_steps = int(total_steps * config.training.warmup_ratio)
        
#         from transformers import get_cosine_schedule_with_warmup
#         self.scheduler = get_cosine_schedule_with_warmup(
#             self.optimizer,
#             num_warmup_steps=warmup_steps,
#             num_training_steps=total_steps
#         )
        
#         # 混合精度
#         self.scaler = torch.cuda.amp.GradScaler() if config.training.fp16 else None
        
#         # 训练状态
#         self.global_step = 0
#         self.best_val_loss = float('inf')
        
#         print(f"✓ 模型加载完成")
#         print(f"✓ 优化器: {config.training.optimizer_type}")
#         print(f"✓ 总训练步数: {total_steps}")
#         print("=" * 60)
    
#     def construct_prompt(self, task: str, label: str = None) -> str:
#         """构造任务提示"""
#         if task == "environment":
#             prompt = (
#                 "Analyze this dashcam image and describe the driving environment. "
#                 "Provide the weather condition, road type, and lighting condition in the format: "
#                 "'Weather: [weather], Road: [road_type], Light: [light_condition]'."
#             )
#         elif task == "accident_detection":
#             prompt = (
#                 "Look at this dashcam image. Is there an accident happening in this frame? "
#                 "Answer only 'Yes' or 'No'."
#             )
#         elif task == "sequence_prediction":
#             prompt = (
#                 "You are viewing a sequence of dashcam frames in chronological order. "
#                 "Based on this sequence, determine if an accident will occur and describe it. "
#                 "Format your answer as: 'Accident: [Yes/No]. Description: [description]'."
#             )
#         else:
#             raise ValueError(f"Unknown task: {task}")
        
#         return prompt
    
#     def prepare_batch_inputs(self, batch: Dict):
#         """准备batch输入"""
#         images_list = []
#         prompts_list = []
#         labels_list = []
        
#         # 处理单帧任务
#         if "single_frame" in batch:
#             sf = batch["single_frame"]
#             for i in range(len(sf["images"])):
#                 # 加载图像
#                 img_tensor = sf["images"][i]  # [3, 224, 224]
#                 img = self.tensor_to_pil(img_tensor)
                
#                 task = sf["task"][i]
#                 label = sf["labels"][i]
#                 prompt = self.construct_prompt(task, label)
                
#                 images_list.append(img)
#                 prompts_list.append(prompt)
#                 labels_list.append(label)
        
#         # 处理序列任务
#         if "sequence" in batch:
#             seq = batch["sequence"]
#             for i in range(len(seq["sequences"])):
#                 # 加载序列图像
#                 seq_tensor = seq["sequences"][i]  # [T, 3, 224, 224]
#                 mask = seq["masks"][i]  # [T]
                
#                 # 只取有效帧
#                 valid_frames = seq_tensor[mask == 1]
#                 img_sequence = [self.tensor_to_pil(frame) for frame in valid_frames]
                
#                 label = seq["labels"][i]
#                 prompt = self.construct_prompt("sequence_prediction", label)
                
#                 images_list.append(img_sequence)
#                 prompts_list.append(prompt)
#                 labels_list.append(label)
        
#         return images_list, prompts_list, labels_list
    
#     def tensor_to_pil(self, tensor):
#         """将tensor转换为PIL Image"""
#         # Denormalize
#         mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
#         std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
#         tensor = tensor * std + mean
#         tensor = torch.clamp(tensor, 0, 1)
        
#         # To PIL
#         import torchvision.transforms as T
#         to_pil = T.ToPILImage()
#         return to_pil(tensor)
    
#     def compute_loss(self, model_outputs, labels):
#         """计算损失"""
#         # 使用模型的标准语言模型损失
#         return model_outputs.loss
    
#     def train_epoch(self, epoch: int):
#         """训练一个epoch"""
#         self.model.train()
#         epoch_loss = 0
#         num_batches = 0
        
#         pbar = tqdm(self.train_loader, desc=f"Epoch {epoch+1}")
        
#         for batch_idx, batch in enumerate(pbar):
#             # 准备输入
#             images, prompts, labels = self.prepare_batch_inputs(batch)
            
#             # 准备标签文本（作为目标）
#             target_texts = labels
            
#             # 构造输入
#             inputs = prepare_model_inputs(
#                 self.processor,
#                 self.config.model.model_type,
#                 images,
#                 prompts,
#                 self.config.training.device
#             )
            
#             # 准备标签（tokenize目标文本）
#             # label_inputs = self.processor.tokenizer(
#             #     target_texts,
#             #     return_tensors="pt",
#             #     padding=True,
#             #     truncation=True,
#             #     max_length=512
#             # )
#             # inputs["labels"] = label_inputs["input_ids"].to(self.config.training.device)
            
#             # ===== 关键改动：把 labels 对齐到 input_ids 的长度 =====
#             tok = self.processor.tokenizer
#             if tok.pad_token_id is None:
#                 tok.pad_token = tok.eos_token

#             input_ids = inputs["input_ids"]
#             B, L = input_ids.shape
#             labels = torch.full_like(input_ids, fill_value=-100)

#             # 计算每个样本的“提示长度”（不含答案），用相同模板再次 tokenize
#             prompt_texts = inputs.pop("__prompt_texts__")
#             prompt_tok = tok(
#                 prompt_texts,
#                return_tensors="pt",
#                 padding=True,
#                 truncation=True,
#                 add_special_tokens=False,
#             )
#             prompt_lens = (prompt_tok["input_ids"] != tok.pad_token_id).sum(dim=1).tolist()

#             # 把答案 tokens 放到 prompt 后面的位置；超长则截断到 L
#             for i in range(B):
#                 ans_ids = tok(
#                     target_texts[i],
#                     return_tensors="pt",
#                     padding=False,
#                     truncation=True,
#                     add_special_tokens=False,
#                 )["input_ids"][0]
#                 start = min(prompt_lens[i], L)
#                 end = min(start + ans_ids.numel(), L)
#                 if end > start:
#                     labels[i, start:end] = ans_ids[: (end - start)]

#             inputs["labels"] = labels.to(self.config.training.device)
            
#             # 前向传播
#             if self.scaler:
#                 with torch.cuda.amp.autocast():
#                     outputs = self.model(**inputs)
#                     loss = outputs.loss / self.config.training.gradient_accumulation_steps
                
#                 self.scaler.scale(loss).backward()
#             else:
#                 outputs = self.model(**inputs)
#                 loss = outputs.loss / self.config.training.gradient_accumulation_steps
#                 loss.backward()
            
#             # 梯度累积
#             if (batch_idx + 1) % self.config.training.gradient_accumulation_steps == 0:
#                 if self.scaler:
#                     self.scaler.unscale_(self.optimizer)
#                     torch.nn.utils.clip_grad_norm_(
#                         self.model.parameters(),
#                         self.config.training.max_grad_norm
#                     )
#                     self.scaler.step(self.optimizer)
#                     self.scaler.update()
#                 else:
#                     torch.nn.utils.clip_grad_norm_(
#                         self.model.parameters(),
#                         self.config.training.max_grad_norm
#                     )
#                     self.optimizer.step()
                
#                 self.scheduler.step()
#                 self.optimizer.zero_grad()
#                 self.global_step += 1
            
#             epoch_loss += loss.item() * self.config.training.gradient_accumulation_steps
#             num_batches += 1
            
#             # 更新进度条
#             pbar.set_postfix({
#                 'loss': f'{loss.item():.4f}',
#                 'lr': f'{self.scheduler.get_last_lr()[0]:.2e}'
#             })
            
#             # 日志
#             if self.global_step > 0 and self.global_step % self.config.training.logging_steps == 0:
#                 avg_loss = epoch_loss / num_batches
#                 print(f"\nStep {self.global_step}: loss={avg_loss:.4f}")
            
#             # 保存checkpoint
#             if self.global_step > 0 and self.global_step % self.config.training.save_steps == 0:
#                 self.save_checkpoint(f"step_{self.global_step}")
            
#             # 验证
#             if self.global_step  > 0 and self.global_step % self.config.training.eval_steps == 0:
#                 val_loss = self.validate()
#                 print(f"\nValidation loss: {val_loss:.4f}")
                
#                 if val_loss < self.best_val_loss:
#                     self.best_val_loss = val_loss
#                     self.save_checkpoint("best")
                
#                 self.model.train()
        
#         return epoch_loss / num_batches
    
#     @torch.no_grad()
#     def validate(self):
#         """验证"""
#         self.model.eval()
#         total_loss = 0
#         num_batches = 0
        
#         for batch in tqdm(self.val_loader, desc="Validating"):
#             images, prompts, labels = self.prepare_batch_inputs(batch)
            
#             inputs = prepare_model_inputs(
#                 self.processor,
#                 self.config.model.model_type,
#                 images,
#                 prompts,
#                 self.config.training.device
#             )
            
#             # ===== 使用与训练相同的标签对齐逻辑 =====
#             tok = self.processor.tokenizer
#             if tok.pad_token_id is None:
#                 tok.pad_token = tok.eos_token

#             input_ids = inputs["input_ids"]
#             B, L = input_ids.shape
#             labels_tensor = torch.full_like(input_ids, fill_value=-100)

#             # 计算提示长度
#             prompt_texts = inputs.pop("__prompt_texts__")
#             prompt_tok = tok(
#                 prompt_texts,
#                 return_tensors="pt",
#                 padding=True,
#                 truncation=True,
#                 add_special_tokens=False,
#             )
#             prompt_lens = (prompt_tok["input_ids"] != tok.pad_token_id).sum(dim=1).tolist()

#             # 对齐答案
#             for i in range(B):
#                 ans_ids = tok(
#                     labels[i],
#                     return_tensors="pt",
#                     padding=False,
#                     truncation=True,
#                     add_special_tokens=False,
#                 )["input_ids"][0]
#                 start = min(prompt_lens[i], L)
#                 end = min(start + ans_ids.numel(), L)
#                 if end > start:
#                     labels_tensor[i, start:end] = ans_ids[: (end - start)]

#             inputs["labels"] = labels_tensor.to(self.config.training.device)
            
#             outputs = self.model(**inputs)
#             total_loss += outputs.loss.item()
#             num_batches += 1
        
#         return total_loss / num_batches
    
#     def save_checkpoint(self, name: str):
#         """保存checkpoint"""
#         save_dir = os.path.join(self.config.training.output_dir, name)
#         os.makedirs(save_dir, exist_ok=True)
        
#         # 保存模型
#         self.model.save_pretrained(save_dir)
#         self.processor.save_pretrained(save_dir)
        
#         # 保存训练状态
#         torch.save({
#             'global_step': self.global_step,
#             'best_val_loss': self.best_val_loss,
#             'optimizer_state': self.optimizer.state_dict(),
#             'scheduler_state': self.scheduler.state_dict(),
#         }, os.path.join(save_dir, "training_state.pt"))
        
#         print(f"✓ Checkpoint saved: {save_dir}")
    
#     def train(self):
#         """完整训练流程"""
#         print("=" * 60)
#         print("开始训练")
#         print("=" * 60)
        
#         for epoch in range(self.config.training.num_epochs):
#             print(f"\nEpoch {epoch+1}/{self.config.training.num_epochs}")
            
#             train_loss = self.train_epoch(epoch)
#             print(f"Epoch {epoch+1} - Train Loss: {train_loss:.4f}")
            
#             # Epoch结束后验证
#             val_loss = self.validate()
#             print(f"Epoch {epoch+1} - Val Loss: {val_loss:.4f}")
            
#             if val_loss < self.best_val_loss:
#                 self.best_val_loss = val_loss
#                 self.save_checkpoint("best")
            
#             # 每个epoch保存一次
#             self.save_checkpoint(f"epoch_{epoch+1}")
        
#         print("\n" + "=" * 60)
#         print("训练完成！")
#         print(f"最佳验证损失: {self.best_val_loss:.4f}")
#         print("=" * 60)

# trainer.py
# -*- coding: utf-8 -*-
"""
多任务 VLM 预训练器（Qwen2.5-VL 系列友好）
- 采用聊天模板对齐：仅对 assistant 回复部分计算损失（labels 其余位置置为 -100）
- 统一处理单帧与序列任务：images 以 list[PIL.Image] 形式传入处理器
- 混合精度：优先 bf16（若可用），否则 fp16；包含 NaN/Inf 保护与梯度裁剪
- 日志/保存/验证：仅在完成一次 optimizer.step() 后按步距触发，避免重复触发

依赖：
- transformers >= 4.44（Qwen2-VL: AutoModelForImageTextToText / Qwen2VLProcessor）
- torch >= 2.1
- 数据管道需提供 batch 结构，见 prepare_batch_inputs() 的说明
"""

import os
import json
from typing import Dict, List, Tuple, Optional

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm
from PIL import Image

from config import PretrainConfig
from model_loader import load_model_and_processor  # 只依赖加载模型/处理器


def _to_pil(t: torch.Tensor) -> Image.Image:
    """
    将张量 [3, H, W] (0~1 标准化前张量) 转为 PIL Image。
    若输入已是 PIL 则直接返回。
    """
    if isinstance(t, Image.Image):
        return t
    assert t.ndim == 3 and t.shape[0] == 3, "expect CHW tensor"
    # 反归一化（若上游做了 ImageNet 标准化）
    mean = torch.tensor([0.485, 0.456, 0.406], dtype=t.dtype, device=t.device).view(3, 1, 1)
    std = torch.tensor([0.229, 0.224, 0.225], dtype=t.dtype, device=t.device).view(3, 1, 1)
    # 兼容：如果数据本身已在 0~1，可关闭下行两句
    x = t * std + mean
    x = x.clamp(0, 1).cpu()
    import torchvision.transforms as T
    return T.ToPILImage()(x)


class MultiTaskTrainer:
    """
    多任务预训练器

    期望 DataLoader yield 的 batch 结构（示例）：
    batch = {
        "single_frame": {
            "images": Tensor[B, 3, H, W] 或 List[PIL]，
            "task":   List[str],  # "environment" | "accident_detection"
            "labels": List[str],  # 文本答案（例如 'Yes'/'No' 或结构化描述）
        },
        "sequence": {
            "sequences": Tensor[B, T, 3, H, W] 或 List[List[PIL]],
            "masks":     Tensor[B, T] (1=有效帧),
            "labels":    List[str],  # 文本答案
        }
    }
    """

    def __init__(self, config: PretrainConfig, train_loader: DataLoader, val_loader: DataLoader):
        self.config = config
        self.train_loader = train_loader
        self.val_loader = val_loader

        print("=" * 60)
        print("初始化模型...")
        self.model, self.processor = load_model_and_processor(config.model)
        self.device = torch.device(config.training.device)
        self.model.to(self.device)

        # tokenizer pad token
        tok = self.processor.tokenizer
        if tok.pad_token_id is None:
            tok.pad_token = tok.eos_token

        # 优化器
        optim_type = getattr(config.training, "optimizer_type", "adamw").lower()
        lr = config.training.learning_rate
        wd = config.training.weight_decay
        if optim_type == "adamw":
            self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr, weight_decay=wd)
        else:
            raise ValueError(f"Unsupported optimizer_type: {optim_type}")

        # 训练步数与调度器
        self.total_steps = (
            len(train_loader) * config.training.num_epochs
        ) // max(1, config.training.gradient_accumulation_steps)
        warmup_steps = int(self.total_steps * config.training.warmup_ratio)

        from transformers import get_cosine_schedule_with_warmup
        self.scheduler = get_cosine_schedule_with_warmup(
            self.optimizer, num_warmup_steps=warmup_steps, num_training_steps=self.total_steps
        )

        # 混合精度
        self.use_bf16 = bool(getattr(config.training, "bf16", False)) and torch.cuda.is_available()
        self.use_fp16 = bool(getattr(config.training, "fp16", False)) and torch.cuda.is_available() and not self.use_bf16
        self.autocast_dtype: Optional[torch.dtype] = torch.bfloat16 if self.use_bf16 else (torch.float16 if self.use_fp16 else None)
        self.scaler = torch.amp.GradScaler("cuda") if self.use_fp16 else None

        # 其他状态
        self.global_step = 0
        self.best_val_loss = float("inf")

        # 可选：开启梯度检查点
        if getattr(config.training, "gradient_checkpointing", False):
            if hasattr(self.model, "gradient_checkpointing_enable"):
                self.model.gradient_checkpointing_enable()

        print(f"✓ 模型加载完成")
        print(f"✓ 优化器: {optim_type}")
        print(f"✓ 总训练步数: {self.total_steps}")
        print("=" * 60)

    # ========= 任务模板 =========
    def construct_prompt(self, task: str) -> str:
        """
        返回 user 提示文本，答案只由 labels 提供。
        """
        if task == "environment":
            return (
                "Analyze this dashcam image and describe the driving environment. "
                "Provide the weather condition, road type, and lighting condition in the format: "
                "'Weather: [weather], Road: [road_type], Light: [light_condition]'."
            )
        elif task == "accident_detection":
            return (
                "Look at this dashcam image. Is there an accident happening in this frame? "
                "Answer only 'Yes' or 'No'."
            )
        elif task == "sequence_prediction":
            return (
                "You are viewing a sequence of dashcam frames in chronological order. "
                "Based on this sequence, determine if an accident will occur and describe it. "
                "Format your answer as: 'Accident: [Yes/No]. Description: [description]'."
            )
        else:
            raise ValueError(f"Unknown task: {task}")

    # ========= 数据整理 =========
    def prepare_batch_inputs(self, batch: Dict) -> Tuple[List[List[Image.Image]], List[str], List[str]]:
        """
        归一化 batch 成 3 个并行列表：
            images_list:  List[ List[PIL.Image] ]   # 每条样本是若干帧（单帧也用长度为1的列表）
            prompts_list: List[str]                 # user 提示
            labels_list:  List[str]                 # assistant 文本答案
        """
        images_list: List[List[Image.Image]] = []
        prompts_list: List[str] = []
        labels_list: List[str] = []

        if "single_frame" in batch:
            sf = batch["single_frame"]
            imgs = sf["images"]
            # 支持张量/列表
            if isinstance(imgs, torch.Tensor):  # [B, 3, H, W]
                for i in range(imgs.shape[0]):
                    images_list.append([_to_pil(imgs[i])])
                    prompts_list.append(self.construct_prompt(sf["task"][i]))
                    labels_list.append(sf["labels"][i])
            else:  # List[PIL]
                for i in range(len(imgs)):
                    images_list.append([imgs[i] if isinstance(imgs[i], Image.Image) else _to_pil(imgs[i])])
                    prompts_list.append(self.construct_prompt(sf["task"][i]))
                    labels_list.append(sf["labels"][i])

        if "sequence" in batch:
            seq = batch["sequence"]
            seqs = seq["sequences"]
            masks = seq.get("masks", None)

            if isinstance(seqs, torch.Tensor):  # [B, T, 3, H, W]
                B, T = seqs.shape[0], seqs.shape[1]
                for i in range(B):
                    if masks is not None:
                        valid_idx = (masks[i] == 1).nonzero(as_tuple=False).flatten().tolist()
                    else:
                        valid_idx = list(range(T))
                    frames = [ _to_pil(seqs[i, j]) for j in valid_idx ]
                    if len(frames) == 0 and T > 0:  # fallback 至第一帧
                        frames = [ _to_pil(seqs[i, 0]) ]
                    images_list.append(frames)
                    prompts_list.append(self.construct_prompt("sequence_prediction"))
                    labels_list.append(seq["labels"][i])
            else:  # List[List[PIL]]
                for i in range(len(seqs)):
                    frames = seqs[i]
                    frames_pil = [ f if isinstance(f, Image.Image) else _to_pil(f) for f in frames ]
                    if masks is not None:
                        m = masks[i]
                        if isinstance(m, torch.Tensor):
                            frames_pil = [frames_pil[j] for j in (m == 1).nonzero(as_tuple=False).flatten().tolist()]
                    if len(frames_pil) == 0 and len(frames) > 0:
                        frames_pil = [frames_pil[0]]
                    images_list.append(frames_pil)
                    prompts_list.append(self.construct_prompt("sequence_prediction"))
                    labels_list.append(seq["labels"][i])

        return images_list, prompts_list, labels_list

    # ========= 编码与标签构建（聊天模板一致） =========
    def _build_texts_for_sample(self, images: List[Image.Image], prompt: str, answer: str) -> Tuple[str, str]:
        """
        基于聊天模板，返回 (prompt_only_text, full_text)
        训练时 full_text = user(msg) + assistant(answer)；计算损失时屏蔽 user 区段。
        """
        # prompt-only（无 assistant）
        msgs_prompt_only = [
            {
                "role": "user",
                "content": [{"type": "image", "image": img} for img in images] + [{"type": "text", "text": prompt}],
            }
        ]
        # full（含 assistant）
        msgs_full = [
            {
                "role": "user",
                "content": [{"type": "image", "image": img} for img in images] + [{"type": "text", "text": prompt}],
            },
            {
                "role": "assistant",
                "content": [{"type": "text", "text": answer}],
            },
        ]

        # 训练：不加 generation_prompt
        t_prompt = self.processor.apply_chat_template(
            msgs_prompt_only, tokenize=False, add_generation_prompt=False
        )
        t_full = self.processor.apply_chat_template(
            msgs_full, tokenize=False, add_generation_prompt=False
        )
        return t_prompt, t_full

    def _encode_batch_with_labels(
        self, images_list: List[List[Image.Image]], prompts_list: List[str], labels_list: List[str]
    ) -> Dict[str, torch.Tensor]:
        """
        对一批样本：
          1) 生成 prompt-only 文本与 full 文本
          2) 分别走 processor(text=..., images=...) 得到两套 input_ids
          3) 根据 prompt-only 的非 pad 长度，构建 full 的 labels（prompt 部分置 -100）
        返回可直接喂给 model(**inputs) 的字典（已放到正确的 device）
        """
        texts_prompt_only, texts_full = [], []
        for imgs, p, a in zip(images_list, prompts_list, labels_list):
            t_p, t_f = self._build_texts_for_sample(imgs, p, a)
            texts_prompt_only.append(t_p)
            texts_full.append(t_f)

        # 编码（注意：两次都要带 images，保证图像占位符 token 一致）
        enc_full = self.processor(
            text=texts_full,
            images=images_list,
            padding=True,
            truncation=True,
            return_tensors="pt",
        )
        enc_prompt = self.processor(
            text=texts_prompt_only,
            images=images_list,
            padding=True,
            truncation=True,
            return_tensors="pt",
        )

        input_ids = enc_full["input_ids"]
        attn_mask = enc_full["attention_mask"]
        B, L = input_ids.shape
        pad_id = self.processor.tokenizer.pad_token_id

        # 计算每条样本的 prompt-only 长度（非 pad token 数量）
        prompt_lens = (enc_prompt["input_ids"] != pad_id).sum(dim=1).tolist()

        # 构建 labels：拷贝 full 的 input_ids，再把 prompt 区段置为 -100
        labels = input_ids.clone()
        for i in range(B):
            plen = min(prompt_lens[i], L)
            labels[i, :plen] = -100
            # 保险：至少保留一个可学习 token，避免全 -100 导致 loss 为 NaN
            if torch.all(labels[i] == -100):
                last_idx = plen - 1 if plen > 0 else L - 1
                labels[i, last_idx] = input_ids[i, last_idx]

        # 转设备
        for k in enc_full.keys():
            if isinstance(enc_full[k], torch.Tensor):
                enc_full[k] = enc_full[k].to(self.device, non_blocking=True)
        labels = labels.to(self.device, non_blocking=True)
        enc_full["labels"] = labels
        return enc_full

    # ========= 训练/验证 =========
    def train_epoch(self, epoch: int) -> float:
        self.model.train()
        running_loss = 0.0
        n_batches = 0

        pbar = tqdm(self.train_loader, desc=f"Epoch {epoch+1}")
        grad_accum = max(1, self.config.training.gradient_accumulation_steps)
        max_norm = getattr(self.config.training, "max_grad_norm", 1.0)

        # 清梯度
        self.optimizer.zero_grad(set_to_none=True)

        for step_in_epoch, batch in enumerate(pbar):
            images_list, prompts_list, labels_list = self.prepare_batch_inputs(batch)
            inputs = self._encode_batch_with_labels(images_list, prompts_list, labels_list)

            # 前向
            loss = None
            if self.autocast_dtype is not None:
                with torch.amp.autocast("cuda", dtype=self.autocast_dtype):
                    outputs = self.model(**inputs)
                    loss = outputs.loss / grad_accum
            else:
                outputs = self.model(**inputs)
                loss = outputs.loss / grad_accum

            # NaN/Inf 保护（前向）
            if not torch.isfinite(loss):
                print(f"[WARN] Non-finite loss detected (forward): {loss.item()}. Skip this micro-batch.")
                self.optimizer.zero_grad(set_to_none=True)
                continue

            # 反向
            if self.scaler is not None:
                self.scaler.scale(loss).backward()
            else:
                loss.backward()

            # 累积步
            do_step = ((step_in_epoch + 1) % grad_accum == 0)
            if do_step:
                # 反 NaN/Inf（反向后 unscale 再裁剪）
                if self.scaler is not None:
                    self.scaler.unscale_(self.optimizer)
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm)

                # 检测梯度异常
                found_inf = False
                for p in self.model.parameters():
                    if p.grad is not None and (torch.isnan(p.grad).any() or torch.isinf(p.grad).any()):
                        found_inf = True
                        break
                if found_inf:
                    print("[WARN] Found NaN/Inf gradients. Skipping step and zeroing grads.")
                    self.optimizer.zero_grad(set_to_none=True)
                    if self.scaler is not None:
                        self.scaler.update()
                    continue

                # step + scheduler
                if self.scaler is not None:
                    self.scaler.step(self.optimizer)
                    self.scaler.update()
                else:
                    self.optimizer.step()
                self.scheduler.step()
                self.optimizer.zero_grad(set_to_none=True)

                # 全局步递增，仅在真正 step 后进行
                self.global_step += 1

                # ---- 日志/保存/验证 均放在这里（只触发一次） ----
                if self.global_step % max(1, self.config.training.logging_steps) == 0:
                    avg_loss = (running_loss + loss.item() * grad_accum) / max(1, (n_batches + 1))
                    print(f"\nStep {self.global_step}: loss={avg_loss:.4f}, lr={self.scheduler.get_last_lr()[0]:.2e}")

                if self.global_step % max(1, self.config.training.save_steps) == 0:
                    self.save_checkpoint(f"step_{self.global_step}")

                if self.global_step % max(1, self.config.training.eval_steps) == 0:
                    val_loss = self.validate()
                    print(f"\nValidation loss: {val_loss:.4f}")
                    if val_loss < self.best_val_loss:
                        self.best_val_loss = val_loss
                        self.save_checkpoint("best")
                    self.model.train()

            # 统计（注意：running_loss 累加的是未除以 grad_accum 的 loss）
            running_loss += loss.item() * grad_accum
            n_batches += 1

            pbar.set_postfix({
                "loss": f"{loss.item():.4f}",
                "lr": f"{self.scheduler.get_last_lr()[0]:.2e}",
            })

        return running_loss / max(1, n_batches)

    @torch.no_grad()
    def validate(self) -> float:
        self.model.eval()
        total_loss = 0.0
        n_batches = 0

        for batch in tqdm(self.val_loader, desc="Validating"):
            images_list, prompts_list, labels_list = self.prepare_batch_inputs(batch)
            inputs = self._encode_batch_with_labels(images_list, prompts_list, labels_list)

            if self.autocast_dtype is not None:
                with torch.amp.autocast("cuda", dtype=self.autocast_dtype):
                    outputs = self.model(**inputs)
                    loss = outputs.loss
            else:
                outputs = self.model(**inputs)
                loss = outputs.loss

            # 守护
            if not torch.isfinite(loss):
                continue
            total_loss += loss.item()
            n_batches += 1

        return total_loss / max(1, n_batches)

    def save_checkpoint(self, name: str):
        save_dir = os.path.join(self.config.training.output_dir, name)
        os.makedirs(save_dir, exist_ok=True)

        # 保存模型/处理器
        self.model.save_pretrained(save_dir)
        self.processor.save_pretrained(save_dir)

        # 保存训练状态
        torch.save(
            {
                "global_step": self.global_step,
                "best_val_loss": self.best_val_loss,
                "optimizer_state": self.optimizer.state_dict(),
                "scheduler_state": self.scheduler.state_dict(),
            },
            os.path.join(save_dir, "training_state.pt"),
        )
        print(f"✓ Checkpoint saved: {save_dir}")

    def train(self):
        print("=" * 60)
        print("开始训练")
        print("=" * 60)

        for epoch in range(self.config.training.num_epochs):
            print(f"\nEpoch {epoch+1}/{self.config.training.num_epochs}")

            train_loss = self.train_epoch(epoch)
            print(f"Epoch {epoch+1} - Train Loss: {train_loss:.4f}")

            # epoch 结束做一次验证
            val_loss = self.validate()
            print(f"Epoch {epoch+1} - Val Loss: {val_loss:.4f}")

            if val_loss < self.best_val_loss:
                self.best_val_loss = val_loss
                self.save_checkpoint("best")

            self.save_checkpoint(f"epoch_{epoch+1}")

        print("\n" + "=" * 60)
        print("训练完成！")
        print(f"最佳验证损失: {self.best_val_loss:.4f}")
        print("=" * 60)