完整的模型训练流程

class UILoRATrainer:
    """UI LoRA训练器 - 核心训练类"""
    
    def __init__(self, config_path: str = "config/model_config.yaml"):
        """初始化训练器"""
        with open(config_path, 'r', encoding='utf-8') as f:
            self.config = yaml.safe_load(f)
        
        self.model = None
        self.tokenizer = None
        self.lora_config = None
    
    def setup_model_and_tokenizer(self, model_name: str):
        """设置模型和分词器"""
        logger.info(f"加载模型: {model_name}")
        
        # 加载分词器
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        
        # 加载模型
        self.model = AutoModelForSeq2SeqLM.from_pretrained(
            model_name,
            torch_dtype=torch.float16 if self.config["training"]["fp16"] else torch.float32,
            device_map="auto" if torch.cuda.is_available() else None
        )
        
        # 设置LoRA配置
        self.lora_config = LoraConfig(
            task_type=TaskType.SEQ_2_SEQ_LM,
            r=self.config["lora"]["r"],
            lora_alpha=self.config["lora"]["lora_alpha"],
            lora_dropout=self.config["lora"]["lora_dropout"],
            target_modules=self.config["lora"]["target_modules"]
        )
        
        # 应用LoRA
        self.model = get_peft_model(self.model, self.lora_config)
        
        # 打印可训练参数
        self.model.print_trainable_parameters()

class UIDataset(Dataset):
    """UI数据集类 - 处理训练数据"""
    
    def __init__(self, data_path: str, tokenizer, max_length: int = 512):
        """初始化数据集"""
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.data = self._load_data(data_path)
    
    def _load_data(self, data_path: str) -> List[Dict[str, str]]:
        """加载JSONL格式的训练数据"""
        data = []
        with open(data_path, 'r', encoding='utf-8') as f:
            for line in f:
                sample = json.loads(line.strip())
                data.append(sample)
        return data
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        """获取单个训练样本"""
        sample = self.data[idx]
        
        # 编码输入和输出
        input_text = sample["instruction"]  # 中文Prompt
        output_text = sample["output_json_minified"]  # UI-DSL JSON
        
        # Tokenize输入
        input_encoding = self.tokenizer(
            input_text,
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )
        
        # Tokenize输出
        output_encoding = self.tokenizer(
            output_text,
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )
        
        return {
            "input_ids": input_encoding["input_ids"].squeeze(),
            "attention_mask": input_encoding["attention_mask"].squeeze(),
            "labels": output_encoding["input_ids"].squeeze()
        }

def prepare_datasets(self, train_path: str, val_path: str) -> tuple:
    """准备训练和验证数据集"""
    logger.info("准备训练和验证数据集")
    
    # 创建数据集
    train_dataset = UIDataset(
        train_path, 
        self.tokenizer, 
        self.config["model"]["max_length"]
    )
    
    val_dataset = UIDataset(
        val_path, 
        self.tokenizer, 
        self.config["model"]["max_length"]
    )
    
    return train_dataset, val_dataset

# LoRA配置参数（来自config/model_config.yaml）
lora_config = {
    "r": 16,                    # LoRA秩，控制适配器大小
    "lora_alpha": 32,           # LoRA缩放参数
    "lora_dropout": 0.1,        # LoRA dropout率
    "target_modules": ["q", "v"] # 目标模块，指定要微调的层
}

# LoRA适配器结构
# 原始权重: W (d × k)
# LoRA分解: W = W₀ + ΔW = W₀ + BA
# 其中: B (d × r), A (r × k), r << min(d,k)

class LoRAAdapter:
    """LoRA适配器实现原理"""
    
    def __init__(self, original_dim: int, rank: int):
        self.rank = rank
        # 低秩矩阵分解
        self.lora_A = nn.Parameter(torch.randn(rank, original_dim) * 0.01)
        self.lora_B = nn.Parameter(torch.zeros(original_dim, rank))
        self.scaling = 1.0 / rank
    
    def forward(self, x):
        # LoRA前向传播
        return x + (x @ self.lora_A.T @ self.lora_B.T) * self.scaling

# 选择Transformer的关键模块进行LoRA微调
target_modules = ["q", "v"]  # Query和Value矩阵

# 模块选择策略：
# 1. Query (q): 控制注意力查询，影响模型对输入的理解
# 2. Value (v): 控制注意力值，影响模型输出生成
# 3. 不选择Key (k): 保持原始知识结构
# 4. 不选择输出层: 保持原始输出分布

