feat(mlx): add KTO preference tuning - #912
Conversation
KTO trains on UNPAIRED data - one completion per row with a binary desirable/undesirable label, not chosen/rejected pairs - mirroring TRL's KTOTrainer (loss_type="kto", Eqn 7 of arXiv:2402.01306) on Apple Silicon. - MLXKTOTrainer / MLXKTOConfig: standalone trainer + config mirroring TRL's API, reusing MLXTrainer's optimizer/schedule/clip/decay helpers. The SFT compile/accumulate loop is left untouched. - Unpaired batch builder with the mismatched-pair KL variant (completions rolled +1 within the batch, TRL _get_kl_dataset), the summed-logp extractor (_kto_sum_logp), the clamped batch KL baseline (_kto_kl_baseline), and the KTO loss (_kto_loss). Reference model is the LoRA-disabled (scale=0) forward; reference and KL forwards are detached and scales restored in a finally. - Guards: per_device_train_batch_size < 2 fails loud (the mismatched-pair roll would self-pair and collapse the KL term); non-PEFT models fail loud (no reference forward without adapters). Correctness (Step 1): _kto_loss matched a torch reimplementation of TRL's kto_loss and a pure-numpy reference to ~1e-8 across weight settings, with the KL baseline matching TRL at 0.31288. Tests: - tests/test_mlx_kto_loss.py: pure-logic under the torch shim (loss vs numpy and pinned TRL values, KL clamp, single-label batches, batch-size guard). - tests/test_mlx_kto_train_metal.py: Metal-gated e2e (finite/decreasing loss, scale restoration incl. the forward-throws finally path, non-PEFT error, sum-logp vs numpy). Limitations: LoRA-only (reference is the adapter-off forward; full-finetune KTO raises); not mx.compile'd (eager, 4 forwards/step); within-batch KL roll vs TRL's global roll (equivalent for the batch-level estimate, not row-identical); kl=0 is expected often (raw estimate frequently negative -> clamped), which is TRL-faithful.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements Kahneman-Tversky Optimization (KTO) for unpaired preference tuning in MLX, introducing the MLXKTOTrainer, configuration classes, batching utilities, and associated unit tests. Feedback on the implementation highlights two key issues: first, the train method should return an MLXTrainOutput object rather than a list of floats to maintain compatibility with the parent MLXTrainer interface; second, _kto_tokenize_row should utilize tokenizer.encode instead of calling the tokenizer directly to prevent crashes with certain tokenizer wrappers, and should properly truncate completion lengths to avoid potential out-of-memory errors.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| step += 1 | ||
|
|
||
| return self._train_loss_history |
There was a problem hiding this comment.
The train method currently returns self._train_loss_history (a list of floats). However, MLXTrainer.train (which this class inherits from) returns an MLXTrainOutput object (a dict subclass with properties like metrics, global_step, and training_loss).
Any downstream pipeline or training script that expects MLXTrainOutput and accesses these properties (e.g., output.metrics or output.global_step) will crash with an AttributeError: 'list' object has no attribute 'metrics'.
We should wrap the return value in MLXTrainOutput to maintain compatibility with the MLXTrainer interface.
| step += 1 | |
| return self._train_loss_history | |
| step += 1 | |
| avg_loss = ( | |
| sum(self._train_loss_history) / len(self._train_loss_history) | |
| if self._train_loss_history else 0.0 | |
| ) | |
| return MLXTrainOutput({ | |
| "train_loss": avg_loss, | |
| "train_steps": step, | |
| "total_train_steps": total_steps, | |
| }) |
| def _kto_tokenize_row(tokenizer, prompt, completion, args): | ||
| """Tokenize one KTO row to (prompt_ids, completion_ids), truncated to the | ||
| configured lengths. Prompt is truncated (kept-completion) if the combined | ||
| length exceeds max_length, matching TRL's keep-completion policy.""" | ||
| p = tokenizer(prompt, add_special_tokens=False)["input_ids"] | ||
| c = tokenizer(completion, add_special_tokens=False)["input_ids"] | ||
| max_completion = args.max_completion_length | ||
| if max_completion is not None and max_completion > 0: | ||
| c = c[:max_completion] | ||
| if args.max_prompt_length and args.max_prompt_length > 0: | ||
| p = p[-args.max_prompt_length:] | ||
| if args.max_length and args.max_length > 0 and len(p) + len(c) > args.max_length: | ||
| # Keep the whole completion; trim the prompt from the left. | ||
| keep = max(args.max_length - len(c), 0) | ||
| p = p[-keep:] if keep > 0 else [] | ||
| return p, c |
There was a problem hiding this comment.
The current implementation of _kto_tokenize_row calls tokenizer(prompt, add_special_tokens=False) directly. If the tokenizer is an unpatched TokenizerWrapper from mlx_lm (which is common if the trainer is instantiated directly without going through FastMLXModel.from_pretrained), it will raise a TypeError: 'TokenizerWrapper' object is not callable because TokenizerWrapper does not implement __call__ by default.
Additionally, if max_completion_length is not set or is larger than max_length, and the completion itself exceeds max_length, the sequence length of the batch will exceed max_length.
Using tokenizer.encode is much safer and more robust as it works for both TokenizerWrapper and Hugging Face tokenizers directly without dict indexing. We should also truncate c to args.max_length if it exceeds it to prevent out-of-memory (OOM) issues.
def _kto_tokenize_row(tokenizer, prompt, completion, args):
"""Tokenize one KTO row to (prompt_ids, completion_ids), truncated to the
configured lengths. Prompt is truncated (kept-completion) if the combined
length exceeds max_length, matching TRL's keep-completion policy."""
try:
p = tokenizer.encode(prompt, add_special_tokens=False)
except TypeError:
p = tokenizer.encode(prompt)
try:
c = tokenizer.encode(completion, add_special_tokens=False)
except TypeError:
c = tokenizer.encode(completion)
max_completion = args.max_completion_length
if max_completion is not None and max_completion > 0:
c = c[:max_completion]
if args.max_length and args.max_length > 0 and len(c) > args.max_length:
c = c[:args.max_length]
if args.max_prompt_length and args.max_prompt_length > 0:
p = p[-args.max_prompt_length:]
if args.max_length and args.max_length > 0 and len(p) + len(c) > args.max_length:
# Keep the whole completion; trim the prompt from the left.
keep = max(args.max_length - len(c), 0)
p = p[-keep:] if keep > 0 else []
return p, c![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3cafce8e9a
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| self._set_optimizer_lr_for_step(optimizer, step) | ||
| optimizer.update(model, grad) |
There was a problem hiding this comment.
When gradient_accumulation_steps > 1 (the inherited default is 4), the KTO loop still advances the LR and calls optimizer.update for every individual batch. That silently trains with only per_device_train_batch_size examples per optimizer step instead of the requested effective batch size, which changes the optimization dynamics for any KTO run using the shared MLX training config defaults or an explicit accumulation setting.
Useful? React with 👍 / 👎.
| ) | ||
| step += 1 | ||
|
|
||
| return self._train_loss_history |
There was a problem hiding this comment.
When a KTO run finishes successfully, this return exits without the final-save/checkpoint path that MLXTrainer uses (save_steps=0 means save at end in the inherited config). In contexts where users rely on the inherited output_dir/save_steps behavior, trainer.train() updates only the in-memory model and leaves no adapters under output_dir, so the trained result is lost once the process exits unless they know to call save_model() manually.
Useful? React with 👍 / 👎.
MLXKTOTrainer.train() returned self._train_loss_history, a list of floats, while the MLXTrainer.train() it overrides returns an MLXTrainOutput. Any caller written against the base contract -- output.metrics, output.global_step, output.training_loss -- hit AttributeError on the KTO path, and downstream code that treats every trainer's result uniformly broke on KTO alone. Return an MLXTrainOutput with the same core keys (train_loss as the mean over the run, train_runtime, train_steps, total_train_steps). The per-step losses remain on self._train_loss_history for anything that wants them. Update the finite-and-decreasing test to read the loss history off the trainer and additionally assert the MLXTrainOutput contract.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c55e435b51
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| "non-empty completions per batch)." | ||
| ) | ||
|
|
||
| total_steps = args.max_steps if args.max_steps and args.max_steps > 0 else len(batches) |
There was a problem hiding this comment.
When callers disable step-based training with max_steps <= 0 and set num_train_epochs > 1, KTO still sets total_steps to just len(batches). The outer epoch loop does not compensate because it stops as soon as step >= total_steps, so an epoch-based KTO run silently trains for one pass instead of the requested number of epochs.
Useful? React with 👍 / 👎.
| self._eval_callbacks = [] | ||
| self._kl_history = [] | ||
|
|
||
| def train(self, resume_from_checkpoint: str | None = None): |
There was a problem hiding this comment.
When a caller passes resume_from_checkpoint, this override accepts the argument but never stores, validates, or loads it, so KTO starts at step 0 with a fresh optimizer even though the caller requested a resumed run. That can silently redo interrupted training or overwrite the intended continuation; this should either load the checkpoint like MLXTrainer.train() or reject the unsupported argument explicitly.
Useful? React with 👍 / 👎.
| if args.max_length and args.max_length > 0 and len(p) + len(c) > args.max_length: | ||
| # Keep the whole completion; trim the prompt from the left. | ||
| keep = max(args.max_length - len(c), 0) | ||
| p = p[-keep:] if keep > 0 else [] |
There was a problem hiding this comment.
When max_completion_length is unset or larger than max_length and a completion alone exceeds max_length, this branch computes keep = 0 and only removes the prompt; the overlong completion is still padded into the batch. That silently ignores the configured sequence cap and can create unexpectedly huge tensors or OOM on a single long completion instead of truncating or rejecting it.
Useful? React with 👍 / 👎.
| self._is_vlm = False | ||
| # Raw KTO rows: do NOT wrap in the SFT tokenized dataset view. | ||
| self.train_dataset = train_dataset | ||
| self.eval_dataset = eval_dataset |
There was a problem hiding this comment.
When a KTO run is configured with eval_dataset plus eval_steps, load_best_model_at_end, or early stopping, the dataset stored here is never read by train(): the loop only builds training batches and returns. Users therefore get no eval metrics, eval callbacks, best-model restore, or early stopping despite passing the standard trainer arguments.
Useful? React with 👍 / 👎.
| p = tokenizer(prompt, add_special_tokens=False)["input_ids"] | ||
| c = tokenizer(completion, add_special_tokens=False)["input_ids"] |
There was a problem hiding this comment.
With the normal tokenizer returned by FastMLXModel.from_pretrained(), this calls the mlx-lm TokenizerWrapper object directly, but that wrapper exposes encode/the inner HF tokenizer rather than being callable. As a result the advertised KTO path fails during batch construction with a TypeError before training starts unless the caller happens to pass a raw Hugging Face tokenizer instead.
Useful? React with 👍 / 👎.
| mx.eval(model.parameters(), optimizer.state, loss) | ||
|
|
||
| train_loss = float(loss) | ||
| self._train_loss_history.append(train_loss) |
There was a problem hiding this comment.
When the same MLXKTOTrainer instance is trained more than once, this append reuses _train_loss_history from the previous run because train() never calls _reset_run_state() at the start. The returned train_loss is then averaged across old and new runs, and callbacks/consumers reading the history see stale losses mixed with the current training session.
Useful? React with 👍 / 👎.
MLXKTOTrainer.train() updated the model in memory and returned without ever writing to output_dir. The KTO config defaults save_steps=0, whose documented meaning is "save at end of training" and which MLXTrainer honors via a final save_model() call. On the KTO path the trained adapters were therefore lost on process exit unless the caller knew to call save_model() by hand. Mirror the base trainer's end-of-run save: on the main process, call save_model() (KTO is LoRA-only, so this takes the adapter path) and report the output directory. A ValueError from save_model is downgraded to a skip message, as in MLXTrainer. Add a Metal regression asserting adapters.safetensors and adapter_config.json land in output_dir after a run.
The KTO loop called optimizer.update() once per micro-batch and never read gradient_accumulation_steps, so a run with the config default of 4 (or any explicit value > 1) silently trained at an effective batch size of per_device_train_batch_size, stepping the optimizer and LR schedule 4x more often than requested -- different optimization dynamics from every other MLX trainer and from TRL's KTOTrainer. Accumulate grad_accum micro-batches, average their gradients, and take one optimizer step, so the effective batch size is per_device_train_batch_size * grad_accum. step/total_steps and the LR schedule now count optimizer steps; when max_steps is unset, one epoch is len(batches) // grad_accum steps, matching MLXTrainer. Gradient clipping, weight decay, and the LR step run once per optimizer step on the mean gradient. A simple mean is used rather than the token-weighted accumulation MLXTrainer applies to SFT: the KTO loss is a per-example preference objective, so each micro-batch contributes equally. Tests: pin the existing tests to grad_accum=1 (unchanged one-step-per-batch behavior) and add coverage that grad_accum=2 halves the optimizer steps per epoch; the new test fails on the pre-fix loop (6 steps instead of 3).
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 61691d1a2f
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if args.max_length and args.max_length > 0 and len(p) + len(c) > args.max_length: | ||
| # Keep the whole completion; trim the prompt from the left. | ||
| keep = max(args.max_length - len(c), 0) | ||
| p = p[-keep:] if keep > 0 else [] |
There was a problem hiding this comment.
When a completion alone is longer than args.max_length (or max_completion_length is set above max_length), this branch sets keep = 0 and drops only the prompt, leaving the overlong completion untouched. _kto_pad will then emit sequences longer than the configured context cap, so KTO can unexpectedly exceed the model's positional limit or OOM despite max_length being set; truncate the completion as well so len(p) + len(c) <= max_length always holds.
Useful? React with 👍 / 👎.
| if args.max_steps and args.max_steps > 0: | ||
| total_steps = args.max_steps | ||
| else: | ||
| total_steps = max(len(batches) // grad_accum, 1) |
There was a problem hiding this comment.
For MLXKTOConfig(max_steps=0, num_train_epochs=N), this computes total_steps from only a single pass over batches, so epoch-based KTO training silently runs one epoch no matter what num_train_epochs is set to. Since the config inherits num_train_epochs from MLXTrainingConfig and the base trainer uses it when max_steps <= 0, users requesting multiple epochs will get substantially under-trained adapters unless this multiplies by the requested epoch count.
Useful? React with 👍 / 👎.
| self._eval_callbacks = [] | ||
| self._kl_history = [] | ||
|
|
||
| def train(self, resume_from_checkpoint: str | None = None): |
There was a problem hiding this comment.
This method accepts resume_from_checkpoint, but the argument is never used before creating a fresh optimizer and starting step = 0, so a caller resuming a failed KTO run silently restarts from scratch and can overwrite the output adapters as if training had resumed. Either load the adapter/optimizer/trainer state like MLXTrainer.train() does or reject the argument loudly to avoid corrupting resumed experiments.
Useful? React with 👍 / 👎.
| self._apply_manual_weight_decay(model, optimizer, grad) | ||
| self._set_optimizer_lr_for_step(optimizer, step) |
There was a problem hiding this comment.
With AdamW/Muon/Lion and a warmup/linear/cosine schedule, _apply_manual_weight_decay reads optimizer.learning_rate before this step's LR is installed, so every step after the first applies decoupled decay using the previous step's LR while the gradient update uses the current one. The default KTO config has weight_decay > 0 and a scheduled LR, so move _set_optimizer_lr_for_step before the decay helper.
Useful? React with 👍 / 👎.
| grad = self._apply_coupled_weight_decay(model, grad) | ||
| self._apply_manual_weight_decay(model, optimizer, grad) | ||
| self._set_optimizer_lr_for_step(optimizer, step) | ||
| optimizer.update(model, grad) |
There was a problem hiding this comment.
When run with mlx.distributed world_size > 1, this update applies each rank's gradient directly; the KTO loop never shards KTO batches or calls _distributed_sum_gradient_tree() / nn.average_gradients like MLXTrainer does, yet only rank 0 saves at the end. Distributed KTO therefore produces independent per-rank training and persists one rank's adapters instead of the global-batch result.
Useful? React with 👍 / 👎.
| for _epoch in range(max_epochs): | ||
| if stop: | ||
| break | ||
| for batch in batches: |
There was a problem hiding this comment.
If the model was left in eval mode before MLXKTOTrainer.train() (for example after generation/evaluation), this loop starts taking optimizer steps without model.train(), so MLX modules gated on training such as LoRA dropout remain disabled for the whole run. The base trainer explicitly enters train mode before its loop; KTO should do the same before the first forward.
Useful? React with 👍 / 👎.
| # Accumulate the mean gradient. Dividing each micro-batch by | ||
| # grad_accum before summing keeps the accumulator at update | ||
| # scale; sum(g_i / N) == mean(g_i). | ||
| contrib = tree_map(lambda g: g / grad_accum, grad) |
There was a problem hiding this comment.
When lora_plus_ratio is set, this accumulation treats every gradient leaf the same (g / grad_accum) and never applies the inherited LoRA+ multiplier to lora_b leaves, so KTO silently ignores a supported MLXTrainingConfig knob and trains with the wrong per-parameter LR. Mirror the base trainer's per-leaf scaling before accumulation/update or reject the option for KTO.
Useful? React with 👍 / 👎.
| self._kl_history = [] | ||
|
|
||
| def train(self, resume_from_checkpoint: str | None = None): | ||
| args = self.args |
There was a problem hiding this comment.
When the same MLXKTOTrainer instance is trained more than once, this method starts a new local step counter but never clears _train_loss_history, _kl_history, or the prior global-step state. The second run's metrics are therefore averaged together with stale losses/KL from the previous run, unlike MLXTrainer.train() which resets per-run state before entering its loop.
Useful? React with 👍 / 👎.
| self._is_vlm = False | ||
| # Raw KTO rows: do NOT wrap in the SFT tokenized dataset view. | ||
| self.train_dataset = train_dataset | ||
| self.eval_dataset = eval_dataset |
There was a problem hiding this comment.
The constructor accepts and stores eval_dataset, but the KTO training loop never builds eval batches or calls _evaluate, so eval_steps, load_best_model_at_end, and early stopping are silently ignored whenever users pass validation data. This can leave KTO experiments with no validation metrics or best-model selection even though the inherited config exposes those controls.
Useful? React with 👍 / 👎.
| configured lengths. Prompt is truncated (kept-completion) if the combined | ||
| length exceeds max_length, matching TRL's keep-completion policy.""" | ||
| p = tokenizer(prompt, add_special_tokens=False)["input_ids"] | ||
| c = tokenizer(completion, add_special_tokens=False)["input_ids"] |
There was a problem hiding this comment.
Completions are tokenized verbatim here and the inherited append_eos=True default is never applied, so common KTO rows whose completion text does not already include an EOS never train the adapter to end the answer. Since MLXTrainingConfig documents append_eos as enabled by default and the SFT path honors it, KTO should append tokenizer.eos_token_id before padding/truncation or explicitly reject the option.
Useful? React with 👍 / 👎.
Three issues in MLXKTOTrainer surfaced by the follow-up review, all on the paths the recent KTO commits touched: - LR/decay ordering: _apply_manual_weight_decay reads optimizer.learning_rate, but the KTO loop installed this step's LR after the decay helpers, so decoupled weight decay used the previous step's LR every step past the first. Move _set_optimizer_lr_for_step ahead of the decay calls, matching the base trainer. - num_train_epochs: when max_steps is unset, total_steps was one pass over the data // grad_accum, ignoring num_train_epochs, so epoch-based KTO ran a single epoch regardless of the requested count. Multiply by num_train_epochs, matching MLXTrainer. - LoRA+: the accumulation scales every gradient leaf by 1/grad_accum and never applies the lora_plus_ratio multiplier to lora_b, so a set lora_plus_ratio was silently ignored. Per-leaf LoRA+ scaling is woven into the base trainer's update path and is not reimplemented here; reject the option loudly instead of training lora_b at the wrong LR. Tests: num_train_epochs=2 doubles the optimizer steps; lora_plus_ratio>0 raises. Both fail on the pre-fix loop.
The standalone KTO loop does not implement checkpoint resume, an eval loop, or MLX distributed training, yet MLXKTOTrainer accepted all three and silently ignored them: resume_from_checkpoint would restart from scratch and overwrite the output adapters, eval_dataset/eval_steps/load_best_model_at_end produced no evaluation, and under world_size > 1 each rank trained independently while only rank 0 saved. Raise a clear error for each at the start of train() rather than return a result that looks trained but ignored the request. These are the pre-existing KTO limitations flagged in review; implementing them (mirroring MLXTrainer's resume/eval/DDP paths) is left to a follow-up. Tests: resume_from_checkpoint and eval_dataset both raise.
Two small correctness gaps in MLXKTOTrainer.train(), matching MLXTrainer: - The loop began taking optimizer steps without model.train(), so if the model was left in eval mode (e.g. after generation/evaluation) training- gated modules like LoRA dropout stayed disabled for the whole run. The reference forward disables adapters by scale rather than eval mode, so entering train mode does not affect it. - Per-run metric state (_train_loss_history, _kl_history, _global_step) was never reset, so re-running the same trainer instance blended the new run's loss/KL with the previous run's -- and train() reports the mean over _train_loss_history in its MLXTrainOutput. Reset the metric state and call model.train() before the loop. Test asserts a second train() call reports 3 steps of history, not 6.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7e1bcc90ed
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| batches.append(dict( | ||
| comp_ids=_kto_pad([p + c for p, c in zip(prompts, comps)], pad_id), | ||
| comp_labels=_kto_pad([[-100] * len(p) + list(c) for p, c in zip(prompts, comps)], -100), | ||
| kl_ids=_kto_pad([p + cr for p, cr in zip(prompts, rolled)], pad_id), |
There was a problem hiding this comment.
When completions in the same KTO batch have different lengths, each prompt is truncated only against its own completion, but this KL path concatenates that prompt with a different rolled completion without applying max_length again. Fresh evidence here is that the rolled KL pair can exceed max_length even when every original prompt+completion row fit, so the KL forward can allocate past the configured context window or OOM.
Useful? React with 👍 / 👎.
| pad_id = 0 | ||
|
|
||
| rows = [] | ||
| for ex in dataset: |
There was a problem hiding this comment.
When args.streaming=True or an HF IterableDataset is passed, this loop tries to exhaust the dataset to build all KTO rows before training starts. Unlike the base MLXTrainer streaming path, an unbounded stream will hang or OOM and finite streams lose streaming semantics, so KTO should either implement an iterator-based batcher or fail loudly for streaming datasets.
Useful? React with 👍 / 👎.
| # only in memory and are lost on process exit unless the caller knows to | ||
| # call save_model() by hand. KTO is LoRA-only, so save_model() takes the | ||
| # adapter path. | ||
| if self.is_main_process: |
There was a problem hiding this comment.
When callers set the inherited save_steps > 0 to keep intermediate checkpoints, the KTO loop never checks it after optimizer steps and only reaches this final adapter save. Interrupted runs therefore leave no checkpoint-* directory, optimizer state, or trainer state despite the standard MLX training config requesting checkpointing; either mirror the base checkpoint path or reject save_steps > 0 as unsupported for KTO.
Useful? React with 👍 / 👎.
| p, c = _kto_tokenize_row(tokenizer, ex["prompt"], ex["completion"], args) | ||
| if len(c) == 0: | ||
| continue # nothing to score | ||
| rows.append((p, c, bool(ex["label"]))) |
There was a problem hiding this comment.
When KTO data comes from a CSV or other source that stores binary labels as strings such as "False" or "0", this truthiness conversion turns every non-empty string into True. That silently moves undesirable examples into the desirable branch and trains on corrupted preference labels; parse accepted boolean/0-1 encodings explicitly or reject unsupported label types.
Useful? React with 👍 / 👎.
| self._train_loss_history.append(train_loss) | ||
| self._kl_history.append(acc_kl) | ||
| self._global_step = step + 1 | ||
| if args.logging_steps and (step % max(int(args.logging_steps), 1) == 0): |
There was a problem hiding this comment.
When a caller registers add_step_callback() to drive progress reporting or callback-based cancellation, the KTO loop records the step and only enters this print branch; it never iterates self._step_callbacks like MLXTrainer.train() does. Those consumers see no progress events and cannot stop from a callback during KTO training, despite the inherited callback API being present.
Useful? React with 👍 / 👎.
| self._eval_callbacks = [] | ||
| self._kl_history = [] | ||
|
|
||
| def train(self, resume_from_checkpoint: str | None = None): |
There was a problem hiding this comment.
Because this override takes over train() without calling the base setup path, a KTO run configured with the inherited neftune_noise_alpha > 0 never calls _install_neftune() and trains without the requested embedding noise. Since the KTO guards already reject unsupported inherited options, this should either install/remove NEFTune around the loop or fail loudly instead of silently ignoring the knob.
Useful? React with 👍 / 👎.
| ref_comp = mx.stop_gradient(_kto_sum_logp(model(batch["comp_ids"]), batch["comp_labels"])) | ||
| ref_kl = mx.stop_gradient(_kto_sum_logp(model(batch["kl_ids"]), batch["kl_labels"])) |
There was a problem hiding this comment.
When the model has any nonzero training-mode dropout, these adapter-off reference forwards still run while the surrounding KTO loop has put the model in training mode. The reference log-probs and KL baseline then include fresh dropout masks instead of the deterministic frozen reference policy, which corrupts the KTO log-ratios; temporarily switch to eval mode for the reference forwards and restore training mode afterwards.
Useful? React with 👍 / 👎.
| acc_loss = 0.0 # running mean of the window's losses | ||
| acc_kl = 0.0 # running mean of the window's KL diagnostics | ||
| stop = False | ||
| max_epochs = 1_000_000 |
There was a problem hiding this comment.
When max_steps is larger than 1_000_000 * len(batches) (for example a tiny smoke dataset repeated for a long adapter run), this fixed epoch cap exits the outer loop before step reaches total_steps and returns a successful MLXTrainOutput with fewer updates than requested. The loop should be driven directly by step < total_steps instead of a silent epoch ceiling.
Useful? React with 👍 / 👎.
| args.beta, args.desirable_weight, args.undesirable_weight, | ||
| ) | ||
|
|
||
| loss, grad = nn.value_and_grad(model, loss_fn)(model) |
There was a problem hiding this comment.
For models that rely on the base MLX training setup to patch missing VJPs or disable fused training paths (for example gated-delta Qwen3.5/qwen3_next/Kimi-style layers or Qwen VL fused MRoPE), KTO jumps straight into value_and_grad without that setup. Those supported model families can therefore fail gradient computation or train through the wrong fused path under KTO; mirror the base pre-train patching or reject those models explicitly.
Useful? React with 👍 / 👎.
| # Accumulate the mean gradient. Dividing each micro-batch by | ||
| # grad_accum before summing keeps the accumulator at update | ||
| # scale; sum(g_i / N) == mean(g_i). | ||
| contrib = tree_map(lambda g: g / grad_accum, grad) |
There was a problem hiding this comment.
When callers LoRA the lm_head/embed_tokens paths and set the inherited embedding_learning_rate, the base trainer scales those gradient leaves relative to the main LR, but KTO averages every gradient uniformly here after only rejecting LoRA+. That silently trains embedding/LM-head adapter leaves at the main learning rate instead of the configured embedding LR; apply the same per-leaf scaling or reject this option for KTO.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7e1bcc90ed
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| p = tokenizer(prompt, add_special_tokens=False)["input_ids"] | ||
| c = tokenizer(completion, add_special_tokens=False)["input_ids"] |
There was a problem hiding this comment.
For boundary-sensitive BPE/SentencePiece tokenizers, concatenating tokens produced by two independent calls is not generally equivalent to tokenizing the actual prompt + completion text. This also bypasses the inherited append_eos=True behavior, so both reward and KL paths optimize synthetic token streams without the completion's EOS probability. The existing _tokenize_mlx_prompt_completion helper in unsloth_zoo/mlx/utils.py handles joint encoding, boundary masking, and EOS; use equivalent processing here.
Useful? React with 👍 / 👎.
| if args.max_length and args.max_length > 0 and len(p) + len(c) > args.max_length: | ||
| # Keep the whole completion; trim the prompt from the left. | ||
| keep = max(args.max_length - len(c), 0) | ||
| p = p[-keep:] if keep > 0 else [] |
There was a problem hiding this comment.
When max_completion_length is unset (the default) and a completion alone exceeds max_length, this branch sets keep to zero and removes the prompt but never truncates c; the resulting sequence therefore still exceeds the configured cap. A long completion can consequently cause unexpectedly large allocations or OOM despite the caller setting max_length, so the completion itself must also be capped.
Useful? React with 👍 / 👎.
| kl_ids=_kto_pad([p + cr for p, cr in zip(prompts, rolled)], pad_id), | ||
| kl_labels=_kto_pad([[-100] * len(p) + list(cr) for p, cr in zip(prompts, rolled)], -100), |
There was a problem hiding this comment.
The prompt is truncated against its own completion before batching, but the KL sequence replaces that completion with another row's potentially much longer one. For example, a long-prompt/short-completion row followed by a short-prompt/long-completion row can make p + cr exceed max_length even though both original pairs fit, inflating memory and computing the KL baseline outside the configured context limit. Retruncate each prompt against its rolled completion before constructing kl_ids and kl_labels.
Useful? React with 👍 / 👎.
| # Enter training mode so training-gated modules (e.g. LoRA dropout) are | ||
| # active, in case the model was left in eval mode by a prior | ||
| # generation/evaluation. The reference forward disables adapters by | ||
| # scale, not by eval mode, so this does not affect it. | ||
| model.train() |
There was a problem hiding this comment.
MLXKTOConfig inherits gradient_checkpointing=True, but this standalone path enters training without calling the base trainer's apply_gradient_checkpointing setup or providing an equivalent implementation. On the larger models for which this default is needed—especially with KTO's multiple forwards—the setting is silently ignored and training can OOM; apply and clean up checkpointing around the loop, or reject the option explicitly.
Useful? React with 👍 / 👎.
| # Honor the documented save_steps=0 contract (save at end of training), | ||
| # matching MLXTrainer.train(). Without this the trained adapters live | ||
| # only in memory and are lost on process exit unless the caller knows to | ||
| # call save_model() by hand. KTO is LoRA-only, so save_model() takes the | ||
| # adapter path. | ||
| if self.is_main_process: |
There was a problem hiding this comment.
When save_steps > 0, the KTO loop still performs only this final save and never creates checkpoint-N directories, so the inherited save_steps and save_total_limit settings are silently ineffective. A long KTO run interrupted before completion loses all persisted progress; add the periodic adapter-save branch used by MLXTrainer, or reject nonzero save_steps rather than accepting it.
Useful? React with 👍 / 👎.
| self._train_loss_history.append(train_loss) | ||
| self._kl_history.append(acc_kl) | ||
| self._global_step = step + 1 | ||
| if args.logging_steps and (step % max(int(args.logging_steps), 1) == 0): |
There was a problem hiding this comment.
This condition tests the zero-based step, so logging_steps=10 logs optimizer steps 1, 11, 21, and so on instead of 10, 20, 30, and it may omit the final step entirely. Test (step + 1) % logging_steps and include the final-step case to match the base trainer's logging contract.
Useful? React with 👍 / 👎.
| args.beta, args.desirable_weight, args.undesirable_weight, | ||
| ) | ||
|
|
||
| loss, grad = nn.value_and_grad(model, loss_fn)(model) |
There was a problem hiding this comment.
This direct value_and_grad path bypasses the architecture setup performed by MLXTrainer.train, including patch_gated_delta for Qwen3.5, Qwen3-Next, and Kimi Linear models and the differentiable fallback for fused MRoPE. On those supported architectures the unpatched fused kernels have no usable VJP, so KTO reaches this line and fails on its first backward pass. Run the same model-specific training setup and cleanup around the KTO loop.
Useful? React with 👍 / 👎.
| rows = [] | ||
| for ex in dataset: |
There was a problem hiding this comment.
MLXKTOConfig inherits the streaming option, but _build_kto_batches always exhausts the entire dataset into rows before training. With a true streaming or unbounded iterable and streaming=True, training never reaches the first optimizer step or grows memory without bound even when max_steps is finite. Reject streaming mode explicitly or consume and batch the iterator incrementally.
Useful? React with 👍 / 👎.
| contrib = tree_map(lambda g: g / grad_accum, grad) | ||
| if acc_grad is None: | ||
| acc_grad = contrib | ||
| else: | ||
| acc_grad = tree_map(lambda a, g: a + g, acc_grad, contrib) | ||
| acc_loss += float(loss) / grad_accum | ||
| acc_kl += float(kl_scalar) / grad_accum |
There was a problem hiding this comment.
_build_kto_batches retains a trailing batch whenever it has at least two rows, so an accumulation window can contain both a full batch and a smaller tail. Because each batch loss is already a per-example mean and this code gives every micro-batch the same 1 / grad_accum weight, examples in the smaller tail receive disproportionately large gradients; for batch sizes four and two, each tail example is weighted twice as much. Accumulate loss sums and example counts, or scale each contribution by its row count.
Useful? React with 👍 / 👎.
| @dataclass | ||
| class MLXKTOConfig(MLXTrainingConfig): |
There was a problem hiding this comment.
Decorating this subclass with @dataclass generates a new initializer and bypasses MLXTrainingConfig.__init__, so _unsloth_mlx_warmup_steps_explicit is never recorded. In particular, when a caller explicitly sets warmup_steps=5 (the class default) together with a positive warmup_ratio, _resolve_warmup_steps treats the step count as implicit and uses the ratio instead, unlike the base config. Delegate through the base initializer or reproduce its explicit-field bookkeeping.
Useful? React with 👍 / 👎.
Bring KTO up to date with main (incl. unslothai#873 HF-style callbacks). Clean auto-merge, 0 conflicts: KTO's additions live in a self-contained trainer.py region main did not touch.
CSV-ingested KTO datasets store the desirable/undesirable label as a string.
bool("false") and bool("0") are both True, so every undesirable row was
silently flipped to desirable. Add _kto_parse_label to parse the common string
spellings (and pass real bool/int/float through), and reject ambiguous values.
Regression test asserts string labels "false"/"0" -> [False, False].
_kto_tokenize_row set keep=0 when a completion was longer than max_length, dropping the prompt but leaving the whole completion, so the returned sequence exceeded max_length whenever max_completion_length was unset. Cap the completion to max_length in that case. Regression test asserts len(p)+len(c) <= max_length.
_build_kto_batches materializes the whole dataset to form the mismatched-pair KL batches, so a streaming/IterableDataset was silently exhausted up front. Reject it (via args.streaming or a bare iterable-without-__len__) with a clear NotImplementedError, mirroring the accepted ORPO/DPO streaming guard. Regression test covers both the config flag and a generator passed directly.
MLXKTOConfig accepts any loss_type, but the training path always calls _kto_loss and never reads the field, so a caller selecting another TRL KTO variant would silently get plain KTO. Fail fast in __init__. Regression test asserts a non-kto loss_type raises.
MLXKTOTrainer accepted arbitrary **kwargs and discarded them, so a caller passing TRL-style ref_model=... had it silently dropped while KTO compared against its own adapter-off forward. Reject ref_model loudly in __init__. Regression test asserts it raises.
_kto_tokenize_row tokenized completions with add_special_tokens=False and never appended EOS, so the inherited append_eos=True was ignored and the model was never trained to emit EOS. Append it (guarded against a double EOS / no-eos tokenizer, gated on append_eos), and -- deviating from TRL, which truncates then appends and can drop it -- preserve the trailing EOS through both the max_completion_length and max_length caps. Regression test covers append, no-double, disabled, and EOS survival under each cap.
The KTO loop bypasses the base trainer's architecture setup, so a gated-delta model (Qwen3.5/Qwen3-Next) fails in the backward with a missing custom-kernel VJP, and a VLM fails at batch-build (no pad_token_id) and forward (LanguageModelOutput, not raw logits). Both were already hard failures; reject them up front (model_has_gated_delta_layers / a processor with image_processor) for UX clarity. Not a correctness fix. Regression test covers both.
KTO's reference is the adapter-off (LoRA scale=0) forward of the same model, so a trainable non-LoRA tensor (directly-trained lm_head/embed_tokens, non-LoRA bias) drifts during training and corrupts the reference. Add a local _kto_model_has_non_lora_trainable_params (verbatim mirror of unslothai#832's helper, kept local to avoid depending on that unmerged branch) and reject unconditionally in train(). Unlike DPO/GRPO there is no reference_free escape -- this is a structural limit, per Daniel's 16:20Z unslothai#832 reversal. Regression test covers the predicate and the guard.
Bare @DataClass regenerated __init__, bypassing MLXTrainingConfig.__init__ and its setup (e.g. _unsloth_mlx_warmup_steps_explicit). init=False inherits the parent init while fields() still registers the KTO-specific fields, which the parent __init__ reads via fields(type(self)). Mirrors unslothai#830's MLXORPOConfig/MLXDPOConfig. Regression test asserts the inherited init and intact KTO fields.
The log condition tested the zero-based step counter, so logging_steps=N logged at steps 1, N+1, 2N+1, ... instead of N, 2N, ... Gate on (step + 1) to match MLXTrainer's one-based cadence. Metal regression test asserts logging_steps=2 over 4 steps logs at [2,4].
_kto_loss is a mean over the batch's rows, but the grad-accum loop weighted each micro-batch equally (grad / grad_accum). A window mixing a full batch and the trailing size-2 batch therefore over-weighted the small one. Weight each micro-batch's gradient/loss/KL by its row count and normalize by the window total, matching the pair/example-count accumulation Daniel adopted for ORPO/DPO/GRPO. Metal regression test (lr=0) asserts the window loss equals the example-weighted mean, not the equal mean.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c2f1057217
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| epochs = max(int(getattr(args, "num_train_epochs", 1) or 1), 1) | ||
| total_steps = max((len(batches) * epochs) // grad_accum, 1) |
There was a problem hiding this comment.
When max_steps is unset and the number of prepared KTO batches is not divisible by gradient_accumulation_steps, this floors away the epoch tail. For example, 5 batches with gradient_accumulation_steps=2 and num_train_epochs=1 reports total_steps=2; after two full windows the loop stops before the 5th batch, so that data never contributes. The preceding int(...) also truncates fractional epoch counts. Use the base trainer's ceil-per-epoch step calculation and flush partial epoch windows instead of dropping them.
Useful? React with 👍 / 👎.
| "trainer's patch_gated_delta setup, so the backward fails with a " | ||
| "missing custom-kernel VJP. Use a non-gated-delta model for KTO." | ||
| ) | ||
| if hasattr(self.tokenizer, "image_processor"): |
There was a problem hiding this comment.
When a VLM caller passes the text tokenizer as tokenizer and the multimodal object via the existing processor= argument, this check misses it because image_processor is on self.processor, not on self.tokenizer. The unsupported KTO VLM run then proceeds into text-only batch construction/model calls instead of raising the intended clear error, ignoring image inputs or failing downstream; include self.processor/model-based VLM detection in this guard.
Useful? React with 👍 / 👎.
| if stop: | ||
| break | ||
| for batch in batches: | ||
| if step >= total_steps or self.stop_requested: |
There was a problem hiding this comment.
When a KTO run is cancelled through trainer.stop_requested = True, the loop breaks here but train() never clears the flag or advances _run_generation on exit. Reusing the same trainer then immediately satisfies this condition at step 0 and returns a successful output after saving the previous adapters, even though the metric state is reset above; mirror the base trainer's prior-generation stop clearing and run-end generation update.
Useful? React with 👍 / 👎.
| else: | ||
| epochs = max(int(getattr(args, "num_train_epochs", 1) or 1), 1) | ||
| total_steps = max((len(batches) * epochs) // grad_accum, 1) | ||
| optimizer = self._build_optimizer(total_steps) |
There was a problem hiding this comment.
On Apple Silicon, the base MLX trainer applies _configure_memory_limits() before training and restores the process-global caps in a finally, honoring both the default guard and explicit memory_limit_gb / cache_limit_gb / wired_limit_gb. This KTO setup goes straight from optimizer construction into the multi-forward training loop, so those inherited guard settings are ignored and large KTO runs can exceed the configured Metal cap instead of failing cleanly; configure and restore the guard around this loop, or reject those options for KTO.
Useful? React with 👍 / 👎.
| def __init__(self, model, tokenizer, train_dataset, args=None, | ||
| eval_dataset=None, processor=None, **kwargs): |
There was a problem hiding this comment.
Because this constructor still accepts arbitrary **kwargs but only inspects ref_model, a ported KTO call with unsupported standard arguments such as data_collator= or callbacks= succeeds while those objects are discarded; the run then uses the built-in raw-row batcher and has no HF callback handler, so custom collation, cancellation, or logging hooks never run. After handling the supported kwargs, reject any leftovers or wire them through explicitly.
Useful? React with 👍 / 👎.
1 participant
Summary
Adds KTO preference tuning to the MLX path on Apple Silicon, mirroring TRL's
KTOTrainer(loss_type="kto", Eqn 7 of arXiv:2402.01306).The selling point is the data. KTO trains on unpaired examples — a single completion per row with a binary desirable/undesirable label — instead of the chosen/rejected pairs DPO and ORPO require. Binary-labelled completions are far cheaper to collect than matched preference pairs (any "was this output good or bad?" signal works: thumbs up/down, task success/failure, filters), which makes KTO the practical choice when you don't have paired data. This is the first unpaired preference method on the MLX trainer.
What's included
MLXKTOTrainer/MLXKTOConfig— standalone trainer + config mirroring TRL's API, reusingMLXTrainer's optimizer / schedule / gradient-clip / weight-decay helpers. The SFT compile/accumulate loop is left completely untouched._get_kl_dataset), the summed-logp extractor (_kto_sum_logp), the clamped batch KL baseline (_kto_kl_baseline), and the KTO loss (_kto_loss).stop_gradient'd; adapter scales are restored in afinallyso a throwing forward can't leave the model adapter-disabled.per_device_train_batch_size < 2raises (the mismatched-pair roll would self-pair and collapse the KL term to the reward); a non-PEFT model raises (no reference forward without adapters).Correctness evidence
_kto_lossmatches both numpy and a torch reimplementation of TRL'skto_lossto ~1e-8, across three weight settings:clamp(min=0)branch (raw negative estimate → 0.0)._kto_sum_logpvs numpy log-softmax reference: max|Δ| = 7e-07.Tests
tests/test_mlx_kto_loss.py— pure-logic, runs under the torch shim in normal CI (no Metal needed): loss vs numpy and vs the pinned TRL values, the KL clamp branch, single-label (all-desirable / all-undesirable) batches, and thebatch_size < 2guard. 5 passed.tests/test_mlx_kto_train_metal.py— Metal-gated e2e (skips cleanly off Apple Silicon): finite & decreasing loss, LoRA-scale restoration including the scales-restored-when-a-forward-throwsfinallypath, the non-PEFT fail-loud error, and_kto_sum_logpvs numpy. 5 passed.(The
_kto_sum_logpnumeric check lives in the Metal file rather than the shim file — its.astype(mx.float32)resolves a real mlx dtype under the torch shim's pytest import order; it runs against real MLX instead. A code comment explains this in place.)Limitations (please read — not buried)
mx.compile'd. The KTO step runs eager with 4 forwards/step (policy, policy-KL, reference, reference-KL). Correct and verified, but slower per step than the compiled SFT path. The blocker is structural: the LoRA-disable reference mutates module scale state, which undermx.compilewould execute at trace time only and freeze the traced scale. Threading scale as a compiled input is a plausible follow-up, but it needs its own benchmark — compiling 1 of 4 forwards may not pay for the added risk.kl = 0is expected often. The raw mismatched-pair estimate is frequently negative early in training and gets clamped to 0. This is TRL-faithful, not a bug (verified: the LoRA is active in the KL forward; the raw estimate was e.g. −4.82 → clamped).