DPO script

train_dpo.py

@@ -0,0 +1,112 @@
+# /// script
+# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "datasets", "transformers", "accelerate", "torch"]
+# ///
+"""DPO on offline compile-ok vs compile-fail preferences (compiler-as-judge)."""
+
+import os
+
+from datasets import load_dataset
+from peft import LoraConfig
+from trl import DPOConfig, DPOTrainer
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+DATASET_ID = os.environ.get("DATASET_ID", "gonzalolinares/cpp-compiler-prefs")
+BASE_MODEL = os.environ.get("BASE_MODEL", "gonzalolinares/qwen25-1.5b-cpp-sft")
+FALLBACK_MODEL = os.environ.get("FALLBACK_MODEL", "Qwen/Qwen2.5-1.5B-Instruct")
+HUB_MODEL_ID = os.environ.get("HUB_MODEL_ID", "gonzalolinares/qwen25-1.5b-cpp-dpo")
+OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "qwen25-1.5b-cpp-dpo")
+
+
+def resolve_model(model_id: str) -> str:
+    try:
+        AutoTokenizer.from_pretrained(model_id)
+        return model_id
+    except Exception:
+        print(f"Could not load {model_id}, falling back to {FALLBACK_MODEL}")
+        return FALLBACK_MODEL
+
+
+def main() -> None:
+    model_id = resolve_model(BASE_MODEL)
+    ds = load_dataset(DATASET_ID, split="train")
+    # Normalize to prompt/chosen/rejected text if needed
+    def to_dpo(example):
+        def as_text(x):
+            if isinstance(x, list):
+                # list of chat messages -> last assistant or join
+                parts = []
+                for m in x:
+                    if isinstance(m, dict):
+                        parts.append(f"{m.get('role', '')}: {m.get('content', '')}")
+                    else:
+                        parts.append(str(m))
+                return "\n".join(parts)
+            return str(x)
+
+        prompt = example.get("prompt")
+        chosen = example.get("chosen")
+        rejected = example.get("rejected")
+        # Prefer conversational: if prompt is message list without assistant
+        if isinstance(prompt, list) and prompt and isinstance(prompt[0], dict):
+            # TRL DPO can take conversational format
+            return {
+                "prompt": prompt,
+                "chosen": chosen if isinstance(chosen, list) else [{"role": "assistant", "content": str(chosen)}],
+                "rejected": rejected
+                if isinstance(rejected, list)
+                else [{"role": "assistant", "content": str(rejected)}],
+            }
+        return {
+            "prompt": as_text(prompt),
+            "chosen": as_text(chosen),
+            "rejected": as_text(rejected),
+        }
+
+    ds = ds.map(to_dpo)
+    split = ds.train_test_split(test_size=0.1, seed=42)
+
+    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto")
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    trainer = DPOTrainer(
+        model=model,
+        processing_class=tokenizer,
+        train_dataset=split["train"],
+        eval_dataset=split["test"],
+        peft_config=LoraConfig(
+            r=16,
+            lora_alpha=32,
+            lora_dropout=0.05,
+            bias="none",
+            task_type="CAUSAL_LM",
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+        ),
+        args=DPOConfig(
+            output_dir=OUTPUT_DIR,
+            num_train_epochs=2,
+            per_device_train_batch_size=1,
+            per_device_eval_batch_size=1,
+            gradient_accumulation_steps=8,
+            learning_rate=5e-5,
+            logging_steps=5,
+            eval_strategy="steps",
+            eval_steps=20,
+            max_length=1024,
+            max_prompt_length=512,
+            bf16=True,
+            push_to_hub=True,
+            hub_model_id=HUB_MODEL_ID,
+            report_to="trackio",
+            project="cpp-compiler-rl",
+            run_name="dpo-compiler-prefs",
+        ),
+    )
+    trainer.train()
+    trainer.push_to_hub()
+    print(f"Pushed to {HUB_MODEL_ID}")
+
+
+if __name__ == "__main__":
+    main()