def setup_training_args(self, output_dir: str) -> TrainingArguments:
    """设置训练参数"""
    return TrainingArguments(
        # 基础配置
        output_dir=output_dir,
        per_device_train_batch_size=self.config["training"]["batch_size"],  # 4
        per_device_eval_batch_size=self.config["training"]["batch_size"],   # 4
        gradient_accumulation_steps=self.config["training"]["gradient_accumulation_steps"],  # 4
        
        # 训练轮数和学习率
        num_train_epochs=self.config["training"]["num_epochs"],  # 3
        learning_rate=float(self.config["training"]["learning_rate"]),  # 2e-4
        
        # 优化器配置
        warmup_steps=self.config["training"]["warmup_steps"],  # 100
        max_grad_norm=self.config["training"]["max_grad_norm"],  # 1.0
        
        # 精度配置
        fp16=self.config["training"]["fp16"],  # True，使用半精度训练
        
        # 数据加载配置
        dataloader_num_workers=self.config["training"]["dataloader_num_workers"],  # 4
        
        # 日志和保存配置
        logging_steps=10,        # 每10步记录一次日志
        eval_steps=100,          # 每100步进行一次验证
        save_steps=500,          # 每500步保存一次模型
        eval_strategy="steps",   # 按步数进行验证
        save_strategy="steps",   # 按步数保存模型
        
        # 模型选择配置
        load_best_model_at_end=True,           # 训练结束时加载最佳模型
        metric_for_best_model="eval_loss",     # 使用验证损失作为最佳模型指标
        greater_is_better=False,               # 损失越小越好
        
        # 其他配置
        report_to=None,                        # 不使用wandb等工具
        remove_unused_columns=False,           # 保留所有列
    )

# 学习率调度策略
learning_rate_schedule = {
    "warmup_steps": 100,        # 预热步数
    "max_learning_rate": 2e-4,  # 最大学习率
    "min_learning_rate": 1e-6,  # 最小学习率
    "decay_type": "cosine"      # 余弦衰减
}

# 学习率调优建议：
# 1. 从2e-4开始，根据收敛情况调整
# 2. 如果训练不稳定，降低到1e-4
# 3. 如果收敛太慢，可以提高到5e-4

# 批次大小配置策略
batch_config = {
    "per_device_batch_size": 4,      # 单设备批次大小
    "gradient_accumulation_steps": 4, # 梯度累积步数
    "effective_batch_size": 16,      # 有效批次大小 = 4 * 4
    "memory_optimization": True      # 内存优化
}

# 批次大小调优建议：
# 1. GPU内存充足时，增加per_device_batch_size
# 2. GPU内存不足时，增加gradient_accumulation_steps
# 3. 保持effective_batch_size在16-32之间

def train(self, train_dataset, val_dataset, output_dir: str):
    """开始训练"""
    logger.info("开始训练")
    
    # 设置训练参数
    training_args = self.setup_training_args(output_dir)
    
    # 数据整理器
    data_collator = DataCollatorForSeq2Seq(
        tokenizer=self.tokenizer,
        model=self.model,
        padding=True
    )
    
    # 创建训练器
    trainer = Trainer(
        model=self.model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        data_collator=data_collator,
        tokenizer=self.tokenizer,
    )
    
    # 开始训练
    trainer.train()
    
    # 保存模型
    trainer.save_model()
    self.tokenizer.save_pretrained(output_dir)
    
    logger.info(f"模型已保存到: {output_dir}")

class TrainingMonitor:
    """训练监控器"""
    
    def __init__(self):
        self.training_history = {
            "train_loss": [],
            "eval_loss": [],
            "learning_rate": [],
            "epoch": []
        }
    
    def log_training_step(self, step: int, loss: float, lr: float):
        """记录训练步骤"""
        self.training_history["train_loss"].append(loss)
        self.training_history["learning_rate"].append(lr)
        
        if step % 10 == 0:
            logger.info(f"Step {step}: Loss={loss:.4f}, LR={lr:.2e}")
    
    def log_evaluation(self, step: int, eval_loss: float):
        """记录验证结果"""
        self.training_history["eval_loss"].append(eval_loss)
        logger.info(f"Evaluation at step {step}: Eval Loss={eval_loss:.4f}")
    
    def plot_training_curves(self, save_path: str):
        """绘制训练曲线"""
        import matplotlib.pyplot as plt
        
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
        
        # 损失曲线
        ax1.plot(self.training_history["train_loss"], label="Train Loss")
        ax1.plot(self.training_history["eval_loss"], label="Eval Loss")
        ax1.set_xlabel("Steps")
        ax1.set_ylabel("Loss")
        ax1.legend()
        ax1.set_title("Training and Validation Loss")
        
        # 学习率曲线
        ax2.plot(self.training_history["learning_rate"])
        ax2.set_xlabel("Steps")
        ax2.set_ylabel("Learning Rate")
        ax2.set_title("Learning Rate Schedule")
        
        plt.tight_layout()
        plt.savefig(save_path)
        plt.close()

