train.py

"""
GPT-2 预训练主脚本
使用 Hugging Face Trainer
"""
import os
import argparse
import torch
import wandb
import swanlab
from transformers import (
    Trainer,
    TrainingArguments,
    get_scheduler
)
from transformers.trainer_callback import TrainerCallback

from config import get_config
from model import create_model, get_tokenizer
from dataset import PretrainDataset, DataCollatorForLM


def setup_swanlab_auth():
    """从环境变量读取 SwanLab API Key 并尝试登录。"""
    api_key = os.environ.get("SWANLAB_API_KEY", "").strip()
    if not api_key:
        return False
    if hasattr(swanlab, "login"):
        try:
            swanlab.login(api_key=api_key)
        except TypeError:
            try:
                swanlab.login(api_key)
            except Exception:
                pass
        except Exception:
            pass
    return True


class LoggingCallback(TrainerCallback):
    """自定义日志回调"""

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is not None and state.is_world_process_zero:
            # 打印关键指标
            if "loss" in logs:
                print(f"Step {state.global_step}: loss={logs['loss']:.4f}", end="")
                if "learning_rate" in logs:
                    print(f", lr={logs['learning_rate']:.2e}", end="")
                print()


def compute_metrics(eval_pred):
    """计算评估指标"""
    import numpy as np

    logits, labels = eval_pred
    # 计算困惑度
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()

    # 使用交叉熵计算
    loss_fct = torch.nn.CrossEntropyLoss()
    loss = loss_fct(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1)
    )
    perplexity = torch.exp(loss).item()

    return {"perplexity": perplexity}


