xiaoli-llm-chat/train_v2.py

#!/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_v2.json')
    if not os.path.exists(vocab_path):
        print("❌ 未找到 vocab_v2.json，请先运行固定vocab生成脚本")
        return
    with open(vocab_path, 'r', encoding='utf-8') as f:
        vocab = json.load(f)
    print(f"词汇表: {len(vocab)} (固定, 来自 vocab_v2.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()