"""
@description: 模型训练主文件,包含模型训练相关的核心功能实现
@author: 唐国梁Tommy（B站，公众号同名，个人网站： www.tgltommy.com）
@date: 2025-02-16
"""

import os
import argparse
import datetime
import json

import torch
from transformers import (
    AutoModelForCausalLM, 
    AutoTokenizer,
    BitsAndBytesConfig,
    set_seed
)
import transformers

# 分布式训练
from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs

from peft import (
    AutoPeftModelForCausalLM,
    LoraConfig, 
    get_peft_model,
    prepare_model_for_kbit_training
)

from datasets import load_dataset
from huggingface_hub import snapshot_download

import wandb
from dotenv import load_dotenv

from config import (
    TrainingConfig,
    ModelConfig, 
    DataConfig,
)


def is_local_path(path):
    """
    检查是否为本地路径
    """
    return os.path.exists(path) and os.path.isdir(path)


def download_model(model_name):
    """
    下载或使用本地模型
    Args:
        model_name: 模型名称或本地路径
    Returns:
        str: 模型路径
    """
    if is_local_path(model_name):
        print(f"使用本地模型: {model_name}")
        return model_name
    
    print(f"从Hugging Face下载模型: {model_name}")
    os.makedirs("/tmp_folder", exist_ok=True)
    
    try:
        snapshot_download(repo_id=model_name, local_dir="/tmp_folder")
        print(f"模型下载完成，保存在: /tmp_folder")
        return "/tmp_folder"
    except Exception as e:
        print(f"模型下载失败: {str(e)}")
        raise


def initialize_tokenizer(model_name):
    """初始化分词器"""
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenizer.pad_token = tokenizer.eos_token
    return tokenizer


def process_datasets(train_ds, test_ds, tokenizer, data_config, accelerator):
    """
    处理训练和测试数据集
    """
    def process_dataset(dataset, is_train=True):
        """统一的数据集处理逻辑"""
        if dataset is None:
            return None
            
        processed_ds = dataset.map(
            lambda sample: tokenizer(
                sample["text"],
                max_length=data_config.max_seq_length if is_train else None,
                padding="max_length" if is_train and data_config.pad_to_max_length else False,
                truncation=data_config.truncation if is_train else True
            ),
            remove_columns=list(dataset.features)
        )
        print(f"Total number of {'train' if is_train else 'test'} samples: {len(processed_ds)}")
        return processed_ds

    with accelerator.main_process_first():
        lm_train_dataset = process_dataset(train_ds, True)
        lm_test_dataset = process_dataset(test_ds, False)
        
        return lm_train_dataset, lm_test_dataset


def get_fsdp_config(gradient_checkpointing: bool) -> tuple[str, dict]:
    """
    获取FSDP(Fully Sharded Data Parallel)分布式训练配置
    
    Args:
        gradient_checkpointing: 是否启用梯度检查点
        
    Returns:
        tuple[str, dict]: FSDP模式和配置字典
    """
    fsdp = "full_shard auto_wrap offload"
    fsdp_config = {
        'backward_prefetch': 'backward_pre',
        'cpu_ram_efficient_loading': True,
        'offload_params': True,
        'forward_prefetch': False,
        'use_orig_params': False,
        'activation_checkpointing': gradient_checkpointing
    }
    return fsdp, fsdp_config


def get_training_configurations(
    is_distributed: bool, 
    torch_dtype: torch.dtype, 
    training_config: TrainingConfig
) -> tuple[BitsAndBytesConfig, dict]:
    """
    获取训练相关配置，包括量化配置和分布式训练配置
    
    Args:
        is_distributed: 是否为分布式训练
        torch_dtype: 训练使用的数据类型
        training_config: 训练相关配置对象
        
    Returns:
        tuple[BitsAndBytesConfig, dict]: BitsAndBytes量化配置和FSDP配置
    """
    if is_distributed:
        # 获取FSDP配置
        fsdp, fsdp_config = get_fsdp_config(training_config.gradient_checkpointing)
        fsdp_configurations = {
            "fsdp": fsdp,
            "fsdp_config": fsdp_config,
            "tf32": True
        }
        
        # 分布式训练时的量化存储配置
        bnb_config_params = {"bnb_4bit_quant_storage": torch_dtype}
    else:
        fsdp_configurations = {}
        bnb_config_params = {}

    # 创建量化配置
    bnb_config_dict = {**training_config.bnb_config, **bnb_config_params}
    bnb_config = BitsAndBytesConfig(**bnb_config_dict)
    
    return bnb_config, fsdp_configurations


