import os import torch import torch.nn.functional as F from torch.utils.data import DataLoader from transformers import Trainer from src.utils import batch_to_device from src.classifier_utils import HomogeneousBatchSampler # 手动实现 Focal Loss (带 alpha/gamma 控制) def sigmoid_focal_loss(inputs, targets, alpha: float = 0.25, gamma: float = 2.0, reduction: str = "mean"): """ Loss = -alpha * (1 - p)^gamma * log(p) """ p = torch.sigmoid(inputs) ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none") p_t = p * targets + (1 - p) * (1 - targets) loss = ce_loss * ((1 - p_t) ** gamma) if alpha >= 0: alpha_t = alpha * targets + (1 - alpha) * (1 - targets) loss = alpha_t * loss if reduction == "mean": return loss.mean() elif reduction == "sum": return loss.sum() return loss 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 # AOP 侧别启停（env）：qry|tgt|both self._aop_apply = os.getenv("AOP_APPLY", "both").strip().lower() self._grad_check_done = False def _rank_of_diagonal(self, sim_mat: torch.Tensor): """ 返回每个样本的正例(diag)在本行中的排名(1-based)，以及 topk 命中率。 """ B = sim_mat.size(0) # argsort 降序 order = torch.argsort(sim_mat, dim=1, descending=True) # [B, B] gt = torch.arange(B, device=sim_mat.device).view(-1, 1) # [B,1] # 找到每行中 diag 的位置 ranks = (order == gt).nonzero(as_tuple=False)[:, 1] + 1 # 1-based top1 = (ranks == 1).float().mean().item() top5 = (ranks <= 5).float().mean().item() if B >= 5 else float('nan') top10 = (ranks <= 10).float().mean().item() if B >= 10 else float('nan') return ranks, top1, top5, top10 # ---------------- Dataloader ---------------- 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, ) # ---------------- Optimizer ---------------- 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 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 # ---------------- helpers: AOP 侧别开关 ---------------- def _enable_for_side(self, side: str) -> bool: side = side.lower() if self._aop_apply == "both": return True return self._aop_apply == side from contextlib import contextmanager @contextmanager def _aop_switch(self, enable: bool): """ 暂时按侧别启停 AOP（仅该 forward），同步 wrapper 与底座。 """ enc = self.backbone.encoder old = getattr(enc, "aop_prune_config", None) def _set_cfg(mod, cfg): setattr(mod, "aop_prune_config", cfg) base = mod.get_base_model() if hasattr(mod, "get_base_model") else None if base is None and hasattr(mod, "model"): base = mod.model if base is not None: setattr(base, "aop_prune_config", cfg) if hasattr(base, "model"): setattr(base.model, "aop_prune_config", cfg) if old is not None and enable is False: cfg = dict(old) if isinstance(old, dict) else None if isinstance(cfg, dict): cfg["enabled"] = False _set_cfg(enc, cfg) try: yield finally: _set_cfg(enc, old) # ---------------- Pooling ---------------- def _perform_pooling(self, hidden_state, attention_mask): """ 修复版 Pooling：增加索引边界检查，防止 CUDA Index Out of Bounds """ pooling_method = self.model_args.pooling batch_size, seq_len, _ = hidden_state.shape if pooling_method in ("last", "eos"): if attention_mask is None: attention_mask = torch.ones(batch_size, seq_len, dtype=torch.long, device=hidden_state.device) # [关键修复] 检测填充方向 # 如果最后一列全是 1，说明没有右边 padding，极大概率是左填充 (Left Padding) # 或者模型本身就是 Left Padding 的 is_left_padding = (attention_mask[:, -1].sum() == batch_size) if is_left_padding: # 左填充：有效内容挤在右边，直接取最后一个 Token reps = hidden_state[:, -1, :] else: # 右填充：有效内容在左边，需要计算长度 if attention_mask.shape[1] > seq_len: attention_mask = attention_mask[:, :seq_len] eos_indices = attention_mask.sum(dim=1) - 1 eos_indices = eos_indices.clamp(min=0, max=seq_len - 1) indices_expanded = eos_indices.unsqueeze(1).unsqueeze(2).expand(batch_size, 1, hidden_state.size(-1)) reps = torch.gather(hidden_state, 1, indices_expanded).squeeze(1) else: # 兜底 mean pooling 等其他逻辑 reps = hidden_state[:, -1, :] if self.model_args.normalize: reps = F.normalize(reps, p=2, dim=-1) return reps def _match_mask(self, h, pre_mask, post_mask): """ 选择与该层 hidden_state 长度匹配的 mask（优先 post，再 pre；否则全1兜底） """ if post_mask is not None and post_mask.size(1) == h.size(1): return post_mask if pre_mask is not None and pre_mask.size(1) == h.size(1): return pre_mask return torch.ones(h.size(0), h.size(1), dtype=torch.long, device=h.device) # ---------------- Loss entry ---------------- def compute_loss(self, model, inputs, return_outputs=False, **kwargs): loss = self._compute_early_exit_loss(model, inputs) return (loss, None) if return_outputs else loss # ---------------- Core loss (V5) ---------------- def _compute_early_exit_loss(self, model, inputs) -> torch.Tensor: self.backbone.eval() model.train() device = self.args.device qry_inputs, tgt_inputs = inputs # === 定义分块大小 (Chunk Size) === CHUNK_SIZE = 128 def forward_chunked(input_batch, side="tgt"): """ 分块执行 Backbone Forward，提取特征后立即释放显存。 """ # 1. 确定样本总数 total_len = input_batch["input_ids"].shape[0] reps_mid_list = [] reps_last_list = [] # 2. 循环切片 for i in range(0, total_len, CHUNK_SIZE): # 2.1 切片: 构造小 batch chunk = {} for k, v in input_batch.items(): if v is None: chunk[k] = None continue if isinstance(v, torch.Tensor): chunk[k] = v[i : i + CHUNK_SIZE] elif isinstance(v, list): chunk[k] = v[i : i + CHUNK_SIZE] else: chunk[k] = v # 2.2 搬运到 GPU chunk = batch_to_device(chunk, device) # 2.3 Backbone Forward (No Grad) with torch.no_grad(): with self._aop_switch(self._enable_for_side(side)): outputs = self.backbone.encoder( **chunk, return_dict=True, output_hidden_states=True ) # 2.4 立即提取需要的层 pre_mask = chunk.get("attention_mask", None) post_mask = getattr(outputs, "attention_mask", None) # Mid Layer 处理 h_mid = outputs.hidden_states[self.target_layer_idx] m_mid = self._match_mask(h_mid, pre_mask, post_mask) r_mid = self._perform_pooling(h_mid, m_mid) # Last Layer 处理 h_last = outputs.hidden_states[-1] m_last = self._match_mask(h_last, pre_mask, post_mask) r_last = self._perform_pooling(h_last, m_last) # 2.5 存入列表 reps_mid_list.append(r_mid) reps_last_list.append(r_last) # 2.6 显式删除引用，辅助 GC 释放显存 del outputs, h_mid, h_last, chunk, pre_mask, post_mask # torch.cuda.empty_cache() # 3. 拼接所有小块 return torch.cat(reps_mid_list, dim=0), torch.cat(reps_last_list, dim=0) # === 1. 使用分块函数提取特征 === tgt_reps_mid, tgt_reps_last = forward_chunked(tgt_inputs, side="tgt") qry_reps_mid, qry_reps_last = forward_chunked(qry_inputs, side="qry") # === 2. 计算相似度与 Loss === 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) # 相似度计算 cos_mid = torch.matmul(qry_reps_mid, tgt_reps_mid.T) cos_last = torch.matmul(qry_reps_last, tgt_reps_last.T) scores_mid = cos_mid / temp probs_mid = torch.softmax(scores_mid, dim=1) # ====== Debug: 仅保留排名打印，移除 Mask 检查 ====== if self.state.global_step < 3 and self.args.local_rank == 0: # 1) 计算正例排名（mid/last） ranks_mid, top1_mid, top5_mid, top10_mid = self._rank_of_diagonal(cos_mid) ranks_last, top1_last, top5_last, top10_last = self._rank_of_diagonal(cos_last) # [已删除] _check_mask 及其调用，因为无法访问 hidden states 了 # 2) 环境/侧别开关 print( f"[DBG][env] AOP_ENABLED={os.getenv('AOP_ENABLED')} " f"APPLY={os.getenv('AOP_APPLY')} " f"LAYER={os.getenv('AOP_LAYER')} " f"SELECTION={os.getenv('AOP_SELECTION')} " f"KEEP_T={os.getenv('AOP_KEEP_RATIO_TEXT')} " f"KEEP_V={os.getenv('AOP_KEEP_RATIO_VISION')} " f"VPOOL_ENABLED={os.getenv('VPOOL_ENABLED')} " f"VPOOL_LAYER={os.getenv('VPOOL_LAYER')}", flush=True ) # 3) 打印正例排名分布与 topk def _brief(ranks): r = ranks.detach().cpu() return { "min": int(r.min().item()), "p25": int(r.kthvalue(max(1, int(0.25*len(r)))).values.item()) if len(r) >= 4 else None, "med": int(r.median().item()), "p75": int(r.kthvalue(max(1, int(0.75*len(r)))).values.item()) if len(r) >= 4 else None, "max": int(r.max().item()) } print(f"[RANK][mid] top1={top1_mid:.2%} top5={top5_mid:.2%} top10={top10_mid:.2%} dist={_brief(ranks_mid)}", flush=True) print(f"[RANK][last] top1={top1_last:.2%} top5={top5_last:.2%} top10={top10_last:.2%} dist={_brief(ranks_last)}", flush=True) if top1_last < 0.4: print("[WARN] last layer top1 < 40%. 建议先 AOP_ENABLED=0/VPOOL_ENABLED=0 进行对照，确认基座检索能力。", flush=True) # === 特征构造 (27维) === # ... (以下代码保持不变) ... diag_cos = cos_mid.max(dim=1)[0] sorted_cos, _ = torch.sort(cos_mid, dim=1, descending=True) s2_cos = sorted_cos[:, 1] if sorted_cos.size(1) > 1 else sorted_cos[:, 0] margin_mid = diag_cos - s2_cos # 统计量 margin_mean = margin_mid.mean() margin_std = margin_mid.std(unbiased=False) + 1e-6 z_margin_mid = (margin_mid - margin_mean) / margin_std margin_median = margin_mid.median() mad = (margin_mid - margin_median).abs().median() + 1e-6 mad_margin_mid = (margin_mid - margin_median) / mad p1_mid = probs_mid.max(dim=1)[0] H_mid = -(probs_mid * torch.log(probs_mid + 1e-6)).sum(dim=1) gini_mid = 1.0 - (probs_mid ** 2).sum(dim=1) TOPK = min(16, probs_mid.size(1)) topk_vals, _ = torch.topk(probs_mid, k=TOPK, dim=1) topk_mean = topk_vals.mean(dim=1) topk_std = topk_vals.std(dim=1, unbiased=False) topk_cv = topk_std / (topk_mean + 1e-6) centered = topk_vals - topk_mean.unsqueeze(1) var = (centered ** 2).mean(dim=1) + 1e-6 m4 = (centered ** 4).mean(dim=1) topk_kurt = m4 / (var ** 2) topk_med = topk_vals.median(dim=1).values row_mean_cos = cos_mid.mean(dim=1) row_med_cos = cos_mid.median(dim=1).values s1_over_mean = diag_cos - row_mean_cos s1_over_med = diag_cos - row_med_cos sorted_probs, _ = torch.sort(probs_mid, dim=1, descending=True) p1 = sorted_probs[:, 0] p2 = sorted_probs[:, 1] if sorted_probs.size(1) > 1 else sorted_probs[:, 0] shape_H = -(sorted_probs * torch.log(sorted_probs + 1e-6)).sum(dim=1) shape_gini = 1.0 - (sorted_probs ** 2).sum(dim=1) R = min(10, sorted_probs.size(1)) x = torch.arange(R, device=device, dtype=sorted_probs.dtype) x_centered = x - x.mean() denom = (x_centered ** 2).sum() y = torch.log(sorted_probs[:, :R] + 1e-6) slope = (x_centered.unsqueeze(0) * y).sum(dim=1) / denom row_mean_p = probs_mid.mean(dim=1) row_std_p = probs_mid.std(dim=1, unbiased=False) + 1e-6 z1 = (p1_mid - row_mean_p) / row_std_p center_p = probs_mid - row_mean_p.unsqueeze(1) m3 = (center_p ** 3).mean(dim=1) skew = m3 / (row_std_p ** 3 + 1e-6) s1_over_sk = p1_mid - skew TAIL_K = min(10, sorted_probs.size(1)) tail_mean = sorted_probs[:, -TAIL_K:].mean(dim=1) HEAD_K = min(5, sorted_probs.size(1)) head5_mean = sorted_probs[:, :HEAD_K].mean(dim=1) mask_ratio = torch.zeros_like(diag_cos) mask_len = torch.zeros_like(diag_cos) mask_runs = torch.zeros_like(diag_cos) scalar_inputs = torch.stack( [ diag_cos, s2_cos, margin_mid, z_margin_mid, mad_margin_mid, p1_mid, H_mid, gini_mid, topk_mean, topk_std, topk_cv, topk_kurt, topk_med, s1_over_mean, s1_over_med, p1, p2, shape_H, shape_gini, slope, z1, s1_over_sk, tail_mean, head5_mean, mask_ratio, mask_len, mask_runs, ], dim=1, ) # Modality Index 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) # Labels gt = torch.arange(batch_size, device=device) mid_top1 = cos_mid.argmax(dim=1) last_top1 = cos_last.argmax(dim=1) mid_hit = mid_top1.eq(gt) last_hit = last_top1.eq(gt) need_last = (~mid_hit) & last_hit labels = need_last.float().unsqueeze(1) both_correct = mid_hit & last_hit both_wrong = (~mid_hit) & (~last_hit) # ======================================================= # 分类器前向（float32） # ======================================================= scalar_inputs_f32 = scalar_inputs.float() qry_reps_mid_f32 = qry_reps_mid.float() logits = model(scalar_inputs_f32, modality_idx, qry_emb=qry_reps_mid_f32) # Loss: V5调整版 Focal Loss (alpha=0.80, gamma=3.0) loss = sigmoid_focal_loss(logits, labels, alpha=0.80, gamma=3.0, reduction="mean") pred_probs = torch.sigmoid(logits) # Logging if self.state.global_step < 10 and self.args.local_rank == 0: pos_ratio = labels.mean().item() neg_ratio = 1.0 - pos_ratio print(f"\n[Probe Step {self.state.global_step}] Loss: {loss.item():.4f}", flush=True) print(f" - Pred Probs (need_last=1): mean={pred_probs.mean().item():.4f}, std={pred_probs.std().item():.4f}", flush=True) print(f" - Labels: need_last={pos_ratio:.4f}, safe={neg_ratio:.4f}", flush=True) print(f" - mid_hit: {mid_hit.float().mean().item():.2%}, last_hit: {last_hit.float().mean().item():.2%}", flush=True) print(f" - both_correct: {both_correct.float().mean().item():.2%}, both_wrong: {both_wrong.float().mean().item():.2%}", flush=True) return loss # ---------------- training_step ---------------- def training_step(self, model, inputs, num_items_in_batch=None) -> torch.Tensor: 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() 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}]", flush=True) 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 total_grad_norm = total_grad_norm ** 0.5 print(f" - Total Grad Norm: {total_grad_norm:.6f}", flush=True) print(f" - Has Gradient: {has_grad}", flush=True) if self.state.global_step >= 2: self._grad_check_done = True return loss.detach() / self.args.gradient_accumulation_steps