""" Offline preference training for E3 (multi-positive, deviation-weighted) and E4 (faithful DivPO, unweighted). Same LoRA config, same beta, same data pool -- only pair construction and loss weighting differ. DDPO-STYLE PER-SAMPLE LOSS WEIGHTING (E3) ----------------------------------------- TRL's DPOConfig.loss_weights is per LOSS TYPE (for blending sigmoid+hinge), not per example, so it cannot express "weight this pair by its chosen story's deviation". Reimplementing TRL's _compute_loss to reach `per_sequence_loss` would mean maintaining a fork of a 300-line method across every loss variant. Instead we exploit an identity: with per_device_train_batch_size == 1, the batch loss IS that single example's loss, so loss_i * w_i accumulated over gradient_accumulation_steps is exactly the weighted-DPO gradient, with no TRL surgery at all. Cost is running at batch size 1, which for ~4k short rows on a 4B LoRA is minutes. The weight rides through the collator (TRL drops unknown columns otherwise) and is popped before the model call. Weights are pre-normalized to mean 1.0 in build_pairs.py so the weighting changes RELATIVE emphasis across rows without also rescaling the effective learning rate -- otherwise "DDPO weighting" and "lower LR" would be confounded. """ from __future__ import annotations import argparse import json import os import sys from pathlib import Path import torch import yaml ROOT = Path(__file__).resolve().parent.parent def load_rows(path: Path) -> list[dict]: return [json.loads(l) for l in open(path) if l.strip()] def build_dpo_dataset(rows: list[dict], weighted: bool): from datasets import Dataset from data import SYSTEM_PROMPT recs = [] for r in rows: recs.append({ "prompt": [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": r["prompt"]}], "chosen": [{"role": "assistant", "content": r["chosen"]}], "rejected": [{"role": "assistant", "content": r["rejected"]}], "weight": float(r.get("weight", 1.0)) if weighted else 1.0, }) return Dataset.from_list(recs) def make_weighted_classes(): from trl import DPOTrainer from trl.trainer.dpo_trainer import DataCollatorForPreference class WeightedCollator(DataCollatorForPreference): def __call__(self, features, return_tensors=None): w = [float(f.pop("weight", 1.0)) for f in features] batch = super().__call__(features, return_tensors) batch["weight"] = torch.tensor(w, dtype=torch.float32) return batch class WeightedDPOTrainer(DPOTrainer): """Exact per-sample weighting, valid because batch size is 1.""" def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): w = inputs.pop("weight", None) out = super().compute_loss(model, inputs, return_outputs=return_outputs, num_items_in_batch=num_items_in_batch) if w is None: return out scale = w.to(out[0].device if return_outputs else out.device).mean() if return_outputs: loss, extra = out return loss * scale, extra return out * scale return WeightedCollator, WeightedDPOTrainer def main(): ap = argparse.ArgumentParser() ap.add_argument("--config", required=True) ap.add_argument("--max-steps", type=int, default=None) ap.add_argument("--smoke", action="store_true") args = ap.parse_args() cfg = yaml.safe_load(open(args.config)) name = cfg["name"] + ("-smoke" if args.smoke else "") import wandb from peft import LoraConfig from transformers import AutoTokenizer from trl import DPOConfig, DPOTrainer from trl.trainer.dpo_trainer import DataCollatorForPreference import logbook weighted = bool(cfg["dpo"].get("weighted", False)) pairs_path = ROOT / cfg["dpo"]["pairs"] rows = load_rows(pairs_path) if args.smoke: rows = rows[:64] print(f"[{name}] {len(rows)} preference rows | weighted={weighted}") run = None if cfg.get("wandb", True) and os.environ.get("WANDB_API_KEY"): run = wandb.init(project=os.environ.get("WANDB_PROJECT", "div-grpo"), name=name, config=cfg, reinit=True) tok = AutoTokenizer.from_pretrained(cfg["model"]) ds = build_dpo_dataset(rows, weighted) lora = LoraConfig( r=cfg["lora"]["r"], lora_alpha=cfg["lora"]["alpha"], lora_dropout=cfg["lora"].get("dropout", 0.0), target_modules=cfg["lora"]["target_modules"], task_type="CAUSAL_LM", bias="none", ) out_dir = ROOT / "outputs" / name bs = 1 if weighted else cfg["dpo"].get("per_device_train_batch_size", 2) dcfg = DPOConfig( output_dir=str(out_dir), num_train_epochs=cfg["dpo"].get("epochs", 1), max_steps=args.max_steps or -1, per_device_train_batch_size=bs, gradient_accumulation_steps=cfg["dpo"].get("gradient_accumulation_steps", 8), learning_rate=cfg["dpo"]["learning_rate"], beta=cfg["dpo"].get("beta", 0.1), loss_type=cfg["dpo"].get("loss_type", "sigmoid"), # TRL 1.10 dropped max_prompt_length; max_length bounds prompt+completion # jointly. Sized so NOTHING truncates: prompt ~180 tok + a story capped # by the 600-word gate (~780 tok) = ~960, well inside 1600. max_length=cfg["dpo"].get("max_length", 1600), truncation_mode=cfg["dpo"].get("truncation_mode", "keep_start"), lr_scheduler_type=cfg["dpo"].get("lr_scheduler_type", "cosine"), warmup_steps=cfg["dpo"].get("warmup_steps", 20), # TRL 1.10: no warmup_ratio bf16=True, gradient_checkpointing=True, logging_steps=10, save_strategy="no", remove_unused_columns=not weighted, # keep `weight` alive when weighting report_to=["wandb"] if run else [], run_name=name, seed=cfg.get("seed", 42), ) if weighted: WCollator, WTrainer = make_weighted_classes() collator = WCollator(pad_token_id=tok.pad_token_id or tok.eos_token_id) trainer = WTrainer(model=cfg["model"], args=dcfg, train_dataset=ds, processing_class=tok, peft_config=lora, data_collator=collator) else: trainer = DPOTrainer(model=cfg["model"], args=dcfg, train_dataset=ds, processing_class=tok, peft_config=lora) logbook.note(f"START {name}", f"rows={len(rows)} weighted={weighted} bs={bs} " f"beta={dcfg.beta} epochs={dcfg.num_train_epochs}\n\n" f"```yaml\n{yaml.safe_dump(cfg, sort_keys=False)}```") trainer.train() final = out_dir / "final" trainer.save_model(str(final)) tok.save_pretrained(str(final)) hist = trainer.state.log_history json.dump(hist, open(out_dir / "log_history.json", "w"), indent=1) last = [h for h in hist if "loss" in h] logbook.note(f"DONE {name}", f"adapter: `{final}`\n\nfinal loss: " f"`{last[-1] if last else 'n/a'}`") logbook.checkpoint(f"after {name}") if run: run.finish() return 0 if __name__ == "__main__": sys.exit(main())