def save_model(
    trainer, 
    tokenizer, 
    output_path, 
    save_merged_model,
    model_config,
    training_config,
    model_name,
    timestamp,
    gpu_mode,
    model_type,
    per_device_train_batch_size,
    learning_rate,
    num_train_epochs,
    gradient_accumulation_steps,
    accelerator
):
    """
    保存模型及相关配置
    """
    def save_model_with_config(model, is_safe=True):
        """统一的模型保存逻辑"""
        model.save_pretrained(
            output_path,
            safe_serialization=is_safe,
            max_shard_size=model_config.model_max_shard_size
        )
        
    if save_merged_model:
        save_model_with_config(trainer.model, model_config.safe_serialization)

        if accelerator.is_main_process:
            # 清理内存
            del trainer.model
            del trainer
            torch.cuda.empty_cache()

            # 加载PEFT模型并合并
            model = AutoPeftModelForCausalLM.from_pretrained(
                output_path,
                torch_dtype=torch.float16,
                low_cpu_mem_usage=True,
                trust_remote_code=True,
            )
            model = model.merge_and_unload()
            
            save_model_with_config(model, model_config.safe_serialization)
            model.config.save_pretrained(output_path)
    else:
        save_model_with_config(trainer.model, True)

    if accelerator.is_main_process:
        # 保存分词器
        tokenizer.save_pretrained(output_path)
        
        # 保存训练配置
        config_info = {
            "timestamp": timestamp,
            "gpu_mode": gpu_mode,
            "model_type": model_type,
            "base_model": model_name,
            "train_batch_size": per_device_train_batch_size,
            "learning_rate": learning_rate,
            "num_epochs": num_train_epochs,
            "gradient_accumulation_steps": gradient_accumulation_steps,
            "lora_r": training_config.lora_r,
            "lora_alpha": training_config.lora_alpha,
            "lora_dropout": training_config.lora_dropout
        }
        
        with open(os.path.join(output_path, "training_config.json"), "w") as f:
            json.dump(config_info, f, indent=4)


def initialize_wandb(accelerator, args, wb_token):
    """
    初始化wandb配置
    
    Args:
        accelerator: 分布式训练加速器
        args: 命令行参数
        wb_token: wandb token
    """
    if wb_token:
        wandb.login(key=wb_token)
    
    if accelerator.is_main_process:
        wandb.init(
            project=args.wandb_project_name,
            name=args.wandb_run_name,
            job_type="training",
            anonymous="allow",
        )


def setup_model_environment(accelerator, model_name, hf_token):
    """
    设置模型环境，包括token和模型下载
    
    Args:
        accelerator: 分布式训练加速器
        model_name: 模型名称或路径
        hf_token: HuggingFace token
    
    Returns:
        str: 模型路径
    """
    if hf_token is not None:
        os.environ.update({"HF_TOKEN": hf_token})
        accelerator.wait_for_everyone()

    # 检查是否为单机训练环境
    if int(os.environ.get("SM_HOST_COUNT", 1)) == 1:
        model_path = None
        if accelerator.is_main_process:
            model_path = download_model(model_name)
        
        if accelerator.num_processes > 1:
            model_path = accelerator.unwrap_model(model_path) if model_path else None
            model_path = [model_path]
            torch.distributed.broadcast_object_list(model_path, src=0)
            model_path = model_path[0]
    else:
        model_path = download_model(model_name)

    accelerator.wait_for_everyone()
    return model_path


