#!/usr/bin/env python3 """ 修复logits_to_keep参数导致的loss计算错误 验证问题并提供解决方案 """ import json import torch import torch.nn.functional as F from transformers import AutoTokenizer from model.LMConfig import LMConfig from model.model_original import MiniMindLM def demonstrate_logits_to_keep_issue(): """ 演示logits_to_keep参数导致的问题 """ print("🔍 验证logits_to_keep参数问题") print("="*60) device = 'cuda' model_path = 'out/experiment_1_4_0/pretrain_512.pth' # 加载模型 config = LMConfig( dim=512, n_layers=8, n_heads=32, vocab_size=6400, max_seq_len=512, dropout=0.0, norm_eps=1e-5, rope_theta=1e6, use_moe=False ) model = MiniMindLM(config) tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer') state_dict = torch.load(model_path, map_location=device) model.load_state_dict(state_dict, strict=False) model.to(device) model.eval() # 加载测试数据 with open('dataset/stable/eval_data_from_train.json', 'r', encoding='utf-8') as f: sample = json.loads(f.readline().strip()) text = sample['text'] tokens = tokenizer.encode(text, add_special_tokens=False) input_tokens = tokens[:100] target_tokens = tokens[100:130] # 30个目标token print(f"测试样本: {len(tokens)} tokens") print(f"输入: {len(input_tokens)} tokens") print(f"目标: {len(target_tokens)} tokens") with torch.no_grad(): full_input = torch.tensor([tokens[:130]], dtype=torch.long).to(device) target_labels = torch.tensor(target_tokens, dtype=torch.long).to(device) print(f"\n🔬 详细对比不同方法:") # 方法1: 标准forward (正确方法) outputs1 = model(full_input) logits1 = outputs1.logits correct_logits = logits1[0, 99:129, :].contiguous() # 取position 99-128 loss1 = F.cross_entropy(correct_logits, target_labels, reduction='mean') print(f"1. 标准forward (正确):") print(f" 完整logits形状: {logits1.shape}") print(f" 用于计算的logits形状: {correct_logits.shape}") print(f" Loss: {loss1.item():.4f}") # 方法2: 使用logits_to_keep=30 (错误方法) outputs2 = model(full_input, logits_to_keep=30) logits2 = outputs2.logits incorrect_logits = logits2[0, -30:, :].contiguous() # 最后30个 loss2 = F.cross_entropy(incorrect_logits, target_labels, reduction='mean') print(f"\n2. logits_to_keep=30 (eval_model.py方法):") print(f" 部分logits形状: {logits2.shape}") print(f" 用于计算的logits形状: {incorrect_logits.shape}") print(f" Loss: {loss2.item():.4f}") # 方法3: 修复后的方法（不使用logits_to_keep） # 这就是方法1，但为了清晰显示修复方案 print(f"\n3. 修复方法 (不使用logits_to_keep):") print(f" 使用完整forward，然后选择正确的logits切片") print(f" 这与方法1相同，Loss: {loss1.item():.4f}") # 分析差异 print(f"\n📊 数值分析:") print(f" Loss差异: {abs(loss2.item() - loss1.item()):.4f}") print(f" Loss增幅: {(loss2.item() / loss1.item() - 1) * 100:.1f}%") # 检查logits的微小差异如何被放大 logits_diff = torch.abs(correct_logits - incorrect_logits).max() print(f" 最大logits差异: {logits_diff.item():.8f}") # 计算softmax概率的差异 prob1 = F.softmax(correct_logits, dim=-1) prob2 = F.softmax(incorrect_logits, dim=-1) prob_diff = torch.abs(prob1 - prob2).max() print(f" 最大概率差异: {prob_diff.item():.8f}") print(f"\n💡 结论:") print(f" 虽然logits差异很小({logits_diff.item():.8f})，") print(f" 但在交叉熵损失中被显著放大，导致loss增加{(loss2.item() / loss1.item() - 1) * 100:.1f}%") def create_fixed_eval_model(): """ 创建修复后的eval_model.py """ print(f"\n🔧 创建修复后的评估脚本") print("="*60) # 读取原始eval_model.py with open('eval_model.py', 'r', encoding='utf-8') as f: content = f.read() # 修复关键部分：移除logits_to_keep的使用 fixed_content = content.replace( """ # 计算loss（使用forward方法） # 准备用于loss计算的输入 loss_input_ids = torch.tensor([tokens[:input_length + predict_length]], dtype=torch.long).to(device) outputs = model(loss_input_ids, logits_to_keep=predict_length) # 计算loss logits = outputs.logits loss = None if logits is not None: # 重塑logits和目标 shift_logits = logits[0, -predict_length:, :].contiguous() shift_labels = torch.tensor(target_tokens, dtype=torch.long).to(device) # 计算交叉熵损失 loss = F.cross_entropy(shift_logits, shift_labels, reduction='mean') loss = loss.item()""", """ # 计算loss（使用forward方法） # 准备用于loss计算的输入 loss_input_ids = torch.tensor([tokens[:input_length + predict_length]], dtype=torch.long).to(device) outputs = model(loss_input_ids) # 移除logits_to_keep参数 # 计算loss logits = outputs.logits loss = None if logits is not None: # 重塑logits和目标 - 修复：使用正确的位置切片 shift_logits = logits[0, input_length:input_length + predict_length, :].contiguous() shift_labels = torch.tensor(target_tokens, dtype=torch.long).to(device) # 计算交叉熵损失 loss = F.cross_entropy(shift_logits, shift_labels, reduction='mean') loss = loss.item()""" ) # 保存修复后的文件 with open('eval_model_fixed.py', 'w', encoding='utf-8') as f: f.write(fixed_content) print(f"✅ 创建了修复版本：eval_model_fixed.py") print(f"主要修复:") print(f" 1. 移除 logits_to_keep 参数") print(f" 2. 使用正确的位置切片: [input_length:input_length + predict_length]") print(f" 3. 而不是错误的 [-predict_length:]") def test_fixed_evaluation(): """ 测试修复后的评估方法 """ print(f"\n🧪 测试修复后的评估方法") print("="*60) device = 'cuda' model_path = 'out/experiment_1_4_0/pretrain_512.pth' # 加载模型 config = LMConfig( dim=512, n_layers=8, n_heads=32, vocab_size=6400, max_seq_len=512, dropout=0.0, norm_eps=1e-5, rope_theta=1e6, use_moe=False ) model = MiniMindLM(config) tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer') state_dict = torch.load(model_path, map_location=device) model.load_state_dict(state_dict, strict=False) model.to(device) model.eval() # 测试多个样本 total_loss_old = 0 total_loss_fixed = 0 valid_samples = 0 with open('dataset/stable/eval_data_from_train.json', 'r', encoding='utf-8') as f: for i, line in enumerate(f): if i >= 10: # 测试前10个样本 break sample = json.loads(line.strip()) text = sample['text'] tokens = tokenizer.encode(text, add_special_tokens=False) if len(tokens) < 130: continue input_length = 100 predict_length = 30 input_tokens = tokens[:input_length] target_tokens = tokens[input_length:input_length + predict_length] with torch.no_grad(): full_input = torch.tensor([tokens[:input_length + predict_length]], dtype=torch.long).to(device) target_labels = torch.tensor(target_tokens, dtype=torch.long).to(device) # 原始错误方法 outputs_old = model(full_input, logits_to_keep=predict_length) logits_old = outputs_old.logits shift_logits_old = logits_old[0, -predict_length:, :].contiguous() loss_old = F.cross_entropy(shift_logits_old, target_labels, reduction='mean') # 修复后方法 outputs_fixed = model(full_input) logits_fixed = outputs_fixed.logits shift_logits_fixed = logits_fixed[0, input_length:input_length + predict_length, :].contiguous() loss_fixed = F.cross_entropy(shift_logits_fixed, target_labels, reduction='mean') total_loss_old += loss_old.item() total_loss_fixed += loss_fixed.item() valid_samples += 1 print(f"样本{i+1}: 原始{loss_old.item():.4f} -> 修复{loss_fixed.item():.4f}") avg_loss_old = total_loss_old / valid_samples avg_loss_fixed = total_loss_fixed / valid_samples print(f"\n📊 测试结果总结:") print(f" 测试样本数: {valid_samples}") print(f" 原始方法平均loss: {avg_loss_old:.4f}") print(f" 修复方法平均loss: {avg_loss_fixed:.4f}") print(f" 差异: {abs(avg_loss_old - avg_loss_fixed):.4f}") print(f" 修复后loss更接近训练时的教师强制loss (~2.4)") if __name__ == "__main__": demonstrate_logits_to_keep_issue() create_fixed_eval_model() test_fixed_evaluation()