当前位置：首页 > news >正文

32、语言模型训练全流程：从数据到模型的工程化实现

news 2025/9/19 5:27:42

第32课：语言模型训练全流程 - 从数据到模型的工程化实现

学习目标：掌握语言模型训练的完整工程流程，理解大规模训练的技术细节，建立稳定高效的训练系统，实现从数据到模型的端到端能力。

语言模型训练是一个复杂的工程系统，需要协调数据处理、训练控制、监控评估等多个环节。本课程将带你构建一个完整的训练流水线，从原始文本到可用模型。

32.1 数据工程化处理：构建高效数据流水线

> 数据流水线详细解析

原始文本数据是整个流程的起点，通常来自网络爬取、图书语料、代码仓库等多源异构数据。这些数据的特点是体量庞大（通常TB级别）、质量参差不齐、格式多样化。

文本预处理包含三个关键子步骤：

数据清洗：去除HTML标签、控制字符、过短或过长文本，检测和过滤垃圾内容。关键参数包括最小文本长度（通常10-50字符）、最大长度阈值、重复内容检测窗口大小
长度分析：统计文本长度分布，确定合适的截断策略。需要分析P50、P90、P95长度分位数，平衡信息保留和计算效率
质量控制：基于困惑度、语言检测、内容分类等维度过滤低质量文本。设置困惑度阈值（如小于1000）、语言置信度（大于0.9）等参数

分词处理将文本转换为模型可理解的token序列，核心参数包括：

词汇表大小：影响模型参数量和表达能力，通常在30K-100K之间
特殊token设置：包括<pad>, <unk>, <bos>, <eos>等
子词算法选择：BPE、SentencePiece等，影响对未见词汇的处理能力

序列打包通过多序列组合提升GPU利用率：

多序列打包：将多个短序列拼接成固定长度，减少padding浪费。关键是选择拼接策略（贪心、最优匹配等）
注意力掩码生成：确保不同序列间不相互影响。需要生成块对角掩码矩阵，时间复杂度O(n²)
内存优化：使用稀疏矩阵存储掩码，减少30-50%内存占用

动态批处理根据序列长度动态调整批大小：

批大小调整：根据GPU内存和序列长度预估最优批大小。公式：batch_size = GPU_memory / (seq_len * model_size * 4)
负载均衡：确保各GPU工作负载相等，避免木桶效应
GPU利用率优化：目标是达到95%以上的GPU利用率

数据加载器是连接存储和计算的桥梁：

内存映射：使用mmap技术处理超大文件，避免全量加载到内存
懒加载机制：按需加载数据，减少内存峰值
分布式采样：在多GPU环境下确保数据不重复不遗漏

32.2 训练循环设计：稳定可控的训练框架

> 训练循环核心组件解析

训练控制器是整个训练流程的大脑，负责协调所有组件：

控制训练节奏：设定epoch数量、步数限制、时间预算等
异常处理：捕获GPU内存溢出、数据损坏、网络中断等异常
状态管理：维护当前epoch、step、学习率等训练状态

数据加载器交互涉及复杂的缓冲和预取策略：

预取队列深度：通常设为2-4个批次，平衡内存占用和等待时间
多进程加载：使用4-8个worker进程并行处理数据
异常重试机制：数据损坏时的重试次数和策略

前向传播控制包含多个性能优化点：

混合精度训练：使用FP16前向传播，FP32累积梯度
梯度检查点：在内存不足时重新计算中间激活值
动态图优化：减少不必要的计算图构建开销

反向传播管理确保梯度计算的正确性：

梯度累积：在内存受限时分多个小批次累积梯度
梯度同步：多GPU环境下的All-Reduce通信
梯度裁剪：防止梯度爆炸，通常设定阈值为1.0-5.0

检查点机制是训练可靠性的保障：

保存频率：平衡磁盘IO和恢复成本，通常每1000-5000步保存一次
状态完整性：包含模型权重、优化器状态、学习率调度器状态、随机数种子
增量保存：只保存变化的参数，减少磁盘占用

32.3 损失计算与优化策略

> 损失计算和优化详细机制

交叉熵损失是语言模型的标准损失函数：

