cnb.cEr0mSxFAFA
a88c47e273
- 用 CharTokenizer.build_vocab 从 977K 数据生成词表 (vocab_size=47725) - train_cnb.py: batch=224 + accum=2 (等效448) + AMP + torch.compile - build_vocab.py: 词表生成脚本 - vocab_full.json: 含特殊token完整词表 - 模型 41.85M参数, L40 GPU训练
204 lines
8.9 KiB
Python
204 lines
8.9 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
小狸AI - CNB GPU训练脚本
|
||
从 merged_clean_v3.json (977K对) 训练，自动GPU检测
|
||
"""
|
||
import sys, os, json, time, glob, re, random, argparse
|
||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1)
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
from torch.utils.data import Dataset, DataLoader
|
||
|
||
from llm_chat_improved import LLMModel, LLMChat, ConversationDataset, LLMTrainer, tokenizer_instance
|
||
|
||
MODEL_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
|
||
def load_pairs(path):
|
||
with open(path, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
pairs = []
|
||
for item in data:
|
||
if isinstance(item, dict):
|
||
q = item.get("question", "") or item.get("input", "")
|
||
a = item.get("answer", "") or item.get("output", "")
|
||
if q and a:
|
||
pairs.append((q, a))
|
||
return pairs
|
||
|
||
def find_latest_ckpt(model_dir):
|
||
ckpts = glob.glob(os.path.join(model_dir, 'llm_model_v2_resume_epoch_*.pth'))
|
||
if not ckpts:
|
||
ckpts = glob.glob(os.path.join(model_dir, 'llm_model_v2_epoch_*.pth'))
|
||
if not ckpts:
|
||
return None, 0
|
||
nums = [(p, max([int(s) for s in re.findall(r'\d+', os.path.basename(p).replace('.pth',''))])) for p in ckpts]
|
||
best = max(nums, key=lambda x: x[1])
|
||
return best[0], best[1]
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description='小狸AI CNB训练')
|
||
parser.add_argument('--epochs', type=int, default=100, help='训练轮数')
|
||
parser.add_argument('--batch_size', type=int, default=8, help='batch大小')
|
||
parser.add_argument('--lr', type=float, default=3e-4, help='学习率')
|
||
parser.add_argument('--accum', type=int, default=2, help='梯度累积步数')
|
||
parser.add_argument('--data', type=str, default='data/merged_clean_v3.json', help='训练数据路径')
|
||
parser.add_argument('--resume', action='store_true', help='断点续训')
|
||
parser.add_argument('--save_every', type=int, default=5, help='每N轮保存checkpoint')
|
||
args = parser.parse_args()
|
||
|
||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||
print("=" * 60)
|
||
print(f"小狸AI CNB训练")
|
||
print(f" GPU: {torch.cuda.get_device_name(0) if device == 'cuda' else '无'}")
|
||
print(f" 设备: {device}")
|
||
print(f" 数据: {args.data}")
|
||
print(f" epochs={args.epochs}, batch={args.batch_size}, lr={args.lr}")
|
||
print("=" * 60)
|
||
|
||
# ====== 加载数据 ======
|
||
data_path = args.data if os.path.isabs(args.data) else os.path.join(MODEL_DIR, args.data)
|
||
if not os.path.exists(data_path):
|
||
print(f"❌ 数据不存在: {data_path}")
|
||
return
|
||
|
||
pairs = load_pairs(data_path)
|
||
random.seed(42)
|
||
random.shuffle(pairs)
|
||
split = int(len(pairs) * 0.95)
|
||
train_set = pairs[:split]
|
||
val_set = pairs[split:]
|
||
print(f"数据: {len(pairs)}对 | 训练: {len(train_set)} | 验证: {len(val_set)}")
|
||
|
||
# ====== 加载词表 ======
|
||
vocab_path = os.path.join(MODEL_DIR, 'vocab.json')
|
||
with open(vocab_path, 'r', encoding='utf-8') as f:
|
||
vocab = json.load(f)
|
||
# vocab.json 不含特殊token(0-6)，需合并进来，否则：
|
||
# 1) tokenizer 编码 <SOS>/<SEP>/<EOS> 等标记时会逐字符错码
|
||
# 2) vocab 的 id 范围是 7..max，超过 len(vocab) 会造成 embedding 索引越界
|
||
special_tokens = {'<PAD>': 0, '<UNK>': 1, '<SOS>': 2, '<EOS>': 3,
|
||
'<START>': 4, '<END>': 5, '<SEP>': 6}
|
||
for tok, tid in special_tokens.items():
|
||
if tok not in vocab:
|
||
vocab[tok] = tid
|
||
real_vocab_size = max(vocab.values()) + 1
|
||
print(f"词表: {len(vocab)}项, vocab_size={real_vocab_size} (字符+BPE, 来自 vocab.json)")
|
||
tokenizer_instance.char_to_id = vocab.copy()
|
||
tokenizer_instance.id_to_char = {v: k for k, v in vocab.items()}
|
||
tokenizer_instance.special_tokens = special_tokens.copy()
|
||
tokenizer_instance._sorted_vocab_cache = None
|
||
for attr in ['_multi_trie', '_multi_keys', '_single_map']:
|
||
if hasattr(tokenizer_instance, attr):
|
||
delattr(tokenizer_instance, attr)
|
||
|
||
# ====== 模型 ======
|
||
model = LLMModel(vocab_size=real_vocab_size, d_model=320, n_heads=8,
|
||
n_layers=6, d_ff=1536, max_len=256, dropout=0.1,
|
||
use_memory=False, use_thinking=False, use_moe=False)
|
||
total = sum(p.numel() for p in model.parameters())
|
||
print(f"模型: {total:,} 参数 ({total/1e6:.2f}M)")
|
||
|
||
start_epoch = 0
|
||
if args.resume:
|
||
ckpt_path, last_epoch = find_latest_ckpt(MODEL_DIR)
|
||
if ckpt_path:
|
||
ckpt = torch.load(ckpt_path, map_location=device)
|
||
old_cfg = ckpt.get('model_config', {})
|
||
if (old_cfg.get('vocab_size', 0) == real_vocab_size
|
||
and old_cfg.get('d_model', 0) == 320
|
||
and old_cfg.get('n_layers', 0) == 6):
|
||
model.load_state_dict(ckpt['model_state_dict'])
|
||
start_epoch = last_epoch
|
||
print(f"断点续训: epoch {start_epoch} ← {os.path.basename(ckpt_path)}")
|
||
else:
|
||
print(f"⚠️ 架构不匹配，从头训练")
|
||
else:
|
||
print("无断点，从头训练")
|
||
|
||
model = model.to(device)
|
||
|
||
# ====== 火力全开：torch.compile (inductor融合, 无CUDA graphs) + 梯度累积 ======
|
||
# 物理batch=192(不OOM), accum=2 → 等效batch=384
|
||
fwd_model = model
|
||
if device == 'cuda':
|
||
print("🔥 启用 torch.compile (mode=default, inductor融合)...")
|
||
fwd_model = torch.compile(model, mode='default')
|
||
print("✓ torch.compile 已就绪（首步触发编译约30-60s）")
|
||
|
||
# ====== DataLoader ======
|
||
n_workers = 16 # CPU火力全开，16核全用
|
||
train_ds = ConversationDataset(train_set, vocab, max_length=256)
|
||
train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True,
|
||
num_workers=n_workers, pin_memory=(device == 'cuda'),
|
||
persistent_workers=True, prefetch_factor=4, drop_last=True)
|
||
val_ds = ConversationDataset(val_set, vocab, max_length=256)
|
||
val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False,
|
||
num_workers=n_workers, pin_memory=(device == 'cuda'),
|
||
persistent_workers=True, prefetch_factor=4, drop_last=True)
|
||
print(f"批次: 训练{len(train_loader)} | 验证{len(val_loader)} | num_workers={n_workers} | prefetch=4 | drop_last=True")
|
||
|
||
# ====== Trainer ======
|
||
config = {
|
||
'batch_size': args.batch_size,
|
||
'gradient_accumulation_steps': args.accum,
|
||
'epochs': args.epochs,
|
||
'd_model': 320, 'n_heads': 8, 'n_layers': 6,
|
||
'd_ff': 1536, 'max_length': 256, 'dropout': 0.1,
|
||
'use_amp': device == 'cuda',
|
||
}
|
||
trainer = LLMTrainer(fwd_model, device, lr=args.lr, use_amp=config['use_amp'])
|
||
|
||
# ====== 训练 ======
|
||
print(f"\n开始训练 epoch {start_epoch+1} → {args.epochs}")
|
||
print("=" * 60)
|
||
t_start = time.time()
|
||
best_val = float('inf')
|
||
|
||
for epoch in range(start_epoch + 1, args.epochs + 1):
|
||
t0 = time.time()
|
||
train_loss = trainer.train_epoch(train_loader, epoch, vocab, config)
|
||
val_loss = trainer.evaluate(val_loader, config)
|
||
dt = time.time() - t0
|
||
|
||
marker = " ★" if val_loss < best_val else ""
|
||
if val_loss < best_val:
|
||
best_val = val_loss
|
||
trainer.save_model(os.path.join(MODEL_DIR, 'xiaoli_cnb_best.pth'), vocab, epoch)
|
||
|
||
elapsed = time.time() - t_start
|
||
eta = (elapsed / (epoch - start_epoch)) * (args.epochs - epoch) if epoch > start_epoch else 0
|
||
|
||
print(f"E{epoch}/{args.epochs} | train={train_loss:.4f} val={val_loss:.4f}{marker} | "
|
||
f"{dt:.0f}s/epoch | {elapsed/60:.1f}min已用 ETA{eta/60:.0f}min")
|
||
|
||
# 保存
|
||
if epoch % args.save_every == 0:
|
||
ckpt_path = os.path.join(MODEL_DIR, f'llm_model_v2_resume_epoch_{epoch}.pth')
|
||
trainer.save_model(ckpt_path, vocab, epoch)
|
||
# 清理旧resume
|
||
for f in glob.glob(os.path.join(MODEL_DIR, 'llm_model_v2_resume_*.pth')):
|
||
en = max([int(s) for s in re.findall(r'\d+', os.path.basename(f).replace('.pth',''))])
|
||
if en <= epoch - 15:
|
||
try: os.remove(f)
|
||
except: pass
|
||
print(f" 💾 已保存 epoch {epoch}")
|
||
|
||
# ====== 完成 ======
|
||
final_path = os.path.join(MODEL_DIR, 'xiaoli_cnb_final.pth')
|
||
trainer.save_model(final_path, vocab, args.epochs)
|
||
total_time = (time.time() - t_start) / 60
|
||
print(f"\n✅ 训练完成! 耗时 {total_time:.1f}min")
|
||
|
||
# 测试
|
||
print("\n测试:")
|
||
chat = LLMChat(model, vocab, device, max_length=256)
|
||
for q in ["你好", "你是谁", "你叫什么名字", "今天天气真好", "我有点难过"]:
|
||
r = chat.generate_response(q, max_length=60, temperature=0.3)
|
||
print(f" Q: {q}")
|
||
print(f" A: {r}\n")
|
||
|
||
if __name__ == "__main__":
|
||
main() |