MLX ORPO/DPO: fix fast-tokenizer unwrap, LoRA reference collection, and pair-count accumulation

Three correctness fixes on the preference (ORPO/DPO) path:

- create_preference_batches unwrapped the tokenizer with
  getattr(tokenizer, "_tokenizer", tokenizer), which on a plain
  PreTrainedTokenizerFast returns the low-level Rust tokenizers.Tokenizer
  (no eos_token_id; encode returns Encoding), crashing before training.
  Add _hf_encoding_tokenizer, which only unwraps non-HF wrappers such as
  mlx-lm's TokenizerWrapper, matching _get_processor_tokenizer.

- 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 were never reached, leaving _lora_mods
  empty and raising "model has none" for ordinary LoRA models. Use
  iter_mlx_lora_modules to walk the module tree, which also covers
  LoRAEmbedding, LoRASwitchLinear and DoRA.

- The ORPO/DPO loss functions returned the response-token count as the
  accumulation weight, but these losses are means over preference pairs.
  With gradient_accumulation_steps > 1 (default 4) that let long-completion
  microbatches dominate and broke equivalence with a single batch of the
  same pairs. Return the pair count instead.

Add regression tests covering all three.

Author: danielhanchen

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

unsloth_zoo/mlx/trainer.py

@@ -50,7 +50,6 @@
 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")
@@ -2043,9 +2042,12 @@ def _main_print(*print_args, **print_kwargs):
             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)]
+                # tree_flatten only recurses dict/list/tuple, so a bare
+                # nn.Module is a single leaf and never reaches nested LoRA
+                # layers. Walk the module tree instead so LoRALinear,
+                # LoRAEmbedding, LoRASwitchLinear and DoRA adapters are all
+                # collected for the reference (adapters-off) pass.
+                _lora_mods = [mod for _, mod in iter_mlx_lora_modules(model)]
                 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 "") + ").")

unsloth_zoo/mlx/utils.py

@@ -897,7 +897,13 @@ def loss_fn(model, batch, lengths, labels=None):
         # (model) dtype so the combined loss dtype is unchanged.
         or_loss = _orpo_odds_ratio_loss(logp_c, logp_r)
         loss = sft + beta * or_loss.astype(sft.dtype)
-        return loss, mask.sum()
+        # Return the PAIR count, not the response-token count. This is a
+        # preference loss reduced as a mean over pairs (the odds-ratio term),
+        # so the trainer's accumulate-then-normalize machinery must weight each
+        # microbatch by its pair count. Weighting by tokens would make
+        # long-completion microbatches dominate and break equivalence with a
+        # single batch of the same pairs under gradient_accumulation_steps > 1.
+        return loss, mx.array(B, dtype=mx.int32)
     return loss_fn
 
 
@@ -942,7 +948,7 @@ def _row_logp_and_mask(model, batch, lengths):
 
     def loss_fn(model, batch, lengths, labels=None):
         B = batch.shape[0] // 2
-        pol, ntoks = _row_logp_and_mask(model, batch, lengths)
+        pol, _ = _row_logp_and_mask(model, batch, lengths)
         pol_c, pol_r = pol[:B], pol[B:]
 
         if reference_free or not _mods:
@@ -962,9 +968,35 @@ def loss_fn(model, batch, lengths, labels=None):
 
         logits = beta * ((pol_c - ref_c) - (pol_r - ref_r))
         loss = -mx.mean(nn.log_sigmoid(logits))
-        return loss, ntoks
+        # Return the PAIR count, not the response-token count. DPO is a mean
+        # over preference pairs, so the trainer must weight each microbatch by
+        # its pair count when accumulating; token weighting would let
+        # long-completion microbatches dominate and break equivalence with a
+        # single batch of the same pairs under gradient_accumulation_steps > 1.
+        return loss, mx.array(B, dtype=mx.int32)
     return loss_fn
 
+
+def _hf_encoding_tokenizer(tokenizer):
+    """Return the Hugging Face tokenizer used for text encoding.
+
+    mlx-lm's ``TokenizerWrapper`` stores the HF tokenizer under ``_tokenizer``,
+    so unwrapping it exposes ``eos_token_id`` and a list-returning ``encode``.
+    HF fast tokenizers ALSO expose ``_tokenizer``, but there it is the
+    low-level Rust ``tokenizers.Tokenizer`` (no ``eos_token_id``; ``encode``
+    returns ``Encoding`` objects). Only unwrap when the object is not already
+    an HF tokenizer, mirroring ``_get_processor_tokenizer`` /
+    ``_resolve_response_mask_tokenizer``.
+    """
+    try:
+        from transformers import PreTrainedTokenizerBase
+        if isinstance(tokenizer, PreTrainedTokenizerBase):
+            return tokenizer
+    except Exception:
+        pass
+    return getattr(tokenizer, "_tokenizer", tokenizer)
+
+
 def create_preference_batches(dataset, tokenizer, batch_size, max_seq_length,
                         prompt_key="prompt", chosen_key="chosen",
                         rejected_key="rejected", pad_to_multiple=32,
@@ -1004,7 +1036,7 @@ def create_preference_batches(dataset, tokenizer, batch_size, max_seq_length,
             f"Unsloth MLX: unsupported preference dataset_order={order_mode!r}. "
             "Expected 'default', 'sequential', or 'torch_randperm'."
         )
-    hf = getattr(tokenizer, "_tokenizer", tokenizer)
+    hf = _hf_encoding_tokenizer(tokenizer)
     eos_id = hf.eos_token_id
     pad_id = eos_id if eos_id is not None else 0