def evaluate(self, test_dataset) -> Dict[str, float]:
    """评估模型"""
    logger.info("开始评估")
    
    # 设置评估参数
    eval_args = TrainingArguments(
        output_dir="./eval_output",
        per_device_eval_batch_size=self.config["training"]["batch_size"],
        dataloader_num_workers=self.config["training"]["dataloader_num_workers"],
        remove_unused_columns=False,
    )
    
    # 数据整理器
    data_collator = DataCollatorForSeq2Seq(
        tokenizer=self.tokenizer,
        model=self.model,
        padding=True
    )
    
    # 创建评估器
    trainer = Trainer(
        model=self.model,
        args=eval_args,
        eval_dataset=test_dataset,
        data_collator=data_collator,
        tokenizer=self.tokenizer,
    )
    
    # 评估
    eval_results = trainer.evaluate()
    
    logger.info(f"评估结果: {eval_results}")
    return eval_results

def compute_metrics(eval_pred):
    """计算自定义评估指标"""
    predictions, labels = eval_pred
    
    # 解码预测结果
    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
    
    # 计算BLEU分数
    bleu_score = calculate_bleu(decoded_preds, decoded_labels)
    
    # 计算ROUGE分数
    rouge_score = calculate_rouge(decoded_preds, decoded_labels)
    
    # 计算JSON格式正确率
    json_accuracy = calculate_json_accuracy(decoded_preds)
    
    return {
        "bleu": bleu_score,
        "rouge": rouge_score,
        "json_accuracy": json_accuracy
    }

def calculate_json_accuracy(predictions: List[str]) -> float:
    """计算JSON格式正确率"""
    correct_count = 0
    total_count = len(predictions)
    
    for pred in predictions:
        try:
            json.loads(pred)
            correct_count += 1
        except:
            continue
    
    return correct_count / total_count if total_count > 0 else 0

# 内存优化配置
memory_optimization = {
    "fp16": True,                    # 半精度训练
    "gradient_checkpointing": True,  # 梯度检查点
    "dataloader_pin_memory": False,  # 禁用pin_memory
    "dataloader_num_workers": 0,     # 单进程数据加载
}

# 梯度检查点实现
def enable_gradient_checkpointing(model):
    """启用梯度检查点以节省内存"""
    if hasattr(model, "gradient_checkpointing_enable"):
        model.gradient_checkpointing_enable()
    elif hasattr(model, "enable_gradient_checkpointing"):
        model.enable_gradient_checkpointing()

# 训练稳定性配置
stability_config = {
    "max_grad_norm": 1.0,           # 梯度裁剪
    "warmup_steps": 100,            # 学习率预热
    "weight_decay": 0.01,           # 权重衰减
    "adam_epsilon": 1e-8,           # Adam优化器epsilon
}

# 梯度裁剪实现
def clip_gradients(model, max_norm: float = 1.0):
    """梯度裁剪"""
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

# 学习率调度器配置
scheduler_config = {
    "scheduler_type": "cosine",     # 余弦调度
    "num_warmup_steps": 100,        # 预热步数
    "num_training_steps": 1000,     # 总训练步数
    "min_lr_ratio": 0.1,            # 最小学习率比例
}

