"""
模型训练的配置参数
"""

class TrainingConfig:
    """训练相关的配置参数"""
    
    def __init__(self):
        # LoRA 配置
        self.lora_r = 8                # LoRA秩,控制可训练参数数量，建议范围4~32
        self.lora_alpha = 16           # LoRA缩放因子，用于缩放LoRA产生的更新值
        self.lora_dropout = 0.1        # LoRA dropout率，防止过拟合，建议范围0.1~0.3
        self.lora_target_modules = "all-linear"  # 指定应用LoRA的模块，应用于所有线性层
        # "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"
        
        # 模型训练配置
        self.torch_dtype = "bf16"      # 训练精度: fp16, bf16, fp32
        self.gradient_checkpointing = True  # 启用梯度检查点
        self.use_flash_attention = True      # 是否使用Flash Attention
        self.seed = 666                      # 随机种子
        
        # 4bit量化配置
        self.bnb_config = {
            "load_in_4bit": True,  # 是否启用4bit量化加载模型
            "bnb_4bit_use_double_quant": True,  # 是否使用双重量化以进一步压缩内存
            "bnb_4bit_quant_type": "nf4",  # 量化类型,nf4表示normal float 4-bit
            "bnb_4bit_compute_dtype": "bfloat16",  # 计算时使用的数据类型
            "llm_int8_skip_modules": None,  # None表示不跳过任何模块
            "llm_int8_threshold": 6.0,  # int8量化的阈值
            "llm_int8_has_fp16_weight": False  # 是否保留fp16权重
        }
        
        # GPU相关配置
        self.gpu_memory_threshold = 24 * 1024 * 1024 * 1024  # 24GB
        self.large_gpu_batch_size = 16       # 默认为 16
        
        # 添加内存优化配置
        self.max_memory = {0: "21GB"}       # 限制 GPU 0 的最大内存使用
        self.offload_folder = "offload"     # 设置模型权重卸载目录
        
        # 优化器配置字典
        self.optimizer_config = {
            "lr": 2e-4, # 学习率
            "betas": (0.9, 0.999), # Adam优化器的beta参数
            "eps": 1e-8, # 数值稳定性的小值
            "weight_decay": 0.01 # 权重衰减系数,用于L2正则化
        }

class ModelConfig:
    """
    模型相关的配置参数类
    """
    
    def __init__(self):
        # 模型缓存目录路径,用于存储下载的模型文件
        self.cache_dir = "/tmp/.cache"
        # 临时模型目录路径,用于存储中间模型文件
        self.temp_model_dir = "/tmp/model"
        # 模型分片的最大大小,超过此大小将分片存储
        self.model_max_shard_size = "2GB"
        
        # 模型保存配置
        self.safe_serialization = True # 是否使用安全序列化方式保存模型 
        self.save_strategy = "no"      # 模型保存策略,"no"表示不保存,"steps"表示按步数保存
        self.save_merged_model = False  # 是否保存合并后的完整模型,False仅保存LoRA权重

class DataConfig:
    """数据处理相关的配置参数"""
    
    def __init__(self):
        # 数据处理配置
        self.max_seq_length = 2048     # 最大序列长度
        self.pad_to_max_length = True  # 是否填充到最大长度
        self.truncation = True         # 是否截断过长序列
