Feat(mlx): Add ORPO (loss_type='orpo') for text models to MLXTrainer

tests/mlx_simulation/mlx_nn_stub.py

@@ -407,6 +407,14 @@ def _init(shape, dtype=torch.float32):
 init_module.uniform = _init_uniform
 
 
+# ---------------------------------------------------------------------------
+# Activations
+# ---------------------------------------------------------------------------
+def log_sigmoid(a, **kw):
+    """mlx.nn.log_sigmoid(x) = log(sigmoid(x)); route to torch's stable impl."""
+    return F.logsigmoid(a)
+
+
 # ---------------------------------------------------------------------------
 # Module-level __getattr__: any unknown nn.X returns _Noop.
 # ---------------------------------------------------------------------------

tests/test_mlx_orpo_odds_ratio_stability.py

@@ -0,0 +1,100 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""Regression test for ORPO odds-ratio numerical stability in float16.
+
+The odds-ratio term uses log(1 - p) = log(-expm1(logp)) with a 1e-12 floor.
+In float16 that floor underflows to 0.0, so a perfectly-predicted response row
+(logp -> 0, e.g. an empty response span) makes mx.log(0.0) = -inf and yields
+NaN gradients. The fix computes the term in float32 (mirroring TRL's
+ORPOTrainer.odds_ratio_loss upcast). This test pins that a float16 logp row
+with an exact 0.0 stays finite.
+"""
+
+from __future__ import annotations
+
+import math
+import sys
+
+import pytest
+
+
+@pytest.fixture(autouse=True, scope="module")
+def _install_shim():
+    shim_prefixes = ("mlx", "mlx_lm", "mlx_vlm")
+    real_mlx_modules = {
+        name: module
+        for name, module in sys.modules.items()
+        if any(name == prefix or name.startswith(f"{prefix}.") for prefix in shim_prefixes)
+    }
+    from mlx_simulation import simulate_mlx_on_torch
+    from mlx_simulation.mlx_stub import _MLXFinder
+    simulate_mlx_on_torch()
+    for name in list(sys.modules):
+        if name == "unsloth_zoo.mlx" or name.startswith("unsloth_zoo.mlx."):
+            sys.modules.pop(name, None)
+    yield
+    for name in list(sys.modules):
+        if (
+            name == "unsloth_zoo.mlx" or name.startswith("unsloth_zoo.mlx.")
+            or any(name == prefix or name.startswith(f"{prefix}.") for prefix in shim_prefixes)
+        ):
+            sys.modules.pop(name, None)
+    sys.meta_path[:] = [
+        finder for finder in sys.meta_path
+        if not isinstance(finder, _MLXFinder)
+    ]
+    sys.modules.update(real_mlx_modules)
+
+
+def test_naive_float16_odds_ratio_underflows_to_inf():
+    """Document the bug: the pre-fix float16 expression is non-finite on logp=0."""
+    import mlx.core as mx
+
+    logp_c = mx.array([0.0], dtype=mx.float16)  # perfectly-predicted chosen row
+    # Pre-fix formulation: floor built in the input (float16) dtype.
+    val_c = mx.maximum(-mx.expm1(logp_c), mx.array(1e-12, logp_c.dtype))
+    assert float(val_c.astype(mx.float32)[0]) == 0.0  # 1e-12 underflowed to 0.0
+    assert not math.isfinite(float(mx.log(val_c).astype(mx.float32)[0]))
+
+
+def test_orpo_odds_ratio_loss_finite_on_perfect_float16_row():
+    """The fixed helper stays finite on a float16 logp row containing 0.0."""
+    import mlx.core as mx
+    from unsloth_zoo.mlx.utils import _orpo_odds_ratio_loss
+
+    # Row 0: perfectly-predicted chosen response (logp -> 0). Row 1: ordinary.
+    logp_c = mx.array([0.0, -0.5], dtype=mx.float16)
+    logp_r = mx.array([-1.0, -2.0], dtype=mx.float16)
+    or_loss = _orpo_odds_ratio_loss(logp_c, logp_r)
+    val = float(or_loss.astype(mx.float32))
+    assert math.isfinite(val), f"ORPO odds-ratio loss not finite: {val}"
+    assert or_loss.dtype == mx.float32
+
+
+def test_orpo_odds_ratio_loss_matches_float32_reference():
+    """float16 inputs give the same result as float32 inputs (upcast parity)."""
+    import mlx.core as mx
+    from unsloth_zoo.mlx.utils import _orpo_odds_ratio_loss
+
+    logp_c16 = mx.array([-0.2, -0.5, -1.5], dtype=mx.float16)
+    logp_r16 = mx.array([-0.9, -2.0, -0.3], dtype=mx.float16)
+    logp_c32 = logp_c16.astype(mx.float32)
+    logp_r32 = logp_r16.astype(mx.float32)
+    a = float(_orpo_odds_ratio_loss(logp_c16, logp_r16).astype(mx.float32))
+    b = float(_orpo_odds_ratio_loss(logp_c32, logp_r32).astype(mx.float32))
+    assert math.isfinite(a) and math.isfinite(b)
+    assert abs(a - b) < 1e-2

unsloth_zoo/mlx/trainer.py

@@ -50,6 +50,7 @@
 import mlx.nn as nn
 import mlx.optimizers as optim
 from mlx.utils import tree_flatten, tree_map, tree_reduce, tree_unflatten
+from mlx_lm.tuner.lora import LoRALinear
 
 _PAD_MULTIPLE = 32
 SUPPORTED_MLX_OPTIMIZERS = ("adafactor", "adamw", "adam", "sgd", "muon", "lion")
@@ -169,6 +170,9 @@ def _with_input_ids(self, item):
     make_vlm_cce_loss_fn,
     make_vlm_baseline_loss_fn,
     create_batches,
+    create_preference_batches,
+    make_orpo_loss_fn,
+    make_dpo_loss_fn,
     create_ordered_batches,
     iterate_training_batches,
     create_vlm_batches,
@@ -577,6 +581,10 @@ class MLXTrainingConfig:
 
     # Eval
     eval_steps: int = 0  # 0 = disabled
+    loss_type: str = "sft"  # "sft" or "orpo"
+    orpo_beta: float = 0.1  # ORPO odds-ratio weight (TRL default)
+    dpo_beta: float = 0.1  # DPO beta (TRL default)
+    reference_free: bool = False  # DPO: drop the reference term if True
     load_best_model_at_end: bool = False
     metric_for_best_model: str = "eval_loss"
     greater_is_better: bool = False
@@ -666,6 +674,20 @@ def __init__(self, *args, **kwargs):
         )
 
 
+@dataclass
+class MLXORPOConfig(MLXTrainingConfig):
+    """ORPO config mirroring TRL's ORPOConfig. Presets loss_type='orpo';
+    tune orpo_beta (inherited). Use with MLXORPOTrainer."""
+    loss_type: str = "orpo"
+
+
+@dataclass
+class MLXDPOConfig(MLXTrainingConfig):
+    """DPO config mirroring TRL's DPOConfig. Presets loss_type='dpo';
+    tune dpo_beta / reference_free (inherited). Use with MLXDPOTrainer."""
+    loss_type: str = "dpo"
+
+
 class MLXTrainer:
     """MLX-native trainer for Apple Silicon, mirroring SFTTrainer's constructor API."""
 
