437 lines
18 KiB
Python
437 lines
18 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
XiaLi 云端 GPU 训练脚本 v1.0
|
|
=============================
|
|
两阶段训练 + ModelScope 自动下载数据 + 断点续训
|
|
|
|
用法（云GPU上）:
|
|
python3 train_cloud.py # 默认训练
|
|
python3 train_cloud.py --resume # 从最新检查点续训
|
|
python3 train_cloud.py --epochs 10 # 指定预训练轮数
|
|
"""
|
|
import sys, os, json, time, glob, re, math, random
|
|
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
|
|
)
|
|
|
|
# ── 配置 ──────────────────────────────────────────────────
|
|
class CloudConfig:
|
|
# 模型架构
|
|
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 # 关掉记忆/思考/MoE 加速
|
|
use_thinking: bool = False
|
|
use_moe: bool = False
|
|
|
|
# 预训练
|
|
pt_batch_size: int = 64 # 32G显存可跑64, 24G改32
|
|
pt_grad_accum: int = 1 # 大batch不需要累积
|
|
pt_epochs: int = 5
|
|
pt_lr: float = 5e-4
|
|
pt_min_lr: float = 5e-5
|
|
pt_warmup_steps: int = 200
|
|
pt_max_length: int = 256 # 长序列学更好
|
|
pt_weight_decay: float = 0.1
|
|
pt_max_norm: float = 1.0
|
|
|
|
# 对齐
|
|
align_batch_size: int = 32
|
|
align_grad_accum: int = 1
|
|
align_epochs: int = 20
|
|
align_lr: float = 1e-4
|
|
align_min_lr: float = 1e-5
|
|
align_warmup_steps: int = 50
|
|
align_weight_decay: float = 0.01
|
|
align_max_norm: float = 1.0
|
|
|
|
# 通用
|
|
device: str = "cuda" if torch.cuda.is_available() else "cpu"
|
|
num_workers: int = 4
|
|
use_amp: bool = True
|
|
seed: int = 42
|
|
save_dir: str = "/root/xiaoli-llm-chat" # AutoDL 默认路径
|
|
log_every: int = 10
|
|
save_interval: int = 1
|
|
keep_checkpoints: int = 3
|
|
|
|
|
|
# ── 词表加载 ──────────────────────────────────────────────
|
|
def load_vocab(vocab_path: str) -> dict:
|
|
special = {'<PAD>': 0, '<UNK>': 1, '<SOS>': 2, '<EOS>': 3,
|
|
'<START>': 4, '<END>': 5, '<SEP>': 6}
|
|
with open(vocab_path, 'r') as f:
|
|
normal_vocab = json.load(f)
|
|
full_vocab = {**special, **normal_vocab}
|
|
# 同步到全局 tokenizer
|
|
tokenizer_instance.char_to_id = full_vocab.copy()
|
|
tokenizer_instance.id_to_char = {v: k for k, v in full_vocab.items()}
|
|
for attr in ['_multi_trie', '_multi_keys', '_single_map']:
|
|
if hasattr(tokenizer_instance, attr): delattr(tokenizer_instance, attr)
|
|
return full_vocab
|
|
|
|
|
|
# ── 预训练数据集（流式加载，不爆内存）────────────────────
|
|
class PretrainDataset(Dataset):
|
|
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)
|
|
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}")
|
|
total_chars = sum(len(t) for t in self.texts)
|
|
print(f" 预训练数据: {len(self.texts):,} 条, {total_chars:,} 字符")
|
|
|
|
def _load_file(self, path: Path):
|
|
try:
|
|
if path.suffix == '.jsonl':
|
|
with open(path, 'r', encoding='utf-8', errors='ignore') as f:
|
|
for line in f:
|
|
try:
|
|
item = json.loads(line)
|
|
t = item.get('text', '').strip()
|
|
if len(t) >= 10: self.texts.append(t)
|
|
except: continue
|
|
except: pass
|
|
|
|
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_ids = torch.tensor(tokens[:-1], dtype=torch.long)
|
|
target_ids = torch.tensor(tokens[1:], dtype=torch.long)
|
|
loss_mask = torch.ones_like(target_ids)
|
|
return {'input_ids': input_ids, 'target_ids': target_ids,
|
|
'loss_mask': loss_mask, 'weights': 1.0}
|
|
|
|
|
|
# ── 训练器 ──────────────────────────────────────────────
|
|
class CloudTrainer:
|
|
def __init__(self, model: nn.Module, config: CloudConfig, 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 用半速
|
|
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, betas=(0.9, 0.95), eps=1e-8)
|
|
self.scaler = torch.cuda.amp.GradScaler(enabled=True)
|
|
self.scheduler = None
|
|
|
|
if torch.cuda.is_available():
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
torch.backends.cudnn.allow_tf32 = True
|
|
torch.backends.cudnn.benchmark = True
|
|
|
|
def build_scheduler(self, total_steps: int):
|
|
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
|
|
base_lr = self.optimizer.param_groups[0]['lr']
|
|
|
|
def lr_lambda(step):
|
|
if step < warmup:
|
|
return step / max(1, warmup)
|
|
progress = (step - warmup) / max(1, total_steps - warmup)
|
|
cosine = 0.5 * (1 + math.cos(math.pi * progress))
|
|
return (1 - min_lr / base_lr) * cosine + min_lr / base_lr
|
|
|
|
self.scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda)
|
|
|
|
def train_epoch(self, dataloader, epoch, total_epochs, stage_name="pretrain"):
|
|
self.model.train()
|
|
total_loss = 0
|
|
n_batches = len(dataloader)
|
|
accum_steps = self.config.pt_grad_accum if stage_name == "pretrain" else self.config.align_grad_accum
|
|
max_norm = self.config.pt_max_norm if stage_name == "pretrain" else self.config.align_max_norm
|
|
|
|
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, non_blocking=True)
|
|
|
|
with torch.amp.autocast('cuda', enabled=True):
|
|
outputs = self.model(input_ids, mask)
|
|
loss = self.criterion(outputs.reshape(-1, outputs.size(-1)), target_ids.reshape(-1))
|
|
loss = loss / accum_steps
|
|
|
|
self.scaler.scale(loss).backward()
|
|
|
|
if (batch_idx + 1) % accum_steps == 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: self.scheduler.step()
|
|
|
|
total_loss += loss.item() * accum_steps
|
|
|
|
if (batch_idx + 1) % max(1, n_batches // 10) == 0:
|
|
lr = self.optimizer.param_groups[0]['lr']
|
|
mem = f" | Mem: {torch.cuda.max_memory_allocated()/1024**3:.1f}GiB" if torch.cuda.is_available() else ""
|
|
print(f" [{stage_name}] Epoch {epoch}/{total_epochs} | "
|
|
f"Batch {batch_idx+1}/{n_batches} | Loss: {total_loss/(batch_idx+1):.4f} | LR: {lr:.2e}{mem}")
|
|
|
|
return total_loss / n_batches
|
|
|
|
def save_checkpoint(self, path: str, vocab: dict, epoch: int, loss: float):
|
|
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,
|
|
'epoch': epoch,
|
|
'loss': loss,
|
|
'config': {k: getattr(self.config, k) for k in dir(self.config)
|
|
if not k.startswith('_') and not callable(getattr(self.config, k))},
|
|
}, path)
|
|
print(f" ✅ 已保存: {path} (epoch {epoch}, loss {loss:.4f})")
|
|
|
|
def load_checkpoint(self, path: str) -> tuple:
|
|
cp = torch.load(path, map_location=self.device, weights_only=False)
|
|
core = getattr(self.model, '_orig_mod', self.model)
|
|
core.load_state_dict(cp['model_state_dict'])
|
|
self.optimizer.load_state_dict(cp['optimizer_state_dict'])
|
|
if cp.get('scheduler_state_dict') and self.scheduler:
|
|
self.scheduler.load_state_dict(cp['scheduler_state_dict'])
|
|
return cp.get('epoch', 0), cp.get('loss', float('inf')), cp.get('vocab', {})
|
|
|
|
|
|
# ── 数据下载 ──────────────────────────────────────────────
|
|
def download_data(save_dir: str):
|
|
"""从 ModelScope 下载预训练数据"""
|
|
data_dir = os.path.join(save_dir, 'pretrain_data')
|
|
if os.path.exists(data_dir) and list(Path(data_dir).glob("*.jsonl")):
|
|
print(" 预训练数据已存在，跳过下载")
|
|
return data_dir
|
|
|
|
os.makedirs(data_dir, exist_ok=True)
|
|
print(" 📥 从 ModelScope 下载预训练数据...")
|
|
try:
|
|
import subprocess
|
|
subprocess.run([
|
|
sys.executable, '-m', 'pip', 'install', 'modelscope', '-q'
|
|
], check=True, capture_output=True)
|
|
subprocess.run([
|
|
sys.executable, '-m', 'modelscope', 'download',
|
|
'--dataset', 'gongjy/minimind_dataset',
|
|
'--include', 'pretrain_t2t_mini*',
|
|
'--local_dir', data_dir,
|
|
], check=True)
|
|
print(f" ✅ 下载完成: {data_dir}")
|
|
except Exception as e:
|
|
print(f" ⚠️ 自动下载失败: {e}")
|
|
print(f" ⚠️ 请手动下载放到 {data_dir}")
|
|
print(f" ⚠️ 命令: modelscope download --dataset gongjy/minimind_dataset --include 'pretrain_t2t_mini*' --local_dir {data_dir}")
|
|
return data_dir
|
|
|
|
|
|
# ── 主函数 ──────────────────────────────────────────────
|
|
def main():
|
|
import argparse
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--resume", action="store_true", help="从最新检查点续训")
|
|
parser.add_argument("--epochs", type=int, default=5, help="预训练轮数")
|
|
parser.add_argument("--align-epochs", type=int, default=20, help="对齐轮数")
|
|
parser.add_argument("--save-dir", type=str, default="/root/xiaoli-llm-chat")
|
|
parser.add_argument("--batch-size", type=int, default=0, help="覆盖batch大小")
|
|
args = parser.parse_args()
|
|
|
|
config = CloudConfig()
|
|
config.pt_epochs = args.epochs
|
|
config.align_epochs = args.align_epochs
|
|
config.save_dir = args.save_dir
|
|
|
|
if args.batch_size > 0:
|
|
config.pt_batch_size = args.batch_size
|
|
config.align_batch_size = args.batch_size // 2
|
|
|
|
if not torch.cuda.is_available():
|
|
print("⚠️ 未检测到 GPU！云端训练建议使用 AutoDL / 矩池云 / 恒源云")
|
|
|
|
print("=" * 65)
|
|
print(f" XiaLi 云端训练")
|
|
print(f" GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'N/A'}")
|
|
print(f" 显存: {torch.cuda.get_device_properties(0).total_mem / 1024**3:.1f} GiB" if torch.cuda.is_available() else "")
|
|
print(f" 模型: {config.d_model}d, {config.n_layers}L, {config.n_heads}H")
|
|
print(f" 预训练: {config.pt_epochs} epochs, batch={config.pt_batch_size}")
|
|
print(f" 对齐: {config.align_epochs} epochs, batch={config.align_batch_size}")
|
|
print("=" * 65)
|
|
|
|
os.makedirs(config.save_dir, exist_ok=True)
|
|
|
|
# 1. 加载词表
|
|
vocab_path = os.path.join(config.save_dir, 'vocab.json')
|
|
if not os.path.exists(vocab_path):
|
|
print(f"❌ 未找到 {vocab_path}")
|
|
print(f" 请从 Gitee 仓库获取: https://gitee.com/shuiyu1123/xiaoli-llm-chat")
|
|
return
|
|
vocab = load_vocab(vocab_path)
|
|
print(f"词表: {len(vocab)} tokens, vocab_size={max(vocab.values())+1}")
|
|
|
|
# 2. 创建模型
|
|
model = LLMModel(
|
|
vocab_size=max(vocab.values())+1,
|
|
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,
|
|
)
|
|
params = sum(p.numel() for p in model.parameters())
|
|
print(f"参数量: {params:,} ({params/1e6:.2f}M)")
|
|
|
|
# 3. 续训检测
|
|
start_epoch = 1
|
|
if args.resume:
|
|
checkpoints = sorted(glob.glob(os.path.join(config.save_dir, 'pretrain_epoch_*.pth')))
|
|
if checkpoints:
|
|
latest = checkpoints[-1]
|
|
trainer_pt = CloudTrainer(model, config, stage="pretrain")
|
|
loaded_epoch, loaded_loss, _ = trainer_pt.load_checkpoint(latest)
|
|
start_epoch = loaded_epoch + 1
|
|
print(f" 从 {latest} 续训 (epoch {loaded_epoch}, loss {loaded_loss:.4f})")
|
|
|
|
# 4. 下载数据
|
|
pt_dir = download_data(config.save_dir)
|
|
|
|
# 5. 预训练
|
|
print(f"\n{'='*65}")
|
|
print(f" 阶段一：预训练 ({config.pt_epochs} epochs)")
|
|
print(f"{'='*65}")
|
|
|
|
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=True,
|
|
)
|
|
|
|
trainer = CloudTrainer(model, config, stage="pretrain")
|
|
total_steps = len(pt_loader) * config.pt_epochs
|
|
trainer.build_scheduler(total_steps)
|
|
|
|
for epoch in range(start_epoch, config.pt_epochs + 1):
|
|
t0 = time.time()
|
|
loss = trainer.train_epoch(pt_loader, epoch, config.pt_epochs, "pretrain")
|
|
epoch_time = time.time() - t0
|
|
lr = trainer.optimizer.param_groups[0]['lr']
|
|
print(f" → Epoch {epoch}/{config.pt_epochs} | Loss: {loss:.4f} | LR: {lr:.2e} | 耗时: {epoch_time:.0f}s")
|
|
|
|
if epoch % config.save_interval == 0:
|
|
ckpt_path = os.path.join(config.save_dir, f'pretrain_epoch_{epoch}.pth')
|
|
trainer.save_checkpoint(ckpt_path, vocab, epoch, loss)
|
|
# 清理旧检查点
|
|
all_ckpts = sorted(glob.glob(os.path.join(config.save_dir, 'pretrain_epoch_*.pth')))
|
|
for old in all_ckpts[:-config.keep_checkpoints]:
|
|
os.remove(old)
|
|
|
|
print(f"\n ✅ 预训练完成")
|
|
|
|
# 6. 对齐
|
|
print(f"\n{'='*65}")
|
|
print(f" 阶段二：角色对齐 ({config.align_epochs} epochs)")
|
|
print(f"{'='*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}")
|
|
print(f" ⚠️ 跳过对齐阶段")
|
|
else:
|
|
with open(merged_path, 'r') as f:
|
|
raw_data = json.load(f)
|
|
pairs = []
|
|
for item in raw_data:
|
|
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.95)
|
|
train_pairs = pairs[:split]
|
|
print(f" 对齐数据: {len(pairs)} 条 (训练 {len(train_pairs)} + 验证 {len(pairs)-split})")
|
|
|
|
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=True,
|
|
)
|
|
|
|
trainer2 = CloudTrainer(model, config, stage="align")
|
|
total_align_steps = len(train_loader) * config.align_epochs
|
|
trainer2.build_scheduler(total_align_steps)
|
|
|
|
for epoch in range(1, config.align_epochs + 1):
|
|
t0 = time.time()
|
|
loss = trainer2.train_epoch(train_loader, epoch, config.align_epochs, "align")
|
|
epoch_time = time.time() - t0
|
|
print(f" → Epoch {epoch}/{config.align_epochs} | Loss: {loss:.4f} | 耗时: {epoch_time:.0f}s")
|
|
|
|
if epoch % 5 == 0:
|
|
ckpt_path = os.path.join(config.save_dir, f'xiaoli_align_epoch_{epoch}.pth')
|
|
trainer2.save_checkpoint(ckpt_path, vocab, epoch, loss)
|
|
|
|
# 保存最终模型
|
|
final_path = os.path.join(config.save_dir, 'xiaoli_final.pth')
|
|
trainer2.save_checkpoint(final_path, vocab, config.align_epochs, loss)
|
|
print(f"\n ✅ 对齐完成！最终模型: {final_path}")
|
|
|
|
print(f"\n{'='*65}")
|
|
print(f" 🎉 训练完成！")
|
|
print(f" 模型保存在: {config.save_dir}")
|
|
print(f" {'='*65}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |