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

tests/test_mlx_trainer_internals.py

@@ -1949,3 +1949,172 @@ def test_quantized_linear_forward():
     # x @ W.T  with W = [[0,1,2,3,4,5,6,7], [0,1,2,3,4,5,6,7]] = [28, 28]
     out = layer(x)
     torch.testing.assert_close(out, torch.tensor([[28.0, 28.0]]))
+
+
+# --- Preference (ORPO / DPO) review-round-3 regressions ----------------
+
+def _build_fast_tokenizer(eos="<eos>"):
+    """A plain HF PreTrainedTokenizerFast (its ``_tokenizer`` is the Rust
+    backend, which exposes neither ``eos_token_id`` nor a list-returning
+    ``encode``)."""
+    from transformers import PreTrainedTokenizerFast
+    from tokenizers import Tokenizer, models, pre_tokenizers
+    vocab = {eos: 0, "a": 1, "b": 2, "c": 3, "prompt": 4, "good": 5, "bad": 6}
+    inner = Tokenizer(models.WordLevel(vocab=vocab, unk_token=eos))
+    inner.pre_tokenizer = pre_tokenizers.Whitespace()
+    return PreTrainedTokenizerFast(tokenizer_object=inner, eos_token=eos)
+
+
+def test_preference_tokenizer_keeps_public_hf_fast_tokenizer_api():
+    # Regression: create_preference_batches used getattr(tok, "_tokenizer", tok),
+    # which unwraps a plain PreTrainedTokenizerFast to its low-level Rust
+    # tokenizers.Tokenizer (no eos_token_id; encode returns an Encoding), so
+    # ORPO/DPO batching crashed before training. _hf_encoding_tokenizer must
+    # only unwrap non-HF wrappers (mlx-lm's TokenizerWrapper).
+    from unsloth_zoo.mlx.utils import _hf_encoding_tokenizer
+
+    ptf = _build_fast_tokenizer()
+    assert ptf._tokenizer is not ptf  # the Rust backend is present...
+    resolved = _hf_encoding_tokenizer(ptf)
+    assert resolved is ptf  # ...but we must NOT unwrap to it
+    assert resolved.eos_token_id == 0
+    assert isinstance(resolved.encode("prompt good"), list)
+
+    # A TokenizerWrapper-style object (not a PreTrainedTokenizerBase) that
+    # proxies the HF tokenizer under _tokenizer IS unwrapped.
+    class _Wrapper:
+        def __init__(self, inner):
+            self._tokenizer = inner
+    inner = object()
+    assert _hf_encoding_tokenizer(_Wrapper(inner)) is inner
+
+
+def test_create_preference_batches_runs_with_plain_fast_tokenizer():
+    # End-to-end: batching a plain PreTrainedTokenizerFast must not raise.
+    from unsloth_zoo.mlx.utils import create_preference_batches
+
+    ptf = _build_fast_tokenizer()
+    dataset = [
+        {"prompt": "prompt", "chosen": "good", "rejected": "bad"},
+        {"prompt": "prompt a", "chosen": "good b", "rejected": "bad c"},
+    ]
+    batches = create_preference_batches(
+        dataset=dataset,
+        tokenizer=ptf,
+        batch_size=2,
+        max_seq_length=16,
+        pad_to_multiple=8,
+    )
+    assert len(batches) == 1
+    batch, lengths, extra = batches[0]
+    assert extra is None
+    # 2 pairs -> 4 rows (chosen block then rejected block).
+    assert batch.shape[0] == 4
+    assert lengths.shape[0] == 4
+
+
+def test_dpo_reference_collects_nested_lora_modules():
+    # Regression: the DPO reference collected adapters with
+    #   tree_flatten(model, is_leaf=lambda x: isinstance(x, LoRALinear)),
+    # but tree_flatten only recurses dict/list/tuple, so a bare nn.Module is a
+    # single leaf and nested adapters are never reached -> _lora_mods empty ->
+    # make_dpo_loss_fn raises "model has none". iter_mlx_lora_modules walks the
+    # module tree instead.
+    import mlx.nn as nn
+    from mlx.utils import tree_flatten
+    from unsloth_zoo.mlx.utils import iter_mlx_lora_modules
+
+    class Adapter(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.lora_a = torch.zeros(2, 4)
+            self.lora_b = torch.zeros(4, 2)
+
+    class Block(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.q_proj = Adapter()
+
+    class Model(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.layers = [Block(), Block()]
+
+    model = Model()
+    # The old leaf-scan treats the whole module as one leaf: no adapter found.
+    flat = tree_flatten(model, is_leaf=lambda x: isinstance(x, Adapter))
+    assert all(not isinstance(m, Adapter) for _, m in flat)
+    # The fix traverses modules and finds every adapter that owns lora_a/lora_b.
+    mods = [m for _, m in iter_mlx_lora_modules(model)]
+    assert len(mods) == 2
+    assert all(hasattr(m, "lora_a") and hasattr(m, "lora_b") for m in mods)
+
+
+class _StubLM:
+    def __init__(self, vocab):
+        self.vocab = vocab
+
+    def __call__(self, x):
+        import mlx.core as mx
+        return mx.zeros((x.shape[0], x.shape[1], self.vocab))
+
+
+def _preference_batch(vocab=8):
+    # 2 pairs (B=2) -> 4 rows; chosen rows have long responses, rejected short,
+    # so the response-token count differs from the pair count.
+    import mlx.core as mx
+    batch = mx.array(
+        [
+            [4, 1, 1, 1, 1],  # chosen pair 0
+            [4, 2, 2, 2, 2],  # chosen pair 1
+            [4, 3, 0, 0, 0],  # rejected pair 0
+            [4, 3, 0, 0, 0],  # rejected pair 1
+        ],
+        dtype=mx.int32,
+    )
+    # [response_start, seq_end): chosen span 4 tokens each, rejected 1 each.
+    lengths = mx.array(
+        [
+            [1, 5],
+            [1, 5],
+            [1, 2],
+            [1, 2],
+        ]
+    )
+    return batch, lengths
+
+
+def _patch_per_token_ce(monkeypatch):
+    # The shim's nn.losses.cross_entropy defaults to reduction="mean" and is
+    # not shape-faithful; real mlx returns a per-token tensor. Patch it so the
+    # loss fns run end-to-end (the pair-count return is what is under test).
+    import mlx.nn as nn
+    monkeypatch.setattr(
+        nn.losses,
+        "cross_entropy",
+        lambda logits, targets, **kw: torch.zeros(tuple(targets.shape)),
+    )
+
+
+def test_dpo_loss_returns_pair_count_not_token_count(monkeypatch):
+    # Regression: DPO is a mean over PAIRS, so the accumulate-then-normalize
+    # trainer must weight each microbatch by its pair count. Returning the
+    # response-token count made long-completion microbatches dominate and broke
+    # equivalence with a single batch under gradient_accumulation_steps > 1.
+    from unsloth_zoo.mlx.utils import make_dpo_loss_fn
+
+    _patch_per_token_ce(monkeypatch)
+    batch, lengths = _preference_batch()
+    loss_fn = make_dpo_loss_fn(beta=0.1, lora_mods=None, reference_free=True)
+    _loss, count = loss_fn(_StubLM(8), batch, lengths)
+    assert int(count) == 2  # pair count (B), not the 10 response tokens
+
+
+def test_orpo_loss_returns_pair_count_not_token_count(monkeypatch):
+    from unsloth_zoo.mlx.utils import make_orpo_loss_fn
+
+    _patch_per_token_ce(monkeypatch)
+    batch, lengths = _preference_batch()
+    loss_fn = make_orpo_loss_fn(beta=0.1)
+    _loss, count = loss_fn(_StubLM(8), batch, lengths)
+    assert int(count) == 2  # pair count (B), not the response-token count