From c369804d3c4a20695d064a7966219c4314d362c1 Mon Sep 17 00:00:00 2001 From: shuiyu1123 Date: Sun, 21 Jun 2026 01:58:19 +0000 Subject: [PATCH] =?UTF-8?q?[=E8=87=AA=E5=8A=A8]=20=E4=BA=91=E7=AB=AFGPU?= =?UTF-8?q?=E8=AE=AD=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_cloud.py | 437 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 437 insertions(+) create mode 100644 train_cloud.py diff --git a/train_cloud.py b/train_cloud.py new file mode 100644 index 0000000..81591cb --- /dev/null +++ b/train_cloud.py @@ -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 = {'': 0, '': 1, '': 2, '': 3, + '': 4, '': 5, '': 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('', 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() \ No newline at end of file