def train(config):
    """训练函数"""
    # 设置随机种子
    torch.manual_seed(config.training.seed)

    # 创建输出目录
    os.makedirs(config.training.output_dir, exist_ok=True)

    # 初始化 SwanLab（仅在主进程）
    if int(os.environ.get("LOCAL_RANK", 0)) == 0:
        authed = setup_swanlab_auth()
        if not authed:
            print("⚠️ 未检测到 SWANLAB_API_KEY，SwanLab 可能无法上传到你的账号")
        swanlab_run = swanlab.init(
            project="GPT2-Dropout-Comparison",
            experiment_name=f"gpt2-dropout-{config.model.resid_pdrop}",
            config={
                "model_size": config.model.model_size,
                "dropout": config.model.resid_pdrop,
                "learning_rate": config.training.learning_rate,
                "batch_size": config.training.per_device_train_batch_size,
                "gradient_accumulation": config.training.gradient_accumulation_steps,
                "warmup_steps": config.training.warmup_steps,
                "weight_decay": config.training.weight_decay,
                "adam_beta1": config.training.adam_beta1,
                "adam_beta2": config.training.adam_beta2,
                "max_grad_norm": config.training.max_grad_norm,
                "num_epochs": config.training.num_train_epochs,
                "block_size": config.data.block_size,
            },
        )

    # 创建模型（如果是resume，只创建架构，权重会从checkpoint加载）
    # 如果resume_from_checkpoint存在，from_scratch应该为False以便加载权重
    should_init_from_scratch = config.model.from_scratch and not config.training.resume_from_checkpoint

    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=should_init_from_scratch
    )

    # 获取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
    )

    # 数据整理器
    data_collator = DataCollatorForLM()

    # 训练参数
    training_args = TrainingArguments(
        output_dir=config.training.output_dir,

        # 训练轮数
        num_train_epochs=config.training.num_train_epochs,
        max_steps=config.training.max_steps,

        # Batch size
        per_device_train_batch_size=config.training.per_device_train_batch_size,
        per_device_eval_batch_size=config.training.per_device_eval_batch_size,
        gradient_accumulation_steps=config.training.gradient_accumulation_steps,

        # 优化器
        learning_rate=config.training.learning_rate,
        weight_decay=config.training.weight_decay,
        adam_beta1=config.training.adam_beta1,
        adam_beta2=config.training.adam_beta2,
        warmup_steps=config.training.warmup_steps,
        max_grad_norm=config.training.max_grad_norm,
        lr_scheduler_type=config.training.lr_scheduler_type,

        # 日志
        logging_dir=os.path.join(config.training.output_dir, "logs"),
        logging_steps=config.training.logging_steps,
        logging_first_step=True,

        # 保存
        save_steps=config.training.save_steps,
        save_total_limit=5,

        # 评估
        eval_strategy="steps",
        eval_steps=config.training.eval_steps,
        load_best_model_at_end=True,
        metric_for_best_model="eval_loss",
        greater_is_better=False,

        # 混合精度
        fp16=config.training.fp16,
        bf16=config.training.bf16,

        # 其他
        seed=config.training.seed,
        dataloader_num_workers=config.training.dataloader_num_workers,
        dataloader_pin_memory=True,
        ddp_find_unused_parameters=False,

        # 报告
        report_to=["tensorboard", "swanlab"],

        # 断点续训
        resume_from_checkpoint=config.training.resume_from_checkpoint,
    )

    # 创建Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        data_collator=data_collator,
        callbacks=[LoggingCallback()],
    )

    # 开始训练
    print("\n" + "=" * 50)
    print("Starting training...")
    if config.training.resume_from_checkpoint:
        print(f"Resuming from checkpoint: {config.training.resume_from_checkpoint}")
        # 列出可用的checkpoints
        import glob
        checkpoints = sorted(glob.glob(os.path.join(config.training.output_dir, "checkpoint-*")))
        if checkpoints:
            print(f"Available checkpoints: {[os.path.basename(c) for c in checkpoints]}")
            print(f"Latest checkpoint: {os.path.basename(checkpoints[-1])}")
        else:
            print("WARNING: No checkpoints found!")
    print("=" * 50)

    trainer.train(resume_from_checkpoint=config.training.resume_from_checkpoint)

    # 保存最终模型
    print("\nSaving final model...")
    trainer.save_model(os.path.join(config.training.output_dir, "final"))
    tokenizer.save_pretrained(os.path.join(config.training.output_dir, "final"))

    # 最终评估
    print("\nFinal evaluation...")
    eval_results = trainer.evaluate()
    print(f"Final eval loss: {eval_results['eval_loss']:.4f}")
    final_ppl = torch.exp(torch.tensor(eval_results['eval_loss'])).item()
    print(f"Final perplexity: {final_ppl:.2f}")

    # 记录最终指标到 SwanLab（仅在主进程）
    if int(os.environ.get("LOCAL_RANK", 0)) == 0:
        swanlab.log({
            "final/eval_loss": eval_results['eval_loss'],
            "final/perplexity": final_ppl,
        })
        # 结束 SwanLab 实验
        swanlab.finish()

    return trainer


def main():
    parser = argparse.ArgumentParser(description="Pretrain GPT-2 with Dropout")

    # 模型参数
    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=False)

    # 数据参数
    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=12)
    parser.add_argument("--gradient_accumulation", type=int, default=40)
    parser.add_argument("--learning_rate", type=float, default=3e-4)
    parser.add_argument("--warmup_steps", type=int, default=2000)
    parser.add_argument("--weight_decay", type=float, default=0.1)
    parser.add_argument("--adam_beta1", type=float, default=0.9)
    parser.add_argument("--adam_beta2", type=float, default=0.95)
    parser.add_argument("--max_grad_norm", type=float, default=1.0)

    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)
    parser.add_argument("--bf16", action="store_true", default=False)

    parser.add_argument("--resume", type=str, default=None)

    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.adam_beta1 = args.adam_beta1
    config.training.adam_beta2 = args.adam_beta2
    config.training.max_grad_norm = args.max_grad_norm
    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
    config.training.bf16 = args.bf16

    # 处理resume参数：支持True(自动查找)或具体路径
    if args.resume and args.resume.lower() == "true":
        config.training.resume_from_checkpoint = True
    else:
        config.training.resume_from_checkpoint = args.resume

    # 打印配置
    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"Resume from checkpoint: {config.training.resume_from_checkpoint}")
    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"Output dir: {config.training.output_dir}")
    print("=" * 50)

    # 开始训练
    train(config)


if __name__ == "__main__":
    main()