feat(mlx): add KTO preference tuning

tests/test_mlx_kto_loss.py

@@ -0,0 +1,140 @@
+# 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/>.
+
+"""Pure-logic regression for the MLX KTO loss primitives.
+
+Covers the unpaired KTO loss (`_kto_loss`), the batch KL baseline clamp
+(`_kto_kl_baseline`), the summed-logp extractor (`_kto_sum_logp`), and the
+batch-size guard in `_build_kto_batches`. Runs under the torch shim so Linux
+CI collection covers it without MLX/Metal (all four primitives are backend-
+agnostic MLX-array math).
+
+Reference values are the ones proven in the KTO Step 1 investigation, where
+`_kto_loss` matched a torch reimplementation of TRL's kto_loss and a pure-numpy
+hand computation to ~1e-8 across desirable/undesirable weight settings.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+
+@pytest.fixture(autouse=True, scope="module")
+def _install_shim():
+    from mlx_simulation import simulate_mlx_on_torch
+    simulate_mlx_on_torch()
+
+
+# --------------------------------------------------------------------------
+# Fixed inputs + independent numpy reference (the KTO Step 1 fixtures).
+# --------------------------------------------------------------------------
+def _fixed_logps():
+    rng = np.random.default_rng(0)
+    return dict(
+        pol_ch=rng.normal(-8, 2, size=3).astype(np.float32),
+        pol_rej=rng.normal(-9, 2, size=2).astype(np.float32),
+        pol_kl=rng.normal(-8.5, 2, size=5).astype(np.float32),
+        ref_ch=rng.normal(-8, 2, size=3).astype(np.float32),
+        ref_rej=rng.normal(-9, 2, size=2).astype(np.float32),
+        ref_kl=rng.normal(-8.5, 2, size=5).astype(np.float32),
+    )
+
+
+def _numpy_kto_loss(pol_ch, pol_rej, ref_ch, ref_rej, kl, beta, wd, wu):
+    sig = lambda z: 1.0 / (1.0 + np.exp(-z))
+    ch = wd * (1 - sig(beta * ((pol_ch - ref_ch) - kl)))
+    rej = wu * (1 - sig(beta * (kl - (pol_rej - ref_rej))))
+    return float(np.concatenate([ch, rej]).mean())
+
+
+# Pinned from Step 1 (MLX == numpy == TRL-torch to ~1e-8).
+_EXPECTED = {(1.0, 1.0): 0.4772096276283264,
+             (1.33, 1.0): 0.566311240196228,
+             (1.0, 1.5): 0.5808119177818298}
+_EXPECTED_KL = 0.31288090348243713
+
+
+def test_kto_loss_matches_reference_across_weights():
+    import mlx.core as mx
+    from unsloth_zoo.mlx.trainer import _kto_loss, _kto_kl_baseline
+    f = _fixed_logps()
+    kl = float(_kto_kl_baseline(mx.array(f["pol_kl"]), mx.array(f["ref_kl"])))
+    assert kl == pytest.approx(_EXPECTED_KL, abs=1e-6)  # matches TRL 0.31288
+    for (wd, wu), expected in _EXPECTED.items():
+        loss = float(_kto_loss(
+            mx.array(f["pol_ch"]), mx.array(f["pol_rej"]),
+            mx.array(f["ref_ch"]), mx.array(f["ref_rej"]), mx.array(kl),
+            0.1, wd, wu,
+        ))
+        npy = _numpy_kto_loss(f["pol_ch"], f["pol_rej"], f["ref_ch"], f["ref_rej"], kl, 0.1, wd, wu)
+        assert loss == pytest.approx(npy, abs=1e-6), f"w=({wd},{wu}) MLX vs numpy"
+        assert loss == pytest.approx(expected, abs=1e-6), f"w=({wd},{wu}) vs pinned"
+
+
+def test_kl_baseline_clamps_negative_to_zero():
+    import mlx.core as mx
+    from unsloth_zoo.mlx.trainer import _kto_kl_baseline
+    # policy far below reference on the mismatched completions -> raw mean < 0.
+    pol_kl = mx.array([-30.0, -28.0, -35.0])
+    ref_kl = mx.array([-20.0, -22.0, -19.0])
+    assert float(_kto_kl_baseline(pol_kl, ref_kl)) == 0.0
+    # positive estimate passes through unclamped.
+    assert float(_kto_kl_baseline(ref_kl, pol_kl)) > 0.0
+
+
+@pytest.mark.parametrize("labels_present", ["desirable_only", "undesirable_only"])
+def test_kto_loss_single_label_batches_are_finite(labels_present):
+    import mlx.core as mx
+    from unsloth_zoo.mlx.trainer import _kto_loss
+    empty = mx.array(np.array([], dtype=np.float32))
+    if labels_present == "desirable_only":
+        pol_ch, ref_ch = mx.array([-8.0, -9.0]), mx.array([-8.3, -9.4])
+        pol_rej = ref_rej = empty
+    else:
+        pol_rej, ref_rej = mx.array([-8.0, -9.0]), mx.array([-8.3, -9.4])
+        pol_ch = ref_ch = empty
+    loss = float(_kto_loss(pol_ch, pol_rej, ref_ch, ref_rej, mx.array(0.5), 0.1, 1.0, 1.0))
+    assert np.isfinite(loss) and 0.0 <= loss <= 1.0
+
+
+# NOTE: _kto_sum_logp's numpy-reference check lives in the Metal-gated
+# test_mlx_kto_train_metal.py. Its .astype(mx.float32) / cross_entropy path
+# resolves a real mlx dtype under the torch shim in the pytest import order, so
+# it is exercised against real MLX instead of the shim.
+
+
+def test_build_kto_batches_requires_batch_size_two():
+    from unsloth_zoo.mlx.trainer import _build_kto_batches, MLXKTOConfig
+
+    class _DummyTokenizer:
+        pad_token_id = 0
+        eos_token_id = 0
+        def __call__(self, text, add_special_tokens=False):
+            return {"input_ids": [1, 2, 3]}
+
+    dataset = [{"prompt": "a", "completion": " b", "label": True},
+               {"prompt": "c", "completion": " d", "label": False}]
+    with pytest.raises(ValueError) as exc:
+        _build_kto_batches(dataset, _DummyTokenizer(), MLXKTOConfig(per_device_train_batch_size=1))
+    msg = str(exc.value)
+    assert "per_device_train_batch_size" in msg and ">= 2" in msg
+
+    # batch_size >= 2 builds batches with the mismatched-pair KL variant.
+    batches = _build_kto_batches(dataset, _DummyTokenizer(), MLXKTOConfig(per_device_train_batch_size=2))
+    assert len(batches) == 1
+    for key in ("comp_ids", "comp_labels", "kl_ids", "kl_labels", "label"):
+        assert key in batches[0]

tests/test_mlx_kto_train_metal.py

@@ -0,0 +1,214 @@
+# 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/>.
+
+"""Real-MLX regression for MLXKTOTrainer on Apple Silicon.
+
+Trains a 4-bit Qwen2.5-0.5B + LoRA with KTO on a small unpaired dataset and
+checks the loss is finite and decreasing, that the LoRA-disable reference path
+restores adapter scales (including when a forward throws mid-step), and that a
+non-PEFT model raises the clear LoRA-only error.
+
+Metal-gated so Linux CI collection skips cleanly.
+"""
+
+import math
+
+import pytest
+
+try:
+    import mlx.core as mx
+    _METAL = mx.metal.is_available()
+except Exception:
+    _METAL = False
+
+if not _METAL:
+    print("NOTICE: Metal unavailable; MLX KTO training tests will be skipped.")
+
+metal_only = pytest.mark.skipif(not _METAL, reason="requires Apple Silicon Metal")
+
+MODEL = "unsloth/Qwen2.5-0.5B"
+
+
+def _dataset(n=24):
+    rows = []
+    for i in range(n):
+        prompt = f"### Question: what is {i} plus {i}?\n### Answer:"
+        if i % 2 == 0:
+            rows.append({"prompt": prompt, "completion": f" {2 * i}.", "label": True})
+        else:
+            rows.append({"prompt": prompt, "completion": f" {2 * i + 7}, wrong.", "label": False})
+    return rows
+
+
+def _load_peft():
+    from unsloth_zoo.mlx.loader import FastMLXModel
+    mx.random.seed(3407)
+    model, tok = FastMLXModel.from_pretrained(MODEL, max_seq_length=256, load_in_4bit=True)
+    model = FastMLXModel.get_peft_model(model, r=8, lora_alpha=16, lora_dropout=0, random_state=3407)
+    return model, tok
+
+
+def _config(**overrides):
+    from unsloth_zoo.mlx.trainer import MLXKTOConfig
+    # gradient_accumulation_steps=1 keeps these tests at one optimizer step per
+    # batch; the accumulation path has its own test below.
+    base = dict(per_device_train_batch_size=4, max_steps=6, warmup_steps=1,
+                gradient_accumulation_steps=1,
+                learning_rate=1e-4, beta=0.1, logging_steps=99, seed=3407, report_to="none")
+    base.update(overrides)
+    return MLXKTOConfig(**base)
+
+
+@metal_only
+def test_kto_trains_finite_and_decreasing(tmp_path):
+    from unsloth_zoo.mlx.trainer import MLXKTOTrainer
+    model, tok = _load_peft()
+    trainer = MLXKTOTrainer(model=model, tokenizer=tok, train_dataset=_dataset(),
+                            args=_config(output_dir=str(tmp_path)))
+    output = trainer.train()
+
+    # train() returns MLXTrainOutput, matching MLXTrainer; the per-step losses
+    # live on the trainer.
+    hist = trainer._train_loss_history
+    assert len(hist) == 6, f"expected 6 steps, got {len(hist)}"
+    assert all(math.isfinite(x) for x in hist), f"non-finite loss: {hist}"
+    assert hist[-1] < hist[0], f"loss did not decrease: {hist}"
+    assert output.global_step == 6, f"expected 6 steps, got {output.global_step}"
+    assert math.isfinite(output.training_loss)
+    assert output["total_train_steps"] == 6
+    # KL baseline recorded per step, finite and clamped >= 0.
+    assert trainer._kl_history and all(math.isfinite(k) and k >= 0.0 for k in trainer._kl_history)
+
+
+@metal_only
+def test_kto_gradient_accumulation_reduces_optimizer_steps(tmp_path):
+    # 24 examples / batch 4 = 6 micro-batches. With max_steps unset, one epoch
+    # is len(batches) // grad_accum optimizer steps: grad_accum=2 -> 3 steps,
+    # each averaging 2 micro-batches. If accumulation were ignored (one update
+    # per micro-batch) this would be 6 steps, so the history length pins that
+    # grad_accum actually groups micro-batches into optimizer steps.
+    from unsloth_zoo.mlx.trainer import MLXKTOTrainer
+    model, tok = _load_peft()
+    trainer = MLXKTOTrainer(
+        model=model, tokenizer=tok, train_dataset=_dataset(),
+        args=_config(output_dir=str(tmp_path), max_steps=0,
+                     gradient_accumulation_steps=2),
+    )
+    output = trainer.train()
+    assert len(trainer._train_loss_history) == 3, trainer._train_loss_history
+    assert len(trainer._kl_history) == 3
+    assert output.global_step == 3
+    assert output["total_train_steps"] == 3
+    assert all(math.isfinite(x) for x in trainer._train_loss_history)
+
+
+@metal_only
+def test_kto_saves_adapters_at_end(tmp_path):
+    # save_steps defaults to 0 (save at end); train() must leave adapters on
+    # disk, not only in memory, matching MLXTrainer.
+    from unsloth_zoo.mlx.trainer import MLXKTOTrainer
+    model, tok = _load_peft()
+    trainer = MLXKTOTrainer(model=model, tokenizer=tok, train_dataset=_dataset(),
+                            args=_config(output_dir=str(tmp_path)))
+    trainer.train()
+    written = {p.name for p in tmp_path.iterdir()}
+    assert "adapters.safetensors" in written, f"no adapters saved: {sorted(written)}"
+    assert "adapter_config.json" in written, f"no adapter config saved: {sorted(written)}"
+
+
+@metal_only
+def test_lora_scales_restored_after_training(tmp_path):
+    from unsloth_zoo.mlx.trainer import MLXKTOTrainer
+    from unsloth_zoo.mlx.utils import iter_mlx_lora_modules
+    model, tok = _load_peft()
+    before = [m.scale for _, m in iter_mlx_lora_modules(model)]
+    assert before and all(s != 0 for s in before)
+
+    trainer = MLXKTOTrainer(model=model, tokenizer=tok, train_dataset=_dataset(),
+                            args=_config(output_dir=str(tmp_path)))
+    trainer.train()
+
+    after = [m.scale for _, m in iter_mlx_lora_modules(model)]
+    assert after == before, "LoRA scales not restored to their pre-training values"
+
+
+@metal_only
+def test_lora_scales_restored_when_forward_throws(tmp_path, monkeypatch):
+    """The reference forward runs with scales set to 0; if it throws, the
+    finally must restore scales so the model is not left adapter-disabled."""
+    import unsloth_zoo.mlx.trainer as T
+    from unsloth_zoo.mlx.trainer import MLXKTOTrainer
+    from unsloth_zoo.mlx.utils import iter_mlx_lora_modules
+    model, tok = _load_peft()
+    before = [m.scale for _, m in iter_mlx_lora_modules(model)]
+
+    # Fail on the 2nd sum-logp call: that is the reference forward, which runs
+    # while adapter scales are 0 (policy-KL is the 1st call, scales still on).
+    real = T._kto_sum_logp
+    calls = {"n": 0}
+    def flaky(logits, labels):
+        calls["n"] += 1
+        if calls["n"] == 2:
+            raise RuntimeError("injected forward failure")
+        return real(logits, labels)
+    monkeypatch.setattr(T, "_kto_sum_logp", flaky)
+
+    trainer = MLXKTOTrainer(model=model, tokenizer=tok, train_dataset=_dataset(),
+                            args=_config(output_dir=str(tmp_path)))
+    with pytest.raises(RuntimeError, match="injected forward failure"):
+        trainer.train()
+
+    after = [m.scale for _, m in iter_mlx_lora_modules(model)]
+    assert after == before, "scales left disabled after a throwing reference forward"
+    assert all(s != 0 for s in after), "adapters left at scale 0"
+
+
+@metal_only
+def test_kto_sum_logp_matches_numpy():
+    """Summed-logp extractor vs a from-scratch numpy log-softmax reference."""
+    import numpy as np
+    from unsloth_zoo.mlx.trainer import _kto_sum_logp
+    rng = np.random.default_rng(1)
+    B, T, V = 3, 6, 11
+    logits = rng.normal(0, 1, size=(B, T, V)).astype(np.float32)
+    labels = rng.integers(0, V, size=(B, T)).astype(np.int64)
+    labels[0, :2] = -100; labels[1, :3] = -100; labels[2, :1] = -100  # masked prompt prefixes
+
+    got = np.array(_kto_sum_logp(mx.array(logits), mx.array(labels)))
+
+    inp, tgt = logits[:, :-1, :], labels[:, 1:]
+    m = inp.max(axis=-1, keepdims=True)
+    logsm = inp - (m + np.log(np.exp(inp - m).sum(axis=-1, keepdims=True)))
+    ref = np.zeros(B)
+    for b in range(B):
+        for t in range(tgt.shape[1]):
+            if tgt[b, t] != -100:
+                ref[b] += logsm[b, t, tgt[b, t]]
+    assert np.abs(got - ref).max() < 1e-4
+
+
+@metal_only
+def test_non_peft_model_raises_lora_only_error(tmp_path):
+    from unsloth_zoo.mlx.loader import FastMLXModel
+    from unsloth_zoo.mlx.trainer import MLXKTOTrainer
+    mx.random.seed(3407)
+    model, tok = FastMLXModel.from_pretrained(MODEL, max_seq_length=256, load_in_4bit=True)
+    # No get_peft_model -> no LoRA adapters -> reference forward is impossible.
+    trainer = MLXKTOTrainer(model=model, tokenizer=tok, train_dataset=_dataset(),
+                            args=_config(output_dir=str(tmp_path)))
+    with pytest.raises(ValueError) as exc:
+        trainer.train()
+    assert "LoRA" in str(exc.value)