diff options
| -rw-r--r-- | balanced_dataloader.py | 126 | ||||
| -rw-r--r-- | train_gee_balanced.py | 219 |
2 files changed, 345 insertions, 0 deletions
diff --git a/balanced_dataloader.py b/balanced_dataloader.py new file mode 100644 index 0000000..a0b89f7 --- /dev/null +++ b/balanced_dataloader.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python3 +""" +平衡的数据加载器 - 确保每个批次包含男女样本 +""" +import torch +from torch.utils.data import Dataset, DataLoader +import random +from typing import List, Dict + +class BalancedGEEDataset(Dataset): + def __init__(self, data: List[Dict]): + self.data = data + # 按性别分组 + self.male_data = [item for item in data if item['gender'] == 'male'] + self.female_data = [item for item in data if item['gender'] == 'female'] + + print(f"📊 数据分布: male={len(self.male_data)}, female={len(self.female_data)}") + + # 确保有足够的数据 + if len(self.male_data) == 0 or len(self.female_data) == 0: + raise ValueError("数据中必须包含男性和女性样本") + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + return self.data[idx] + + def create_balanced_batches(self, batch_size: int, num_batches: int = None): + """创建平衡的批次""" + if batch_size < 2: + raise ValueError("batch_size必须>=2才能保证性别平衡") + + # 每个批次中男女样本的数量 + male_per_batch = batch_size // 2 + female_per_batch = batch_size - male_per_batch + + batches = [] + max_batches = num_batches or (len(self.data) // batch_size) + + for i in range(max_batches): + batch = [] + + # 随机选择男性样本 + male_samples = random.sample(self.male_data, + min(male_per_batch, len(self.male_data))) + batch.extend(male_samples) + + # 随机选择女性样本 + female_samples = random.sample(self.female_data, + min(female_per_batch, len(self.female_data))) + batch.extend(female_samples) + + # 打乱批次内的顺序 + random.shuffle(batch) + batches.append(batch) + + print(f"批次 {i+1}: male={len(male_samples)}, female={len(female_samples)}") + + return batches + +def balanced_collate(batch): + """平衡的collate函数""" + inputs = [item["input"] for item in batch] + genders = [item["gender"] for item in batch] + + # 检查批次平衡性 + male_count = sum(1 for g in genders if g == 'male') + female_count = sum(1 for g in genders if g == 'female') + + print(f"🔍 批次检查: male={male_count}, female={female_count}") + + if male_count == 0: + print("⚠️ 警告: 批次中没有男性样本！") + if female_count == 0: + print("⚠️ 警告: 批次中没有女性样本！") + + return { + "input": inputs, + "gender": genders + } + +def create_balanced_dataloader(data: List[Dict], batch_size: int, num_batches: int = 10): + """创建平衡的数据加载器""" + dataset = BalancedGEEDataset(data) + + if batch_size < 2: + print("⚠️ 警告: batch_size < 2，无法保证性别平衡") + return DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=balanced_collate) + + # 创建平衡批次 + balanced_batches = dataset.create_balanced_batches(batch_size, num_batches) + + # 展平批次为单个数据点 + flat_data = [] + for batch in balanced_batches: + flat_data.extend(batch) + + # 创建新的数据集 + balanced_dataset = BalancedGEEDataset(flat_data) + + return DataLoader(balanced_dataset, batch_size=batch_size, shuffle=False, collate_fn=balanced_collate) + +# 测试函数 +if __name__ == "__main__": + # 测试平衡数据加载器 + import sys + sys.path.append('.') + from dataset.gee_processor import GEEProcessor + + class MockTokenizer: + def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=True): + return messages[0]["content"] + + processor = GEEProcessor(MockTokenizer()) + test_data = processor.create_test_data(num_samples=20) + + print("🧪 测试平衡数据加载器") + dataloader = create_balanced_dataloader(test_data, batch_size=4, num_batches=3) + + for i, batch in enumerate(dataloader): + print(f"\n批次 {i+1}:") + print(f" 输入数量: {len(batch['input'])}") + print(f" 性别: {batch['gender']}") + if i >= 2: # 只测试前3个批次 + break
\ No newline at end of file diff --git a/train_gee_balanced.py b/train_gee_balanced.py new file mode 100644 index 0000000..e9ad2a4 --- /dev/null +++ b/train_gee_balanced.py @@ -0,0 +1,219 @@ +#!/usr/bin/env python3 +""" +GEE训练脚本 - 使用平衡数据加载器 +确保每个批次都包含男女样本 +""" +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 + +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.gee_loss import GEELoss, gender_to_label +from balanced_dataloader import create_balanced_dataloader, balanced_collate + +os.environ.setdefault("NCCL_TIMEOUT", "2700") +os.environ.setdefault("TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC", "2700") + +def parse_args(): + parser = argparse.ArgumentParser() + # GEE相关参数 + parser.add_argument('--lambda_weight', type=float, default=1.0, help='GEE lambda weight') + parser.add_argument('--use_l1', action='store_true', help='Use L1 loss instead of L2') + parser.add_argument('--auto_anneal', action='store_true', help='Use automatic annealing') + + # 简化参数 + 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=2e-5, help='Learning rate') + parser.add_argument('--max_steps', type=int, default=5, 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='gee_balanced', help='Run name') + parser.add_argument('--wandb_project', type=str, default='one-shot-gee', 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') + + 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") + + # 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"🚀 开始平衡版GEE训练 - {args.run_name}") + print(f"📊 批次配置: micro_batch={args.micro_batch_size}, effective_batch={args.effective_batch}") + + # 加载模型 + 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处理器和损失函数 + gee_processor = GEEProcessor(tokenizer) + gee_loss_fn = GEELoss(lambda_weight=args.lambda_weight, use_l1=args.use_l1) + + 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=50) + + # 检查数据平衡性 + 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_balanced_dataloader( + train_data, + batch_size=args.micro_batch_size, + num_batches=args.max_steps * 2 # 确保有足够的批次 + ) + else: + print("❌ 请使用 --use_test_data 进行测试") + return + + optimizer = AdamW(model.parameters(), lr=args.learning_rate) + model, optimizer, train_loader = accelerator.prepare(model, optimizer, train_loader) + + print(f"🎯 训练配置:") + print(f" Lambda权重: {args.lambda_weight}") + print(f" 最大步数: {args.max_steps}") + print(f" 学习率: {args.learning_rate}") + + # 开始训练 + model.train() + initial_entropy_gap = None + + 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: + # 准备输入 + 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 = gee_loss_fn.compute_token_entropy(logits, pad_mask) + H_i = gee_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) + + # 计算GEE损失 + loss, metrics = gee_loss_fn.compute_gee_loss(H_i, gender_labels) + + # 自动退火（可选） + if args.auto_anneal and initial_entropy_gap is None: + initial_entropy_gap = metrics['entropy_gap'] + + if args.auto_anneal and initial_entropy_gap > 0: + current_gap = metrics['entropy_gap'] + anneal_factor = current_gap / initial_entropy_gap + new_lambda = args.lambda_weight * anneal_factor + gee_loss_fn.update_lambda(new_lambda) + metrics['lambda_weight'] = new_lambda + + # 反向传播 + accelerator.backward(loss) + accelerator.clip_grad_norm_(model.parameters(), 1.0) + optimizer.step() + optimizer.zero_grad() + + # 日志记录 + if accelerator.is_main_process: + if step % args.log_steps == 0: + print(f"📈 Step {step} | loss={loss.item():.6f} | " + f"entropy_gap={metrics['entropy_gap']:.6f} | " + f"H_male={metrics['H_male']:.6f} | " + f"H_female={metrics['H_female']:.6f} | " + f"批次性别={batch['gender']}") + wandb.log({"step": step, **metrics}) + + except Exception as e: + print(f"❌ 训练步骤 {step} 出错: {e}") + continue + + if accelerator.is_main_process: + print("🎉 平衡版GEE训练完成!") + wandb.finish() + +if __name__ == "__main__": + main()
\ No newline at end of file |