#!/usr/bin/env python3 """ 纯偏见减少训练脚本 目标：只最小化男女间熵差，不进行整体熵最小化 """ import argparse import os import torch import torch.nn.functional as F from torch.optim import AdamW import pandas as pd import numpy as np from pathlib import Path import wandb from accelerate import Accelerator, DeepSpeedPlugin from accelerate.utils import set_seed from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM # 导入自定义模块 import sys sys.path.append('.') from dataset.gee_processor import GEEProcessor from losses.debiasing_loss import DebiasingLoss, gender_to_label from smart_balanced_dataloader import create_smart_balanced_dataloader os.environ.setdefault("NCCL_TIMEOUT", "2700") os.environ.setdefault("TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC", "2700") def parse_args(): parser = argparse.ArgumentParser() # Debiasing相关参数 parser.add_argument('--scale_factor', type=float, default=1.0, help='Debiasing loss scale factor') parser.add_argument('--use_l1', action='store_true', help='Use L1 loss instead of L2') parser.add_argument('--target_gap', type=float, default=0.01, help='Target entropy gap for early stopping') # 模型参数 parser.add_argument('--model_name', type=str, default='Qwen2.5-Math-1.5B-Instruct', help='Model name') parser.add_argument('--model_path', type=str, required=True, help='Model path') parser.add_argument('--effective_batch', type=int, default=4, help='Global batch size') parser.add_argument('--micro_batch_size', type=int, default=2, help='Micro batch size (must >=2 for balance)') parser.add_argument('--learning_rate', type=float, default=1e-5, help='Learning rate (lower for debiasing)') parser.add_argument('--max_steps', type=int, default=20, help='Maximum training steps') parser.add_argument('--sample_temp', type=float, default=0.7, help='Generation temperature') # 运行参数 parser.add_argument('--run_name', type=str, default='pure_debiasing', help='Run name') parser.add_argument('--wandb_project', type=str, default='debiasing-only', help='W&B project name') parser.add_argument('--use_test_data', action='store_true', help='Use synthetic test data') parser.add_argument('--seed', type=int, default=42, help='Random seed') parser.add_argument('--log_steps', type=int, default=1, help='Logging frequency') parser.add_argument('--save_steps', type=int, default=10, help='Save frequency') return parser.parse_args() def main(): args = parse_args() set_seed(args.seed) # 强制检查batch_size if args.micro_batch_size < 2: print("❌ 错误: micro_batch_size必须>=2才能保证性别平衡！") print("请使用: --micro_batch_size 2 或更大") return # DeepSpeed配置 ds_config = { "train_micro_batch_size_per_gpu": args.micro_batch_size, "train_batch_size": args.effective_batch, "gradient_accumulation_steps": max(1, args.effective_batch // args.micro_batch_size), "bf16": {"enabled": True}, "zero_optimization": { "stage": 2, "offload_optimizer": {"device": "cpu"}, "offload_param": {"device": "none"} }, "gradient_clipping": 1.0, } accelerator = Accelerator( mixed_precision="bf16", gradient_accumulation_steps=max(1, args.effective_batch // args.micro_batch_size), deepspeed_plugin=DeepSpeedPlugin(hf_ds_config=ds_config) ) print = accelerator.print print(f"🎯 开始纯偏见减少训练 - {args.run_name}") print(f"📊 配置信息:") print(f" 批次大小: micro={args.micro_batch_size}, effective={args.effective_batch}") print(f" 缩放因子: {args.scale_factor}") print(f" 目标熵差: {args.target_gap}") print(f" 学习率: {args.learning_rate}") print(f" ⚠️ 注意: 只进行偏见减少，不进行熵最小化") # 加载模型 model_path = args.model_path config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) config.use_cache = False model = AutoModelForCausalLM.from_pretrained(model_path, config=config, trust_remote_code=True) model.gradient_checkpointing_enable() tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token # 初始化纯偏见减少损失函数 gee_processor = GEEProcessor(tokenizer) debiasing_loss_fn = DebiasingLoss(use_l1=args.use_l1, scale_factor=args.scale_factor) if accelerator.is_main_process: wandb.init(project=args.wandb_project, name=args.run_name, config=vars(args)) # 准备数据 if args.use_test_data: print("📊 使用合成测试数据...") train_data = gee_processor.create_test_data(num_samples=100) # 检查数据平衡性 male_count = sum(1 for item in train_data if item['gender'] == 'male') female_count = sum(1 for item in train_data if item['gender'] == 'female') print(f"原始数据: male={male_count}, female={female_count}") # 创建智能平衡的数据加载器 train_loader = create_smart_balanced_dataloader( train_data, batch_size=args.micro_batch_size, num_batches=args.max_steps + 5 ) else: print("❌ 请使用 --use_test_data 进行测试") return optimizer = AdamW(model.parameters(), lr=args.learning_rate) model, optimizer = accelerator.prepare(model, optimizer) print(f"🎯 开始训练...") print(f" 目标: 最小化 |H_female - H_male|") print(f" 理想结果: entropy_gap → {args.target_gap}") # 开始训练 model.train() best_gap = float('inf') converged_steps = 0 for step, batch in enumerate(train_loader, start=1): if step > args.max_steps: print(f"🛑 达到最大步数 {args.max_steps}，训练结束") break with accelerator.accumulate(model): try: # 验证批次平衡性 male_count = sum(1 for g in batch["gender"] if g == 'male') female_count = sum(1 for g in batch["gender"] if g == 'female') if male_count == 0 or female_count == 0: print(f"💥 Step {step}: 批次不平衡！male={male_count}, female={female_count}") continue # 准备输入 inputs = tokenizer( batch["input"], return_tensors="pt", padding="longest", truncation=True, max_length=1024 ).to(accelerator.device) # 生成回答 with torch.no_grad(): gen_ids = accelerator.unwrap_model(model).generate( **inputs, max_new_tokens=128, do_sample=True, top_p=0.95, temperature=args.sample_temp, synced_gpus=True, repetition_penalty=1.15, pad_token_id=tokenizer.pad_token_id, use_cache=False ) # 准备完整序列 seq = torch.cat([inputs.input_ids, gen_ids[:, inputs.input_ids.shape[1]:]], dim=1) pad_mask = seq.ne(tokenizer.pad_token_id) prompt_lengths = pad_mask[:, :inputs.input_ids.shape[1]].sum(-1) # 计算logits和熵 logits = model(seq, attention_mask=pad_mask).logits H_tok = debiasing_loss_fn.compute_token_entropy(logits, pad_mask) H_i = debiasing_loss_fn.compute_sample_entropy(H_tok, prompt_lengths) # 准备性别标签 gender_labels = torch.tensor([ gender_to_label(g) for g in batch["gender"] ], device=accelerator.device) # 计算纯偏见减少损失 loss, metrics = debiasing_loss_fn.compute_debiasing_loss(H_i, gender_labels) # 反向传播 accelerator.backward(loss) accelerator.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() optimizer.zero_grad() # 检查收敛 current_gap = metrics['entropy_gap'] if current_gap < best_gap: best_gap = current_gap converged_steps = 0 else: converged_steps += 1 # 早停检查 if current_gap <= args.target_gap: print(f"🎉 达到目标熵差 {args.target_gap}！当前: {current_gap:.6f}") print(f"在第 {step} 步实现目标") break # 日志记录 if accelerator.is_main_process: if step % args.log_steps == 0: gap_direction = "📉" if current_gap <= best_gap else "📈" print(f"{gap_direction} Step {step} | loss={loss.item():.6f} | " f"gap={current_gap:.6f} | " f"H_male={metrics['H_male']:.6f} | " f"H_female={metrics['H_female']:.6f} | " f"批次[{male_count}M,{female_count}F]") # 添加进度指标 progress = max(0, min(1, (0.1 - current_gap) / 0.1)) * 100 # 假设从0.1开始 metrics['progress_percent'] = progress metrics['best_gap'] = best_gap wandb.log({"step": step, **metrics}) # 保存检查点 if accelerator.is_main_process and step % args.save_steps == 0: ckpt = Path(f"checkpoints/{args.model_name}/{args.run_name}") / f"step_{step}" ckpt.mkdir(parents=True, exist_ok=True) accelerator.unwrap_model(model).save_pretrained(ckpt, safe_serialization=True) tokenizer.save_pretrained(ckpt) print(f"💾 检查点已保存: {ckpt}") except Exception as e: print(f"❌ 训练步骤 {step} 出错: {e}") continue if accelerator.is_main_process: print(f"\n🎉 纯偏见减少训练完成!") print(f"📊 最终结果:") print(f" 最佳熵差: {best_gap:.6f}") print(f" 目标熵差: {args.target_gap}") if best_gap <= args.target_gap: print("✅ 成功达到偏见减少目标！") elif best_gap <= 0.05: print("⚠️ 偏见显著减少，接近目标") else: print("❌ 偏见减少效果有限，可能需要更多训练步数或调整参数") # 保存最终模型 final = Path(f"checkpoints/{args.model_name}/{args.run_name}") / "final" final.mkdir(parents=True, exist_ok=True) accelerator.unwrap_model(model).save_pretrained(final, safe_serialization=True) tokenizer.save_pretrained(final) print(f"💾 最终模型已保存: {final}") wandb.finish() if __name__ == "__main__": main()