[自动] 云端GPU训练脚本
This commit is contained in:
shuiyu1123 authored and Gitee committed
1 parent
0ccce55ab0
commit
c369804d3c
1 file changed
+437
+437
@@ -0,0 +1,437 @@
|
||||
#!/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()
|
||||
Reference in new issue
Block a user