# 自定义学习率调度器
class CosineSchedulerWithWarmup:
    """带预热的余弦学习率调度器"""
    
    def __init__(self, optimizer, num_warmup_steps, num_training_steps, min_lr_ratio=0.1):
        self.optimizer = optimizer
        self.num_warmup_steps = num_warmup_steps
        self.num_training_steps = num_training_steps
        self.min_lr_ratio = min_lr_ratio
        self.base_lrs = [group['lr'] for group in optimizer.param_groups]
    
    def step(self, step):
        if step < self.num_warmup_steps:
            # 预热阶段
            lr = self.base_lrs[0] * step / self.num_warmup_steps
        else:
            # 余弦衰减阶段
            progress = (step - self.num_warmup_steps) / (self.num_training_steps - self.num_warmup_steps)
            lr = self.base_lrs[0] * (self.min_lr_ratio + (1 - self.min_lr_ratio) * 0.5 * (1 + math.cos(math.pi * progress)))
        
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="LoRA微调训练")
    parser.add_argument("--model_name", type=str, default="google/flan-t5-base",
                       help="基础模型名称")
    parser.add_argument("--train_file", type=str, required=True,
                       help="训练数据文件路径")
    parser.add_argument("--val_file", type=str, required=True,
                       help="验证数据文件路径")
    parser.add_argument("--test_file", type=str, default=None,
                       help="测试数据文件路径")
    parser.add_argument("--output_dir", type=str, required=True,
                       help="输出目录")
    parser.add_argument("--config", type=str, default="config/model_config.yaml",
                       help="配置文件路径")
    
    args = parser.parse_args()
    
    # 检查CUDA可用性
    if torch.cuda.is_available():
        logger.info(f"使用GPU: {torch.cuda.get_device_name()}")
        logger.info(f"GPU内存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f}GB")
    else:
        logger.warning("未检测到GPU，将使用CPU训练（速度较慢）")
    
    # 创建训练器
    trainer = UILoRATrainer(args.config)
    
    # 设置模型和分词器
    trainer.setup_model_and_tokenizer(args.model_name)
    
    # 准备数据集
    train_dataset, val_dataset = trainer.prepare_datasets(args.train_file, args.val_file)
    
    # 开始训练
    trainer.train(train_dataset, val_dataset, args.output_dir)
    
    # 如果有测试集，进行评估
    if args.test_file:
        test_dataset = UIDataset(
            args.test_file, 
            trainer.tokenizer, 
            trainer.config["model"]["max_length"]
        )
        eval_results = trainer.evaluate(test_dataset)
        
        # 保存评估结果
        eval_path = Path(args.output_dir) / "eval_results.json"
        with open(eval_path, 'w', encoding='utf-8') as f:
            json.dump(eval_results, f, indent=2)
        
        logger.info(f"评估结果已保存到: {eval_path}")

if __name__ == "__main__":
    main()

def analyze_training_logs(log_file: str):
    """分析训练日志"""
    train_losses = []
    eval_losses = []
    
    with open(log_file, 'r') as f:
        for line in f:
            if "train_loss" in line:
                # 提取训练损失
                loss = extract_loss_from_log(line)
                train_losses.append(loss)
            elif "eval_loss" in line:
                # 提取验证损失
                loss = extract_loss_from_log(line)
                eval_losses.append(loss)
    
    # 分析训练趋势
    if len(train_losses) > 0:
        print(f"最终训练损失: {train_losses[-1]:.4f}")
        print(f"训练损失下降趋势: {train_losses[0] - train_losses[-1]:.4f}")
    
    if len(eval_losses) > 0:
        print(f"最终验证损失: {eval_losses[-1]:.4f}")
        print(f"验证损失下降趋势: {eval_losses[0] - eval_losses[-1]:.4f}")

def extract_loss_from_log(log_line: str) -> float:
    """从日志行中提取损失值"""
    import re
    match = re.search(r'loss[=:]\s*([\d.]+)', log_line)
    return float(match.group(1)) if match else 0.0

class TrainingDiagnostics:
    """训练问题诊断器"""
    
    def diagnose_training_issues(self, trainer, eval_results):
        """诊断训练问题"""
        issues = []
        
        # 检查过拟合
        if self.check_overfitting(trainer, eval_results):
            issues.append("检测到过拟合，建议增加正则化或减少训练轮数")
        
        # 检查欠拟合
        if self.check_underfitting(trainer, eval_results):
            issues.append("检测到欠拟合，建议增加训练轮数或提高学习率")
        
        # 检查梯度问题
        if self.check_gradient_issues(trainer):
            issues.append("检测到梯度问题，建议检查学习率或梯度裁剪")
        
        return issues
    
    def check_overfitting(self, trainer, eval_results):
        """检查过拟合"""
        train_loss = trainer.state.log_history[-1].get("train_loss", 0)
        eval_loss = eval_results.get("eval_loss", 0)
        
        # 如果验证损失明显高于训练损失，可能存在过拟合
        return eval_loss > train_loss * 1.5
    
    def check_underfitting(self, trainer, eval_results):
        """检查欠拟合"""
        train_loss = trainer.state.log_history[-1].get("train_loss", 0)
        
        # 如果训练损失仍然很高，可能存在欠拟合
        return train_loss > 1.0
    
    def check_gradient_issues(self, trainer):
        """检查梯度问题"""
        # 检查梯度范数
        total_norm = 0
        for p in trainer.model.parameters():
            if p.grad is not None:
                param_norm = p.grad.data.norm(2)
                total_norm += param_norm.item() ** 2
        total_norm = total_norm ** (1. / 2)
        
        # 如果梯度范数过大或过小，可能存在梯度问题
        return total_norm > 10.0 or total_norm < 0.01