def create_training_args(args, wb_token, fsdp_configurations, is_distributed, batch_size):
    """
    创建训练参数配置
    
    Args:
        args: 命令行参数
        wb_token: wandb token
        fsdp_configurations: FSDP配置
        is_distributed: 是否分布式训练
        batch_size: 根据GPU显存调整后的batch_size
    
    Returns:
        TrainingArguments: 训练参数配置
    """
    return transformers.TrainingArguments(
        # 训练和评估的批次大小
        per_device_train_batch_size=batch_size,  # 使用调整后的batch_size
        per_device_eval_batch_size=args.eval_batch_size,
        
        # 梯度累积相关配置
        gradient_accumulation_steps=args.grad_accum_steps,  # 梯度累积步数
        max_grad_norm=1.0,  # 梯度裁剪的最大范数
        
        # 日志记录相关配置
        report_to="wandb" if wb_token else None,  # 是否使用wandb记录日志
        logging_strategy="steps",  # 按步数记录日志
        logging_steps=args.logging_steps,  # 日志记录间隔步数
        log_on_each_node=False,  # 是否在每个节点记录日志
        
        # 训练相关配置
        num_train_epochs=args.epochs,  # 训练轮数
        learning_rate=args.learning_rate,  # 学习率
        lr_scheduler_type="cosine",  # 学习率调度器类型
        warmup_ratio=0.03,  # 预热比例
        bf16=True,  # 是否使用bf16精度
        
        # 分布式训练配置
        ddp_find_unused_parameters=False,  # DDP是否查找未使用的参数
        
        # 模型保存和评估策略
        save_strategy="epoch",  # 按epoch保存模型
        evaluation_strategy="epoch",  # 按epoch进行评估
        save_steps=1,  # 保存间隔步数
        eval_steps=1,  # 评估间隔步数
        save_total_limit=3,  # 最多保存的检查点数量
        output_dir="outputs",  # 输出目录
        
        # 早停相关配置
        load_best_model_at_end=True if args.early_stopping_patience > 0 else False,  # 是否加载最佳模型
        metric_for_best_model="eval_loss",  # 用于选择最佳模型的指标
        greater_is_better=False,  # 指标是否越大越好
        
        # 优化器配置
        optim="adamw_torch",  # 优化器类型
        optim_args="weight_decay=0.01",  # 优化器参数
        
        # 分布式训练额外配置
        **(fsdp_configurations if is_distributed else {}),  # FSDP配置
    )


