import math import os import time import random from typing import Iterator, List import argparse import concurrent.futures import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import IterableDataset, DataLoader from torch import amp import sentencepiece as spm from tqdm import tqdm import itertools print(f"--- PYTHON EXECUTING THIS FILE: {__file__} ---") try: from torch.utils.tensorboard import SummaryWriter except ImportError: SummaryWriter = None # --- 配置 --- class Config: # --- 模型架构 --- vocab_size = 72000 d_model = 1024 n_heads = 4 n_layers = 8 d_ff = 4096 seq_len = 1024 dropout = 0.1 # --- 训练硬件与精度 --- dtype = torch.float16 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # --- 学习率热重启配置 --- lr_max = 7e-6 warmup_steps = 1000 lr_min = 1e-6 restart_warmup_steps = 200 lr_restart_cycles = [45000, 2000, 3000] # --- 批次与步数 --- batch_size = 5 grad_accum_steps = 6 max_steps = warmup_steps + sum(lr_restart_cycles) + len(lr_restart_cycles) * restart_warmup_steps save_every = 500 out_dir = './checkpoints' # --- 其他训练参数 --- max_grad_norm = 1.0 # --- 高级正则化超参数 --- label_smoothing = 0.05 distill_temp = 1.2 distill_alpha = 0 # --- 分词器类 --- class SPTokenizer: def __init__(self, model_file: str, seq_len: int): self.sp = spm.SentencePieceProcessor(model_file=model_file) self.seq_len = seq_len def encode(self, text: str, pad=True): ids = self.sp.encode(text, out_type=int) if pad: if len(ids) > self.seq_len: ids = ids[:self.seq_len] else: pad_id = self.sp.pad_id() if self.sp.pad_id() != -1 else 0 ids += [pad_id] * (self.seq_len - len(ids)) return ids def decode(self, ids: List[int]): valid_ids = [id for id in ids if id >= 0 and id < self.sp.vocab_size()] return self.sp.decode(valid_ids) # --- 数据集类 --- class WeightedTextLineDataset(IterableDataset): def __init__(self, file_weights: dict[str, int], tokenizer: SPTokenizer, skip_lines: int = 0): super().__init__() self.tokenizer = tokenizer self.skip_lines = skip_lines if not file_weights: raise ValueError("文件权重字典不能为空。") self.filepaths = list(file_weights.keys()) self.weights = list(file_weights.values()) print("加权数据集已初始化。") def _create_line_iterator(self, filepath): """为单个文件创建一个无限循环的行生成器""" while True: try: with open(filepath, 'r', encoding='utf-8', errors='ignore') as f: for line in f: yield line except Exception as e: print(f"警告：读取文件 {filepath} 时出错: {e}") time.sleep(1) # 稍作等待再重试 def __iter__(self) -> Iterator[torch.Tensor]: # 为每个文件创建一个独立的、无限循环的行迭代器 iterators = [self._create_line_iterator(fp) for fp in self.filepaths] # 创建一个总的、加权的行生成器 def weighted_line_generator(): while True: # 根据权重，随机选择一个文件的迭代器 chosen_iterator = random.choices(iterators, weights=self.weights, k=1)[0] try: # 从被选中的迭代器中获取下一行 yield next(chosen_iterator) except StopIteration: # 由于 _create_line_iterator 是无限的，这里理论上不会发生 continue # 在总的加权迭代器上，执行一次性的、高效的跳过 line_it = weighted_line_generator() if self.skip_lines > 0: print(f"数据集：正在快速跳过前 {self.skip_lines} 行（按权重分布）...") line_it = itertools.islice(line_it, self.skip_lines, None) print("跳过完成。") self.skip_lines = 0 # 确保在后续的数据集循环中不再跳过 # 从正确的位置开始，处理并 yield 数据 for line in line_it: line = line.strip() if line: tokens = self.tokenizer.encode(line) yield torch.tensor(tokens, dtype=torch.long) # --- 模型定义 --- class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() self.eps = eps self.scale = nn.Parameter(torch.ones(dim)) def forward(self, x): norm = x.pow(2).mean(-1, keepdim=True).add(self.eps).sqrt() return x / norm * self.scale class MultiHeadSelfAttention(nn.Module): def __init__(self, d_model, n_heads, seq_len, dropout=0.0): super().__init__() assert d_model % n_heads == 0 self.n_heads = n_heads self.head_dim = d_model // n_heads self.scale = 1.0 / math.sqrt(self.head_dim) self.qkv = nn.Linear(d_model, d_model * 3, bias=False) self.out = nn.Linear(d_model, d_model) self.register_buffer('mask', torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0)) self.dropout = nn.Dropout(dropout) def forward(self, x): B, T, C = x.size() qkv = self.qkv(x) q, k, v = qkv.chunk(3, dim=-1) q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) att = torch.matmul(q, k.transpose(-2, -1)) * self.scale # 注意: 你的旧版本注意力mask逻辑 mask = self.mask[:, :, :T, :T] if self.mask.dim() == 4 else self.mask[:, :T, :T] att = att.masked_fill(mask == 0, float('-inf')) att = F.softmax(att, dim=-1) att = self.dropout(att) out = torch.matmul(att, v) out = out.transpose(1, 2).contiguous().reshape(B, T, C) return self.out(out) class FeedForward(nn.Module): def __init__(self, d_model, d_ff, dropout=0.0): super().__init__() self.net = nn.Sequential( nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model), nn.Dropout(dropout) ) def forward(self, x): return self.net(x) class TransformerBlock(nn.Module): def __init__(self, d_model, n_heads, d_ff, seq_len, dropout=0.0): super().__init__() self.attn = MultiHeadSelfAttention(d_model, n_heads, seq_len, dropout) self.norm1 = RMSNorm(d_model) self.ff = FeedForward(d_model, d_ff, dropout) self.norm2 = RMSNorm(d_model) def forward(self, x): x = x + self.attn(self.norm1(x)) x = x + self.ff(self.norm2(x)) return x class TinyDecoderModel(nn.Module): def __init__(self, cfg: Config): super().__init__() self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model) self.pos_emb = nn.Parameter(torch.zeros(cfg.seq_len, cfg.d_model)) self.layers = nn.ModuleList([ TransformerBlock(cfg.d_model, cfg.n_heads, cfg.d_ff, cfg.seq_len, cfg.dropout) for _ in range(cfg.n_layers) ]) self.ln_f = RMSNorm(cfg.d_model) self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False) nn.init.normal_(self.pos_emb, mean=0.0, std=0.02) def forward(self, idx): B, T = idx.size() x = self.tok_emb(idx) + self.pos_emb[:T] for l in self.layers: x = l(x) x = self.ln_f(x) logits = self.head(x) return logits # --- 学习率调度器 --- def get_lr_with_cyclical_warmup(step: int, cfg: Config) -> float: if step < cfg.warmup_steps: return cfg.lr_max * (step + 1) / cfg.warmup_steps effective_step = step - cfg.warmup_steps for cycle_len in cfg.lr_restart_cycles: full_cycle_len = cfg.restart_warmup_steps + cycle_len if effective_step < full_cycle_len: step_in_this_cycle = effective_step if step_in_this_cycle < cfg.restart_warmup_steps: return cfg.lr_min + (cfg.lr_max - cfg.lr_min) * (step_in_this_cycle / cfg.restart_warmup_steps) else: step_in_decay = step_in_this_cycle - cfg.restart_warmup_steps decay_ratio = step_in_decay / cycle_len coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) return cfg.lr_min + coeff * (cfg.lr_max - cfg.lr_min) effective_step -= full_cycle_len return cfg.lr_min # --- 检查点函数 --- def save_checkpoint(model, optimizer, scaler, step, out_dir): os.makedirs(out_dir, exist_ok=True) path = os.path.join(out_dir, f'ckpt_step{step}.pt') torch.save({ 'step': step, 'model_state': model.state_dict(), 'opt_state': optimizer.state_dict(), 'scaler_state': scaler.state_dict(), }, path) print(f"\n检查点已保存: {path}") # --- 主训练函数 --- def train(file_weights: dict[str, int], cfg: Config, resume_from_ckpt=None, skip_lines=0): tokenizer = SPTokenizer(model_file=r"D:\a\uyghur-dictionary\mymodel.model", seq_len=cfg.seq_len) writer = SummaryWriter(log_dir=os.path.join('./logs', time.strftime("%Y%m%d-%H%M%S"))) if SummaryWriter else None model = TinyDecoderModel(cfg).to(cfg.device) optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr_max) scaler = amp.GradScaler(enabled=(cfg.dtype == torch.float16)) step = 0 if resume_from_ckpt: print(f"正在从检查点恢复: {resume_from_ckpt}") checkpoint = torch.load(resume_from_ckpt, map_location=cfg.device) model.load_state_dict(checkpoint['model_state']) optimizer.load_state_dict(checkpoint['opt_state']) if 'scaler_state' in checkpoint: scaler.load_state_dict(checkpoint['scaler_state']) step = checkpoint['step'] print(f"恢复成功，将从 step {step + 1} 继续训练。") if step > 0 and skip_lines == 0: lines_per_step = cfg.batch_size * cfg.grad_accum_steps skip_lines = step * lines_per_step dataset = WeightedTextLineDataset(file_weights, tokenizer, skip_lines=skip_lines) dataloader = DataLoader(dataset, batch_size=cfg.batch_size, num_workers=0) data_iter = iter(dataloader) model.train() optimizer.zero_grad(set_to_none=True) start_time = time.time() initial_step = step pbar = tqdm(range(step, cfg.max_steps), desc="Training Steps", initial=step, total=cfg.max_steps) for step in pbar: lr = get_lr_with_cyclical_warmup(step, cfg) for g in optimizer.param_groups: g['lr'] = lr accumulated_loss = 0.0 for i in range(cfg.grad_accum_steps): pbar.set_description(f"Step {step+1}/{cfg.max_steps} [Accum. {i+1}/{cfg.grad_accum_steps}]") try: batch = next(data_iter) except StopIteration: data_iter = iter(dataloader) batch = next(data_iter) batch = batch.to(cfg.device) with amp.autocast(device_type=cfg.device.type.replace(':', ''), dtype=cfg.dtype): logits = model(batch) shift_logits = logits[:, :-1, :].contiguous() shift_labels = batch[:, 1:].contiguous() loss_hard = F.cross_entropy( shift_logits.view(-1, cfg.vocab_size), shift_labels.view(-1), ignore_index=0, label_smoothing=cfg.label_smoothing ) if cfg.distill_temp > 1.0 and cfg.distill_alpha > 0.0: with torch.no_grad(): logits_teacher = shift_logits / cfg.distill_temp q_soft_target = F.softmax(logits_teacher, dim=-1) log_p_cold = F.log_softmax(shift_logits, dim=-1) loss_distill = -torch.sum(q_soft_target * log_p_cold, dim=-1).mean() loss = cfg.distill_alpha * loss_distill + (1.0 - cfg.distill_alpha) * loss_hard else: loss = loss_hard accumulated_loss += loss.item() scaler.scale(loss / cfg.grad_accum_steps).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.max_grad_norm) scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_none=True) if writer: avg_loss_this_step = accumulated_loss / cfg.grad_accum_steps if math.isfinite(avg_loss_this_step): writer.add_scalar('Loss/Total_Loss', avg_loss_this_step, step) writer.add_scalar('Meta/Learning_Rate', lr, step) writer.flush() pbar.set_postfix(loss=f"{avg_loss_this_step:.4f}", lr=f"{lr:.2e}") if (step + 1) % cfg.save_every == 0: if writer: writer.flush() save_checkpoint(model, optimizer, scaler, step + 1, cfg.out_dir) if writer: writer.flush() save_checkpoint(model, optimizer, scaler, step + 1, cfg.out_dir) if writer: writer.close() if __name__ == "__main__": if torch.cuda.is_available(): torch.cuda.empty_cache() parser = argparse.ArgumentParser(description="使用高级正则化和热重启学习率进行模型训练。") parser.add_argument('--resume_from_ckpt', type=str, default=None, help='指定要恢复训练的检查点文件路径。') parser.add_argument('--skip_lines', type=int, default=0, help='手动指定数据集要跳过的初始行数。') args = parser.parse_args() # --- 【核心】使用字典来定义文件和权重 --- file_weights_map = { } if args.resume_from_ckpt and not os.path.exists(args.resume_from_ckpt): print(f"错误: 检查点文件 '{args.resume_from_ckpt}' 不存在！") exit() cfg = Config() # --- 【核心】将字典和正确的参数名传入 train 函数 --- train(file_weights=file_weights_map, cfg=cfg, resume_from_ckpt=args.resume_from_ckpt, skip_lines=args.skip_lines)