def save_model(self, output_dir: str):
    """保存训练好的模型"""
    # 保存LoRA权重
    self.model.save_pretrained(output_dir)
    
    # 保存分词器
    self.tokenizer.save_pretrained(output_dir)
    
    # 保存训练配置
    config_path = Path(output_dir) / "training_config.json"
    with open(config_path, 'w', encoding='utf-8') as f:
        json.dump(self.config, f, indent=2)
    
    # 保存LoRA配置
    lora_config_path = Path(output_dir) / "lora_config.json"
    with open(lora_config_path, 'w', encoding='utf-8') as f:
        json.dump(self.lora_config.to_dict(), f, indent=2)
    
    logger.info(f"模型已保存到: {output_dir}")

def load_model(self, model_path: str):
    """加载训练好的模型"""
    # 加载基础模型
    base_model = AutoModelForSeq2SeqLM.from_pretrained(
        "google/flan-t5-base",
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        device_map="auto" if torch.cuda.is_available() else None
    )
    
    # 加载LoRA权重
    self.model = PeftModel.from_pretrained(base_model, model_path)
    
    # 加载分词器
    self.tokenizer = AutoTokenizer.from_pretrained(model_path)
    
    # 设置为评估模式
    self.model.eval()
    
    logger.info(f"模型已从 {model_path} 加载")

# 超参数调优建议
hyperparameter_recommendations = {
    "learning_rate": {
        "range": [1e-5, 5e-4],
        "recommended": 2e-4,
        "description": "学习率过高可能导致训练不稳定，过低可能导致收敛缓慢"
    },
    "batch_size": {
        "range": [2, 16],
        "recommended": 4,
        "description": "批次大小影响训练稳定性和内存使用"
    },
    "lora_rank": {
        "range": [8, 64],
        "recommended": 16,
        "description": "LoRA秩影响模型容量和训练效率"
    },
    "num_epochs": {
        "range": [1, 10],
        "recommended": 3,
        "description": "训练轮数需要根据数据集大小和模型收敛情况调整"
    }
}

# 训练环境配置建议
training_environment = {
    "hardware": {
        "gpu": "NVIDIA RTX 4090 (24GB) 或更高",
        "memory": "32GB RAM 或更高",
        "storage": "100GB SSD 或更高"
    },
    "software": {
        "python": "3.10+",
        "pytorch": "2.0+",
        "cuda": "11.8+",
        "transformers": "4.30+"
    },
    "optimization": {
        "mixed_precision": True,
        "gradient_checkpointing": True,
        "dataloader_workers": 4
    }
}

# 训练监控检查点
training_checkpoints = {
    "step_100": "检查初始收敛情况",
    "step_500": "检查训练稳定性",
    "step_1000": "检查过拟合情况",
    "epoch_end": "检查整体训练效果"
}

def training_checkpoint(trainer, step: int):
    """训练检查点"""
    if step == 100:
        # 检查初始收敛
        current_loss = trainer.state.log_history[-1].get("train_loss", 0)
        if current_loss > 2.0:
            logger.warning("初始损失过高，可能需要调整学习率")
    
    elif step == 500:
        # 检查训练稳定性
        recent_losses = [log.get("train_loss", 0) for log in trainer.state.log_history[-10:]]
        if max(recent_losses) - min(recent_losses) > 0.5:
            logger.warning("训练不稳定，建议降低学习率")
    
    elif step == 1000:
        # 检查过拟合
        if "eval_loss" in trainer.state.log_history[-1]:
            train_loss = trainer.state.log_history[-1].get("train_loss", 0)
            eval_loss = trainer.state.log_history[-1].get("eval_loss", 0)
            if eval_loss > train_loss * 1.3:
                logger.warning("可能存在过拟合，建议早停或增加正则化")