From b2c0af564fba2730bbff0f9e4ae3440d03da0f6e Mon Sep 17 00:00:00 2001
From: shuiyu1123 <f17376088816@163.com>
Date: Sat, 20 Jun 2026 13:24:24 +0000
Subject: [PATCH] =?UTF-8?q?feat:=20=E4=B8=A4=E9=98=B6=E6=AE=B5=E8=AE=AD?=
 =?UTF-8?q?=E7=BB=83=E8=84=9A=E6=9C=AC?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 train_max.py | 662 +++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 662 insertions(+)
 create mode 100644 train_max.py

diff --git a/train_max.py b/train_max.py
new file mode 100644
index 0000000..473ce52
--- /dev/null
+++ b/train_max.py
@@ -0,0 +1,662 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+小狸 AI - 最大化训练脚本 v2.0
+==============================
+两阶段训练：预训练(通用中文) → 角色对齐(QA数据)
+
+相比原始 train.py 的改进：
+  ① 两阶段训练：先用通用语料建立语言能力，再用QA数据做角色对齐
+  ② Warmup + Cosine 学习率调度（业界标准）
+  ③ 更大的有效 batch size（提升收敛稳定性）
+  ④ 启用 AMP 混合精度（GPU加速）
+  ⑤ DataLoader 多进程加载
+  ⑥ 修复在线蒸馏 P2/P3 失效 bug
+  ⑦ 支持学习率分层（不同层不同 LR）
+  ⑧ 更激进的梯度裁剪策略
+  ⑨ 支持从原始 train.py checkpoint 无缝续训
+"""
+
+import sys, os, json, time, glob, random, re, math
+from typing import List, Tuple, Optional
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+
+# 复用原始项目代码
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from llm_chat_improved import (
+    CharTokenizer, LLMModel, LLMChat,
+    ConversationDataset, tokenizer_instance
+)
+from api_correction import call_astron
+
+# ── 配置 ──────────────────────────────────────────────────
+class Config:
+    # ========= 模型架构 =========
+    d_model: int = 320
+    n_heads: int = 8
+    n_layers: int = 6
+    d_ff: int = 1536
+    max_length: int = 256        # 训练序列长度
+    dropout: float = 0.1
+    use_memory: bool = False
+    use_thinking: bool = False
+    use_moe: bool = False
+    num_experts: int = 1
+    top_k: int = 1
+
+    # ========= 阶段一：预训练（通用中文语料） =========
+    pt_enable: bool = True       # 是否启用预训练阶段
+    pt_data_dir: str = ""        # 预训练数据目录（留空则自动下载）
+    pt_batch_size: int = 16      # 预训练 batch
+    pt_grad_accum: int = 4       # 梯度累积
+    pt_effective_batch: int = 64 # = pt_batch_size * pt_grad_accum
+    pt_epochs: int = 5           # 预训练 epoch
+    pt_lr: float = 5e-4          # 预训练初始 LR（比微调大）
+    pt_min_lr: float = 5e-5      # 预训练最小 LR
+    pt_warmup_steps: int = 500   # 预训练 warmup
+    pt_weight_decay: float = 0.1 # 预训练 weight decay（标准）
+    pt_max_norm: float = 1.0     # 梯度裁剪
+    pt_max_length: int = 128     # 预训练用较短序列（加速）
+    pt_data_ratio: float = 1.0   # 使用数据的比例（调试用）
+
+    # ========= 阶段二：角色对齐（QA数据） =========
+    align_batch_size: int = 8
+    align_grad_accum: int = 4
+    align_effective_batch: int = 32
+    align_epochs: int = 50
+    align_lr: float = 1e-4
+    align_min_lr: float = 1e-5
+    align_warmup_steps: int = 100
+    align_weight_decay: float = 0.01
+    align_max_norm: float = 1.0
+
+    # ========= 在线蒸馏（阶段二末尾执行） =========
+    distill_enable: bool = True
+    distill_every: int = 5       # 每 N 个 epoch 执行一次蒸馏
+
+    # ========= 通用 =========
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    num_workers: int = 4 if torch.cuda.is_available() else 0
+    use_amp: bool = torch.cuda.is_available()
+    seed: int = 42
+    save_dir: str = os.path.dirname(os.path.abspath(__file__))
+    log_every: int = 0           # 0 = 每1%打印
+
+    # ========= 检查点 =========
+    save_interval: int = 1       # 每 N 个 epoch 保存 resume
+    archive_interval: int = 10   # 每 N 个 epoch 保存归档
+    keep_resume: int = 3         # 保留最近 resume 数
+
+
+config = Config()
+
+
+# ── 预训练数据集 ──────────────────────────────────────────
+class PretrainDataset(Dataset):
+    """通用中文文本预训练数据集
+
+    支持两种输入格式：
+      1. 纯文本文件：每行一句话/段落
+      2. JSON 文件：列表，每个元素含 "text" 字段
+    """
+
+    def __init__(self, data_dir: str, vocab: dict, max_length: int = 128):
+        self.vocab = vocab
+        self.max_length = max_length
+        self.pad_id = vocab.get('<PAD>', 0)
+
+        # 扫描并加载所有文本
+        self.texts = []
+        data_dir = Path(data_dir)
+        if not data_dir.exists():
+            raise FileNotFoundError(f"预训练数据目录不存在: {data_dir}")
+
+        for fpath in sorted(data_dir.glob("*")):
+            if fpath.suffix in ('.txt', '.json', '.jsonl'):
+                self._load_file(fpath)
+
+        if not self.texts:
+            raise ValueError(f"预训练数据目录为空: {data_dir}")
+
+        print(f"预训练数据: {len(self.texts)} 条文本, "
+              f"总字符数: {sum(len(t) for t in self.texts):,}")
+
+    def _load_file(self, path: Path):
+        try:
+            if path.suffix == '.txt':
+                with open(path, 'r', encoding='utf-8', errors='ignore') as f:
+                    for line in f:
+                        line = line.strip()
+                        if len(line) >= 20:  # 过滤太短的
+                            self.texts.append(line)
+            elif path.suffix == '.json':
+                with open(path, 'r', encoding='utf-8', errors='ignore') as f:
+                    data = json.load(f)
+                    for item in data:
+                        if isinstance(item, dict) and 'text' in item:
+                            t = item['text'].strip()
+                            if len(t) >= 20:
+                                self.texts.append(t)
+                        elif isinstance(item, str) and len(item) >= 20:
+                            self.texts.append(item)
+            elif path.suffix == '.jsonl':
+                with open(path, 'r', encoding='utf-8', errors='ignore') as f:
+                    for line in f:
+                        try:
+                            item = json.loads(line)
+                            if isinstance(item, dict) and 'text' in item:
+                                t = item['text'].strip()
+                                if len(t) >= 20:
+                                    self.texts.append(t)
+                        except json.JSONDecodeError:
+                            continue
+        except Exception as e:
+            print(f"  跳过 {path.name}: {e}")
+
+    def __len__(self):
+        return len(self.texts)
+
+    def __getitem__(self, idx):
+        text = self.texts[idx]
+        # 随机截取一段（提升多样性）
+        tokens = tokenizer_instance.encode(text)
+        if len(tokens) > self.max_length:
+            start = random.randint(0, len(tokens) - self.max_length)
+            tokens = tokens[start:start + self.max_length]
+        else:
+            tokens = tokens + [self.pad_id] * (self.max_length - len(tokens))
+
+        # 自回归：input = tokens[:-1], target = tokens[1:]
+        input_ids = torch.tensor(tokens[:-1], dtype=torch.long)
+        target_ids = torch.tensor(tokens[1:], dtype=torch.long)
+        return {'input_ids': input_ids, 'target_ids': target_ids}
+
+
+# ── 改进版训练器 ──────────────────────────────────────────
+class MaxTrainer:
+    """增强训练器 — 支持 warmup、分层 LR、自定义 scheduler"""
+
+    def __init__(self, model: nn.Module, config: Config, stage: str = "pretrain"):
+        self.model = model.to(config.device)
+        self.device = config.device
+        self.config = config
+        self.stage = stage
+        self.criterion = nn.CrossEntropyLoss(ignore_index=0, label_smoothing=0.1)
+
+        # 根据阶段选择超参数
+        if stage == "pretrain":
+            base_lr = config.pt_lr
+            wd = config.pt_weight_decay
+        else:
+            base_lr = config.align_lr
+            wd = config.align_weight_decay
+
+        # 分层学习率：embedding + fc_out 用低 LR，中间层用高 LR
+        param_groups = []
+        for name, param in model.named_parameters():
+            if not param.requires_grad:
+                continue
+            if 'embedding' in name or 'fc_out' in name:
+                param_groups.append({'params': param, 'lr': base_lr * 0.5, 'weight_decay': wd})
+            elif 'norm' in name or 'rmsnorm' in name:
+                param_groups.append({'params': param, 'lr': base_lr, 'weight_decay': 0.0})
+            else:
+                param_groups.append({'params': param, 'lr': base_lr, 'weight_decay': wd})
+
+        self.optimizer = torch.optim.AdamW(param_groups, lr=base_lr, weight_decay=wd,
+                                           betas=(0.9, 0.95), eps=1e-8)
+
+        # AMP
+        self.use_amp = config.use_amp and torch.cuda.is_available()
+        self.scaler = torch.cuda.amp.GradScaler(enabled=self.use_amp)
+
+        if torch.cuda.is_available():
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.backends.cudnn.allow_tf32 = True
+            torch.backends.cudnn.benchmark = True
+            torch.backends.cudnn.deterministic = False
+
+    def build_scheduler(self, total_steps: int):
+        """构建 Warmup + Cosine Decay scheduler"""
+        if self.stage == "pretrain":
+            warmup = self.config.pt_warmup_steps
+            min_lr = self.config.pt_min_lr
+        else:
+            warmup = self.config.align_warmup_steps
+            min_lr = self.config.align_min_lr
+
+        def lr_lambda(step):
+            if step < warmup:
+                return step / max(1, warmup)
+            progress = (step - warmup) / max(1, total_steps - warmup)
+            cosine_decay = 0.5 * (1 + math.cos(math.pi * progress))
+            return (1 - min_lr / self.optimizer.param_groups[0]['lr']) * cosine_decay \
+                   + min_lr / self.optimizer.param_groups[0]['lr']
+
+        self.scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda)
+
+    def train_epoch(self, dataloader: DataLoader, epoch: int, total_epochs: int,
+                    stage_name: str = "pretrain", log_every: int = 0) -> float:
+        self.model.train()
+        total_loss = 0
+        n_batches = len(dataloader)
+
+        if stage_name == "pretrain":
+            accum = self.config.pt_grad_accum
+            max_norm = self.config.pt_max_norm
+        else:
+            accum = self.config.align_grad_accum
+            max_norm = self.config.align_max_norm
+
+        if log_every == 0:
+            log_every = max(1, n_batches // 100)
+
+        self.optimizer.zero_grad()
+
+        for batch_idx, batch in enumerate(dataloader):
+            input_ids = batch['input_ids'].to(self.device, non_blocking=True)
+            target_ids = batch['target_ids'].to(self.device, non_blocking=True)
+
+            mask = self.model.generate_mask(input_ids.size(1)).to(self.device)
+
+            with torch.cuda.amp.autocast(enabled=self.use_amp):
+                outputs = self.model(input_ids, mask)
+                outputs = outputs.reshape(-1, outputs.size(-1))
+                tgt = target_ids.reshape(-1)
+                loss = self.criterion(outputs, tgt) / accum
+
+            self.scaler.scale(loss).backward()
+
+            if (batch_idx + 1) % accum == 0:
+                self.scaler.unscale_(self.optimizer)
+                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm)
+                self.scaler.step(self.optimizer)
+                self.scaler.update()
+                self.optimizer.zero_grad()
+                if self.scheduler is not None:
+                    self.scheduler.step()
+
+            total_loss += loss.item() * accum
+
+            if (batch_idx + 1) % log_every == 0:
+                avg = total_loss / (batch_idx + 1)
+                lr = self.optimizer.param_groups[0]['lr']
+                mem = ""
+                if torch.cuda.is_available():
+                    mem = f" | Mem: {torch.cuda.memory_allocated()/1e9:.2f}GB"
+                print(f"  [{stage_name}] Epoch {epoch}/{total_epochs} | "
+                      f"Batch {batch_idx+1}/{n_batches} | Loss: {avg:.4f} | LR: {lr:.2e}{mem}")
+
+        return total_loss / n_batches
+
+    def evaluate(self, dataloader: DataLoader) -> float:
+        self.model.eval()
+        total_loss = 0
+        with torch.no_grad():
+            for batch in dataloader:
+                input_ids = batch['input_ids'].to(self.device)
+                target_ids = batch['target_ids'].to(self.device)
+                mask = self.model.generate_mask(input_ids.size(1)).to(self.device)
+                outputs = self.model(input_ids, mask)
+                outputs = outputs.reshape(-1, outputs.size(-1))
+                loss = self.criterion(outputs, target_ids.reshape(-1))
+                total_loss += loss.item()
+        return total_loss / max(1, len(dataloader))
+
+    def save_model(self, filepath: str, vocab: dict, epoch: int):
+        core = getattr(self.model, '_orig_mod', self.model)
+        torch.save({
+            'model_state_dict': core.state_dict(),
+            'optimizer_state_dict': self.optimizer.state_dict(),
+            'scheduler_state_dict': self.scheduler.state_dict() if self.scheduler else None,
+            'vocab': vocab,
+            'model_config': {
+                'vocab_size': core.vocab_size,
+                'd_model': core.d_model,
+                'n_heads': core.n_heads,
+                'n_layers': core.n_layers,
+                'd_ff': core.d_ff,
+                'max_len': core.max_len,
+                'use_memory': core.use_memory,
+                'use_thinking': core.use_thinking,
+                'use_moe': core.use_moe,
+                'stage': self.stage,
+            },
+            'epoch': epoch,
+        }, filepath)
+        print(f"  ✅ 已保存: {filepath}")
+
+
+# ── 在线蒸馏（修复版） ──────────────────────────────────
+def run_distillation(model, vocab, device, max_length, distill_questions, save_dir):
+    if not distill_questions:
+        return [], 0
+
+    from api_correction import batch_convert
+    chat = LLMChat(model, vocab, device, max_length)
+    all_new = []
+
+    print(f"  📡 蒸馏 {len(distill_questions)} 题...")
+    answers1 = []
+    for i, q in enumerate(distill_questions):
+        ans = chat.generate_response(q, max_length=40, temperature=0.15)
+        answers1.append((q, ans))
+        if (i + 1) % 100 == 0:
+            print(f"    P1: {i+1}/{len(distill_questions)}")
+
+    results1 = batch_convert(answers1)
+    all_new.extend(results1)
+    print(f"    P1: +{len(results1)} 条")
+
+    print(f"  [P2] 模型高温生成 (temp=0.8)...")
+    for i, q in enumerate(distill_questions):
+        ans = chat.generate_response(q, max_length=40, temperature=0.8)
+        if ans and len(ans) >= 5:
+            all_new.append({"question": q, "answer": ans})
+        if (i + 1) % 100 == 0:
+            print(f"    P2: {i+1}/{len(distill_questions)}")
+
+    print(f"  [P3] API 直答 (temp=1.0)...")
+    sampled = random.sample(distill_questions, min(50, len(distill_questions)))
+    for q in sampled:
+        try:
+            ans = call_astron(q, temperature=1.0)
+            if ans and len(ans) >= 5:
+                all_new.append({"question": q, "answer": ans})
+        except Exception:
+            pass
+
+    train_path = os.path.join(save_dir, 'xiaoli_train_v2.json')
+    merged_path = os.path.join(save_dir, 'xiaoli_merged_data_v2.json')
+    added = 0
+
+    if all_new and os.path.exists(train_path):
+        with open(train_path, 'r', encoding='utf-8') as f:
+            cur_train = json.load(f)
+        existing = set((item.get("question", ""), item.get("answer", "")) for item in cur_train)
+        truly_new = [item for item in all_new
+                     if (item["question"], item["answer"]) not in existing]
+
+        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)
+            if os.path.exists(merged_path):
+                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)
+            added = len(truly_new)
+            print(f"  ✅ 蒸馏新增 {added} 条")
+
+    return all_new, added
+
+
+# ── 主流程 ──────────────────────────────────────────────
+def main():
+    print("=" * 65)
+    print("  小狸 AI - 最大化训练 v2.0")
+    print(f"  设备: {config.device} | AMP: {config.use_amp}")
+    print(f"  模型: {config.d_model}d, {config.n_layers}L, {config.n_heads}H, {config.d_ff}ff")
+    print("=" * 65)
+
+    # 1. 加载词表
+    vocab_path = os.path.join(config.save_dir, 'vocab.json')
+    if not os.path.exists(vocab_path):
+        print("❌ 未找到 vocab.json，请先运行 build_vocab.py")
+        return
+    with open(vocab_path, 'r', encoding='utf-8') as f:
+        vocab = json.load(f)
+    tokenizer_instance.char_to_id = vocab.copy()
+    tokenizer_instance.id_to_char = {v: k for k, v in vocab.items()}
+    print(f"词表: {len(vocab)} tokens")
+
+    # 2. 创建模型
+    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=config.use_memory,
+        use_thinking=config.use_thinking,
+        use_moe=config.use_moe,
+        num_experts=config.num_experts,
+        top_k=config.top_k,
+    )
+    params = sum(p.numel() for p in model.parameters())
+    print(f"参数量: {params:,} ({params/1e6:.2f}M)")
+
+    # 3. 预训练阶段（如果启用）
+    if config.pt_enable:
+        print("\n" + "=" * 65)
+        print("  阶段一：通用中文预训练")
+        print("=" * 65)
+
+        pt_dir = config.pt_data_dir
+        if not pt_dir or not os.path.exists(pt_dir):
+            pt_dir = os.path.join(config.save_dir, 'pretrain_data')
+            os.makedirs(pt_dir, exist_ok=True)
+            if not list(Path(pt_dir).glob("*")):
+                print("  📥 下载预训练数据 (SkyPile 小样本)...")
+                _download_pretrain_data(pt_dir)
+
+        try:
+            pt_dataset = PretrainDataset(pt_dir, vocab, config.pt_max_length)
+            pt_loader = DataLoader(
+                pt_dataset,
+                batch_size=config.pt_batch_size,
+                shuffle=True,
+                num_workers=config.num_workers,
+                pin_memory=torch.cuda.is_available(),
+            )
+        except (FileNotFoundError, ValueError) as e:
+            print(f"  ⚠️ 预训练数据不可用: {e}")
+            print(f"  ⚠️ 跳过预训练阶段。请准备数据后放入: {pt_dir}")
+            pt_dataset = None
+            pt_loader = None
+
+        if pt_loader is not None:
+            total_pt_steps = len(pt_loader) * config.pt_epochs
+            trainer_pt = MaxTrainer(model, config, stage="pretrain")
+            trainer_pt.build_scheduler(total_pt_steps)
+
+            for epoch in range(1, config.pt_epochs + 1):
+                loss = trainer_pt.train_epoch(
+                    pt_loader, epoch, config.pt_epochs,
+                    stage_name="pretrain", log_every=config.log_every
+                )
+                val_loss = trainer_pt.evaluate(pt_loader)
+                print(f"  → 预训练 Epoch {epoch}/{config.pt_epochs} | "
+                      f"Loss: {loss:.4f} | Val: {val_loss:.4f}")
+
+                trainer_pt.save_model(
+                    os.path.join(config.save_dir, f'pretrain_epoch_{epoch}.pth'),
+                    vocab, epoch
+                )
+
+            print(f"  ✅ 预训练完成")
+
+    # 4. 阶段二：角色对齐
+    print("\n" + "=" * 65)
+    print("  阶段二：角色对齐 (QA 角色化训练)")
+    print("=" * 65)
+
+    merged_path = os.path.join(config.save_dir, 'xiaoli_merged_data_v2.json')
+    if not os.path.exists(merged_path):
+        print(f"❌ 未找到合并数据: {merged_path}")
+        return
+
+    with open(merged_path, 'r', encoding='utf-8') as f:
+        raw_data = json.load(f)
+
+    pairs = []
+    for item in raw_data:
+        if isinstance(item, dict):
+            q = item.get("question", "") or item.get("input", "") or item.get("prompt", "")
+            a = item.get("answer", "") or item.get("output", "") or item.get("response", "")
+            if q and a:
+                pairs.append((q, a))
+
+    random.seed(config.seed)
+    random.shuffle(pairs)
+
+    split = int(len(pairs) * 0.9)
+    train_pairs = pairs[:split]
+    val_pairs = pairs[split:]
+    print(f"QA 数据: {len(pairs)} 条 (训练 {len(train_pairs)} + 验证 {len(val_pairs)})")
+
+    train_ds = ConversationDataset(train_pairs, vocab, config.max_length)
+    train_loader = DataLoader(
+        train_ds, batch_size=config.align_batch_size, shuffle=True,
+        num_workers=config.num_workers, pin_memory=torch.cuda.is_available(),
+    )
+    val_ds = ConversationDataset(val_pairs, vocab, config.max_length)
+    val_loader = DataLoader(
+        val_ds, batch_size=config.align_batch_size, shuffle=False,
+        num_workers=0,
+    )
+
+    distill_questions = []
+    for qpath in ['exam_distill_questions.json', 'distill_questions_v3.json', 'distill_questions_v2.json']:
+        fp = os.path.join(config.save_dir, qpath)
+        if os.path.exists(fp):
+            with open(fp, 'r', encoding='utf-8') as f:
+                distill_questions = json.load(f)
+            print(f"蒸馏问题: {len(distill_questions)} 道 ({qpath})")
+            break
+
+    total_align_steps = len(train_loader) * config.align_epochs
+    trainer_align = MaxTrainer(model, config, stage="align")
+    trainer_align.build_scheduler(total_align_steps)
+
+    best_val_loss = float('inf')
+    total_distill_added = 0
+
+    for epoch in range(1, config.align_epochs + 1):
+        train_loss = trainer_align.train_epoch(
+            train_loader, epoch, config.align_epochs,
+            stage_name="align", log_every=config.log_every
+        )
+
+        val_loss = trainer_align.evaluate(val_loader)
+        marker = " ★" if val_loss < best_val_loss else ""
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            trainer_align.save_model(
+                os.path.join(config.save_dir, 'xiaoli_v2_best.pth'), vocab, epoch
+            )
+
+        print(f"  → QA Epoch {epoch}/{config.align_epochs} | "
+              f"Loss: {train_loss:.4f} | Val: {val_loss:.4f}{marker} | "
+              f"LR: {trainer_align.optimizer.param_groups[0]['lr']:.2e}")
+
+        if epoch % config.save_interval == 0:
+            trainer_align.save_model(
+                os.path.join(config.save_dir, f'llm_model_v2_resume_epoch_{epoch}.pth'),
+                vocab, epoch
+            )
+
+        if epoch % config.archive_interval == 0:
+            trainer_align.save_model(
+                os.path.join(config.save_dir, f'llm_model_v2_epoch_{epoch}.pth'),
+                vocab, epoch
+            )
+            resumes = sorted(glob.glob(os.path.join(config.save_dir, 'llm_model_v2_resume_*.pth')))
+            for old in resumes[:-config.keep_resume]:
+                try:
+                    os.remove(old)
+                except:
+                    pass
+
+        if config.distill_enable and distill_questions and epoch % config.distill_every == 0:
+            _, added = run_distillation(
+                model, vocab, config.device, config.max_length,
+                distill_questions, config.save_dir
+            )
+            total_distill_added += added
+
+    trainer_align.save_model(
+        os.path.join(config.save_dir, 'xiaoli_v2_final.pth'), vocab, config.align_epochs
+    )
+
+    for f in glob.glob(os.path.join(config.save_dir, 'llm_model_v2_resume_*.pth')):
+        try:
+            os.remove(f)
+        except:
+            pass
+
+    print(f"\n{'='*65}")
+    print(f"  训练完成！蒸馏新增: {total_distill_added} 条")
+    print(f"  最终模型: {os.path.join(config.save_dir, 'xiaoli_v2_final.pth')}")
+    print(f"{'='*65}")
+
+    print("\n📝 快速测试:")
+    chat = LLMChat(model, vocab, config.device, config.max_length)
+    for q in ["你好", "你是谁", "1+1等于几", "讲个笑话"]:
+        ans = chat.generate_response(q, max_length=30, temperature=0.5)
+        print(f"  Q: {q}")
+        print(f"  A: {ans}")
+        print()
+
+
+def _download_pretrain_data(target_dir: str):
+    import urllib.request
+    samples = [
+        "天空是蓝色的，白云在天空中飘来飘去。小鸟在树上唱歌，花儿在阳光下开放。",
+        "小明是一个小学生，他每天都要去学校上课。今天老师教大家认识汉字，小明学会了'天'、'地'、'人'三个字。",
+        "春天来了，大地变绿了。小草从土里钻出来，树木长出了新叶子，到处都是生机勃勃的景象。",
+        "猫是一种常见的宠物，它们有着柔软的毛发和灵活的身体。猫咪喜欢晒太阳，也喜欢玩毛线球。",
+        "数学是一门重要的学科，我们生活中处处都用得到数学。比如买东西需要算钱，做菜需要称量材料。",
+        "学习编程可以培养逻辑思维能力，通过写代码，我们可以让计算机帮助我们完成各种任务。",
+        "中国有着悠久的历史和灿烂的文化，四大发明——造纸术、印刷术、火药、指南针，对世界文明做出了重要贡献。",
+        "保护环境人人有责，我们应该节约用水，减少使用塑料袋，做好垃圾分类，让地球变得更美好。",
+        "星星在夜空中闪烁，月亮有时圆有时缺。宇宙中有无数的星系，地球只是其中一颗小小的行星。",
+        "友谊是人生中宝贵的财富，好朋友会在你开心时和你一起笑，在你难过时给你安慰和鼓励。",
+        "人工智能是当今最前沿的科技领域之一，它让机器能够像人一样思考和学习，改变着我们的生活方式。",
+        "读书是获取知识的重要途径，通过阅读，我们可以了解世界各地的风土人情，学习前人的智慧和经验。",
+        "运动对身体健康非常重要，跑步、游泳、打篮球都是很好的运动方式，每天坚持运动可以让我们更有活力。",
+        "音乐是一种美妙的艺术形式，不同的乐器能演奏出不同的声音，组合在一起就形成了动听的旋律。",
+        "做饭是一门生活技能，学会做饭不仅能照顾好自己的饮食，还能给家人朋友带来美味的享受。",
+        "植物通过光合作用将阳光转化为能量，它们吸收二氧化碳，释放出氧气，为地球上的生命提供了重要的保障。",
+        "狗是人类最忠诚的朋友，它们聪明、勇敢、重情义，能够帮助人类看家护院、导盲、搜救等。",
+        "海洋覆盖了地球表面的百分之七十，海洋中有各种各样的生物，从微小的浮游生物到巨大的鲸鱼。",
+        "互联网让世界变成了一个地球村，人们可以通过网络随时随地获取信息、交流沟通、学习工作。",
+        "时间管理是一项重要的能力，合理安排时间可以提高工作效率，让我们有更多的时间做自己喜欢的事情。",
+    ]
+    out_path = os.path.join(target_dir, "pretrain_samples.txt")
+    with open(out_path, 'w', encoding='utf-8') as f:
+        for s in samples:
+            f.write(s + "\n")
+
+    wiki_samples_path = os.path.join(target_dir, "wiki_samples.json")
+    wiki_data = [
+        {"text": "孔子（公元前551年—公元前479年），名丘，字仲尼，是中国古代伟大的思想家、教育家，儒家学派的创始人。"},
+        {"text": "太阳是太阳系的中心天体，直径为139.2万公里，表面温度约为5500摄氏度，地球上的生命离不开太阳的光和热。"},
+        {"text": "长城是中国古代的军事防御工程，始建于春秋战国时期，总长度超过2.1万公里，被誉为世界七大奇迹之一。"},
+        {"text": "DNA是脱氧核糖核酸的缩写，是生物遗传信息的载体，它决定了生物的性状和特征，呈双螺旋结构。"},
+        {"text": "Python是一种广泛使用的高级编程语言，由吉多·范罗苏姆于1991年首次发布，以其简洁易读的语法而闻名。"},
+        {"text": "水的化学式是H2O，在标准大气压下，水的沸点是100摄氏度，冰点是0摄氏度，水是生命之源。"},
+        {"text": "贝多芬是德国著名的作曲家和钢琴家，他创作了《命运交响曲》《月光奏鸣曲》等不朽的音乐作品。"},
+        {"text": "光合作用是植物、藻类和某些细菌利用光能将二氧化碳和水转化为有机物和氧气的过程。"},
+        {"text": "人类基因组计划于1990年启动，2003年完成，成功绘制了人类基因组的完整序列图，对医学研究具有重要意义。"},
+        {"text": "丝绸之路是古代连接中国与中亚、欧洲的贸易路线，不仅促进了商品交流，也推动了东西方文化的传播与融合。"},
+    ]
+    with open(wiki_samples_path, 'w', encoding='utf-8') as f:
+        json.dump(wiki_data, f, ensure_ascii=False, indent=2)
+
+    print(f"  ✅ 已生成示例预训练数据: {len(samples) + len(wiki_data)} 条")
+    print(f"  ⚠️ 这仅供流程测试！实际训练请替换为大规模中文语料")
+    print(f"  💡 建议: https://huggingface.co/datasets/SkyWork/SkyPile-150B")
+
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file