src/train_dpo.py

"""
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())