#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 小狸完整数据集训练 - 支持断点续训 + batch=4防OOM """ import sys, os sys.path.insert(0, '/workspace/xiaoli-llm-chat') sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1) sys.stderr = os.fdopen(sys.stderr.fileno(), 'w', buffering=1) import torch import torch.nn as nn import json import time import glob import random from torch.utils.data import Dataset, DataLoader from llm_chat_improved import ( CharTokenizer, build_vocabulary, LLMModel, LLMChat, ConversationDataset, LLMTrainer, tokenizer_instance ) from api_correction import batch_convert, batch_generate_multi_temp def dict_to_tuple(data_list): result = [] for item in data_list: if isinstance(item, dict): q = item.get("question", "") or item.get("input", "") or item.get("prompt", "") or item.get("instruction", "") a = item.get("answer", "") or item.get("output", "") or item.get("response", "") if q and a: result.append((q, a)) elif isinstance(item, (list, tuple)) and len(item) >= 2: result.append((str(item[0]), str(item[1]))) return result def load_multiturn_data(path): with open(path, 'r', encoding='utf-8') as f: data = json.load(f) single_turn = [] for item in data: if isinstance(item, dict) and "conversations" in item: convs = item["conversations"] for i in range(0, len(convs) - 1, 2): if i + 1 < len(convs): c1, c2 = convs[i], convs[i+1] if isinstance(c1, dict): h = c1.get("value", "") or c1.get("content", "") a = c2.get("value", "") or c2.get("content", "") elif isinstance(c1, list) and len(c1) >= 2: h, a = str(c1[1]), str(c2[1]) else: continue if h and a: single_turn.append((h, a)) return single_turn def find_latest_checkpoint(model_dir): """查找最新的checkpoint，返回 (路径, epoch数) 或 (None, 0)""" checkpoints = glob.glob(os.path.join(model_dir, 'llm_model_v2_resume_epoch_*.pth')) if not checkpoints: checkpoints = glob.glob(os.path.join(model_dir, 'llm_model_v2_epoch_*.pth')) if not checkpoints: return None, 0 best_path = None best_epoch = 0 for path in checkpoints: import re nums = [int(s) for s in re.findall(r'\d+', os.path.basename(path).replace('.pth',''))] if nums: epoch = max(nums) if epoch > best_epoch: best_epoch = epoch best_path = path return best_path, best_epoch def migrate_weights(old_sd, new_sd_init, old_d_model, new_d_model, old_n_layers, new_n_layers, old_d_ff, new_d_ff, old_vocab_size, new_vocab_size): """权重扩展迁移：将旧模型权重扩展到新架构 支持 d_model / n_layers / d_ff / vocab_size 全部变化 """ import torch new_sd = {} old_keys = list(old_sd.keys()) new_keys = list(new_sd_init.keys()) for key in new_keys: if 'freq' in key: # RoPE频率表，新模型自动生成，跳过 new_sd[key] = new_sd_init[key].clone() continue # 解析层索引（如 transformer_blocks.3.attention.wq.weight） old_key = key new_w = new_sd_init[key].clone() old_w_shape = None if key in old_sd: old_w = old_sd[key] old_w_shape = old_w.shape new_shape = new_w.shape # 1. 严格匹配 = 直接复制 if old_w.shape == new_shape: new_w = old_w.clone() # 2. 二维权重（线性层 + embedding） elif old_w.dim() == 2: o0, o1 = old_w.shape n0, n1 = new_shape # embedding.weight 和 fc_out.weight: 词汇表维度变化 if 'embedding' in key and old_vocab_size != new_vocab_size: # 词汇表扩展：复制旧词嵌入，新词随机 min_v = min(old_vocab_size, new_vocab_size) min_d = min(o1, n1) new_w[:min_v, :min_d] = old_w[:min_v, :min_d] print(f" 迁移 {key}: ({o0},{o1})→({n0},{n1}) vocab扩展") elif 'fc_out' in key and old_vocab_size != new_vocab_size: min_vocab = min(old_vocab_size, new_vocab_size) min_feat = min(o1, n1) new_w[:min_vocab, :min_feat] = old_w[:min_vocab, :min_feat] print(f" 迁移 {key}: ({o0},{o1})→({n0},{n1}) vocab扩展") elif o0 <= n0 and o1 <= n1: # 权重扩展（old 是 new 的子集，放在左上角） new_w[:o0, :o1] = old_w if o0 < n0 or o1 < n1: print(f" 扩展 {key}: ({o0},{o1})→({n0},{n1})") elif n0 <= o0 and n1 <= o1: # 权重缩小（new 是 old 的子集，取左上角） new_w = old_w[:n0, :n1].clone() print(f" 缩小 {key}: ({o0},{o1})→({n0},{n1})") else: # 特殊形状 → 尽量复制重合部分 min0 = min(o0, n0) min1 = min(o1, n1) new_w[:min0, :min1] = old_w[:min0, :min1] print(f" 部分迁移 {key}: ({o0},{o1})→({n0},{n1})") # 3. 一维权重（LayerNorm, bias等） elif old_w.dim() == 1: o0 = old_w.shape[0] n0 = new_w.shape[0] min_d = min(o0, n0) new_w[:min_d] = old_w[:min_d] if o0 != n0: print(f" 迁移 {key}: ({o0})→({n0}) norm扩展") else: new_w = old_w.clone() else: # 旧模型中没有这个key（可能是不存在的层） # 检查是否是额外的新层 # 如 transformer_blocks.6.xxx → 旧模型只有0-5 new_sd[key] = new_w # 保持随机初始化 continue new_sd[key] = new_w return new_sd def main(): print("=" * 60) print("小狸训练 - 从头训练版 v2 (13.2M架构)") print("=" * 60) device = 'cpu' model_dir = '/workspace/xiaoli-llm-chat' # 加载数据 - 从合并文件加载 merged_path = os.path.join(model_dir, 'xiaoli_merged_data_v2.json') if not os.path.exists(merged_path): print("❌ 未找到合并数据文件 xiaoli_merged_data.json") print("请先运行合并脚本") return with open(merged_path, 'r', encoding='utf-8') as f: data = json.load(f) unique = dict_to_tuple(data) print(f"加载合并数据: {len(unique)}条") # ====== 划分训练/验证集（保存到独立文件，避免问题重复问题 ====== val_split = 0.1 train_path = os.path.join(model_dir, 'xiaoli_train_v2.json') val_path = os.path.join(model_dir, 'xiaoli_valid_v2.json') if os.path.exists(train_path) and os.path.exists(val_path): # 从独立文件加载（精确的position-based划分） with open(train_path, 'r', encoding='utf-8') as f: train_set = dict_to_tuple(json.load(f)) with open(val_path, 'r', encoding='utf-8') as f: val_set = dict_to_tuple(json.load(f)) # 检查是否有新数据（来自外部脚本） train_qs = set(q for q, _ in train_set) val_qs = set(q for q, _ in val_set) new_entries = [(q, a) for q, a in unique if q not in train_qs and q not in val_qs] if new_entries: train_set.extend(new_entries) with open(train_path, 'w', encoding='utf-8') as f: json.dump([{"question": q, "answer": a} for q, a in train_set], f, ensure_ascii=False, indent=2) print(f" 新增外部数据: {len(new_entries)}条 → 加入训练集") print(f" 训练集: {len(train_set)}条 | 验证集: {len(val_set)}条 (来自文件)") else: # 首次划分：随机打乱后按位置切分 shuffled = unique.copy() random.seed(42) random.shuffle(shuffled) split_idx = int(len(shuffled) * val_split) val_set = shuffled[:split_idx] train_set = shuffled[split_idx:] # 保存两份独立文件 with open(train_path, 'w', encoding='utf-8') as f: json.dump([{"question": q, "answer": a} for q, a in train_set], f, ensure_ascii=False, indent=2) with open(val_path, 'w', encoding='utf-8') as f: json.dump([{"question": q, "answer": a} for q, a in val_set], f, ensure_ascii=False, indent=2) print(f" 首次划分: 训练集{len(train_set)}条 + 验证集{len(val_set)}条 (已保存)") print(f" 训练集: {len(train_set)}条 | 验证集: {len(val_set)}条") # ====== 加载固定vocab ====== vocab_path = os.path.join(model_dir, 'vocab.json') if not os.path.exists(vocab_path): print("❌ 未找到 vocab.json，请先运行BPE词表生成脚本") return with open(vocab_path, 'r', encoding='utf-8') as f: vocab = json.load(f) print(f"词汇表: {len(vocab)} (BPE, 来自 vocab.json)") # 同步tokenizer（新字符映射为UNK） tokenizer_instance.char_to_id = vocab.copy() tokenizer_instance.id_to_char = {v: k for k, v in vocab.items()} tokenizer_instance._sorted_vocab_cache = None for attr in ['_multi_trie', '_multi_keys', '_single_map']: if hasattr(tokenizer_instance, attr): delattr(tokenizer_instance, attr) # ====== 加载蒸馏问题(v3 + 真题) ====== distill_q_path = os.path.join(model_dir, 'exam_distill_questions.json') if not os.path.exists(distill_q_path): distill_q_path = os.path.join(model_dir, 'distill_questions_v3.json') if not os.path.exists(distill_q_path): distill_q_path = os.path.join(model_dir, 'distill_questions_v2.json') if os.path.exists(distill_q_path): with open(distill_q_path, 'r', encoding='utf-8') as f: distill_questions = json.load(f) print(f"蒸馏问题: {len(distill_questions)}道") else: distill_questions = [] print("蒸馏问题: 无") config = { 'batch_size': 4, 'gradient_accumulation_steps': 2, 'epochs': 200, 'd_model': 320, 'n_heads': 8, 'n_layers': 6, 'd_ff': 1536, 'max_length': 256, 'dropout': 0.1, 'use_amp': False, } model = LLMModel( vocab_size=len(vocab), d_model=config['d_model'], n_heads=config['n_heads'], n_layers=config['n_layers'], d_ff=config['d_ff'], max_len=config['max_length'], dropout=config['dropout'], use_memory=False, use_thinking=False, use_moe=False, num_experts=1, top_k=1, ) total_params = sum(p.numel() for p in model.parameters()) print(f"参数量: {total_params:,} ({total_params/1e6:.2f}M)") # 训练集/验证集 DataLoader train_dataset = ConversationDataset(train_set, vocab, config['max_length']) train_loader = DataLoader(train_dataset, batch_size=config['batch_size'], shuffle=True, num_workers=0) print(f"训练批次: {len(train_loader)}") val_dataset = ConversationDataset(val_set, vocab, config['max_length']) val_loader = DataLoader(val_dataset, batch_size=config['batch_size'], shuffle=False, num_workers=0) print(f"验证批次: {len(val_loader)}") # 检查断点 ckpt_path, start_epoch = find_latest_checkpoint(model_dir) if ckpt_path and start_epoch > 0: print(f"\n发现断点: epoch {start_epoch} ({ckpt_path})") checkpoint = torch.load(ckpt_path, map_location=device) old_cfg = checkpoint.get('model_config', {}) old_d_model = old_cfg.get('d_model', 256) old_n_layers = old_cfg.get('n_layers', 6) old_d_ff = old_cfg.get('d_ff', 1024) # 检测架构变化，自动做权重扩展迁移 new_sd_init = model.state_dict() # 随机初始化的新权重 old_sd = checkpoint['model_state_dict'] old_vocab_size = old_cfg.get('vocab_size', 0) new_vocab_size = len(vocab) # 如果架构没变化，直接加载 if (old_d_model == config['d_model'] and old_n_layers == config['n_layers'] and old_d_ff == config['d_ff'] and old_vocab_size == new_vocab_size): model.load_state_dict(old_sd) print("✅ 架构一致，直接加载权重") else: print(f" 架构变化: d_model={old_d_model}→{config['d_model']}, " f"layers={old_n_layers}→{config['n_layers']}, " f"d_ff={old_d_ff}→{config['d_ff']}, " f"vocab={old_vocab_size}→{new_vocab_size}") print(" 执行权重扩展迁移（宽度扩展 + 深度复制）...") new_sd = migrate_weights(old_sd, new_sd_init, old_d_model, config['d_model'], old_n_layers, config['n_layers'], old_d_ff, config['d_ff'], old_vocab_size, new_vocab_size) model.load_state_dict(new_sd) print(" ✅ 权重扩展迁移完成！") # 记录旧checkpoint已使用，从下一epoch开始 start_epoch = start_epoch # 尝试加载优化器（仅当架构vocab完全一致时） if (old_d_model == config['d_model'] and old_n_layers == config['n_layers'] and old_d_ff == config['d_ff'] and old_vocab_size == new_vocab_size and 'optimizer_state_dict' in checkpoint): try: trainer = LLMTrainer(model, device, use_amp=config['use_amp']) # 前向一步恢复优化器（不做完整epoch，太慢） batch = next(iter(train_loader)) input_ids = batch['input_ids'].to(device) target_ids = batch['target_ids'].to(device) mask = model.generate_mask(input_ids.size(1)).to(device) outputs = model(input_ids, mask) loss = trainer.criterion(outputs.reshape(-1, outputs.size(-1)), target_ids.reshape(-1)) loss.backward() trainer.optimizer.step() trainer.optimizer.zero_grad() # 加载优化器状态 trainer.optimizer.load_state_dict(checkpoint['optimizer_state_dict']) print("优化器状态已加载（1步热身）") except Exception as e: print(f"优化器加载失败（将使用新的）: {e}") print(f"\n开始训练: epoch {start_epoch + 1} -> {config['epochs']}") print("=" * 60) start = time.time() # 确保trainer已创建（可能因为架构变化没有在提前创建） if 'trainer' not in dir() or trainer is None: trainer = LLMTrainer(model, device, use_amp=config['use_amp']) best_val_loss = float('inf') total_distill_added = 0 for epoch in range(start_epoch + 1, config['epochs'] + 1): # ====== 每轮从train_path重建dataloader（蒸馏数据已追加至此）====== with open(train_path, 'r', encoding='utf-8') as f: epoch_train = dict_to_tuple(json.load(f)) print(f"\n 训练数据: {len(epoch_train)}条 (含{total_distill_added}条蒸馏新增)") train_dataset = ConversationDataset(epoch_train, vocab, config['max_length']) train_loader = DataLoader(train_dataset, batch_size=config['batch_size'], shuffle=True, num_workers=0) # ====== 训练 ====== t0 = time.time() avg_loss = trainer.train_epoch(train_loader, epoch, vocab, config) dt = time.time() - t0 # ====== 验证 ====== t1 = time.time() val_loss = trainer.evaluate(val_loader, config) dt_val = time.time() - t1 # 打印train/val loss marker = "" if val_loss < best_val_loss: best_val_loss = val_loss marker = " ★" best_path = os.path.join(model_dir, 'xiaoli_v2_best.pth') trainer.save_model(best_path, vocab, epoch) print(f"Epoch {epoch}/{config['epochs']} | Train Loss: {avg_loss:.4f} | Val Loss: {val_loss:.4f}{marker} | {dt:.0f}s+{dt_val:.0f}s") # 每轮保存断点 resume_path = os.path.join(model_dir, f'llm_model_v2_resume_epoch_{epoch}.pth') trainer.save_model(resume_path, vocab, epoch) if epoch % 20 == 0: save_path = os.path.join(model_dir, f'llm_model_v2_epoch_{epoch}.pth') trainer.save_model(save_path, vocab, epoch) print(f" 已保存正式checkpoint: epoch {epoch}") resume_files = sorted(glob.glob(os.path.join(model_dir, 'llm_model_v2_resume_*.pth'))) for old in resume_files[:-5]: try: os.remove(old) except: pass # ====== 每轮末尾：在线蒸馏 ====== if distill_questions: distill_t0 = time.time() try: print(f" 📡 在线蒸馏 ({len(distill_questions)}题 × 多种方式)...") chat = LLMChat(model, vocab, device, config['max_length']) all_new_pairs = [] # ====== Pass 1: 模型生成(temperature=0.15) → API转换 ====== print(" [P1] 模型生成中...") model_answers = [] for i, q in enumerate(distill_questions): ans = chat.generate_response(q, max_length=40, temperature=0.15) model_answers.append((q, ans)) if (i + 1) % 100 == 0: print(f" P1: {i+1}/{len(distill_questions)}") p1_gen = time.time() print(f" [P1] 生成完毕 ({(p1_gen-distill_t0):.0f}s)，API转换中...") def progress(done, total): if done % 50 == 0 or done == total: print(f" API: {done}/{total}", end="\r") p1_results = batch_convert(model_answers, progress) print(f" API: {len(model_answers)}/{len(model_answers)}") all_new_pairs.extend(p1_results) print(f" [P1] +{len(p1_results)}条") # ====== Pass 2: 模型高温生成(temperature=0.7) → 更多多样性 ====== print(f" [P2] 模型高温生成中 (temp=0.7)...") model_answers2 = [] for i, q in enumerate(distill_questions): ans = chat.generate_response(q, max_length=40, temperature=0.7) model_answers2.append((q, ans)) if (i + 1) % 100 == 0: print(f" P2: {i+1}/{len(distill_questions)}") p2_gen = time.time() print(f" [P2] 生成完毕 ({(p2_gen-p1_gen):.0f}s)，API转换中...") p2_results = batch_convert(model_answers2, progress) print(f" API: {len(model_answers2)}/{len(model_answers2)}") all_new_pairs.extend(p2_results) print(f" [P2] +{len(p2_results)}条") # ====== Pass 3: API直接生成多样回答(temp=1.0) ====== print(f" [P3] API直接生成 (temp=1.0, 全部直答)...") p3_t0 = time.time() p3_results = batch_generate_multi_temp( distill_questions, temperatures=[1.0], max_workers=8 ) all_new_pairs.extend(p3_results) print(f" [P3] +{len(p3_results)}条") # ====== 合并去重后追加 ====== if all_new_pairs: with open(train_path, 'r', encoding='utf-8') as f: cur_train = json.load(f) existing_pairs = set((item.get("question",""), item.get("answer","")) for item in cur_train) truly_new = [item for item in all_new_pairs if (item["question"], item["answer"]) not in existing_pairs] if truly_new: cur_train.extend(truly_new) with open(train_path, 'w', encoding='utf-8') as f: json.dump(cur_train, f, ensure_ascii=False, indent=2) # 同步到merged with open(merged_path, 'r', encoding='utf-8') as f: cur_merged = json.load(f) cur_merged.extend(truly_new) with open(merged_path, 'w', encoding='utf-8') as f: json.dump(cur_merged, f, ensure_ascii=False, indent=2) total_distill_added += len(truly_new) distill_total = time.time() - distill_t0 print(f" ✅ 蒸馏完成: +{len(truly_new)}条 (新增{len(set(q for q,a in [(i['question'],i['answer']) for i in truly_new]))}道独特问题) | 总耗时{distill_total:.0f}s") else: print(f" ℹ️ 蒸馏完成: 本轮无新内容（全部已存在）({time.time()-distill_t0:.0f}s)") else: print(f" ℹ️ 蒸馏完成: 无有效转换 ({time.time()-distill_t0:.0f}s)") except Exception as e: print(f" ⚠️ 蒸馏异常: {e}") # 保存最终模型 trainer.save_model(os.path.join(model_dir, 'xiaoli_v2_final.pth'), vocab, config['epochs']) # 清理所有resume文件 for f in glob.glob(os.path.join(model_dir, 'llm_model_v2_resume_*.pth')): try: os.remove(f) except: pass total = time.time() - start print(f"\n训练完成！总耗时: {total/60:.1f}分钟") # 测试 print("\n测试:") chat = LLMChat(model, vocab, device, config['max_length']) for q in ["你好", "你是谁", "你叫什么名字", "你好可爱", "我喜欢你", "1+1等于几"]: r = chat.generate_response(q, max_length=60, temperature=0.8) print(f" Q: {q}") print(f" A: {r}\n") if __name__ == "__main__": main()