""" GPT-2 预训练脚本 (PyTorch原生版本) 不依赖Hugging Face Trainer，更灵活可控 """ import os import math import time import argparse from contextlib import nullcontext import torch import torch.nn as nn from torch.utils.data import DataLoader from torch.cuda.amp import GradScaler, autocast from torch.utils.tensorboard import SummaryWriter from config import get_config from model import create_model, get_tokenizer from dataset import PretrainDataset, StreamingDataset, DataCollatorForLM def get_lr(step, warmup_steps, max_steps, max_lr, min_lr): """ 学习率调度：线性warmup + cosine decay """ # Warmup阶段 if step < warmup_steps: return max_lr * step / warmup_steps # Decay阶段 if step > max_steps: return min_lr # Cosine decay decay_ratio = (step - warmup_steps) / (max_steps - warmup_steps) coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) return min_lr + coeff * (max_lr - min_lr) @torch.no_grad() def evaluate(model, val_loader, device, ctx): """评估函数""" model.eval() total_loss = 0 total_tokens = 0 for batch in val_loader: input_ids = batch["input_ids"].to(device) labels = batch["labels"].to(device) with ctx: outputs = model(input_ids, labels=labels) loss = outputs.loss total_loss += loss.item() * input_ids.numel() total_tokens += input_ids.numel() avg_loss = total_loss / total_tokens perplexity = math.exp(avg_loss) model.train() return avg_loss, perplexity def train(config): """训练主函数""" # 设置设备 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") # 设置随机种子 torch.manual_seed(config.training.seed) if torch.cuda.is_available(): torch.cuda.manual_seed(config.training.seed) # 创建输出目录 os.makedirs(config.training.output_dir, exist_ok=True) # 创建模型 model, model_config = create_model( model_size=config.model.model_size, resid_pdrop=config.model.resid_pdrop, attn_pdrop=config.model.attn_pdrop, embd_pdrop=config.model.embd_pdrop, from_scratch=config.model.from_scratch ) model = model.to(device) # 获取tokenizer tokenizer = get_tokenizer(config.model.model_size) # 创建数据集 print("\nLoading datasets...") train_dataset = PretrainDataset( config.data.train_file, config.data.block_size ) val_dataset = PretrainDataset( config.data.val_file, config.data.block_size ) train_loader = DataLoader( train_dataset, batch_size=config.training.per_device_train_batch_size, shuffle=True, num_workers=config.training.dataloader_num_workers, pin_memory=True, drop_last=True ) val_loader = DataLoader( val_dataset, batch_size=config.training.per_device_eval_batch_size, shuffle=False, num_workers=config.training.dataloader_num_workers, pin_memory=True ) # 计算训练步数 steps_per_epoch = len(train_loader) // config.training.gradient_accumulation_steps if config.training.max_steps > 0: max_steps = config.training.max_steps else: max_steps = steps_per_epoch * config.training.num_train_epochs print(f"Steps per epoch: {steps_per_epoch}") print(f"Max steps: {max_steps}") # 优化器 (使用nanoGPT验证过的参数) optimizer = torch.optim.AdamW( model.parameters(), lr=config.training.learning_rate, weight_decay=config.training.weight_decay, betas=(config.training.adam_beta1, config.training.adam_beta2), eps=config.training.adam_epsilon ) # 混合精度 use_amp = config.training.fp16 or config.training.bf16 if config.training.bf16 and torch.cuda.is_bf16_supported(): dtype = torch.bfloat16 elif config.training.fp16: dtype = torch.float16 else: dtype = torch.float32 scaler = GradScaler(enabled=config.training.fp16) # bf16不需要scaler ctx = autocast(dtype=dtype) if use_amp else nullcontext() # TensorBoard writer = SummaryWriter(os.path.join(config.training.output_dir, "logs")) # 训练循环 print("\n" + "=" * 50) print("Starting training...") print("=" * 50) model.train() global_step = 0 best_val_loss = float("inf") start_time = time.time() for epoch in range(config.training.num_train_epochs): print(f"\nEpoch {epoch + 1}/{config.training.num_train_epochs}") for micro_step, batch in enumerate(train_loader): # 获取学习率 (使用min_lr配置) lr = get_lr( global_step, config.training.warmup_steps, max_steps, config.training.learning_rate, config.training.min_lr ) for param_group in optimizer.param_groups: param_group["lr"] = lr # 前向传播 input_ids = batch["input_ids"].to(device) labels = batch["labels"].to(device) with ctx: outputs = model(input_ids, labels=labels) loss = outputs.loss / config.training.gradient_accumulation_steps # 反向传播 scaler.scale(loss).backward() # 梯度累积 if (micro_step + 1) % config.training.gradient_accumulation_steps == 0: # 梯度裁剪 scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_( model.parameters(), config.training.max_grad_norm ) # 更新参数 scaler.step(optimizer) scaler.update() optimizer.zero_grad() global_step += 1 # 日志 if global_step % config.training.logging_steps == 0: elapsed = time.time() - start_time tokens_per_sec = ( global_step * config.training.per_device_train_batch_size * config.training.gradient_accumulation_steps * config.data.block_size ) / elapsed print( f"Step {global_step}/{max_steps} | " f"Loss: {loss.item() * config.training.gradient_accumulation_steps:.4f} | " f"LR: {lr:.2e} | " f"Tokens/s: {tokens_per_sec:.0f}" ) writer.add_scalar("train/loss", loss.item() * config.training.gradient_accumulation_steps, global_step) writer.add_scalar("train/lr", lr, global_step) writer.add_scalar("train/tokens_per_sec", tokens_per_sec, global_step) # 评估 if global_step % config.training.eval_steps == 0: val_loss, val_ppl = evaluate(model, val_loader, device, ctx) print(f" Eval - Loss: {val_loss:.4f}, PPL: {val_ppl:.2f}") writer.add_scalar("eval/loss", val_loss, global_step) writer.add_scalar("eval/perplexity", val_ppl, global_step) # 保存最佳模型 if val_loss < best_val_loss: best_val_loss = val_loss save_path = os.path.join(config.training.output_dir, "best") os.makedirs(save_path, exist_ok=True) model.save_pretrained(save_path) tokenizer.save_pretrained(save_path) print(f" Saved best model to {save_path}") # 定期保存 if global_step % config.training.save_steps == 0: save_path = os.path.join( config.training.output_dir, f"checkpoint-{global_step}" ) os.makedirs(save_path, exist_ok=True) model.save_pretrained(save_path) tokenizer.save_pretrained(save_path) # 保存优化器状态 torch.save({ "optimizer": optimizer.state_dict(), "scaler": scaler.state_dict(), "global_step": global_step, "best_val_loss": best_val_loss, }, os.path.join(save_path, "training_state.pt")) print(f" Saved checkpoint to {save_path}") # 检查是否达到最大步数 if global_step >= max_steps: break if global_step >= max_steps: break # 保存最终模型 print("\nSaving final model...") save_path = os.path.join(config.training.output_dir, "final") os.makedirs(save_path, exist_ok=True) model.save_pretrained(save_path) tokenizer.save_pretrained(save_path) # 最终评估 val_loss, val_ppl = evaluate(model, val_loader, device, ctx) print(f"\nFinal Results:") print(f" Val Loss: {val_loss:.4f}") print(f" Val PPL: {val_ppl:.2f}") print(f" Best Val Loss: {best_val_loss:.4f}") writer.close() print("\nTraining complete!") def main(): parser = argparse.ArgumentParser(description="Pretrain GPT-2 (PyTorch native)") # 模型参数 parser.add_argument("--model_size", type=str, default="gpt2-medium", choices=["gpt2", "gpt2-medium", "gpt2-large", "gpt2-xl"]) parser.add_argument("--dropout", type=float, default=0.1) parser.add_argument("--from_scratch", action="store_true", default=True) # 数据参数 parser.add_argument("--train_file", type=str, default="./data/train.bin") parser.add_argument("--val_file", type=str, default="./data/val.bin") parser.add_argument("--block_size", type=int, default=1024) # 训练参数 parser.add_argument("--output_dir", type=str, default="./output/gpt2-dropout") parser.add_argument("--num_epochs", type=int, default=1) parser.add_argument("--max_steps", type=int, default=-1) parser.add_argument("--batch_size", type=int, default=4) parser.add_argument("--gradient_accumulation", type=int, default=8) parser.add_argument("--learning_rate", type=float, default=5e-4) parser.add_argument("--warmup_steps", type=int, default=5000) parser.add_argument("--weight_decay", type=float, default=0.01) parser.add_argument("--logging_steps", type=int, default=100) parser.add_argument("--save_steps", type=int, default=5000) parser.add_argument("--eval_steps", type=int, default=1000) parser.add_argument("--fp16", action="store_true", default=True) args = parser.parse_args() # 创建配置 config = get_config( model_size=args.model_size, dropout=args.dropout, from_scratch=args.from_scratch, ) # 更新配置 config.data.train_file = args.train_file config.data.val_file = args.val_file config.data.block_size = args.block_size config.training.output_dir = args.output_dir config.training.num_train_epochs = args.num_epochs config.training.max_steps = args.max_steps config.training.per_device_train_batch_size = args.batch_size config.training.per_device_eval_batch_size = args.batch_size config.training.gradient_accumulation_steps = args.gradient_accumulation config.training.learning_rate = args.learning_rate config.training.warmup_steps = args.warmup_steps config.training.weight_decay = args.weight_decay config.training.logging_steps = args.logging_steps config.training.save_steps = args.save_steps config.training.eval_steps = args.eval_steps config.training.fp16 = args.fp16 # 打印配置 print("=" * 50) print("Training Configuration") print("=" * 50) print(f"Model: {config.model.model_size}") print(f"Dropout: {config.model.resid_pdrop}") print(f"From scratch: {config.model.from_scratch}") print(f"Block size: {config.data.block_size}") print(f"Batch size: {config.training.per_device_train_batch_size}") print(f"Gradient accumulation: {config.training.gradient_accumulation_steps}") print(f"Effective batch size: {config.training.per_device_train_batch_size * config.training.gradient_accumulation_steps}") print(f"Learning rate: {config.training.learning_rate}") print(f"FP16: {config.training.fp16}") print(f"Output dir: {config.training.output_dir}") print("=" * 50) # 开始训练 train(config) if __name__ == "__main__": main()