计算公式：loss = -log(softmax(logits)[target])
数值稳定性：使用log_softmax避免exp溢出，关键参数temperature控制概率分布尖锐程度
忽略索引：设置ignore_index=-1跳过padding token的损失计算

标签平滑减少过拟合和增强泛化能力：

平滑参数：通常设为0.1，将目标概率从1.0调整为0.9
计算方式：smooth_target = (1-α) * target + α / vocab_size
作用机制：防止模型对预测过于自信，提高输出概率分布的熵

损失掩码确保只计算有效位置的损失：

掩码生成：标识哪些位置需要计算损失，padding位置设为0
损失聚合：使用掩码进行加权平均，避免padding影响结果
序列级损失：可以根据序列长度进行归一化

AdamW优化器结合了Adam和权重衰减：

权重衰减：直接在参数上应用L2正则化，系数通常为0.01-0.1
动量参数β1：一阶矩估计的衰减率，通常为0.9，控制梯度的历史记忆
二阶矩估计β2：二阶矩估计的衰减率，通常为0.999，控制梯度方差的估计

学习率调度控制优化过程的收敛速度：

预热阶段：从0线性增长到目标学习率，通常持续1000-5000步，避免初期梯度过大
余弦退火：按余弦函数曲线衰减学习率，公式：lr = min_lr + (max_lr - min_lr) * (1 + cos(π * step / max_steps)) / 2
学习率衰减：指数或多项式衰减，在特定里程碑处降低学习率

32.4 训练监控与可视化

> 监控系统各组件功能详解

实时指标收集是监控系统的数据基础：

损失曲线：记录训练损失、验证损失的变化趋势，采样频率通常为每步或每10步。关键参数包括移动平均窗口大小（通常100-1000步）
学习率变化：跟踪学习率调度器的执行情况，确保按预期变化
梯度统计：包括梯度范数、梯度分布、各层梯度大小，用于诊断训练问题
GPU使用率：监控计算资源利用率，目标是保持90%以上的利用率
内存占用：跟踪GPU显存和系统内存使用情况，防止内存泄漏

异常检测机制自动识别训练过程中的问题：

梯度异常检测：检测梯度爆炸（范数>10）或消失（范数<1e-7）
损失突变检测：识别损失突然上升或下降，可能表示学习率过大或数据异常
硬件故障检测：监控GPU温度、功耗、错误率等硬件指标

可视化展示帮助研究者直观理解训练过程：

TensorBoard集成：标准的深度学习可视化工具，支持标量、直方图、图像等多种图表
Web仪表盘：实时显示关键指标，支持多实验对比
实时图表：损失曲线、学习率变化等关键指标的实时更新

32.5 模型验证与评估

> 模型评估的多维度指标体系

困惑度计算是语言模型最直接的评估指标：

验证集困惑度：在验证集上计算平均困惑度，公式为perplexity = exp(loss)，通常每1000步计算一次
困惑度趋势分析：跟踪困惑度的下降趋势，识别过拟合信号（验证困惑度开始上升）
分层困惑度分析：分别计算不同文本类型、长度区间的困惑度，识别模型弱点

生成质量评估从人类使用角度评估模型：

文本生成采样：使用不同采样策略（greedy、top-k、top-p）生成文本样本
生成质量打分：基于流畅性、连贯性、创意性等维度的主观评分
多样性评估：计算生成文本的词汇多样性、句式多样性，避免重复生成

收敛性判断决定何时停止训练：

损失平稳性检查：检查最近N步的损失变化是否小于阈值（如0.001）
早停机制：当验证损失连续N个评估周期没有改善时停止训练
最优模型选择：基于验证集性能选择最佳checkpoint作为最终模型

32.6 完整训练系统实现

