用TRL和LoRA在HH-RLHF上审计偏好偏差并微调DPO模型

import dataclasses
import importlib.util
import inspect
import os
import re
import subprocess
import sys
import warnings
warnings.filterwarnings("ignore", category=UserWarning)
REQUIRED = ["trl>=0.12", "transformers>=4.45", "accelerate", "datasets", "peft", "scikit-learn"]
def ensure_deps():
   """Install in ONE pip call so the resolver picks a mutually compatible set."""
   try:
       import trl
       import transformers
       return False
   except ImportError:
       print("Installing dependencies...")
       subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "-U", *REQUIRED])
       return True
def drop_broken_torchao():
   """Colab ships torchao 0.10.0; peft demands >0.16 and raises rather than skipping.
   Nothing here uses torchao, so removing it is safer than upgrading (an upgrade can
   drag in a torch build that does not match this runtime)."""
   if importlib.util.find_spec("torchao") is None:
       return False
   try:
       from peft.import_utils import is_torchao_available
       is_torchao_available()
       return False
   except ImportError:
       print("Removing incompatible torchao (unused, but peft raises on it)...")
       subprocess.call([sys.executable, "-m", "pip", "uninstall", "-y", "-q", "torchao"])
       return True
   except Exception:
       return False
_installed = ensure_deps()
_removed = drop_broken_torchao() if not _installed else False
if _installed or _removed:
   print("\nEnvironment changed. RESTART THE RUNTIME (Runtime > Restart session), "
         "then run this cell again.")
   raise SystemExit(0)
import numpy as np
import pandas as pd
import torch
import matplotlib.pyplot as plt
from datasets import load_dataset, concatenate_datasets
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, classification_report, roc_auc_score
import transformers
import trl
from trl import DPOConfig, DPOTrainer
def patch_peft_torchao():
   """Belt and braces: if torchao survived the uninstall, stop peft raising on it."""
   try:
       from peft import import_utils
       from peft.tuners.lora import torchao as lora_torchao
   except ImportError:
       return
   try:
       import_utils.is_torchao_available()
   except ImportError as exc:
       print(f"  neutralising peft's torchao check ({exc})")
       import_utils.is_torchao_available = lambda: False
       lora_torchao.is_torchao_available = lambda: False
MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
SUBSETS = ["helpful-base", "helpful-rejection-sampled", "helpful-online", "harmless-base"]
N_TRAIN_PER_SUBSET = 120
N_TEST_PER_SUBSET = 30
MAX_LENGTH = 512
MAX_PROMPT_LENGTH = 256
BETA = 0.1
MAX_STEPS = 30
BATCH_SIZE = 1
GRAD_ACCUM = 8
LEARNING_RATE = 5e-6
WARMUP_RATIO = 0.1
LOGGING_STEPS = 5
USE_LORA = True
N_REWARD_EVAL = 40
SEED = 17
OUTPUT_DIR = "/content/dpo-hh" if os.path.isdir("/content") else "./dpo-hh"
set_seed(SEED)
rng = np.random.default_rng(SEED)
def report_environment():
   from transformers import TrainingArguments
   cuda = torch.cuda.is_available()
   bf16 = bool(cuda and torch.cuda.is_bf16_supported())
   fp16 = bool(cuda and not bf16)
   device = "cuda" if cuda else "cpu"
   print(f"python       : {sys.version.split()[0]}")
   print(f"torch        : {torch.__version__}")
   print(f"transformers : {transformers.__version__}")
   print(f"trl          : {trl.__version__}")
   print(f"Device: {device} | bf16={bf16} | fp16={fp16}")
   if not cuda:
       print("CPU fallback is enabled; training is intentionally shortened.")
   cfg_fields = {f.name for f in dataclasses.fields(DPOConfig)}
   trainer_params = set(inspect.signature(DPOTrainer.__init__).parameters)
   print(f"DPOConfig subclasses TrainingArguments : {issubclass(DPOConfig, TrainingArguments)}")
   print(f"DPOConfig fields                       : {len(cfg_fields)}")
   for probe in ("warmup_ratio", "warmup_steps", "beta", "max_length", "max_prompt_length"):
       where = [c for c, s in (("DPOConfig", cfg_fields), ("DPOTrainer", trainer_params))
                if probe in s]
       print(f"  {probe:<20} -> {', '.join(where) if where else 'NOT ACCEPTED ANYWHERE'}")
   if not issubclass(DPOConfig, TrainingArguments) or "per_device_train_batch_size" not in cfg_fields:
       print("\n!! DPOConfig looks broken. Reinstall in one command, then restart:")
       print("   pip install -U trl transformers accelerate datasets peft")
   return device, bf16, fp16, cfg_fields, trainer_params
DEVICE, BF16, FP16, CFG_FIELDS, TRAINER_PARAMS = report_environment()