def main(args):
    """
    主函数
    
    第1步：加载环境配置
    """
    load_dotenv() 
    training_config = TrainingConfig()
    model_config = ModelConfig()
    data_config = DataConfig()
    
    
    """
    第2步：初始化分布式训练环境
    """
    # 梯度累积步数
    gradient_accumulation_steps = args.grad_accum_steps
    # 初始化Accelerator对象用于分布式训练
    # - gradient_accumulation_steps: 梯度累积步数,用于增大等效batch size
    # - mixed_precision: 使用bf16混合精度训练,可以减少显存占用并提升训练速度
    # - kwargs_handlers: DDP相关配置,禁用未使用参数检查以提升性能
    accelerator = Accelerator(
        gradient_accumulation_steps=gradient_accumulation_steps,
        mixed_precision="bf16", 
        kwargs_handlers=[
            DistributedDataParallelKwargs(find_unused_parameters=False)
        ]
    )
    # 获取当前进程的本地rank,用于区分不同GPU上的进程
    local_rank = accelerator.local_process_index
    # 判断是否为分布式训练环境(进程数>1表示分布式训练)
    is_distributed = accelerator.num_processes > 1

    """
    第3步：初始化wandb日志记录
    """
    wb_token = os.getenv("WANDB_TOKEN")
    initialize_wandb(accelerator, args, wb_token)

    """
    第4步：设置模型环境并下载模型
    """
    model_name = args.model_name_or_path
    hf_token = os.getenv("HUGGINGFACE_TOKEN")
    model_path = setup_model_environment(accelerator, model_name, hf_token)
    model_name = model_path

    """
    第5步：初始化tokenizer和加载数据集
    """
    tokenizer = initialize_tokenizer(model_name)
    train_ds = load_dataset('json', data_files=args.train_file, split='train')  
    test_ds = load_dataset('json', data_files=args.eval_file, split='train')   

    """
    第6步：设置随机种子并处理数据集
    """
    set_seed(training_config.seed)
    lm_train_dataset, lm_test_dataset = process_datasets(
        train_ds, 
        test_ds, 
        tokenizer,  
        data_config,
        accelerator
    )

    """
    第7步: 根据GPU显存大小调整训练参数并创建training_args
    """
    # 根据GPU显存大小动态调整训练参数
    # 获取第一块GPU的总显存大小,与阈值比较
    if torch.cuda.get_device_properties(0).total_memory >= training_config.gpu_memory_threshold:
        # 如果显存充足,使用较大的batch size
        batch_size = training_config.large_gpu_batch_size
        # 启用flash attention加速训练
        use_flash_attention = training_config.use_flash_attention
    else:
        # 显存不足时使用命令行指定的较小batch size
        batch_size = args.train_batch_size 
        # 关闭flash attention以节省显存
        use_flash_attention = False

    # 设置模型训练的数据类型和量化配置
    # 使用bfloat16混合精度训练,可以在保持数值稳定性的同时减少显存占用
    torch_dtype = torch.bfloat16
    # 获取8bit量化和FSDP分布式训练的配置
    bnb_config, fsdp_configurations = get_training_configurations(
        is_distributed,  # 是否为分布式训练
        torch_dtype,     # 训练使用的数据类型
        training_config  # 训练相关配置
    )

    # 创建Trainer需要的training_args对象
    # 包含了训练所需的各种超参数和配置
    training_args = create_training_args(
        args=args,                           # 命令行参数
        wb_token=wb_token,                  # wandb用于日志记录
        fsdp_configurations=fsdp_configurations,  # FSDP分布式训练配置
        is_distributed=is_distributed,           # 是否为分布式训练
        batch_size=batch_size                    # 使用动态确定的batch size
    )

    """
    第8步：加载预训练模型
    """
    # 从预训练模型加载基础模型
    model = AutoModelForCausalLM.from_pretrained(
        model_name, 
        trust_remote_code=True,   
        quantization_config=bnb_config,  # 4bit量化配置,用于减少显存占用
        # 根据显存大小选择注意力实现方式:
        # - flash_attention_2: 显存充足时使用,可加速训练
        # - sdpa: 显存不足时的默认选项
        attn_implementation="flash_attention_2" if use_flash_attention else "sdpa",
        use_cache=False,           # 关闭KV缓存以节省显存
        max_memory=training_config.max_memory,  # 限制每个GPU的最大显存使用量
        offload_folder=training_config.offload_folder,  # 模型权重卸载目录
        torch_dtype=torch_dtype,   # 使用bfloat16混合精度训练
        low_cpu_mem_usage=True,    # 启用CPU内存优化
        # 设备映射策略:
        # - 分布式训练时为None,由DDP控制
        # - 单卡训练时为"auto",自动分配到可用设备
        device_map="auto" if not is_distributed else None,
        cache_dir=model_config.cache_dir  # 模型缓存目录
    )

    """
    第9步：配置分布式训练设备
    """
    if is_distributed:
        device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")
        model = model.to(device)

    """
    第10步：单GPU训练时的内存优化配置
    """
    if not is_distributed:
        # 启用梯度检查点以减少显存占用
        training_config.gradient_checkpointing = True
        # 为4bit量化训练准备模型,并启用梯度检查点
        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
        # 开启模型的梯度检查点功能
        model.gradient_checkpointing_enable()

    """
    第11步：配置LoRA参数并转换模型
    """
    # 创建LoRA配置对象
    config = LoraConfig(
        r=training_config.lora_r, # LoRA的秩,控制可训练参数数量,建议范围4-32
        lora_alpha=training_config.lora_alpha, # LoRA缩放因子,用于缩放LoRA产生的更新值
        target_modules=training_config.lora_target_modules,  # 指定应用LoRA的目标模块
        lora_dropout=training_config.lora_dropout, # LoRA的dropout率,防止过拟合,建议范围0.1-0.3
        bias="none", # 不对偏置项进行LoRA微调
        task_type="CAUSAL_LM"  # 指定任务类型为因果语言模型
    )

    # 将原始模型转换为PEFT模型,添加LoRA参数
    # get_peft_model会为指定的目标层添加LoRA权重矩阵,使模型可以进行参数高效微调
    model = get_peft_model(model, config)

    """
    第12步：创建训练器并开始训练
    """
    # 创建Trainer对象进行模型训练
    trainer = transformers.Trainer(
        model=model,
        train_dataset=lm_train_dataset,  # 训练数据集
        eval_dataset=lm_test_dataset if lm_test_dataset is not None else None,  # 评估数据集
        args=training_args,  # 训练参数配置
        # 配置early_stopping
        callbacks=[transformers.EarlyStoppingCallback(early_stopping_patience=args.early_stopping_patience)] 
        if args.early_stopping_patience > 0 else None,
        # 数据整理器,用于将数据批次转换为模型输入格式
        # mlm=False 表示不使用掩码语言模型预训练,而是使用因果语言模型
        data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
    )

    # 在主进程中打印可训练参数信息
    # 使用print_trainable_parameters()方法显示模型中可训练参数的数量和比例
    if trainer.accelerator.is_main_process:
        trainer.model.print_trainable_parameters()

    # 开始训练
    trainer.train()

    """
    第13步：保存模型
    """
    # 生成当前时间戳,格式为"年月日_时分秒"，用于区分不同时间保存的模型
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    # 获取输出目录路径
    base_output_dir = args.model_dir
    # 根据是否分布式训练确定GPU模式
    gpu_mode = "multi_gpu" if is_distributed else "single_gpu"
    # 根据是否保存合并模型确定模型类型
    model_type = "merged" if args.save_merged_model else "lora"
    
    # 构建完整的输出路径
    output_path = os.path.join(
        base_output_dir,
        gpu_mode, 
        model_type,
        timestamp
    )
    # 创建输出目录,exist_ok=True表示目录已存在时不报错
    os.makedirs(output_path, exist_ok=True)

    # 保存模型
    save_model(
        trainer=trainer,
        tokenizer=tokenizer,
        output_path=output_path,
        save_merged_model=args.save_merged_model,
        model_config=model_config,
        training_config=training_config,
        model_name=model_name,
        timestamp=timestamp,
        gpu_mode=gpu_mode,
        model_type=model_type,
        per_device_train_batch_size=training_args.per_device_train_batch_size,
        learning_rate=training_args.learning_rate,
        num_train_epochs=training_args.num_train_epochs,
        gradient_accumulation_steps=training_args.gradient_accumulation_steps,
        accelerator=accelerator
    )

    """
    第14步：清理和同步
    """
    # 如果是主进程且启用了wandb,结束wandb记录
    if accelerator.is_main_process and wb_token:
        wandb.finish()

    # 等待所有进程完成同步
    accelerator.wait_for_everyone()