@@ -2014,7 +2036,20 @@ def _main_print(*print_args, **print_kwargs):
                 )
                 _main_print("Unsloth: Using VLM standard cross-entropy loss.")
         else:
-            if use_cce:
+            if getattr(args, "loss_type", "sft") == "orpo":
+                _ob = getattr(args, "orpo_beta", 0.1)
+                loss_fn = make_orpo_loss_fn(beta=_ob)
+                print("Unsloth: Using ORPO loss (beta=" + str(_ob) + ").")
+            elif getattr(args, "loss_type", "sft") == "dpo":
+                _db = getattr(args, "dpo_beta", 0.1)
+                _rf = bool(getattr(args, "reference_free", False))
+                _lora_mods = [mod for _, mod in tree_flatten(
+                    model, is_leaf=lambda x: isinstance(x, LoRALinear))
+                    if isinstance(mod, LoRALinear)]
+                loss_fn = make_dpo_loss_fn(beta=_db, lora_mods=_lora_mods, reference_free=_rf)
+                print("Unsloth: Using DPO loss (beta=" + str(_db) +
+                      (", reference_free" if _rf else "") + ").")
+            elif use_cce:
                 loss_fn = make_cce_loss_fn(model)
                 cce_backend = getattr(loss_fn, "_unsloth_cce_backend", "unknown")
                 _main_print(
@@ -2584,7 +2619,10 @@ def _ddp_compiled_step_fn(batch_data, prev_state, do_update):
         eval_batches = None
         text_completion_only_loss = _text_completion_only_loss_arg(args)
         text_assistant_only_loss = _text_assistant_only_loss_arg(args)
-        if args.eval_steps > 0 and self.eval_dataset is not None:
+        if (getattr(args, "loss_type", "sft") in ("orpo", "dpo")
+                and args.eval_steps > 0 and self.eval_dataset is not None):
+            print(f"Unsloth: eval is not yet supported for {args.loss_type}; skipping eval.")
+        elif args.eval_steps > 0 and self.eval_dataset is not None:
             eval_batch_size = (
                 getattr(args, "per_device_eval_batch_size", None)
                 or args.per_device_train_batch_size
@@ -3274,6 +3312,27 @@ def _prepare_data(self, is_vlm):
         text_assistant_only_loss = _text_assistant_only_loss_arg(args)
         comm_group = self.distributed_world
 
+        if getattr(args, "loss_type", "sft") in ("orpo", "dpo"):
+            if is_vlm:
+                raise ValueError(
+                    f"{args.loss_type.upper()} is not yet supported for VLM models on MLX."
+                )
+            batches = create_preference_batches(
+                dataset=self.train_dataset,
+                tokenizer=self.tokenizer,
+                batch_size=args.per_device_train_batch_size,
+                max_seq_length=args.max_seq_length,
+                num_batches=total_batches_needed,
+                dataset_order=(
+                    "sequential"
+                    if getattr(args, "preserve_dataset_order", False)
+                    else getattr(args, "dataset_order", "default")
+                ),
+                seed=getattr(args, "seed", None),
+                append_eos=bool(getattr(args, "append_eos", True)),
+            )
+            return batches, None
+
         if is_vlm:
             if text_assistant_only_loss:
                 raise ValueError(
@@ -3541,6 +3600,36 @@ def save_model(self, output_dir=None):
             save_merged_model(self.model, self.tokenizer, output_dir)
 
 
+class MLXORPOTrainer(MLXTrainer):
+    """ORPO trainer mirroring TRL's ORPOTrainer. Forces loss_type='orpo' so
+    the class is authoritative regardless of the config passed."""
+    def __init__(self, model, tokenizer, train_dataset, eval_dataset=None,
+                 dataset_text_field=None, max_seq_length=None, packing=None,
+                 data_collator=None, args=None, formatting_func=None, processor=None):
+        if args is None:
+            args = MLXORPOConfig()
+        elif getattr(args, "loss_type", "sft") != "orpo":
+            args.loss_type = "orpo"
+        super().__init__(model, tokenizer, train_dataset, eval_dataset,
+                         dataset_text_field, max_seq_length, packing,
+                         data_collator, args, formatting_func, processor)
+
+
+class MLXDPOTrainer(MLXTrainer):
+    """DPO trainer mirroring TRL's DPOTrainer. Forces loss_type='dpo' so
+    the class is authoritative regardless of the config passed."""
+    def __init__(self, model, tokenizer, train_dataset, eval_dataset=None,
+                 dataset_text_field=None, max_seq_length=None, packing=None,
+                 data_collator=None, args=None, formatting_func=None, processor=None):
+        if args is None:
+            args = MLXDPOConfig()
+        elif getattr(args, "loss_type", "sft") != "dpo":
+            args.loss_type = "dpo"
+        super().__init__(model, tokenizer, train_dataset, eval_dataset,
+                         dataset_text_field, max_seq_length, packing,
+                         data_collator, args, formatting_func, processor)
+
+
 def _create_labeled_batches(dataset, tokenizer, mask_fn, batch_size,
                             max_seq_length, formatting_func=None,
                             dataset_text_field="text", num_batches=None,