import argparse import os import random import time from pathlib import Path import psutil import torch import torch.nn.functional as F from torch.optim import AdamW import pandas as pd import numpy as np from torch.utils.data import Dataset, DataLoader 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 os.environ.setdefault("NCCL_TIMEOUT", "2700") os.environ.setdefault("TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC", "2700") class GEEDataset(Dataset): def __init__(self, data): self.data = data def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx] def custom_collate(batch): return { "input": [item["input"] for item in batch], "gender": [item["gender"] for item in batch] } def parse_args(): parser = argparse.ArgumentParser() # GEE相关参数 parser.add_argument('--lambda_weight', type=float, default=3.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('--bias_eval_steps', type=int, default=10, help='Bias evaluation frequency') parser.add_argument('--balance_dataset', action='store_true', default=True, help='Balance dataset by gender') parser.add_argument('--target_size', type=int, default=None, help='Target dataset size for balancing') # 保留原有参数 parser.add_argument('--model_name', type=str, default='Qwen2.5-Math-7B', help='Model name') parser.add_argument('--model_path', type=str, default=None, help='Local model path') parser.add_argument('--train_data', type=str, default='dataset/1shot_rlvr/pi1_r1280.parquet', help='Training data file path') parser.add_argument('--save_root', type=str, default=None, help='Checkpoint save root directory') parser.add_argument('--effective_batch', type=int, default=64, help='Global batch size') parser.add_argument('--micro_batch_size', type=int, default=2, help='Micro batch size') parser.add_argument('--temperature', type=float, default=0.5, help='Temperature coefficient') parser.add_argument('--learning_rate', type=float, default=2e-5, help='Learning rate') parser.add_argument('--log_steps', type=int, default=1, help='Logging step interval') parser.add_argument('--save_steps', type=int, default=1, help='Checkpoint saving step interval') parser.add_argument('--max_steps', type=int, default=50, help='Maximum training steps') parser.add_argument('--sample_temp', type=float, default=0.5, help='Generation temperature parameter') parser.add_argument('--run_name', type=str, default='one_shot_gee', help='Experiment run name') parser.add_argument('--wandb_project', type=str, default='one-shot-gee', help='W&B project name') parser.add_argument('--wandb_name', type=str, default=None, help='W&B run name') parser.add_argument('--seed', type=int, default=15, help='Random seed') parser.add_argument('--use_test_data', action='store_true', help='Use synthetic test data instead of real data') return parser.parse_args() def apply_chat_template(tokenizer, problem: str) -> str: return tokenizer.apply_chat_template( [{"role": "user", "content": problem}], tokenize=False, add_generation_prompt=True ) def main(): args = parse_args() set_seed(args.seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False world_size = int(os.getenv("WORLD_SIZE", "1")) micro_bs = args.micro_batch_size eff_bs = args.effective_batch accum_steps = max(1, eff_bs // (micro_bs * world_size)) temp = args.temperature lr = args.learning_rate save_root = args.save_root or (f"checkpoints/{args.model_name}/{args.run_name}" if args.run_name else f"checkpoints/{args.model_name}") ds_config = { "train_micro_batch_size_per_gpu": micro_bs, "train_batch_size": eff_bs, "gradient_accumulation_steps": accum_steps, "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=accum_steps, deepspeed_plugin=DeepSpeedPlugin(hf_ds_config=ds_config)) print = accelerator.print model_path = args.model_path or f"/volume/pt-train/models/{args.model_name}" config = AutoConfig.from_pretrained(model_path) config.use_cache = False model = AutoModelForCausalLM.from_pretrained(model_path, config=config) model.gradient_checkpointing_enable() tokenizer = AutoTokenizer.from_pretrained(model_path, padding_side="left") tokenizer.pad_token = tokenizer.eos_token if tokenizer.pad_token is None else tokenizer.pad_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 or args.wandb_name or args.model_name, config=vars(args)) # 准备数据 if args.use_test_data: print("使用合成测试数据...") train_data = gee_processor.create_test_data(num_samples=200) else: print("使用真实数据...") train_data = gee_processor.prepare_gee_data( args.train_data, balance=args.balance_dataset, target_size=args.target_size ) train_loader = DataLoader( GEEDataset(train_data), batch_size=micro_bs, shuffle=True, collate_fn=custom_collate ) optimizer = AdamW(model.parameters(), lr=lr) model, optimizer, train_loader = accelerator.prepare(model, optimizer, train_loader) initial_entropy_gap = None model.train() for step, batch in enumerate(train_loader, start=1): if step > args.max_steps: print(f"Exceed max step {args.max_steps}, training stopped.") break with accelerator.accumulate(model): # 准备输入 inputs = tokenizer( batch["input"], return_tensors="pt", padding="longest", truncation=True, max_length=2048 ).to(accelerator.device) # 生成回答 with torch.no_grad(): gen_ids = accelerator.unwrap_model(model).generate( **inputs, max_new_tokens=512, 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)[:, :4096] 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}") wandb.log({"step": step, **metrics}) if step % args.save_steps == 0: ckpt = Path(save_root) / 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"Checkpoint saved to {ckpt}") if accelerator.is_main_process: final = Path(save_root) / "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 checkpoint saved to {final}") wandb.finish() if __name__ == "__main__": main()