src/trainer_early_exit.py

import torch
import torch.nn.functional as F
import torch.distributed as dist
from torch.utils.data import DataLoader
from transformers import Trainer
from src.utils import batch_to_device
from src.classifier_utils import HomogeneousBatchSampler

class EarlyExitTrainer(Trainer):
    def __init__(self, backbone_model, target_layer_idx, model_args, *args, **kwargs):
        self.max_length = kwargs.pop('max_length', 512)
        super().__init__(*args, **kwargs)
        
        self.backbone = backbone_model.to(self.args.device)
        self.backbone.eval()
        
        self.target_layer_idx = target_layer_idx
        self.model_args = model_args
        
        # 添加梯度检查标志
        self._grad_check_done = False

    def get_train_dataloader(self) -> DataLoader:
        if self.train_dataset is None:
            raise ValueError("Trainer: training requires a train_dataset.")
        train_sampler = HomogeneousBatchSampler(
            self.train_dataset, 
            batch_size=self._train_batch_size,
            drop_last=self.args.dataloader_drop_last
        )
        return DataLoader(
            self.train_dataset,
            batch_sampler=train_sampler,
            collate_fn=self.data_collator,
            num_workers=self.args.dataloader_num_workers,
            pin_memory=self.args.dataloader_pin_memory,
        )

    def create_optimizer(self):
        if self.optimizer is None:
            print(f"\n[Debug Rank {self.args.local_rank}] Creating Optimizer...")
            decay_parameters = []
            no_decay_parameters = []
            trainable_count = 0
            
            # 注意：self.model 可能是 DDP 包装后的，所以用 self.model.named_parameters()
            for name, param in self.model.named_parameters():
                if not param.requires_grad:
                    continue
                trainable_count += 1
                if "bias" in name or "LayerNorm" in name or "BatchNorm" in name:
                    no_decay_parameters.append(param)
                else:
                    decay_parameters.append(param)
            
            print(f"[Debug] Found {trainable_count} trainable parameters.")
            
            self.optimizer = torch.optim.AdamW(
                [
                    {"params": decay_parameters, "weight_decay": self.args.weight_decay},
                    {"params": no_decay_parameters, "weight_decay": 0.0},
                ],
                lr=self.args.learning_rate,
                eps=self.args.adam_epsilon,
            )
        return self.optimizer

    def _perform_pooling(self, hidden_state, attention_mask):
        pooling_method = self.model_args.pooling
        batch_size = hidden_state.shape[0]
        if pooling_method == 'last' or pooling_method == 'eos':
            left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
            if left_padding:
                reps = hidden_state[torch.arange(batch_size), -1, :]
            else:
                eos_indices = attention_mask.sum(dim=1) - 1
                reps = hidden_state[torch.arange(batch_size, device=hidden_state.device), eos_indices]
        else:
            reps = hidden_state[:, -1, :]
        if self.model_args.normalize:
            reps = F.normalize(reps, p=2, dim=-1)
        return reps

    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        """
        覆盖 Trainer 的 compute_loss 方法
        这是 Trainer 真正调用的损失计算入口
        """
        loss = self._compute_early_exit_loss(model, inputs)
        return (loss, None) if return_outputs else loss

    def _compute_early_exit_loss(self, model, inputs) -> torch.Tensor:
        """
        计算 Early Exit 分类器的损失
        """
        self.backbone.eval()
        model.train()
        
        device = self.args.device
        qry_inputs, tgt_inputs = inputs
        qry_inputs = batch_to_device(qry_inputs, device)
        tgt_inputs = batch_to_device(tgt_inputs, device)
        
        with torch.no_grad():
            # Backbone Forward
            tgt_outputs = self.backbone.encoder(**tgt_inputs, return_dict=True, output_hidden_states=True)
            tgt_reps = self._perform_pooling(tgt_outputs.hidden_states[-1], tgt_inputs['attention_mask'])
            
            qry_outputs = self.backbone.encoder(**qry_inputs, return_dict=True, output_hidden_states=True)
            q_hidden_mid = qry_outputs.hidden_states[self.target_layer_idx]
            qry_reps_mid = self._perform_pooling(q_hidden_mid, qry_inputs['attention_mask'])
            
            batch_size = qry_reps_mid.size(0)
            
            # 特征工程
            backbone_ptr = self.backbone.module if hasattr(self.backbone, 'module') else self.backbone
            temp = getattr(backbone_ptr, 'temperature', 0.02)
            
            sim_matrix = torch.matmul(qry_reps_mid, tgt_reps.T) / temp
            diag_mask = torch.eye(batch_size, dtype=torch.bool, device=device)
            sim_matrix_no_diag = sim_matrix.masked_fill(diag_mask, -1e9)
            
            all_topk_vals, all_topk_inds = torch.topk(sim_matrix, k=2, dim=1)
            feat_s1 = torch.diag(sim_matrix) # 使用对角线作为正样本分数
            # 【关键修复】Margin = Top1 - Top2，而不是 Top1 - Top1
            # all_topk_vals[:, 0] 是最大值（通常就是对角线）
            # all_topk_vals[:, 1] 是第二大值（hard negative）
            feat_margin = feat_s1 - all_topk_vals[:, 1]
            
            probs = torch.softmax(sim_matrix, dim=1)
            feat_entropy = -(probs * torch.log(probs + 1e-6)).sum(dim=1)
            
            # Norm/Var 近似
            left_padding = (qry_inputs['attention_mask'][:, -1].sum() == batch_size)
            if left_padding:
                q_raw_pooled = q_hidden_mid[:, -1, :]
            else:
                eos_indices = qry_inputs['attention_mask'].sum(dim=1) - 1
                q_raw_pooled = q_hidden_mid[torch.arange(batch_size, device=device), eos_indices]
            feat_norm = torch.norm(q_raw_pooled, p=2, dim=1)
            feat_var = torch.var(q_raw_pooled, dim=1)
            
            scalar_inputs = torch.stack([feat_s1, feat_margin, feat_entropy, feat_norm, feat_var], dim=1)
            
            # 模态判断
            modality_idx = torch.zeros(batch_size, dtype=torch.long, device=device)
            if 'pixel_values' in qry_inputs and qry_inputs['pixel_values'] is not None:
                pv = qry_inputs['pixel_values']
                if isinstance(pv, list):
                    for i, item in enumerate(pv):
                        if item is not None: modality_idx[i] = 1
                elif isinstance(pv, torch.Tensor) and pv.numel() > 0:
                    modality_idx.fill_(1)

            # 标签生成：衡量中间层检索质量
            # 方法1：检查Top-1是否正确（对角线）
            preds = torch.argmax(sim_matrix, dim=1)
            ground_truth = torch.arange(batch_size, device=device)
            is_correct = (preds == ground_truth).float()
            
            # 方法2：使用归一化的相似度分数作为连续标签
            diag_scores = torch.diag(sim_matrix)  # 正样本得分
            max_scores = sim_matrix.max(dim=1)[0]  # 最高得分
            
            # 混合标签：如果预测正确给1.0，否则用相对分数
            # 这样可以给模型提供更丰富的监督信号
            relative_scores = torch.sigmoid((diag_scores - max_scores) * 2)  # 放大差异
            labels = torch.where(
                is_correct.bool(),
                torch.ones_like(is_correct),
                relative_scores
            ).unsqueeze(1)

        # Classifier Forward
        # scalar_inputs 已经 detach，不需要梯度
        # 模型参数自动有 requires_grad=True
        pred_probs = model(scalar_inputs, modality_idx)
        loss = F.binary_cross_entropy(pred_probs, labels)

        # ========================================================
        # 【梯度探针】 仅在前 10 步检查梯度流
        # ========================================================
        if self.state.global_step < 10 and self.args.local_rank == 0:
            print(f"\n[Probe Step {self.state.global_step}] Loss: {loss.item():.4f}")
            print(f"  - Loss has grad_fn: {loss.grad_fn is not None}")
            print(f"  - Pred Probs: Mean={pred_probs.mean().item():.4f}, Std={pred_probs.std().item():.4f}")
            print(f"  - Pred Probs Range: [{pred_probs.min().item():.4f}, {pred_probs.max().item():.4f}]")
            print(f"  - Labels: Mean={labels.mean().item():.4f}, Std={labels.std().item():.4f}")
            print(f"  - Labels Range: [{labels.min().item():.4f}, {labels.max().item():.4f}]")
            print(f"  - Correct Rate: {is_correct.mean().item():.4f}")
            
            # 检查输入特征的统计
            print(f"  - Scalar Inputs: Mean={scalar_inputs.mean().item():.4f}, Std={scalar_inputs.std().item():.4f}")
            print(f"  - Modality: Text={((modality_idx==0).sum().item())}, Image={((modality_idx==1).sum().item())}")
        # ========================================================

        return loss
    
    def training_step(self, model, inputs, num_items_in_batch=None) -> torch.Tensor:
        """
        覆盖 Trainer 的 training_step 以添加梯度监控
        注意：这个方法在新版 Transformers 中被调用，负责完整的前向+反向过程
        
        Args:
            model: 要训练的模型
            inputs: 输入数据
            num_items_in_batch: batch 中的样本数（新版 Transformers 会传入）
        """
        model.train()
        inputs = self._prepare_inputs(inputs)
        
        # 计算损失
        with self.compute_loss_context_manager():
            loss = self.compute_loss(model, inputs)
        
        if self.args.n_gpu > 1:
            loss = loss.mean()
        
        # 反向传播（使用 Accelerator）
        self.accelerator.backward(loss)
        
        # 在前几步检查梯度
        if not self._grad_check_done and self.args.local_rank == 0:
            print(f"\n[Gradient Check After Backward - Step {self.state.global_step}]")
            inner_model = model.module if hasattr(model, 'module') else model
            has_grad = False
            total_grad_norm = 0.0
            for name, param in inner_model.named_parameters():
                if param.grad is not None:
                    has_grad = True
                    grad_norm = param.grad.norm().item()
                    total_grad_norm += grad_norm ** 2
                    if self.state.global_step < 3:
                        print(f"  - {name}: grad_norm={grad_norm:.6f}")
            
            total_grad_norm = total_grad_norm ** 0.5
            print(f"  - Total Grad Norm: {total_grad_norm:.6f}")
            print(f"  - Has Gradient: {has_grad}")
            
            if self.state.global_step >= 2:
                self._grad_check_done = True
        
        return loss.detach() / self.args.gradient_accumulation_steps