def parse_args():
    """
    解析命令行参数

    Args:
        model_name_or_path: 模型名称或路径,默认使用 deepseek-ai/DeepSeek-R1-Distill-Llama-8B
        train_file: 训练数据文件路径
        eval_file: 评估数据文件路径
        epochs: 训练轮数,默认1轮
        train_batch_size: 每个GPU的训练批次大小,默认2
        eval_batch_size: 每个GPU的评估批次大小,默认4
        grad_accum_steps: 梯度累积步数,默认4步
        learning_rate: 学习率,默认2e-4
        save_merged_model: 是否保存合并后的模型,默认False只保存LoRA权重
        model_dir: 模型保存目录,默认当前目录
        wandb_project_name: W&B项目名称,默认"LLM-Finetuning"
        wandb_run_name: W&B实验名称,默认使用时间戳
        logging_steps: 日志记录间隔步数,默认10步
        early_stopping_patience: 早停的耐心值,默认0表示不使用早停
    """
    parser = argparse.ArgumentParser()
    # 模型参数
    parser.add_argument("--model_name_or_path", default="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B", type=str, help="Input directory for training")    
    # 数据参数
    parser.add_argument("--train_file", type=str, help="Input data for training")
    parser.add_argument("--eval_file", type=str, help="Input data for eval")
    # 训练超参数
    parser.add_argument("--epochs", default=1, type=int, help="number of epochs")

    parser.add_argument("--train_batch_size", default=2, type=int, help="training - mini batch size for each gpu/process")
    # 评估批次大小,默认4
    parser.add_argument("--eval_batch_size", default=4, type=int, help="evaluation - mini batch size for each gpu/process")
    # 梯度累积步数,默认4步
    parser.add_argument("--grad_accum_steps", default=4, type=int, help="gradient accumulation steps")
    # 学习率,默认2e-4
    parser.add_argument("--learning_rate", default=2e-4, type=float, help="learning rate")
    # 模型保存参数
    parser.add_argument("--save_merged_model", action="store_true",
                        help="是否保存合并后的模型。指定时保存合并模型，不指定时仅保存LoRA模型")
    # 模型保存目录,默认当前目录
    parser.add_argument("--model_dir", type=str, default="./", help="output directory for model")
    # wandb相关参数
    parser.add_argument("--wandb_project_name", 
                       default="r1-qlora-finetuning",
                       type=str,
                       help="wandb项目名称")
    # wandb实验名称,默认None
    parser.add_argument("--wandb_run_name",
                       default=None,
                       type=str, 
                       help="wandb实验名称")
    # 日志记录间隔步数,默认10步
    parser.add_argument("--logging_steps",
                       default=10,
                       type=int,
                       help="每多少步记录一次日志")
    # 早停的耐心值,默认0表示不使用早停
    parser.add_argument("--early_stopping_patience", 
                       type=int,
                       default=0,
                       help="早停的耐心值，0表示不使用早停")
    args = parser.parse_args()
    
    # 验证数据文件是否存在    
    if not os.path.exists(args.train_file):
        raise FileNotFoundError(f"训练数据文件不存在: {args.train_file}")
    if not os.path.exists(args.eval_file):
        raise FileNotFoundError(f"评估数据文件不存在: {args.eval_file}")
    
    # 如果未指定wandb实验名称,使用时间戳作为默认名称    
    if args.wandb_run_name is None:
        args.wandb_run_name = f"run_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}"
    
    return args


if __name__ == "__main__":
    """
    程序入口点
    解析命令行参数并调用main函数开始训练
    """
    args = parse_args()
    main(args)