""" dpo_trainer.py - DPO (Direct Preference Optimization) fine-tuning. Turns human-edited outputs into (chosen, rejected) preference pairs and trains a policy model to prefer the human-edited versions. MPS limitations apply identically to trainer.py: - float32, no fp16/bf16, gradient_checkpointing=False, dataloader_num_workers=0, device_map not used. """ from __future__ import annotations import logging from pathlib import Path from typing import Optional import torch from datasets import Dataset from peft import LoraConfig, TaskType, get_peft_model from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments from trl import DPOTrainer logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Preference pair construction # --------------------------------------------------------------------------- def build_preference_pairs(corpus_rows: list[dict]) -> list[dict]: """ Build (prompt, chosen, rejected) triples from human-edited corpus rows. Filters: - human_edited must be True - edited_output must be non-empty and differ from output_text Returns a list of dicts with keys: prompt, chosen, rejected, task_type. Immutable input — corpus_rows is never mutated. """ pairs: list[dict] = [] skipped = 0 for row in corpus_rows: if not row.get("human_edited"): skipped += 1 continue prompt = (row.get("input_text") or "").strip() chosen = (row.get("edited_output") or "").strip() rejected = (row.get("output_text") or "").strip() if not prompt or not chosen or not rejected: skipped += 1 continue if chosen == rejected: skipped += 1 continue pairs.append( { "prompt": prompt, "chosen": chosen, "rejected": rejected, "task_type": row.get("task_type", "general"), } ) logger.info( "build_preference_pairs: %d valid pairs, %d skipped", len(pairs), skipped ) return pairs def _pairs_to_dataset(pairs: list[dict]) -> Dataset: """ Convert preference pair dicts to a HuggingFace Dataset. DPOTrainer expects columns: prompt, chosen, rejected. """ records = [ { "prompt": p["prompt"], "chosen": p["chosen"], "rejected": p["rejected"], } for p in pairs ] return Dataset.from_list(records) # --------------------------------------------------------------------------- # Device selection (mirrors trainer.py) # --------------------------------------------------------------------------- def _select_device() -> str: if torch.backends.mps.is_available() and torch.backends.mps.is_built(): return "mps" if torch.cuda.is_available(): return "cuda" return "cpu" def _load_model_and_tokenizer( base_model_path: str, device: str ) -> tuple: logger.info("DPO: loading tokenizer from %s", base_model_path) tokenizer = AutoTokenizer.from_pretrained( base_model_path, trust_remote_code=True, padding_side="right", ) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token logger.info("DPO: loading model from %s on device=%s", base_model_path, device) if device == "cuda": model = AutoModelForCausalLM.from_pretrained( base_model_path, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True, ) else: model = AutoModelForCausalLM.from_pretrained( base_model_path, torch_dtype=torch.float32, trust_remote_code=True, ) model.config.use_cache = False return model, tokenizer # --------------------------------------------------------------------------- # Main DPO training entry point # --------------------------------------------------------------------------- def run_dpo_training( base_model_path: str, pairs: list[dict], output_dir: str, beta: float = 0.1, num_epochs: int = 1, batch_size: int = 1, gradient_accumulation_steps: int = 4, learning_rate: float = 5e-5, max_length: int = 2048, max_prompt_length: int = 512, lora_r: int = 16, lora_alpha: int = 32, lora_dropout: float = 0.05, ) -> dict: """ Run DPO preference learning. Uses LoRA adapters on the base model to keep memory footprint small. The reference model is the frozen base model; the policy model trains on top of the LoRA adapter. Returns: { "train_loss": float, "adapter_path": str, "n_pairs": int, "device": str, "beta": float, } Raises on fatal errors. """ if len(pairs) < 10: raise ValueError( f"Insufficient preference pairs: need >= 10, got {len(pairs)}" ) device = _select_device() logger.info( "run_dpo_training: device=%s beta=%.2f pairs=%d output=%s", device, beta, len(pairs), output_dir, ) dataset = _pairs_to_dataset(pairs) # Split: 90% train, 10% eval split = dataset.train_test_split(test_size=0.1, seed=42) train_dataset = split["train"] eval_dataset = split["test"] model, tokenizer = _load_model_and_tokenizer(base_model_path, device) # Separate reference model (frozen copy of base) ref_model, _ = _load_model_and_tokenizer(base_model_path, device) for param in ref_model.parameters(): param.requires_grad = False # LoRA adapter on the policy model lora_config = LoraConfig( r=lora_r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, bias="none", task_type=TaskType.CAUSAL_LM, target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], inference_mode=False, ) model = get_peft_model(model, lora_config) model.print_trainable_parameters() if device in ("mps", "cpu"): model = model.to(device) ref_model = ref_model.to(device) output_path = Path(output_dir) output_path.mkdir(parents=True, exist_ok=True) training_args = TrainingArguments( output_dir=str(output_path), num_train_epochs=num_epochs, per_device_train_batch_size=batch_size, per_device_eval_batch_size=batch_size, gradient_accumulation_steps=gradient_accumulation_steps, learning_rate=learning_rate, warmup_ratio=0.1, eval_strategy="steps", eval_steps=25, save_strategy="steps", save_steps=50, load_best_model_at_end=True, metric_for_best_model="eval_loss", greater_is_better=False, logging_steps=5, report_to="none", dataloader_num_workers=0, fp16=False, bf16=False, optim="adamw_torch", gradient_checkpointing=False, remove_unused_columns=False, label_names=["labels"], ) trainer = DPOTrainer( model=model, ref_model=ref_model, args=training_args, beta=beta, train_dataset=train_dataset, eval_dataset=eval_dataset, tokenizer=tokenizer, max_length=max_length, max_prompt_length=max_prompt_length, ) logger.info( "Starting DPO training: %d train pairs, %d eval pairs", len(train_dataset), len(eval_dataset), ) train_result = trainer.train() eval_metrics = trainer.evaluate() logger.info("DPO eval metrics: %s", eval_metrics) adapter_path = str(output_path / "dpo_adapter") model.save_pretrained(adapter_path) tokenizer.save_pretrained(adapter_path) logger.info("Saved DPO adapter to %s", adapter_path) return { "train_loss": round(train_result.training_loss, 4), "eval_loss": round(eval_metrics.get("eval_loss", -1.0), 4), "train_runtime": round(train_result.metrics.get("train_runtime", 0.0), 1), "n_pairs": len(pairs), "train_pairs": len(train_dataset), "eval_pairs": len(eval_dataset), "adapter_path": adapter_path, "device": device, "beta": beta, }