Support loss_type/loss_weights DPO

Author: BrownianNotion

src/axolotl/core/trainers/dpo/__init__.py

@@ -22,6 +22,13 @@ def set_training_args_kwargs(cls, cfg):
         training_args_kwargs = {}
         if cfg.rl is RLType.IPO:
             training_args_kwargs["loss_type"] = ["ipo"]
+        else:
+            if cfg.dpo_loss_type is not None:
+                training_args_kwargs["loss_type"] = cfg.dpo_loss_type
+
+            if cfg.dpo_loss_weights is not None:
+                training_args_kwargs["loss_weights"] = cfg.dpo_loss_weights
+
         # Label smoothing is not compatible with IPO
         if cfg.rl is RLType.DPO and cfg.dpo_label_smoothing:
             training_args_kwargs["label_smoothing"] = cfg.dpo_label_smoothing

src/axolotl/utils/schemas/config.py

@@ -302,6 +302,16 @@ class AxolotlInputConfig(
 
     dpo_padding_free: bool | None = None
 
+    dpo_loss_type: list[str] | None = Field(
+        default=None,
+        json_schema_extra={"description": "List of DPO losses to use."},
+    )
+
+    dpo_loss_weights: list[float] | None = Field(
+        default=None,
+        json_schema_extra={"description": "Weights for each DPO loss."},
+    )
+
     datasets: (
         Annotated[
             list[

src/axolotl/utils/schemas/validation.py

@@ -770,6 +770,20 @@ def check_gdpo(cls, data):
             )
         return data
 
+    @model_validator(mode="before")
+    @classmethod
+    def check_dpo(cls, data):
+        if data.get("rl") == "dpo":
+            loss_types = data.get("dpo_loss_type")
+            loss_weights = data.get("dpo_loss_weights")
+
+            if loss_types and loss_weights and len(loss_types) != len(loss_weights):
+                raise ValueError(
+                    f"`dpo_loss_type` and `dpo_loss_weights` must be the same length, "
+                    f"but got {len(loss_types)} losses and {len(loss_weights)} weights"
+                )
+        return data
+
 
 class OptimizationValidationMixin:
     """Validation methods related to optimization and performance."""

tests/core/test_builders.py

@@ -96,6 +96,8 @@ def fixture_dpo_cfg(base_cfg):
             "dpo_use_weighting": True,
             "dpo_label_smoothing": 0.1,
             "beta": 0.1,  # DPO beta
+            "dpo_loss_type": ["sigmoid", "sft"],
+            "dpo_loss_weights": [1.0, 0.5],
         }
     )
     return cfg
@@ -298,6 +300,8 @@ def test_dpo_training_arguments(self, dpo_cfg, model, tokenizer):
         assert hasattr(training_arguments, "use_weighting")
         assert training_arguments.use_weighting is True
         assert training_arguments.label_smoothing == 0.1
+        assert training_arguments.loss_type == ["sigmoid", "sft"]
+        assert training_arguments.loss_weights == [1.0, 0.5]
 
     def test_orpo_training_arguments(self, orpo_cfg, model, tokenizer):
         builder = HFRLTrainerBuilder(orpo_cfg, model, tokenizer)

tests/e2e/test_dpo.py

@@ -116,6 +116,58 @@ def test_dpo_use_weighting(self, temp_dir):
         train(cfg=cfg, dataset_meta=dataset_meta)
         check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
 
+    @with_temp_dir
+    def test_rpo(self, temp_dir):
+        # For TRL >= 0.29, loss_type=["sigmoid", "sft"], loss_weights=[1, alpha]
+        # replaces loss_type="rpo", rpo_alpha=alpha.
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
+                "sequence_len": 1024,
+                "load_in_8bit": True,
+                "adapter": "lora",
+                "lora_r": 64,
+                "lora_alpha": 32,
+                "lora_dropout": 0.1,
+                "lora_target_linear": True,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "rl": "dpo",
+                "dpo_loss_type": ["sigmoid", "sft"],
+                "dpo_loss_weights": [1.0, 1.0],
+                "datasets": [
+                    {
+                        "path": "arcee-ai/distilabel-intel-orca-dpo-pairs-binarized",
+                        "type": "chatml.ultra",
+                        "split": "train",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "paged_adamw_8bit",
+                "lr_scheduler": "cosine",
+                "max_steps": 20,
+                "save_steps": 10,
+                "warmup_steps": 5,
+                "gradient_checkpointing": True,
+                "gradient_checkpointing_kwargs": {"use_reentrant": True},
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        cli_args = TrainerCliArgs()
+        dataset_meta = load_preference_datasets(cfg=cfg, cli_args=cli_args)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
+
     @pytest.mark.skip("kto_pair no longer supported in trl")
     @with_temp_dir
     def test_kto_pair_lora(self, temp_dir):