在这里插入图片描述

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import numpy as np
import json
import os
import time
import logging
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Iterator
from dataclasses import dataclass
import matplotlib.pyplot as plt
from collections import defaultdict# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)@dataclass
class TrainingConfig:"""训练配置类，包含所有超参数和设置"""# 模型参数vocab_size: int = 10000          # 词汇表大小max_seq_len: int = 512           # 最大序列长度n_layers: int = 6                # transformer层数n_heads: int = 8                 # 注意力头数n_embd: int = 512                # 嵌入维度dropout: float = 0.1             # dropout概率# 训练参数batch_size: int = 32             # 批大小learning_rate: float = 1e-4      # 初始学习率weight_decay: float = 0.01       # 权重衰减max_epochs: int = 10             # 最大训练轮数warmup_steps: int = 1000         # 预热步数eval_interval: int = 500         # 评估间隔save_interval: int = 2000        # 保存间隔# 优化器参数beta1: float = 0.9               # Adam beta1参数beta2: float = 0.999             # Adam beta2参数grad_clip: float = 1.0           # 梯度裁剪阈值label_smoothing: float = 0.1     # 标签平滑参数# 系统参数device: str = 'cuda' if torch.cuda.is_available() else 'cpu'num_workers: int = 4             # 数据加载器工作进程数mixed_precision: bool = True     # 是否使用混合精度训练compile_model: bool = False      # 是否使用torch.compile# 路径参数data_path: str = './data'        # 数据路径output_dir: str = './output'     # 输出目录checkpoint_dir: str = './checkpoints'  # 检查点目录log_dir: str = './logs'          # 日志目录class SimpleTextDataset(Dataset):"""简单的文本数据集类，支持动态padding和序列打包"""def __init__(self, texts: List[str], tokenizer, max_length: int = 512):"""初始化数据集Args:texts: 文本列表tokenizer: 分词器（这里简化为字符级）max_length: 最大序列长度"""self.texts = textsself.max_length = max_length# 创建字符级词汇表（简化实现）all_chars = set(''.join(texts))self.vocab = {char: i+4 for i, char in enumerate(sorted(all_chars))}self.vocab.update({'<pad>': 0, '<unk>': 1, '<bos>': 2, '<eos>': 3})self.vocab_size = len(self.vocab)self.idx_to_char = {v: k for k, v in self.vocab.items()}# 预处理文本数据self.encoded_texts = []for text in texts:# 添加开始和结束标记encoded = [self.vocab['<bos>']]for char in text[:max_length-2]:  # 为bos和eos预留空间encoded.append(self.vocab.get(char, self.vocab['<unk>']))encoded.append(self.vocab['<eos>'])self.encoded_texts.append(encoded)logger.info(f"数据集初始化完成，包含{len(self.encoded_texts)}个样本，词汇表大小{self.vocab_size}")def __len__(self):return len(self.encoded_texts)def __getitem__(self, idx):"""获取单个样本，返回输入序列和目标序列对于语言模型，目标序列是输入序列右移一位"""encoded = self.encoded_texts[idx]# 创建输入序列（去除最后一个token）和目标序列（去除第一个token）if len(encoded) > 1:input_ids = encoded[:-1]target_ids = encoded[1:]else:input_ids = encodedtarget_ids = encodedreturn {'input_ids': torch.tensor(input_ids, dtype=torch.long),'target_ids': torch.tensor(target_ids, dtype=torch.long),'length': len(input_ids)}def collate_fn(batch):"""自定义的批处理函数，支持动态padding将不同长度的序列打包成同一批次，使用padding对齐"""# 获取批次中的最大长度max_len = max([item['length'] for item in batch])input_ids = []target_ids = []attention_mask = []for item in batch:seq_len = item['length']pad_len = max_len - seq_len# 对输入序列进行paddingpadded_input = F.pad(item['input_ids'], (0, pad_len), value=0)padded_target = F.pad(item['target_ids'], (0, pad_len), value=-100)  # -100会被忽略# 创建注意力掩码（1表示有效token，0表示padding）mask = torch.cat([torch.ones(seq_len), torch.zeros(pad_len)])input_ids.append(padded_input)target_ids.append(padded_target)attention_mask.append(mask)return {'input_ids': torch.stack(input_ids),'target_ids': torch.stack(target_ids),'attention_mask': torch.stack(attention_mask)}class SimpleTransformer(nn.Module):"""简化的Transformer模型，专注于训练流程演示"""def __init__(self, config: TrainingConfig):super().__init__()self.config = config# Token嵌入层self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)# 位置嵌入层self.position_embedding = nn.Embedding(config.max_seq_len, config.n_embd)# Transformer层self.layers = nn.ModuleList([nn.TransformerDecoderLayer(d_model=config.n_embd,nhead=config.n_heads,dim_feedforward=4 * config.n_embd,dropout=config.dropout,batch_first=True) for _ in range(config.n_layers)])# 层归一化self.layer_norm = nn.LayerNorm(config.n_embd)# 输出投影层self.output_projection = nn.Linear(config.n_embd, config.vocab_size, bias=False)# 初始化权重self.apply(self._init_weights)# 计算参数量total_params = sum(p.numel() for p in self.parameters())logger.info(f"模型初始化完成，总参数量: {total_params:,}")def _init_weights(self, module):"""权重初始化"""if isinstance(module, nn.Linear):nn.init.normal_(module.weight, mean=0.0, std=0.02)if module.bias is not None:nn.init.zeros_(module.bias)elif isinstance(module, nn.Embedding):nn.init.normal_(module.weight, mean=0.0, std=0.02)def forward(self, input_ids, attention_mask=None):"""前向传播Args:input_ids: 输入token序列 [batch_size, seq_len]attention_mask: 注意力掩码 [batch_size, seq_len]Returns:logits: 输出logits [batch_size, seq_len, vocab_size]"""batch_size, seq_len = input_ids.shapedevice = input_ids.device# Token嵌入token_emb = self.token_embedding(input_ids)# 位置嵌入positions = torch.arange(seq_len, device=device).unsqueeze(0).expand(batch_size, -1)pos_emb = self.position_embedding(positions)# 组合嵌入x = token_emb + pos_emb# 创建因果掩码（下三角矩阵）causal_mask = torch.triu(torch.full((seq_len, seq_len), float('-inf'), device=device), diagonal=1)# 通过Transformer层for layer in self.layers:x = layer(x, x, tgt_mask=causal_mask, tgt_key_padding_mask=~attention_mask.bool() if attention_mask is not None else None)# 层归一化x = self.layer_norm(x)# 输出投影logits = self.output_projection(x)return logitsclass TrainingMetrics:"""训练指标收集和统计类"""def __init__(self):self.reset()def reset(self):"""重置所有指标"""self.losses = []self.perplexities = []self.learning_rates = []self.grad_norms = []self.step_times = []self.step = 0def update(self, loss: float, lr: float, grad_norm: float, step_time: float):"""更新指标"""self.losses.append(loss)self.perplexities.append(np.exp(loss))  # 困惑度是损失的指数self.learning_rates.append(lr)self.grad_norms.append(grad_norm)self.step_times.append(step_time)self.step += 1def get_recent_stats(self, window: int = 100) -> Dict:"""获取最近N步的统计信息"""if not self.losses:return {}recent_losses = self.losses[-window:]recent_times = self.step_times[-window:]return {'avg_loss': np.mean(recent_losses),'avg_perplexity': np.mean(self.perplexities[-window:]),'current_lr': self.learning_rates[-1] if self.learning_rates else 0,'avg_grad_norm': np.mean(self.grad_norms[-window:]),'avg_step_time': np.mean(recent_times),'steps_per_sec': 1.0 / np.mean(recent_times) if recent_times else 0}def plot_metrics(self, save_path: str = None):"""绘制训练指标曲线"""if not self.losses:returnfig, axes = plt.subplots(2, 2, figsize=(15, 10))fig.suptitle('Training Metrics', fontsize=16)# 损失曲线axes[0, 0].plot(self.losses)axes[0, 0].set_title('Training Loss')axes[0, 0].set_xlabel('Step')axes[0, 0].set_ylabel('Loss')axes[0, 0].grid(True)# 困惑度曲线axes[0, 1].plot(self.perplexities)axes[0, 1].set_title('Perplexity')axes[0, 1].set_xlabel('Step')axes[0, 1].set_ylabel('Perplexity')axes[0, 1].grid(True)# 学习率变化axes[1, 0].plot(self.learning_rates)axes[1, 0].set_title('Learning Rate')axes[1, 0].set_xlabel('Step')axes[1, 0].set_ylabel('Learning Rate')axes[1, 0].grid(True)# 梯度范数axes[1, 1].plot(self.grad_norms)axes[1, 1].set_title('Gradient Norm')axes[1, 1].set_xlabel('Step')axes[1, 1].set_ylabel('Grad Norm')axes[1, 1].grid(True)plt.tight_layout()if save_path:plt.savefig(save_path, dpi=150, bbox_inches='tight')logger.info(f"训练指标图表已保存到: {save_path}")plt.show()class LanguageModelTrainer:"""语言模型训练器类，封装完整的训练流程"""def __init__(self, config: TrainingConfig):self.config = configself.metrics = TrainingMetrics()# 创建输出目录os.makedirs(config.output_dir, exist_ok=True)os.makedirs(config.checkpoint_dir, exist_ok=True)os.makedirs(config.log_dir, exist_ok=True)# 设置设备self.device = torch.device(config.device)logger.info(f"使用设备: {self.device}")# 初始化混合精度训练self.scaler = torch.cuda.amp.GradScaler() if config.mixed_precision and self.device.type == 'cuda' else None# 训练状态self.global_step = 0self.epoch = 0self.best_loss = float('inf')def setup_model_and_optimizer(self, dataset):"""设置模型和优化器"""# 更新配置中的词汇表大小self.config.vocab_size = dataset.vocab_size# 创建模型self.model = SimpleTransformer(self.config).to(self.device)# 编译模型（如果支持）if self.config.compile_model and hasattr(torch, 'compile'):logger.info("编译模型以提升性能...")self.model = torch.compile(self.model)# 创建优化器self.optimizer = torch.optim.AdamW(self.model.parameters(),lr=self.config.learning_rate,betas=(self.config.beta1, self.config.beta2),weight_decay=self.config.weight_decay)# 学习率调度器（带预热的余弦退火）def lr_lambda(step):if step < self.config.warmup_steps:return step / self.config.warmup_stepselse:progress = (step - self.config.warmup_steps) / max(1, self.total_steps - self.config.warmup_steps)return 0.5 * (1 + np.cos(np.pi * progress))self.scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda)logger.info("模型和优化器设置完成")def compute_loss(self, logits, targets, attention_mask=None):"""计算交叉熵损失，支持标签平滑和掩码Args:logits: 模型输出 [batch_size, seq_len, vocab_size]targets: 目标token [batch_size, seq_len]attention_mask: 注意力掩码 [batch_size, seq_len]Returns:loss: 标量损失值"""# 展平张量以便计算损失flat_logits = logits.view(-1, logits.size(-1))flat_targets = targets.view(-1)# 计算交叉熵损失if self.config.label_smoothing > 0:# 标签平滑log_probs = F.log_softmax(flat_logits, dim=-1)smooth_loss = -log_probs.mean(dim=-1)nll_loss = F.nll_loss(log_probs, flat_targets, ignore_index=-100, reduction='none')loss = (1 - self.config.label_smoothing) * nll_loss + self.config.label_smoothing * smooth_loss# 应用掩码if attention_mask is not None:mask = (flat_targets != -100).float()loss = (loss * mask).sum() / mask.sum()else:loss = loss.mean()else:loss = F.cross_entropy(flat_logits, flat_targets, ignore_index=-100)return lossdef training_step(self, batch):"""执行一个训练步骤Args:batch: 批次数据Returns:loss: 损失值grad_norm: 梯度范数"""step_start_time = time.time()# 将数据移到设备上input_ids = batch['input_ids'].to(self.device)target_ids = batch['target_ids'].to(self.device)attention_mask = batch['attention_mask'].to(self.device)# 前向传播if self.scaler is not None:# 混合精度训练with torch.cuda.amp.autocast():logits = self.model(input_ids, attention_mask)loss = self.compute_loss(logits, target_ids, attention_mask)# 反向传播self.scaler.scale(loss).backward()# 梯度裁剪self.scaler.unscale_(self.optimizer)grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.grad_clip)# 优化器步骤self.scaler.step(self.optimizer)self.scaler.update()else:# 标准精度训练logits = self.model(input_ids, attention_mask)loss = self.compute_loss(logits, target_ids, attention_mask)# 反向传播loss.backward()# 梯度裁剪grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.grad_clip)# 优化器步骤self.optimizer.step()# 清除梯度self.optimizer.zero_grad(set_to_none=True)# 更新学习率self.scheduler.step()# 记录指标step_time = time.time() - step_start_timecurrent_lr = self.scheduler.get_last_lr()[0]self.metrics.update(loss.item(), current_lr, grad_norm.item(), step_time)self.global_step += 1return loss.item(), grad_norm.item()def evaluate(self, dataloader):"""在验证集上评估模型"""self.model.eval()total_loss = 0num_batches = 0with torch.no_grad():for batch in dataloader:input_ids = batch['input_ids'].to(self.device)target_ids = batch['target_ids'].to(self.device)attention_mask = batch['attention_mask'].to(self.device)if self.scaler is not None:with torch.cuda.amp.autocast():logits = self.model(input_ids, attention_mask)loss = self.compute_loss(logits, target_ids, attention_mask)else:logits = self.model(input_ids, attention_mask)loss = self.compute_loss(logits, target_ids, attention_mask)total_loss += loss.item()num_batches += 1avg_loss = total_loss / max(num_batches, 1)perplexity = np.exp(avg_loss)self.model.train()return avg_loss, perplexitydef save_checkpoint(self, is_best=False):"""保存检查点"""checkpoint = {'epoch': self.epoch,'global_step': self.global_step,'model_state_dict': self.model.state_dict(),'optimizer_state_dict': self.optimizer.state_dict(),'scheduler_state_dict': self.scheduler.state_dict(),'best_loss': self.best_loss,'config': self.config,'metrics': {'losses': self.metrics.losses,'learning_rates': self.metrics.learning_rates}}if self.scaler is not None:checkpoint['scaler_state_dict'] = self.scaler.state_dict()# 保存最新检查点latest_path = os.path.join(self.config.checkpoint_dir, 'latest.pt')torch.save(checkpoint, latest_path)# 保存最佳检查点if is_best:best_path = os.path.join(self.config.checkpoint_dir, 'best.pt')torch.save(checkpoint, best_path)logger.info(f"保存最佳模型检查点: {best_path}")logger.info(f"保存检查点: {latest_path}")def generate_sample(self, dataset, prompt="Hello", max_length=100, temperature=0.8):"""生成文本样本用于质量检查"""self.model.eval()# 编码输入promptinput_ids = [dataset.vocab.get(char, dataset.vocab['<unk>']) for char in prompt]input_ids = torch.tensor([input_ids], device=self.device)generated = input_ids.clone()with torch.no_grad():for _ in range(max_length):if generated.size(1) >= self.config.max_seq_len:break# 前向传播获取下一个token的概率logits = self.model(generated)next_token_logits = logits[0, -1, :] / temperature# 采样下一个tokenprobs = F.softmax(next_token_logits, dim=-1)next_token = torch.multinomial(probs, 1)# 如果生成结束符则停止if next_token.item() == dataset.vocab['<eos>']:breakgenerated = torch.cat([generated, next_token.unsqueeze(0)], dim=1)# 解码生成的文本generated_text = ''.join([dataset.idx_to_char.get(idx.item(), '') for idx in generated[0]])self.model.train()return generated_textdef train(self, train_dataset, val_dataset=None):"""主训练循环"""logger.info("开始训练...")# 设置模型和优化器self.setup_model_and_optimizer(train_dataset)# 创建数据加载器train_loader = DataLoader(train_dataset,batch_size=self.config.batch_size,shuffle=True,collate_fn=collate_fn,num_workers=self.config.num_workers,pin_memory=True if self.device.type == 'cuda' else False)val_loader = Noneif val_dataset is not None:val_loader = DataLoader(val_dataset,batch_size=self.config.batch_size,shuffle=False,collate_fn=collate_fn,num_workers=self.config.num_workers)# 计算总训练步数self.total_steps = len(train_loader) * self.config.max_epochslogger.info(f"训练数据: {len(train_dataset)} 样本")logger.info(f"验证数据: {len(val_dataset) if val_dataset else 0} 样本")logger.info(f"总训练步数: {self.total_steps}")# 训练循环self.model.train()for epoch in range(self.config.max_epochs):self.epoch = epochepoch_start_time = time.time()logger.info(f"开始训练 Epoch {epoch + 1}/{self.config.max_epochs}")for batch_idx, batch in enumerate(train_loader):# 执行训练步骤loss, grad_norm = self.training_step(batch)# 定期输出训练信息if self.global_step % 50 == 0:stats = self.metrics.get_recent_stats()logger.info(f"Step {self.global_step:6d} | "f"Loss: {stats['avg_loss']:.4f} | "f"PPL: {stats['avg_perplexity']:.2f} | "f"LR: {stats['current_lr']:.2e} | "f"GradNorm: {stats['avg_grad_norm']:.3f} | "f"Steps/sec: {stats['steps_per_sec']:.1f}")# 定期评估if val_loader is not None and self.global_step % self.config.eval_interval == 0:val_loss, val_perplexity = self.evaluate(val_loader)logger.info(f"验证 - Loss: {val_loss:.4f}, Perplexity: {val_perplexity:.2f}")# 保存最佳模型if val_loss < self.best_loss:self.best_loss = val_lossself.save_checkpoint(is_best=True)# 生成样本检查质量sample_text = self.generate_sample(train_dataset, prompt="The", max_length=50)logger.info(f"生成样本: {sample_text}")# 定期保存检查点if self.global_step % self.config.save_interval == 0:self.save_checkpoint()# 检查是否需要提前停止if self.global_step >= self.total_steps:breakepoch_time = time.time() - epoch_start_timelogger.info(f"Epoch {epoch + 1} 完成，用时: {epoch_time:.1f}秒")# 每个epoch结束后保存检查点self.save_checkpoint()logger.info("训练完成！")# 绘制训练指标metrics_plot_path = os.path.join(self.config.output_dir, 'training_metrics.png')self.metrics.plot_metrics(metrics_plot_path)return self.model, self.metrics# 演示完整训练流程
def main():"""主函数：演示完整的语言模型训练流程"""logger.info("=== 语言模型训练全流程演示 ===")# 创建训练配置config = TrainingConfig(# 模型配置（小规模便于演示）vocab_size=1000,  # 会根据实际数据更新max_seq_len=128,n_layers=4,n_heads=4,n_embd=256,# 训练配置batch_size=8,learning_rate=1e-3,max_epochs=5,warmup_steps=100,eval_interval=200,save_interval=500,# 系统配置mixed_precision=False,  # 小模型不需要混合精度num_workers=0,  # 简化演示)# 准备演示数据（莎士比亚风格文本）sample_texts = ["To be or not to be that is the question","Whether tis nobler in the mind to suffer","The slings and arrows of outrageous fortune","Or to take arms against a sea of troubles","And by opposing end them to die to sleep","No more and by a sleep to say we end","The heartache and the thousand natural shocks","That flesh is heir to tis a consummation","Devoutly to be wished to die to sleep","To sleep perchance to dream aye theres the rub","For in that sleep of death what dreams may come","When we have shuffled off this mortal coil","Must give us pause theres the respect","That makes calamity of so long life","For who would bear the whips and scorns of time"] * 20  # 重复以增加训练数据# 创建数据集logger.info("创建数据集...")full_dataset = SimpleTextDataset(sample_texts, None, max_length=config.max_seq_len)# 分割训练集和验证集train_size = int(0.8 * len(full_dataset))val_size = len(full_dataset) - train_sizetrain_dataset, val_dataset = torch.utils.data.random_split(full_dataset, [train_size, val_size],generator=torch.Generator().manual_seed(42))logger.info(f"训练集大小: {len(train_dataset)}")logger.info(f"验证集大小: {len(val_dataset)}")# 创建训练器trainer = LanguageModelTrainer(config)# 开始训练try:model, metrics = trainer.train(full_dataset, val_dataset)logger.info("训练成功完成！")# 展示最终统计final_stats = metrics.get_recent_stats(window=100)logger.info("=== 最终训练统计 ===")logger.info(f"最终损失: {final_stats['avg_loss']:.4f}")logger.info(f"最终困惑度: {final_stats['avg_perplexity']:.2f}")logger.info(f"平均训练速度: {final_stats['steps_per_sec']:.1f} steps/sec")# 生成几个样本展示效果logger.info("=== 生成样本展示 ===")prompts = ["To be", "The", "And", "For"]for prompt in prompts:sample = trainer.generate_sample(full_dataset, prompt=prompt, max_length=30)logger.info(f"Prompt: '{prompt}' -> Generated: '{sample}'")except Exception as e:logger.error(f"训练过程中出现错误: {e}")raiseif __name__ == "__main__":main()