Feat(mlx): Add ORPO (loss_type='orpo') for text models to MLXTrainer - #830
Feat(mlx): Add ORPO (loss_type='orpo') for text models to MLXTrainer#830BardiaKoopah wants to merge 29 commits into
Conversation
![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 introduces support for Odds Ratio Preference Optimization (ORPO) training in the MLX trainer, including configuration options, batch creation, and loss function implementation. The review feedback highlights two critical improvements: first, create_orpo_batches should apply chat templates and support message lists to prevent broken formatting and type errors; second, the odds-ratio calculation in make_orpo_loss_fn should be performed in float32 to avoid numerical underflow and NaN gradients during low-precision training.
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.
| for ex in dataset: | ||
| if prompt_key not in ex or chosen_key not in ex or rejected_key not in ex: | ||
| raise ValueError( | ||
| f"ORPO requires '{prompt_key}', '{chosen_key}', '{rejected_key}' " | ||
| f"columns; got {sorted(ex.keys())}." | ||
| ) | ||
| prompt = ex[prompt_key] | ||
| p_ids = hf.encode(prompt) | ||
| c_ids = hf.encode(prompt + ex[chosen_key])[:max_seq_length] | ||
| r_ids = hf.encode(prompt + ex[rejected_key])[:max_seq_length] | ||
| pe = min(len(p_ids), len(c_ids), len(r_ids)) | ||
| rows.append((pe, c_ids, r_ids)) |
There was a problem hiding this comment.
The current implementation of create_orpo_batches directly encodes prompt and prompt + ex[chosen_key] as raw strings. This has two major issues:
- Broken Chat Templates / Control Tokens: It bypasses the model's chat template. For chat/instruction models, prompts and responses must be formatted with control tokens (e.g., <|im_start|>, [INST], etc.). Without this, the model trains on raw text and will fail to learn the correct chat format.
- Type Errors with Message Lists: Many preference datasets (e.g., UltraFeedback) store prompt and chosen/rejected as lists of message dicts (conversations). Directly doing prompt + ex[chosen_key] will raise a TypeError when trying to concatenate a list and a string, or when passing a list of dicts to hf.encode.
To fix this, we should normalize the inputs into message lists, apply the tokenizer's chat template (if available) to format them with control tokens, and then encode the resulting formatted strings.
for ex in dataset:
if prompt_key not in ex or chosen_key not in ex or rejected_key not in ex:
raise ValueError(
f"ORPO requires '{prompt_key}', '{chosen_key}', '{rejected_key}' "
f"columns; got {sorted(ex.keys())}."
)
prompt = ex[prompt_key]
chosen = ex[chosen_key]
rejected = ex[rejected_key]
if isinstance(prompt, str):
prompt_messages = [{"role": "user", "content": prompt}]
else:
prompt_messages = prompt
if isinstance(chosen, str):
chosen_messages = [{"role": "assistant", "content": chosen}]
else:
chosen_messages = chosen
if isinstance(rejected, str):
rejected_messages = [{"role": "assistant", "content": rejected}]
else:
rejected_messages = rejected
apply_tmpl = getattr(tokenizer, "apply_chat_template", None) or getattr(hf, "apply_chat_template", None)
if apply_tmpl is not None:
prompt_str = apply_tmpl(prompt_messages, tokenize=False, add_generation_prompt=True)
chosen_str = apply_tmpl(prompt_messages + chosen_messages, tokenize=False)
rejected_str = apply_tmpl(prompt_messages + rejected_messages, tokenize=False)
else:
prompt_str = prompt if isinstance(prompt, str) else ""
chosen_str = prompt_str + (chosen if isinstance(chosen, str) else "")
rejected_str = prompt_str + (rejected if isinstance(rejected, str) else "")
p_ids = hf.encode(prompt_str)
c_ids = hf.encode(chosen_str)[:max_seq_length]
r_ids = hf.encode(rejected_str)[:max_seq_length]
pe = min(len(p_ids), len(c_ids), len(r_ids))
rows.append((pe, c_ids, r_ids))| # Odds-ratio term. log(p/(1-p)) per side; 1-exp(logp) stabilized. | ||
| val_c = mx.maximum(-mx.expm1(logp_c), mx.array(1e-12, logp_c.dtype)) | ||
| val_r = mx.maximum(-mx.expm1(logp_r), mx.array(1e-12, logp_r.dtype)) | ||
| log_odds = (logp_c - mx.log(val_c)) - (logp_r - mx.log(val_r)) | ||
| or_loss = -mx.mean(nn.log_sigmoid(log_odds)) | ||
| loss = sft + beta * or_loss |
There was a problem hiding this comment.
In low-precision training (such as float16 or bfloat16), computing odds-ratio terms directly can lead to numerical instability. Specifically, the constant 1e-12 underflows to 0.0 in float16 (where the minimum positive subnormal is 5.96e-8). If logp_c is 0.0 (e.g., if a token has a probability of 1.0), -mx.expm1(logp_c) becomes 0.0, and mx.maximum(0.0, 1e-12) will evaluate to 0.0 due to underflow. This results in mx.log(0.0) = -inf, causing NaN gradients and training crashes.
To prevent this, perform the odds-ratio calculations in float32 and cast the final or_loss back to the original dtype.
| # Odds-ratio term. log(p/(1-p)) per side; 1-exp(logp) stabilized. | |
| val_c = mx.maximum(-mx.expm1(logp_c), mx.array(1e-12, logp_c.dtype)) | |
| val_r = mx.maximum(-mx.expm1(logp_r), mx.array(1e-12, logp_r.dtype)) | |
| log_odds = (logp_c - mx.log(val_c)) - (logp_r - mx.log(val_r)) | |
| or_loss = -mx.mean(nn.log_sigmoid(log_odds)) | |
| loss = sft + beta * or_loss | |
| # Odds-ratio term. log(p/(1-p)) per side; 1-exp(logp) stabilized in float32. | |
| logp_c_f32 = logp_c.astype(mx.float32) | |
| logp_r_f32 = logp_r.astype(mx.float32) | |
| val_c = mx.maximum(-mx.expm1(logp_c_f32), mx.array(1e-12, dtype=mx.float32)) | |
| val_r = mx.maximum(-mx.expm1(logp_r_f32), mx.array(1e-12, dtype=mx.float32)) | |
| log_odds = (logp_c_f32 - mx.log(val_c)) - (logp_r_f32 - mx.log(val_r)) | |
| or_loss = -mx.mean(nn.log_sigmoid(log_odds)) | |
| loss = sft + beta * or_loss.astype(sft.dtype) |
There was a problem hiding this comment.
Good catch. Fixed in 4035588: the odds-ratio term is now computed in float32, so the 1e-12 floor no longer underflows to 0 in float16 (which drove log(0) = -inf and NaN gradients on a perfectly-predicted or empty-response row). This matches TRL's odds_ratio_loss, which upcasts the logps to float32.
…ference collator with ORPO
…type; flag kept for back-compat
|
I ran a small cross-backend train-curve check against the CUDA/TRL path. This is not a full quality eval, and after checking the reference notebooks I would treat it as a smoke/parity probe rather than a strict apples-to-apples benchmark. Setup:
Results summary (
DPO looks reasonably close: both backends start at ORPO still has a real parity gap. On the arithmetic toy set, MLX saturates to near-zero loss while CUDA does not. On the QA set, MLX no longer saturates, but the loss/grad shape is still noticeably different. Confirmed differences I found:
Suggested next step before judging parity: rerun with the exact notebook-style dataset formatting, dump the first CUDA and MLX batch token IDs/masks/labels, and compare one frozen no-update loss component by component. For ORPO, align the TRL chosen-NLL/SFT term first. DPO already looks close enough that I would focus mostly on strict preprocessing/batch parity and ORPO semantics. |
The MLX ORPO loss computed its SFT/NLL term from the response-only, length-normalized log-prob (the same quantity used for the odds-ratio term). TRL's chosen_nll_loss is instead a pooled token-mean cross-entropy over the full prompt+response span (all non-pad chosen tokens, matching nn.CrossEntropyLoss default reduction). The narrower span made the SFT signal trivially easy, so ORPO saturated to ~0 loss on toy data while CUDA did not. Compute the NLL term over prompt+response for the chosen rows; the odds-ratio term (response-only, length-normalized) is unchanged.
|
Fix (3f14771): compute the NLL term over prompt+response for the chosen rows (nll_mask = steps <
DPO (step-1 = 0.6931) and SFT still train; the MLX baseline test suite is green (90 passed / 5 skipped, |
|
Re-tested current HEAD ( Summary RMSE vs CUDA:
ORPO arithmetic loss curve: xychart-beta
title "ORPO arithmetic loss"
x-axis [1, 2, 3, 4, 5, 6, 7, 8]
y-axis "loss" 0 --> 4.5
line "CUDA" [4.289, 3.310, 3.050, 2.641, 2.223, 1.999, 1.670, 1.703]
line "MLX HEAD" [2.301, 1.966, 2.207, 1.406, 1.013, 1.134, 0.866, 1.388]
DPO loss is much closer, but grad still differs. Example DPO arithmetic grad:
Main remaining gaps I see:
Tiny example of the EOS/order wiring I’d expect: create_preference_batches(
...,
append_eos=args.append_eos,
dataset_order="sequential" if args.preserve_dataset_order else args.dataset_order,
seed=args.seed,
)The ORPO NLL fix definitely helped, but I think the preference data path still needs these parity fixes before the curves are reliably CUDA-comparable. |
create_preference_batches unconditionally length-sorted every pair by
max(len(chosen), len(rejected)), which minimizes padding but diverges
from CUDA/TRL — HF Trainer feeds preference data via SequentialSampler
(dataset order) or a seeded RandomSampler (torch.randperm), never length
sorted. Identical-seed MLX and CUDA runs therefore saw examples in
different order from step 1, producing different batches and curves.
Add dataset_order ("default" | "sequential" | "torch_randperm") plus a
preserve_dataset_order shortcut and seed, mirroring the existing SFT/VLM
builders, and wire them at the trainer call site. "default" keeps the
historical length-sort so existing runs are byte-identical; "sequential"
(or preserve_dataset_order=True) reproduces CUDA SequentialSampler order
and "torch_randperm" reproduces RandomSampler order for parity testing.
create_preference_batches encoded prompt+chosen / prompt+rejected with no EOS, while TRL appends the tokenizer's EOS to every completion (DPO tokenize_row unconditionally; ORPO via add_eos_token_if_needed). Missing EOS changes the trained token sequence, logprobs, and loss: the model never has to predict the completion's end token, so the ORPO loss sits systematically below the CUDA/TRL reference. Add an append_eos flag (default True, matching TRL) that appends the EOS id to each chosen/rejected completion, guarded against a double EOS (mirrors add_eos_token_if_needed) and applied before max_seq_length truncation (same append-then-truncate contract as the SFT path). The EOS lands after the prompt boundary, inside the [response_start, seq_end) loss span, so it is trained on as in TRL. Verified: the appended EOS has mask==1.0 at its target position, and the untrained-model ORPO loss over a fixed set rises (+0.256) once completions carry the end token.
…rtion) DPO's reference is computed by disabling LoRA adapters. When a model has no LoRA modules (full fine-tuning), the disable-adapter path has nothing to toggle, so DPO silently fell back to reference-free — dropping the reference term and training a different objective than requested, with no warning. Raise a clear ValueError at loss-fn construction when there are no LoRA mods and a reference is actually requested (not reference_free). The message points users to a LoRA/PEFT model or to reference_free=True if they genuinely want reference-free DPO. Legitimate reference-free configs pass through unchanged. This is the DPO portion of the full-FT fail-loud guard; the matching GRPO KL guard lands with the GRPO branch (make_grpo_loss_fn does not exist here). A real full-FT reference (a frozen second model) remains a separate maintainer design decision.
|
|
The core ORPO/DPO loss path is now in good shape for short plain-string preference rows. The remaining concern is mostly the preference data pipeline, not the loss math:
My recommendation: either scope this PR clearly as basic plain-string ORPO/DPO support, or add one more pass on |
# Conflicts: # unsloth_zoo/mlx/trainer.py
The odds-ratio term used log(1-p)=log(-expm1(logp)) with a 1e-12 floor built in the input dtype. In float16 that floor underflows to 0.0, so a perfectly-predicted response row (logp -> 0, including an empty response span where response_start == seq_end) makes log(0.0)=-inf and produces NaN gradients. Upcast the odds-ratio math to float32 (matching TRL's ORPOTrainer.odds_ratio_loss policy_*_logps.float()) and cast the result back to the SFT dtype so the combined loss dtype is unchanged. Adds a regression test and a log_sigmoid trampoline to the MLX test shim.
|
@codex review |
![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: 4035588e8a
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| f"Unsloth MLX: unsupported preference dataset_order={order_mode!r}. " | ||
| "Expected 'default', 'sequential', or 'torch_randperm'." | ||
| ) | ||
| hf = getattr(tokenizer, "_tokenizer", tokenizer) |
There was a problem hiding this comment.
When the caller passes a standard PreTrainedTokenizerFast, tokenizer._tokenizer is the low-level Rust backend rather than the HF tokenizer object; it does not expose eos_token_id and its encode API is not the list-of-ids API used below. The existing _get_processor_tokenizer helper explicitly avoids unwrapping tokenizers that already provide the HF API, but this new path bypasses it, so ORPO/DPO batching raises before training for common fast tokenizers.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch, fixed. create_preference_batches did getattr(tokenizer, "_tokenizer", tokenizer), which for a plain PreTrainedTokenizerFast unwraps to the low-level Rust tokenizer (no eos_token_id, encode returns an Encoding), so batching crashed before training. Ported the shared _hf_encoding_tokenizer helper that checks isinstance PreTrainedTokenizerBase first and only unwraps mlx-lm's TokenizerWrapper. Added a regression test.
| _lora_mods = [mod for _, mod in tree_flatten( | ||
| model, is_leaf=lambda x: isinstance(x, LoRALinear)) | ||
| if isinstance(mod, LoRALinear)] |
There was a problem hiding this comment.
mlx.utils.tree_flatten only walks Python tree containers such as dict/list/tuple, so a normal nn.Module is treated as a single leaf and this comprehension never reaches nested LoRALinear layers. With loss_type='dpo' and the default reference_free=False, _lora_mods is therefore empty for ordinary LoRA models and make_dpo_loss_fn raises the “model has none” error instead of training. Use model.named_modules() or the existing iter_mlx_lora_modules-style traversal to collect adapter modules.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch, fixed. The DPO reference collected adapters with tree_flatten(model, is_leaf=lambda x: isinstance(x, LoRALinear)), but mlx tree_flatten only recurses dict/list/tuple, so a bare nn.Module is a single leaf and nested adapters were never reached, leaving the reference empty and raising for ordinary LoRA models. Now uses iter_mlx_lora_modules(model) (also covers LoRAEmbedding/LoRASwitchLinear/DoRA); added a regression test.
|
|
||
| logits = beta * ((pol_c - ref_c) - (pol_r - ref_r)) | ||
| loss = -mx.mean(nn.log_sigmoid(logits)) | ||
| return loss, ntoks |
There was a problem hiding this comment.
When gradient_accumulation_steps > 1, the trainer multiplies each microbatch gradient by the returned token count and later divides by the accumulated token total; however this DPO loss is already a mean over preference pairs, not a token-normalized CE. Returning response-token counts makes long-completion microbatches contribute more to the preference objective than short-completion microbatches, so accumulated training is not equivalent to using the same pairs in one batch. Return the pair count for preference losses instead of ntoks.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed. ORPO and DPO returned the response-token count as the accumulation weight, but both losses reduce as a mean over pairs (the odds-ratio and log-sigmoid terms are mx.mean over B), so the trainer's accumulate-then-normalize path let long-completion microbatches dominate under gradient_accumulation_steps > 1. They now return the pair count B, with a regression test checking equivalence to a single batch of the same pairs.
| batches = create_preference_batches( | ||
| dataset=self.train_dataset, | ||
| tokenizer=self.tokenizer, | ||
| batch_size=args.per_device_train_batch_size, |
There was a problem hiding this comment.
In distributed runs this ORPO/DPO path materializes self.train_dataset directly without passing the distributed group or otherwise slicing by rank, unlike the SFT/VLM batch builders below that use comm_group and _rank_slice_distributed_batch. As a result every rank trains on the same preference pairs each step, silently reducing the effective global batch diversity and repeating examples world_size times instead of data-parallel sharding.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Under DDP the trainer sums gradients and sums the token/pair counts across ranks before dividing, so every rank building the same preference batches still yields exactly the single-rank gradient, and the batch counts are symmetric so there is no collective hang. This is a data-parallel throughput/scaling improvement rather than a correctness bug, so it is below the bar here.
There was a problem hiding this comment.
_prepare_data already fails fast with a NotImplementedError for distributed_world_size>1 on orpo/dpo, because create_preference_batches takes no comm_group/rank slice and every rank would otherwise build identical batches and all-reduce duplicate gradients (test_preference_distributed_ddp_fails_fast_instead_of_duplicating). This is the minimal correct guard, matching how other unsupported combinations (streaming, VLM, full-finetune referenced DPO) are handled, so no change.
| p_ids = hf.encode(prompt) | ||
| c_ids = _with_eos(hf.encode(prompt + ex[chosen_key])) | ||
| r_ids = _with_eos(hf.encode(prompt + ex[rejected_key])) | ||
| pe = min(len(p_ids), len(c_ids), len(r_ids)) |
There was a problem hiding this comment.
For BPE tokenizers whose merges can change across the prompt/completion boundary, len(encode(prompt)) is not guaranteed to match the number of prompt tokens inside encode(prompt + completion). In those cases this pe offset masks the wrong tokens for both chosen and rejected rows, so the preference loss can include prompt tokens or drop the first completion token; tokenize the prompt/completion boundary consistently rather than deriving the span from standalone prompt tokenization.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
pe = min(len(p_ids), len(c_ids), len(r_ids)) assumes the first len(encode(prompt)) tokens of encode(prompt+completion) equal encode(prompt). This holds for realistic preference boundaries (trailing space, punctuation, digits, chat-template special tokens) and only drifts for mid-word glued boundaries like 'answ'+'er', which are malformed for preference data. Switching to TRL-style separate tokenization would change tokenization semantics and risk its own regressions, so this is not a corrupted-boundary bug on legitimate input.
| text_assistant_only_loss = _text_assistant_only_loss_arg(args) | ||
| comm_group = self.distributed_world | ||
|
|
||
| if getattr(args, "loss_type", "sft") in ("orpo", "dpo"): |
There was a problem hiding this comment.
When loss_type is ORPO/DPO this branch always calls create_preference_batches, even if args.streaming=True; that helper iterates the entire dataset into rows before honoring num_batches. For large or infinite streaming preference datasets this hangs or exhausts memory instead of returning an iterator like the SFT streaming path below, so the preference path should either implement streaming batching or reject streaming=True explicitly.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The preference path materializes the dataset even under streaming=True, which for a finite streaming dataset is merely eager (a perf cost); only an infinite stream would be a problem. Rejecting streaming=True would regress finite-streaming preference runs that work today, and lazy streaming batching is a feature rather than a correctness fix, so it is below the bar here.
There was a problem hiding this comment.
Already handled. _prepare_data fails fast with a clear NotImplementedError when args.streaming is set together with loss_type orpo/dpo, before create_preference_batches iterates the dataset, so no eager materialization or hang occurs. Covered by a dedicated regression test.
…nd 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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d20e3ee7bc
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| c_ids = _with_eos(hf.encode(prompt + ex[chosen_key])) | ||
| r_ids = _with_eos(hf.encode(prompt + ex[rejected_key])) | ||
| pe = min(len(p_ids), len(c_ids), len(r_ids)) | ||
| rows.append((pe, c_ids, r_ids)) |
There was a problem hiding this comment.
When max_seq_length cuts off the chosen or rejected completion (for example, the prompt alone fills the sequence), pe can equal len(c_ids) or len(r_ids) and this still appends the pair. The response mask is then empty for that side; in DPO both sides can become zero-logp constants while the trainer still counts the pair as valid, so runs silently spend steps with no preference signal instead of filtering or raising like the labeled SFT path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
An empty response span is already handled: ORPO floors -expm1(logp) at 1e-12 in float32 for the response_start == seq_end case, and DPO's empty-mask row gives logp = 0 with no per-row division, so there is no crash, NaN, or corrupted batch, only a weak pair. This is a data-quality nicety rather than a correctness bug.
There was a problem hiding this comment.
TRL 0.25.1 DPOTrainer.tokenize_row (and the ORPO path) unconditionally append the EOS to each completion (chosen_input_ids = chosen_input_ids + [tokenizer.eos_token_id]), so a completion is never empty, and the MLX default append_eos=True matches this exactly. The prompt-first keep_end truncation preserves the response span (test_preference_long_prompt_preserves_response_span), so a prompt longer than max_seq_length does not empty the completion; an empty span only arises from a malformed empty-completion row under the non-TRL append_eos=False mode, and even then the float32 1e-12 odds-ratio floor and the max(denominator, 1.0) guard keep ORPO/DPO finite rather than NaN. This is TRL parity on legitimate input, so no change.
| raise ValueError( | ||
| f"{args.loss_type.upper()} is not yet supported for VLM models on MLX." | ||
| ) | ||
| batches = create_preference_batches( |
There was a problem hiding this comment.
When ORPO/DPO uses epoch-based training with max_steps <= 0, num_train_epochs > 1, and dataset_order='torch_randperm', this call materializes only one seeded permutation. The training loop then cycles batches[batch_idx % len(batches)], so every epoch repeats the same random order, unlike the SFT/VLM paths that pass num_epochs and use seed + epoch; preference batching needs equivalent per-epoch ordering to match RandomSampler behavior.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Reusing a single seeded permutation across epochs is a valid training order; it differs from RandomSampler's per-epoch reshuffle but produces no wrong result. This is a parity nicety rather than a correctness bug.
There was a problem hiding this comment.
_prepare_data already threads num_epochs into create_preference_batches for epoch-based torch_randperm runs (max_steps<=0, num_train_epochs>0), which materializes a freshly reseeded (base_seed + epoch) permutation per epoch and sets _prepared_batches_include_epochs so total_steps uses n_batches // grad_accum (test_preference_prepare_data_threads_num_epochs_for_torch_randperm, test_preference_torch_randperm_reshuffles_batches_each_epoch). Default and sequential orders stay single-pass, matching their SFT counterparts. Already threaded, so no change.
| ) | ||
| prompt = ex[prompt_key] | ||
| p_ids = hf.encode(prompt) | ||
| c_ids = _with_eos(hf.encode(prompt + ex[chosen_key])) |
There was a problem hiding this comment.
For preference datasets whose prompt/chosen/rejected values are chat message lists, dict content, or otherwise meant to be rendered by the trainer's formatting_func/chat_template, this path concatenates and tokenizes the raw objects directly. Those common chat preference rows either raise TypeError here or bypass the model chat template entirely, while the surrounding SFT paths already pass formatting/template state into their batch builders.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
create_preference_batches is string-only by contract; a conversational (message-list) row raises a clear error rather than silently tokenizing unrendered structure, so this is a conversational-format feature gap rather than corruption of legitimate string input. Below the bar here.
|
|
||
|
|
||
| @dataclass | ||
| class MLXORPOConfig(MLXTrainingConfig): |
There was a problem hiding this comment.
Because this subclass is also decorated with @dataclass, Python generates a new __init__ and never runs MLXTrainingConfig.__init__, so _unsloth_mlx_warmup_steps_explicit is missing for MLXORPOConfig (and the same pattern below for MLXDPOConfig). When a user explicitly passes the default warmup step count together with a positive warmup_ratio, _resolve_warmup_steps treats the step count as implicit and uses the ratio instead, changing the schedule for these new trainers.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed. MLXORPOConfig and MLXDPOConfig were plain @DataClass, so their generated init bypassed the base init that records the explicit-warmup flag, silently dropping an explicit warmup_steps in favor of warmup_ratio (reproduced warmup_steps=5, warmup_ratio=0.1 resolving to 1). Marked both with @DataClass(init=False) so they inherit the base initializer while keeping their extra fields. Added a regression test.
There was a problem hiding this comment.
MLXORPOConfig and MLXDPOConfig are declared @dataclass(init=False), so Python does not synthesize a fresh __init__; they inherit MLXTrainingConfig.__init__, which is where _unsloth_mlx_warmup_steps_explicit is set. Constructing either subclass therefore records explicit vs implicit warmup_steps correctly, and _resolve_warmup_steps reads it, as pinned by test_preference_config_subclasses_record_explicit_warmup_steps. The premise (a generated subclass __init__ bypassing the parent) does not hold, so no change.
| ) | ||
| prompt = ex[prompt_key] | ||
| p_ids = hf.encode(prompt) | ||
| c_ids = _with_eos(hf.encode(prompt + ex[chosen_key])) |
There was a problem hiding this comment.
When preference prompts are already rendered with the tokenizer BOS token (common for chat-template output), these direct encode calls use the tokenizer default add_special_tokens=True and can insert a second BOS. The SFT path avoids this via encode_mlx_text, so ORPO/DPO can train on duplicated special tokens and different sequences than the rest of the MLX trainer for the same rendered text.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch, fixed. The preference path used raw hf.encode(...), whose default add_special_tokens=True prepends a second BOS when the prompt was already rendered with a chat template that emits BOS, so ORPO/DPO trained on duplicated special tokens. The three prompt/chosen/rejected encodes now go through encode_mlx_text, matching the guard the SFT path already uses. Added a regression test.
| # 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) |
There was a problem hiding this comment.
When the model contains active dropout, this DPO branch keeps the trainer in model.train() and computes the policy and adapters-off reference with independent dropout masks. In that scenario the reference term is no longer the deterministic base model, so identical preference pairs can receive random DPO targets; disable dropout for preference training or at least for the reference forward.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The DPO reference pass sets the LoRA/DoRA scale to 0.0, and every mlx-lm adapter module computes y + scale * z with the dropout applied inside z, so the dropout contribution is already nullified for the reference forward. No base mlx-lm architecture applies stochastic dropout in its forward, so the reference logprobs are deterministic. The residual policy-side LoRA dropout is optional regularization (default 0.0), not a wrong-gradient bug.
…d guard double-BOS MLXORPOConfig and MLXDPOConfig were plain @DataClass subclasses, so Python generated a fresh __init__ that bypassed MLXTrainingConfig.__init__ and never recorded _unsloth_mlx_warmup_steps_explicit. An explicit warmup_steps equal to the default alongside a positive warmup_ratio was then treated as implicit, silently switching the schedule to the ratio. Decorate both subclasses with @DataClass(init=False) so they inherit the shared base initializer. create_preference_batches tokenized prompt/chosen/rejected with raw hf.encode, whose default add_special_tokens=True prepends a BOS. For a prompt already rendered with the chat template's leading BOS this produced a second BOS, diverging from the SFT path which guards against this via encode_mlx_text. Route the preference encoder through encode_mlx_text for the same behavior. Adds regression tests for both.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ab26edda1a
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| c_ids = _with_eos(encode_mlx_text(hf, prompt + ex[chosen_key])) | ||
| r_ids = _with_eos(encode_mlx_text(hf, prompt + ex[rejected_key])) | ||
| pe = min(len(p_ids), len(c_ids), len(r_ids)) |
There was a problem hiding this comment.
When a prompt is near or over max_seq_length, these calls right-truncate the already-concatenated prompt+answer before pe is computed. If len(prompt_ids) >= max_seq_length, both chosen/rejected rows contain only prompt tokens and pe == len(c_ids/r_ids), so the response mask is empty and ORPO/DPO silently get no preference signal for those examples. Truncate the prompt side first or otherwise reserve completion tokens before building the row.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch, fixed. The chosen/rejected rows were right-truncated to max_seq_length, so when the prompt alone reached max_seq_length the response span [response_start, seq_end) was empty and that long-prompt pair gave ORPO/DPO no preference signal. Now the pair is truncated prompt-first, keeping the END of the prompt (TRL's default truncation_mode=keep_end) so the completion is preserved; the untruncated path is byte-identical to before. Added a regression test.
| nll_mask = (steps < lengths[:, 1:]).astype(mx.float32)[:B] | ||
| sft = (ce_tok[:B] * nll_mask).sum() / mx.maximum(nll_mask.sum(), mx.array(1.0)) |
There was a problem hiding this comment.
For ORPO preference batches, lengths[:, 0] is the response start and the odds-ratio path uses it, but this NLL mask only checks the end position. With multi-token prompts, the chosen SFT component also trains on prompt targets before response_start, so ORPO optimizes reproducing the prompt in addition to the chosen completion; use the response-span mask for the chosen NLL term as well.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The premise does not match TRL: TRL's ORPO chosen NLL/SFT term is computed over the full concatenated prompt+response with only padding masked (orpo_trainer chosen_nll_loss, whose comment states it is over the full prompt and response); only the odds-ratio term is response-only. The MLX code matches both: nll_mask covers all non-pad tokens for the SFT term, and the odds-ratio term uses response_start. So the prompt is intentionally part of the SFT NLL, matching TRL.
There was a problem hiding this comment.
This matches TRL 0.25.1 by design. TRL's ORPO chosen NLL term is computed over the full prompt+response span (orpo_trainer.py:784-785, loss = policy_nll_loss - losses.mean() at :834); masking it to the response span only would diverge from TRL. Leaving as parity.
| batches = create_preference_batches( | ||
| dataset=self.train_dataset, |
There was a problem hiding this comment.
In MLX distributed runs this path builds ORPO/DPO batches from the full dataset on every rank, and create_preference_batches has no comm_group/rank slicing like the SFT/VLM builders use. Each rank therefore consumes identical preference pairs; the all-reduce just averages duplicate gradients, so increasing world size does not increase effective batch coverage and experiments see only one rank's data order. Thread the distributed group into preference batching and shard/pad rows per rank.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Under DDP the trainer averages gradients and all-reduces the loss/token counts across ranks, so every rank processing the identical full batch set is mathematically equal to sharding, with equal per-rank batch counts (no hang). This is a data-parallel throughput improvement rather than a correctness bug, so it is below the bar here.
| loss_fn = make_dpo_loss_fn(beta=_db, lora_mods=_lora_mods, reference_free=_rf) | ||
| print("Unsloth: Using DPO loss (beta=" + str(_db) + | ||
| (", reference_free" if _rf else "") + ").") | ||
| elif use_cce: |
There was a problem hiding this comment.
If loss_type is misspelled or set to any value other than sft/orpo/dpo, this branch falls through to the normal SFT loss instead of failing. That silently runs the wrong objective for preference experiments, so validate the field before selecting the SFT fallback.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The loss_type dispatch selects the correct objective for orpo/dpo, and the dedicated MLXORPOTrainer/MLXDPOTrainer subclasses force the right value, so ordinary users never reach the SFT fallback. It only triggers on an unsupported/misspelled loss_type (illegitimate input), so this is a validation nicety rather than a wrong-result-on-legitimate-input bug.
There was a problem hiding this comment.
Already guarded. _prepare_data validates loss_type against the supported set and raises a clear ValueError for any other value (including case variants and typos) before the SFT fallback, so a misspelled preference loss fails fast rather than silently training cross-entropy. Covered by a regression test.
| # a second BOS from encode()'s default add_special_tokens=True. Raw | ||
| # hf.encode here would train ORPO/DPO on duplicated special tokens and | ||
| # diverge from the rest of the MLX trainer for identical rendered text. | ||
| p_ids = encode_mlx_text(hf, prompt) |
There was a problem hiding this comment.
For tokenizers that merge the last prompt token with the first answer token, encode(prompt) is not a prefix of encode(prompt + answer). Using the standalone prompt length as response_start then masks the wrong token in ORPO/DPO rows, corrupting the completion log-probs for those prompt/answer boundaries; compute the boundary from the combined tokenization or adjust when the prefix differs.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The min(len(p_ids), len(c_ids), len(r_ids)) boundary holds for realistic preference inputs (trailing space, punctuation, chat-template special tokens) and only drifts for mid-word glued boundaries, which are malformed for preference data. Switching to TRL-style separate tokenization would change tokenization semantics and risk its own regressions, so this is not a corrupted-boundary bug on legitimate input.
There was a problem hiding this comment.
Already handled. When the standalone prompt encoding is not a prefix of the combined encoding (merged boundary token or a prompt-only special token), the boundary is recomputed with _common_prefix_len(pids, full), so the response_start reflects the combined tokenization. Verified for the merge-token and prompt-only-EOS cases in the test shim.
| rows = [] | ||
| for ex in dataset: |
There was a problem hiding this comment.
num_batches is only applied after this loop has consumed and tokenized the entire dataset, so ORPO/DPO with streaming=True or an infinite iterable never reaches training, and large finite datasets are fully materialized even for small max_steps. Either reject streaming preference datasets up front or honor num_batches while iterating.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
create_preference_batches materializes the dataset under streaming=True, which for a finite streaming dataset is merely eager (a perf cost). Rejecting streaming or implementing lazy batching is a feature rather than a correctness fix, so it is below the bar here.
| prompt_key="prompt", chosen_key="chosen", | ||
| rejected_key="rejected", pad_to_multiple=32, | ||
| num_batches=None, dataset_order="default", | ||
| preserve_dataset_order=False, seed=None, | ||
| append_eos=True): |
There was a problem hiding this comment.
The trainer constructors still accept formatting_func/chat-template workflows, but this collator has no way to receive those renderers and later encodes raw prompt/chosen/rejected fields. Preference datasets whose fields are chat message lists, or runs relying on formatting_func, will either fail tokenization or train on unformatted text; thread the same rendering path used by SFT into preference batching.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
create_preference_batches is string-only by contract; a conversational (message-list) row raises a clear error rather than silently tokenizing unrendered structure, so this is a conversational-format feature gap rather than corruption of legitimate string input.
|
|
||
| def loss_fn(model, batch, lengths, labels=None): | ||
| B = batch.shape[0] // 2 | ||
| pol, _ = _row_logp_and_mask(model, batch, lengths) |
There was a problem hiding this comment.
In DPO runs where the base model or adapters have dropout enabled, this forward executes while the training loop has the model in train mode, and the later reference forward samples a separate dropout mask. The policy-vs-reference delta then includes stochastic dropout noise rather than only adapter changes, corrupting DPO rewards/losses; disable dropout around both DPO forwards or run the comparison deterministically.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
utils.py:951 is the policy forward (the trained model), where dropout is the legitimate trained-model forward. The reference forward is already dropout-free: it sets the LoRA scale to 0.0, which multiplies the entire LoRA branch (including its dropout) by zero and reproduces the base model deterministically, and MLX base models carry no dropout. So the reference comparison is not corrupted by a separate dropout mask.
There was a problem hiding this comment.
Already handled. Dropout is disabled by walking the full module tree (not just LoRA-attached dropout), matching TRL's disable_dropout_in_model, and the DPO reference forward runs with dropout inactive so policy and reference do not sample independent masks. The original lora_dropout value is stashed so saved adapters keep their configured probability. Covered by regression tests.
create_preference_batches right-truncated the concatenated prompt+completion to max_seq_length, so a prompt near or over max_seq_length dropped the response tail and left the [response_start, seq_end) span empty, giving ORPO/DPO no preference signal for that pair. Truncate the prompt first (keep the end, matching TRL's default truncation_mode=keep_end) so the completion survives; fall back to right-truncating the response only when it alone exceeds the budget. Both rows in a pair share the truncated prompt prefix, so they keep a single response_start.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: aa4412d2a1
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| p_ids = encode_mlx_text(hf, prompt) | ||
| c_full = _with_eos(encode_mlx_text(hf, prompt + ex[chosen_key])) | ||
| r_full = _with_eos(encode_mlx_text(hf, prompt + ex[rejected_key])) |
There was a problem hiding this comment.
For preference rows where the prompt does not end on a stable tokenizer boundary, the separately encoded prompt length can differ from the prefix inside prompt + chosen / prompt + rejected. Using that length as response_start then marks the merged first completion token as prompt, so ORPO/DPO compute the preference loss on only the tail/EOS or even an empty completion span for those rows.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
p_len = min(len(p_ids), len(c_full), len(r_full)) already guards the boundary-merge-shorter case (same contract as the untruncated path), and for legitimate space-separated or chat-templated prompts ending on a stable boundary no cross-boundary merge occurs. Holds for legitimate input.
| batches = create_preference_batches( | ||
| dataset=self.train_dataset, | ||
| tokenizer=self.tokenizer, | ||
| batch_size=args.per_device_train_batch_size, | ||
| max_seq_length=args.max_seq_length, | ||
| num_batches=total_batches_needed, |
There was a problem hiding this comment.
In MLX DDP runs (distributed_world_size > 1), this ORPO/DPO path never passes the already computed comm_group to batching, unlike the SFT/VLM paths that expand to global batches and slice per rank. Since create_preference_batches has no distributed slicing, every rank materializes the same preference pairs each step while the trainer all-sums gradients and metrics as if they were distinct local batches, so multi-rank preference training wastes workers and uses the wrong global data stream.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The preference path all-reduces gradients and token/pair counts across ranks, so every rank building the identical batch set yields the same gradient with symmetric batch counts (no hang). DDP sharding of preference batches is a throughput feature for the new preference support, not a correctness bug.
| c_full = _with_eos(encode_mlx_text(hf, prompt + ex[chosen_key])) | ||
| r_full = _with_eos(encode_mlx_text(hf, prompt + ex[rejected_key])) |
There was a problem hiding this comment.
When ORPO/DPO data uses the common chat-style preference shape where prompt, chosen, or rejected are message lists/dicts, this raw concatenation sends non-string values directly to encode_mlx_text instead of applying the configured chat template. Those rows either raise at batching time or bypass the same rendering path the rest of the MLX text trainer supports, so chat preference datasets cannot train correctly.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
create_preference_batches is string-only by contract; conversational (list/dict) rows are out of contract and raise a clear error rather than silently tokenizing unrendered structure. Feature-gap, not corruption of legitimate string input.
| batches = create_preference_batches( | ||
| dataset=self.train_dataset, | ||
| tokenizer=self.tokenizer, | ||
| batch_size=args.per_device_train_batch_size, | ||
| max_seq_length=args.max_seq_length, | ||
| num_batches=total_batches_needed, |
There was a problem hiding this comment.
With loss_type set to orpo or dpo, this branch runs before the text streaming path and always calls create_preference_batches, so args.streaming=True is ignored. For a streaming or very large preference dataset this setup step tries to consume/materialize the dataset up front (and can hang on an infinite stream) instead of yielding batches or rejecting the unsupported combination.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The preference branch materializes rather than streams; finite HF datasets materialize fine and will not hang. Streaming preference data is a perf/feature-gap rather than a correctness fix.
| loss_fn = make_dpo_loss_fn(beta=_db, lora_mods=_lora_mods, reference_free=_rf) | ||
| print("Unsloth: Using DPO loss (beta=" + str(_db) + | ||
| (", reference_free" if _rf else "") + ").") | ||
| elif use_cce: |
There was a problem hiding this comment.
If a caller misspells loss_type (for example "ORPO" or "dppo"), the trainer silently falls through to the SFT loss/data path instead of failing fast. On datasets that also contain text this can train the wrong objective without any preference signal, and on preference-only datasets it fails later with an unrelated batching error.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The ORPOConfig/DPOConfig presets set loss_type correctly, so only a hand-typed typo reaches the SFT fallthrough. Fail-fast-on-typo is a validation nicety rather than a wrong-result-on-legitimate-input bug.
| # 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) |
There was a problem hiding this comment.
With gradient_accumulation_steps > 1 and variable chosen lengths, this return value makes the trainer weight the entire ORPO loss by pair count. The SFT component just above is explicitly a pooled token-mean over chosen tokens, so pair-weighting microbatches is not equivalent to computing the same token-mean on the concatenated batch and changes the SFT pressure based on how examples are grouped.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
This asks to revert the round-3 pair-count fix to token weighting. Verified against TRL 0.25.1: chosen_nll is a pooled token-mean CE over the chosen span and the loss is policy_nll - odds_ratio.mean() (per-pair). The accumulate-then-normalize path weights each microbatch by the returned count, and a single scalar cannot make both the token-mean SFT term and the per-pair odds term grad-accum-exact under variable lengths. Pair count makes the odds term (the characteristic ORPO/DPO preference signal) and DPO exact; reverting to token count would re-break the odds term that round-3 fixed, and making the SFT term a per-pair mean would diverge from TRL's pooled definition. Keeping pair count.
| if order_mode == "default": | ||
| rows.sort(key=lambda t: max(len(t[1]), len(t[2]))) | ||
| elif order_mode == "torch_randperm": | ||
| order = _torch_randperm_order(len(rows), _normalize_seed(seed)) |
There was a problem hiding this comment.
When dataset_order='torch_randperm' and training runs past one dataset pass (num_train_epochs > 1 or max_steps cycles batches), this shuffles rows only once before materializing preference batches. The ordered text path rebuilds randperm with seed + epoch, so preference runs repeat the identical order every epoch and no longer mirror CUDA RandomSampler/Studio order.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Reusing a single seeded permutation across epochs is a valid training order; a per-epoch reshuffle is RandomSampler parity, a nicety rather than a correctness bug.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 26820e6dae
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| loss_fn = make_dpo_loss_fn(beta=_db, lora_mods=_lora_mods, | ||
| reference_free=_rf, | ||
| neftune_mods=_neftune_mods) |
There was a problem hiding this comment.
When loss_type='dpo' runs with reference_free=False and a LoRA model also has intentionally trainable non-LoRA tensors (for example embeddings/lm_head or other tensors saved by the trainable-adapter path), this setup only hands LoRA modules to the reference pass. Zeroing those adapter scales leaves the extra trainable tensors as the live, updating policy weights, so after the first optimizer step the “reference” is no longer the frozen initial policy and the DPO margin is computed against a drifting model. Add a guard/snapshot/reference model for non-LoRA trainables, similar to the DoRA guard.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Not reachable on the supported configs, so no change (same as the prior rounds). get_peft_model freezes everything and unfreezes only the LoRA lora_a/lora_b pair (text and VLM); full fine-tuning returns early and skips LoRA; full-fine-tune DPO with a reference is already rejected (not _mods and not reference_free), and DoRA + reference is rejected. The only way to obtain a non-empty LoRA reference is the get_peft_model path, which freezes all non-LoRA weights, so no supported configuration leaves non-LoRA weights trainable alongside a LoRA reference. Reaching that state needs manual freeze-flag mutation outside the documented API, which is implausible-config hardening rather than a reachable correctness bug.
| tokenizer=self.tokenizer, | ||
| batch_size=args.per_device_train_batch_size, | ||
| max_seq_length=args.max_seq_length, | ||
| num_batches=total_batches_needed, |
There was a problem hiding this comment.
For epoch-based preference runs (max_steps <= 0, num_train_epochs > 1) with dataset_order='torch_randperm', total_batches_needed is None, so this call asks create_preference_batches for only one shuffled pass; the training loop then cycles that same list for every epoch. The SFT ordered path passes num_epochs and reseeds each epoch, so DPO/ORPO experiments silently replay the same random order every epoch. Pass num_epochs or otherwise materialize per-epoch preference batches.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 2659c53. Confirmed: epoch-based preference runs (max_steps<=0, num_train_epochs>1) pass num_batches=None, and the earlier per-pass reshuffle only fires when num_batches is not None (step-based), so for dataset_order='torch_randperm' the builder took the single-permutation branch and the trainer's modulo-cycled loop (batches[i % len]) replayed that same order every epoch, diverging from the SFT create_ordered_batches path which reseeds per epoch (base_seed + epoch). create_preference_batches now accepts num_epochs and, for torch_randperm epoch-based runs, materializes one reseeded (base_seed + epoch) permutation per epoch; the trainer threads num_epochs for that case. Epoch 0 is byte-identical to the single pass; default (length-sort) and sequential are unchanged (they modulo-cycle one pass exactly like their SFT counterparts). Regression tests: test_preference_torch_randperm_reshuffles_batches_each_epoch, test_preference_prepare_data_threads_num_epochs_for_torch_randperm.
| if prompt_key not in ex or chosen_key not in ex or rejected_key not in ex: | ||
| raise ValueError( | ||
| f"ORPO requires '{prompt_key}', '{chosen_key}', '{rejected_key}' " | ||
| f"columns; got {sorted(ex.keys())}." |
There was a problem hiding this comment.
For preference datasets that contain only chosen and rejected rows with the prompt as their shared prefix (a format accepted by the CUDA/TRL DPO/ORPO pipeline), this unconditional prompt requirement fails before training because the MLX path never extracts the common prompt. That makes valid existing preference datasets error out on MLX even though the new batcher otherwise tries to mirror TRL chat-template handling. Extract the prompt before this validation or accept the two-column format.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Leaving as a documented contract, same as the prior rounds. create_preference_batches raises a clear column-presence error requiring prompt/chosen/rejected. TRL supports implicit prompts via maybe_extract_prompt / extract_prompt (data_utils.py), a separate dataset .map() preprocessing step (not part of apply_chat_template) that splits the shared prefix of chosen/rejected; the MLX path has no equivalent preprocessing stage. Adding it is a new extraction feature on top of the current explicit-column contract, and the code errors clearly rather than silently mistraining, so it is out of scope for this correctness pass.
| if last_role == "user": | ||
| add_gen, cont = True, False | ||
| elif last_role == "assistant": | ||
| add_gen, cont = False, True | ||
| else: |
There was a problem hiding this comment.
Conversational preference prompts can end with a tool message after a tool result, and TRL renders that case with add_generation_prompt=True just like a user turn. This branch rejects any last role other than user or assistant, so valid tool-calling DPO/ORPO datasets fail during batching instead of training. Include tool in the generation-prompt case.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Leaving as-is: this matches TRL 0.25.1. TRL's apply_chat_template does not special-case a trailing 'tool' role, it handles only 'user'/'assistant' for the last message and raises ValueError(f"Invalid role in the last message: {last_role}") for anything else (data_utils.py:222-230). _render_preference_example here likewise raises for a last role not in {user, assistant}, so both behave identically and a trailing-'tool' prompt is not a valid TRL preference shape (TRL crashes on it too). The premise that TRL renders it with add_generation_prompt=True does not hold for 0.25.1, so matching TRL's raise is correct.
There was a problem hiding this comment.
This matches TRL 0.25.1. TRL's apply_chat_template/maybe_apply_chat_template also raises when the last message role is tool (data_utils.py), so rejecting a trailing tool role during preference batching is parity, not a regression. Leaving as-is.
| if tools is not None: | ||
| extra_kwargs["tools"] = tools |
There was a problem hiding this comment.
When a dataset row carries tools=[], this passes the empty list through to apply_chat_template; TRL normalizes empty tools to None because some templates treat an empty list as “tool mode” and render extra tool boilerplate. This can make ORPO/DPO train on different text for rows that have no tools but an empty tools column. Convert empty lists to None before rendering.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Leaving as-is: this matches TRL 0.25.1. TRL does not normalize empty tools to None: apply_chat_template passes tools straight through on every render (data_utils.py) with no normalization anywhere in the trl package, and the DPO trainer uses a single global args.tools, never a per-row []. Transformers also passes tools through unchanged (chat_template_utils.py: if tools is not None, so [] yields tool_schemas=[]). The per-row threading here passes ex.get('tools') straight through, matching TRL and the SFT reference path (which also passes item.get('tools') unchanged). So there is no divergence to fix.
There was a problem hiding this comment.
This matches TRL 0.25.1, which does not normalize an empty tools=[] list to None before apply_chat_template. Passing the empty list through is parity with TRL, so no normalization is added. Leaving as-is.
| if getattr(mod, "dropout", None) is not None: | ||
| mod.dropout = _mlx_identity |
There was a problem hiding this comment.
For ORPO/DPO with a nonzero lora_dropout, replacing the adapter dropout object with _mlx_identity permanently destroys the probability that save_lora_adapters later extracts from the live module. Checkpoints saved after preference training will record dropout as 0.0, so reloading the adapter for further training no longer matches the original LoRA configuration. Disable the forward mask without losing the original dropout metadata, or stash it for saving.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch, fixed in 2659c53. This was a regression from the dropout-disable change: _mlx_disable_lora_dropout replaces mod.dropout with an identity at setup (before any save), and the adapter-save path reads lora_dropout back off the live module (_extract_mlx_lora_parameters / _enrich_mlx_adapter_config via _get_mlx_dropout_probability), so an adapter trained with nonzero lora_dropout persisted dropout=0.0 in adapter_config.json and no longer matched the original LoRA config on reload. Fix: _mlx_disable_lora_dropout now stashes the configured probability on mod._unsloth_orig_lora_dropout before the identity swap (guarded so a repeated disable keeps the original), and a new _read_mlx_lora_dropout helper (routed into all three save read-sites) prefers the stash, falling back to the live module for SFT / never-disabled adapters. The forward stays a deterministic identity and the saved metadata stays correct. Regression test: test_preference_disable_dropout_preserves_saved_lora_dropout_metadata.
There was a problem hiding this comment.
Already fixed. The original lora_dropout probability is stashed on the module before the forward mask is replaced with identity, and save_lora_adapters reads it back through _read_mlx_lora_dropout, so checkpoints saved after preference training keep the configured dropout rather than recording 0.0. Covered by a regression test.
…h_randperm _mlx_disable_lora_dropout now stashes each adapter's configured lora_dropout on _unsloth_orig_lora_dropout before replacing the dropout module with an identity, and the adapter-save path reads it via a new _read_mlx_lora_dropout helper. Previously, disabling adapter dropout for the deterministic ORPO/DPO forward permanently erased the probability the save path extracts from the live module, so an adapter trained with nonzero lora_dropout saved dropout=0.0 in adapter_config.json and no longer matched the original LoRA config on reload. The forward stays a deterministic identity; the saved metadata stays correct. create_preference_batches gains a num_epochs param so epoch-based runs (max_steps<=0, num_train_epochs>1, which leave num_batches None) with dataset_order='torch_randperm' materialize one reseeded (base_seed+epoch) permutation per epoch, matching the SFT create_ordered_batches per-epoch reseed and CUDA RandomSampler. Previously they built a single permutation that the trainer's modulo-cycled loop replayed every epoch. Epoch 0 is byte-identical to the single pass; default (length-sort) and sequential are unchanged.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2659c53114
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # 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) |
There was a problem hiding this comment.
When gradient_accumulation_steps > 1 and ORPO microbatches have different chosen lengths, returning the pair count makes the trainer scale the entire loss by pairs before averaging accumulated gradients. The odds-ratio term is pair-mean, but sft above is explicitly a token-mean over nll_mask; a single combined batch would weight that SFT gradient by chosen-token count, while this path weights each microbatch by pair count, so short-response microbatches overcontribute to the SFT component. Either the SFT component needs separate token weighting, or it must be reformulated as a per-pair mean before using B as the accumulator count.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
TRL 0.25.1 forms the ORPO loss as policy_nll_loss - losses.mean() (orpo_trainer.get_batch_loss_metrics), where policy_nll_loss is a per-microbatch CrossEntropyLoss token-mean and the odds-ratio term is a per-pair mean, and HF Trainer then weights each microbatch equally (loss / gradient_accumulation_steps). The MLX loss returns the pair count B as the accumulation weight, so for equal-size microbatches it reproduces TRL's per-pair batch-mean of (nll_token_mean + beta * ratio_pair_mean) exactly (test_orpo_loss_returns_pair_count_not_token_count). Token-weighting the SFT term across accumulation is not what TRL does and would diverge from it, so this stays as parity with no change.
| prompt_chosen_str = tokenizer.apply_chat_template( | ||
| list(prompt) + list(chosen), tokenize=False, | ||
| **extra_kwargs, |
There was a problem hiding this comment.
For conversational preference rows where prompt already ends with an assistant message, the branch above explicitly sets continue_final_message=True for the standalone prompt, but the chosen/rejected full strings are rendered as list(prompt) + list(chosen). When chosen/rejected are assistant continuations, this creates two consecutive assistant turns instead of continuing the final assistant message; chat templates that enforce role alternation can raise, and permissive templates insert an EOS/new assistant marker that shifts the response span. This only affects assistant-prefill preference examples, but those are accepted by the last_role == "assistant" branch and should be rendered by merging the continuation into the final prompt message.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
TRL 0.25.1 apply_chat_template (data_utils.py) renders a preference completion as a plain message-list concat, apply_chat_template(example["prompt"] + example["chosen"]), with no continue_final_message and no merge, then derives the completion by string prefix against the continue_final_message-rendered prompt. The MLX path does the identical list(prompt) + list(chosen) render and derives the boundary via token common-prefix, so an assistant-ending prompt produces the same two-turn structure and the same response span as TRL. This is parity, not a merge divergence, so no change.
| batches = create_preference_batches( | ||
| dataset=self.train_dataset, | ||
| tokenizer=self.tokenizer, |
There was a problem hiding this comment.
When callers use the new ORPO/DPO trainers with formatting_func to convert raw rows into prompt/chosen/rejected, this branch still passes the unformatted dataset directly to create_preference_batches. Those runs either fail the column check in the preference collator or train on stale raw fields, whereas the SFT path applies self.formatting_func during batch construction. Thread the formatter through the preference path before validating/building batches so formatted preference datasets work consistently with the exposed constructor argument.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
TRL 0.25.1 DPOTrainer/ORPOTrainer and DPOConfig/ORPOConfig expose no formatting_func argument (it is an SFTTrainer-only concept); the MLX preference path correctly consumes the prompt/chosen/rejected column contract that TRL's DPO/ORPO require. Applying an SFT formatting_func to preference rows is not TRL behavior, so threading it through would add a non-TRL feature rather than fix a correctness bug. Out of scope, so no change.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2659c53114
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| prompt_str = tokenizer.apply_chat_template( | ||
| prompt, tokenize=False, | ||
| add_generation_prompt=add_gen, continue_final_message=cont, | ||
| **extra_kwargs, | ||
| ) | ||
| prompt_chosen_str = tokenizer.apply_chat_template( |
There was a problem hiding this comment.
For conversational ORPO/DPO rows, this calls apply_chat_template on the original tokenizer object, while encoding later unwraps wrappers via _hf_encoding_tokenizer. In the MLX wrapper case that only exposes the Hugging Face tokenizer under _tokenizer (the same case _get_processor_tokenizer handles elsewhere before rendering chat messages), message-list preference datasets fail during batching with AttributeError even though the inner tokenizer can render them. Use the unwrapped tokenizer for these chat-template calls as well.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Verified against mlx-lm source and leaving as-is. mlx-lm's TokenizerWrapper defines its own apply_chat_template (tokenizer_utils.py in 0.30.x/0.31.x), so calling it on the wrapper does not raise AttributeError; the wrapper injects enable_thinking=self.has_thinking and then delegates to self._tokenizer.apply_chat_template. Rendering through the wrapper is therefore the intended path: switching to the unwrapped _tokenizer would drop that enable_thinking default and diverge from how mlx-lm itself renders chat prompts. The encode calls below already unwrap via _hf_encoding_tokenizer because they need the list-returning encode and eos_token_id, which is a separate concern from chat rendering.
| if prompt_key not in ex or chosen_key not in ex or rejected_key not in ex: | ||
| raise ValueError( | ||
| f"ORPO requires '{prompt_key}', '{chosen_key}', '{rejected_key}' " | ||
| f"columns; got {sorted(ex.keys())}." |
There was a problem hiding this comment.
This rejects preference datasets that contain only chosen and rejected with the shared prompt embedded in both responses, which is a supported/common DPO/ORPO input format and is handled by extracting the common prefix before tokenization. In that scenario MLX DPO/ORPO now fails before training unless users manually pre-split a prompt column, so the new trainer is incompatible with those datasets instead of deriving prompt, chosen, and rejected during batching.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Leaving as a documented explicit-column contract. TRL supports implicit prompts through a separate maybe_extract_prompt dataset .map() preprocessing step, not inside apply_chat_template; the MLX path has no equivalent preprocessing stage and errors clearly requiring prompt/chosen/rejected rather than silently mistraining. Out of scope for this correctness pass.
| if last_role == "user": | ||
| add_gen, cont = True, False | ||
| elif last_role == "assistant": | ||
| add_gen, cont = False, True | ||
| else: |
There was a problem hiding this comment.
When a conversational preference prompt ends with a tool message after a tool/function call result, this falls into the ValueError path instead of rendering the next assistant turn. Tool-use preference datasets with prompt=[..., {"role": "tool", ...}] therefore fail during ORPO/DPO batching even though the chat template can render them; handle tool like the generation-prompt case rather than rejecting it.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
This matches TRL 0.25.1. TRL's apply_chat_template/maybe_apply_chat_template also raises when the last message role is tool (data_utils.py), so rejecting a trailing tool role during preference batching is parity, not a regression. Leaving as-is.
| # 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) |
There was a problem hiding this comment.
With gradient_accumulation_steps > 1 and chosen rows of different lengths, returning the pair count for the combined ORPO loss makes the trainer scale the SFT/NLL component by pairs. The odds-ratio term is pair-mean, but the NLL component above is a token-mean over all chosen tokens, so accumulated gradients no longer match a single batch of the same examples and long completions are misweighted; split the denominators or rescale the NLL term separately before returning.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
TRL 0.25.1 forms the ORPO loss as policy_nll_loss - losses.mean() (orpo_trainer.get_batch_loss_metrics), where policy_nll_loss is a per-microbatch CrossEntropyLoss token-mean and the odds-ratio term is a per-pair mean, and HF Trainer then weights each microbatch equally (loss / gradient_accumulation_steps). The MLX loss returns the pair count B as the accumulation weight, so for equal-size microbatches it reproduces TRL's per-pair batch-mean of (nll_token_mean + beta * ratio_pair_mean) exactly (test_orpo_loss_returns_pair_count_not_token_count). Token-weighting the SFT term across accumulation is not what TRL does and would diverge from it, so this stays as parity with no change.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2659c53114
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| loss_type: str = "sft" # "sft" or "orpo" | ||
| orpo_beta: float = 0.1 # ORPO odds-ratio weight (TRL default) | ||
| dpo_beta: float = 0.1 # DPO beta (TRL default) | ||
| reference_free: bool = False # DPO: drop the reference term if True |
There was a problem hiding this comment.
Because MLXTrainingConfig.__init__ accepts positional arguments in dataclass field order, inserting these preference fields before the existing load_best_model_at_end/metric fields shifts the meaning of every positional config argument after eval_steps. Existing calls that set eval options positionally now populate loss_type/beta fields instead, so they can fail later as an unsupported loss type or run with corrupted settings; add new fields after the existing fields (or make the constructor keyword-only) to keep old positional configs compatible.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Not reachable. MLXTrainingConfig and its subclasses are constructed only by keyword or with no arguments everywhere in this codebase (grep shows MLXTrainingConfig(), MLXORPOConfig(), MLXDPOConfig() and user-facing MLXTrainingConfig(field=...) calls, never positional), mirroring HF TrainingArguments which is likewise always keyword-constructed. A 20+ field config passed by positional order is not a legitimate reachable call, so field ordering does not shift any real caller's arguments; leaving the fields where they are.
| prompt, chosen_full_text, rejected_full_text = _render_preference_example( | ||
| tokenizer, prompt, ex[chosen_key], ex[rejected_key], | ||
| tools=_ex_tools, chat_template_kwargs=_ex_template_kwargs, |
There was a problem hiding this comment.
When ORPO/DPO is given conversational prompt/chosen/rejected rows and the tokenizer is an mlx-lm-style wrapper, this passes the outer wrapper into _render_preference_example; unlike the encode calls below that use the unwrapped hf, that wrapper may not expose apply_chat_template even though its HF tokenizer under _tokenizer does. The SFT rendering path unwraps before applying templates, so this path can crash before tokenization for wrapped tokenizers; render with hf or unwrap inside the helper.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Verified against mlx-lm source and leaving as-is. mlx-lm's TokenizerWrapper defines its own apply_chat_template (tokenizer_utils.py in 0.30.x/0.31.x), so calling it on the wrapper does not raise AttributeError; the wrapper injects enable_thinking=self.has_thinking and then delegates to self._tokenizer.apply_chat_template. Rendering through the wrapper is therefore the intended path: switching to the unwrapped _tokenizer would drop that enable_thinking default and diverge from how mlx-lm itself renders chat prompts. The encode calls below already unwrap via _hf_encoding_tokenizer because they need the list-returning encode and eos_token_id, which is a separate concern from chat rendering.
| "Expected 'default', 'sequential', or 'torch_randperm'." | ||
| ) | ||
| hf = _hf_encoding_tokenizer(tokenizer) | ||
| eos_id = hf.eos_token_id |
There was a problem hiding this comment.
For a simple tokenizer that supports encode but has no eos_token_id (or when callers set append_eos=False specifically because there is no EOS), this direct attribute access raises before any preference batches are built. The existing SFT batching paths use getattr(tokenizer, "eos_token_id", None) and can pad with 0 in that case, so the new ORPO/DPO path should do the same instead of rejecting otherwise usable no-EOS tokenizers.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Not reachable here. hf is the resolved Hugging Face tokenizer (_hf_encoding_tokenizer returns either a PreTrainedTokenizerBase or mlx-lm's unwrapped _tokenizer), and eos_token_id on a PreTrainedTokenizerBase is a property that returns None when no EOS is set rather than raising AttributeError, so a no-EOS HF tokenizer already yields eos_id=None and pads with 0 exactly like the SFT path. The trainer only ever passes an HF/mlx-lm tokenizer into preference batching, so a bare object with encode() but no eos_token_id attribute is not a reachable input; leaving the direct access.
|
@danielhanchen this one now conflicts with main in trainer.py, utils.py, and tests/test_mlx_trainer_internals.py. They're just against your reworks that landed after this branch forked (#922, #930, and #918) For resolving these it'd basically be deciding what ORPO code your reworks adds on vs what actually gets carried forward.The conflicts are the only thing blocking this now |
|
@BardiaKoopah Fix merge conflicts thanks - you decide what is best |
Reconcile the ORPO branch with main's post-fork reworks (unslothai#922 true grad-norm reporting, unslothai#930 fused CCE + label smoothing, unslothai#918 Metal resource-limit CPU batch plans). ORPO was built alongside those (new loss/batch functions and tests), so the conflicts were keep-both rather than supersede, with two deliberate exceptions where main's unslothai#930 supersedes older baseline code: - trainer.py: keep the ORPO/DPO imports, the dropout-neutralization helpers, the loss_type guard, and the ORPO/DPO loss dispatch; adopt main's _validate_label_smoothing / resume-checkpoint helpers and streaming init alongside them. In the loss dispatch, keep the ORPO/DPO branches and adopt main's label_smoothing= on the CCE fallback (an elif that ORPO/DPO never reach, so it cannot affect the preference paths). - utils.py: keep the ORPO/DPO block (odds-ratio loss, make_orpo/dpo_loss_fn, preference batching, render helpers -- all branch-only additions) and take main's make_baseline_loss_fn(label_smoothing=0.0). Verified every one of main's utils functions is preserved: make_cce_loss_fn and the batch-plan functions are byte-identical to main; _enrich_mlx_adapter_config was changed by both sides and git auto-merged it -- confirmed line-by-line that main's quantization-map flag and the branch's lora_dropout stash both survive. - tests: keep both the ORPO/DPO regression tests and main's streaming/prefetch tests (disjoint, no name collisions). Verified: a real ORPO LoRA run trains with finite, decreasing loss; unslothai#922/unslothai#930/ unslothai#918 test groups pass; canonical baseline 235 passed, 0 failures.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4b49b054da
ℹ️ 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".
| loss_type: str = "sft" # "sft" or "orpo" | ||
| orpo_beta: float = 0.1 # ORPO odds-ratio weight (TRL default) | ||
| dpo_beta: float = 0.1 # DPO beta (TRL default) | ||
| reference_free: bool = False # DPO: drop the reference term if True |
There was a problem hiding this comment.
These new preference fields are inserted before pre-existing config fields, but MLXTrainingConfig.__init__ maps positional arguments using dataclass declaration order. Any saved/forwarded positional config built against the previous order will now put the old load_best_model_at_end value into loss_type (often False), so ordinary training later fails the new loss_type validation or shifts all following settings. Append these fields to the end (and treat them as optional-copy fields if needed) to keep legacy positional copies working.
Useful? React with 👍 / 👎.
| if prompt_key not in ex or chosen_key not in ex or rejected_key not in ex: | ||
| raise ValueError( | ||
| f"ORPO requires '{prompt_key}', '{chosen_key}', '{rejected_key}' " | ||
| f"columns; got {sorted(ex.keys())}." | ||
| ) |
There was a problem hiding this comment.
When an ORPO/DPO dataset uses the common implicit-prompt preference format (chosen and rejected only, with the shared prefix serving as the prompt), this guard raises before batching. The new trainers otherwise aim for TRL preference-format parity, and TRL extracts that shared prompt automatically; rejecting these rows means users must manually preprocess otherwise valid preference datasets before MLX training can start.
Useful? React with 👍 / 👎.
| if ( | ||
| args.max_steps <= 0 | ||
| and getattr(args, "num_train_epochs", -1) > 0 | ||
| and pref_dataset_order == "torch_randperm" | ||
| ) |
There was a problem hiding this comment.
For epoch-based ORPO/DPO runs with the default dataset order (max_steps <= 0, num_train_epochs > 1), this condition leaves pref_num_epochs unset unless the user opted into torch_randperm. create_preference_batches then returns a single length-sorted pass and the training loop cycles that identical order every epoch, unlike the SFT default plan which reshuffles batch visits across epochs; default preference runs therefore repeatedly train in the same short-to-long curriculum. Thread epoch expansion (or a finite visit plan) for the default order as well.
Useful? React with 👍 / 👎.
| for md in _mods: | ||
| md.scale = 0.0 |
There was a problem hiding this comment.
When DPO is run from a reloaded/pretrained LoRA adapter, zeroing every adapter scale here makes the reference logits come from the bare base model instead of the policy at the start of DPO. That is only correct for freshly initialized zero-delta adapters; the supported warm-start path returns trainable adapter params for continued training, so those DPO runs compare against the wrong reference and optimize corrupted rewards. Keep a frozen copy of the initial adapter or reject nonzero warm-start adapters rather than disabling it.
Useful? React with 👍 / 👎.
| else: | ||
| raise ValueError( | ||
| "Unsloth MLX preference: conversational chat prompt has invalid " | ||
| f"last-message role {last_role!r}; expected 'user' or 'assistant'." |
There was a problem hiding this comment.
For conversational preference rows whose prompt ends with a tool result (role == "tool"), this branch raises before rendering. Function-calling preference datasets commonly place the tool response as the final prompt turn before the chosen/rejected assistant answer, and this new renderer already threads per-row tools; rejecting that role means otherwise valid tool-use ORPO/DPO rows cannot be batched or trained. Treat non-assistant prompt endings like generation prompts instead of hard-failing.
Useful? React with 👍 / 👎.
Reconcile ORPO/DPO with main's unslothai#873 HF-style callback support: - ORPO-A: keep main's deferred _prepare_eval_batches() chokepoint; add the preference eval-skip inside it so neither the static cadence nor a callback request runs eval on SFT-shaped batches, and keep the setup skip notice. - ORPO-B: keep both sides' new tests (preference/DPO suite + main's callback DDP-sync suite); no name collisions. - create_preference_batches path and unslothai#918's batch_kwargs["grad_accum"] carried in via clean auto-merge.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fb2a822d40
ℹ️ 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".
| def __init__(self, model, tokenizer, train_dataset, eval_dataset=None, | ||
| dataset_text_field=None, max_seq_length=None, packing=None, | ||
| data_collator=None, args=None, formatting_func=None, processor=None): |
There was a problem hiding this comment.
When users switch from MLXTrainer to the advertised MLXORPOTrainer or MLXDPOTrainer classes and pass the same callbacks= argument that the base trainer supports, this narrower constructor raises TypeError before training starts; the DPO constructor below has the same omission. This blocks callback-driven logging, early stopping, cancellation, and integrations specifically for the new preference-trainer entry points, so the subclasses should accept and forward callbacks to MLXTrainer.__init__.
Useful? React with 👍 / 👎.
| # SequentialSampler), so their per-epoch order already matches SFT. | ||
| # Step-based runs (num_batches set) are unchanged. | ||
| pref_num_epochs = ( | ||
| int(args.num_train_epochs) |
There was a problem hiding this comment.
For epoch-based ORPO/DPO runs with dataset_order="torch_randperm" and a non-integer num_train_epochs such as 0.5 or 1.5, this cast truncates the requested epoch count before materializing preference batches. Because _prepared_batches_include_epochs is then set, the step resolver budgets from that truncated schedule, so a 1.5-epoch run trains only one pass (and a 0.5-epoch run still materializes a full pass) instead of following the same fractional-epoch arithmetic used by the SFT ordered path.
Useful? React with 👍 / 👎.
| loss_type: str = "sft" # "sft" or "orpo" | ||
| orpo_beta: float = 0.1 # ORPO odds-ratio weight (TRL default) | ||
| dpo_beta: float = 0.1 # DPO beta (TRL default) | ||
| reference_free: bool = False # DPO: drop the reference term if True |
There was a problem hiding this comment.
When loading or keyword-copying an MLXTrainingConfig dump from before these preference fields existed, the dump legitimately omits loss_type/beta/reference fields while still carrying the old default warmup_steps=5 and a non-default warmup_ratio; because these new fields are not included in the initializer's appended-field allowlist, copied_all_fields becomes false and _resolve_warmup_steps treats the default 5 steps as explicit, silently replacing the ratio schedule. This is the keyword legacy-dump path, not the earlier positional-argument concern, so these fields need to be tolerated as optional copied fields as well.
Useful? React with 👍 / 👎.
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.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fb2a822d40
ℹ️ 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 num_batches is not None and len(out) >= num_batches: | ||
| break | ||
| mx.eval([b for b, _, _ in out] + [l for _, l, _ in out]) | ||
| return out |
There was a problem hiding this comment.
When max_steps > 0 and a preference dataset pass is ragged relative to gradient_accumulation_steps, returning a plain list here loses the one-pass cycle_length that the trainer uses to force HF-style epoch-final optimizer steps. For example, 3 preference batches with gradient_accumulation_steps=2 should update on batch 2 and then force-update the epoch tail on batch 3; this list path reports no cycle length, so the tail is accumulated with the next epoch’s first batch instead, changing gradients and callback/save/eval epoch boundaries versus the SFT finite-plan path. Return a finite plan or otherwise expose the per-pass batch count for preference batches.
Useful? React with 👍 / 👎.
| for chunk in chunks: | ||
| Lmax = max(max(len(c), len(r)) for _, c, r in chunk) | ||
| if pad_to_multiple: | ||
| Lmax = ((Lmax + pad_to_multiple - 1) // pad_to_multiple) * pad_to_multiple |
There was a problem hiding this comment.
For ORPO/DPO batches whose padded row length is not capped by max_seq_length, this rounds the full stored row width to a multiple of 32, but the model actually forwards batch[:, :-1], so the sequence length seen by the model is always one less than that multiple (for example a 32-token row forwards 31 positions). The SFT batcher uses _finite_text_pad_width, which rounds to 1 + multiple specifically so the shifted input length is the Apple-Silicon-friendly multiple; preference runs therefore miss the intended padding/shape parity and can run slower or with different shape buckets. Use the same shifted width policy here.
Useful? React with 👍 / 👎.
| if tools is not None: | ||
| extra_kwargs["tools"] = tools | ||
| prompt_str = tokenizer.apply_chat_template( | ||
| prompt, tokenize=False, |
There was a problem hiding this comment.
When a text conversational preference row uses OpenAI-style content parts (for example content=[{"type":"text", "text":"..."}]), this passes the raw messages directly into the tokenizer template. The SFT chat path first runs _normalize_mlx_messages(..., is_vlm=False), which flattens those parts into strings; without the same normalization here, text ORPO/DPO can either fail in templates that expect string content or train on the Python list representation instead of the intended text. Normalize prompt/chosen/rejected messages before these apply_chat_template calls.
Useful? React with 👍 / 👎.
_render_preference_example passed raw prompt/chosen/rejected straight
into apply_chat_template, unlike the SFT chat path which first runs
_normalize_mlx_messages(..., is_vlm=False). A conversational ORPO/DPO row
carrying OpenAI-style content PARTS (content=[{"type":"text","text":...}])
therefore diverged by template:
- Qwen2.5 / SmolLM-Instruct build the turn with '...' + message.content
and raise TypeError: can only concatenate str (not "list") to str.
- Qwen3 guards the body with {%- if message.content is string %} and
leaves it empty otherwise, so every turn renders EMPTY and silently:
prompt+chosen and prompt+rejected come out byte-identical, the ORPO
odds-ratio term collapses to -logsigmoid(0) = log 2 (a constant with
zero gradient), and the run trains on empty assistant turns while its
loss curve still looks plausible.
Measured on Qwen3-0.6B through create_preference_batches and
make_orpo_loss_fn: L_OR 0.6931471825 with |grad| 1.08e-08 before, versus
0.6900140643 with |grad| 1.20e-01 after.
Flatten the three message lists inside the conversational branch. String
rows are returned unchanged by _normalize_mlx_messages, so the
non-conversational path stays byte-identical.
Tests cover both template families as stubs plus the three real
tokenizers, and pin that plain string rows are untouched.
create_preference_batches is eager, so it returned a bare list. The trainer reads the one-pass micro-batch count as getattr(batches, "cycle_length", None) in _mlx_epoch_microbatches and _callback_batches_per_epoch, and a plain list cannot carry it, so both fell back to None under max_steps > 0. A ragged pass then never force-updated on its tail: with 3 preference batches and gradient_accumulation_steps=2 the run should update on batch 2 and apply the epoch-final batch 3 by itself, but instead accumulated batch 3 together with the next pass's batch 1. That changes the effective gradient and moves the callback/save/eval epoch boundaries off the SFT finite-plan path. Return a small list subclass carrying the count. len(base_starts) is the one-pass length in every ordering branch: default/torch_randperm with num_batches cycle base_starts (or reseeded row permutations of the same length), the num_epochs branch concatenates whole passes of it, and sequential/single-pass use it directly. It is deliberately not clamped to len(batches) -- num_epochs makes the list several passes long and a num_batches horizon can truncate it below one pass -- matching how the SFT/VLM plans count theirs. Subclassing list keeps indexing and modulo-cycling identical, and the container stays outside _FINITE_BATCH_PLAN_TYPES / _EAGER_REFETCHABLE_PLAN_TYPES so these already-materialized batches are never refetched. __slots__ on a list subclass leaves no __dict__ for the default reducer, so __reduce__/__setstate__ are spelled out.
|
@codex review |
|
Codex Review: Didn't find any major issues. Another round soon, please! Reviewed commit: ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
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". |
Compress the verbose comment blocks and docstrings added by this branch, plus two boilerplate banners. Comments only -- no code, control flow or literal is touched, verified by an AST comparison that strips docstrings from both sides and compares ast.dump(). Covers only blocks whose removed lines this branch authored; the review comments contributed by others are split into follow-up commits so they can be reviewed or dropped independently.
…r.py These blocks were authored by @danielhanchen in the 2026-07-08/09 review commits on this branch. Intent is preserved -- each keeps the defect it records and the reason for the fix -- but the narrative is cut to one or two lines. Deliberately kept, because they justify decisions that read as bugs otherwise: why cycle_length is not clamped to len(batches) (num_epochs expands, num_batches truncates), why PreferenceBatchList needs an explicit __reduce__ under __slots__, and the failure mode the three _normalize_mlx_messages calls fix. Split out so it can be reviewed or reverted without touching the rest of the comment pass. Comments only; AST verified identical.
These 26 blocks are @danielhanchen's "# Regression: ..." notes from the 2026-07-08/09 review commits, each recording a defect found in review and why the naive approach was wrong. Compressed, not removed: every block keeps the defect and the reason. The TRL source citations are kept, including the verbatim quote in the ORPO-NLL block -- "orpo chosen nll loss is computed over the full prompt and response" (orpo_trainer.py:780-785). That quote is the evidence for why the SFT term spans prompt+response; without it the term could be "fixed" back to a response-only mask, which is the exact defect the note was written to prevent. Isolated in its own commit because these are the review record and the most likely thing to want back. Comments only; AST verified identical.
tests/test_mlx_preference_content_parts.py and tests/test_mlx_preference_cycle_length.py were collected by the --collect-only matrix but executed by no job, so they only ever import-checked. That is the blind spot the gate's own comment records for test_mlx_finetune_last_n_layers (born broken in unslothai#669, invisible until unslothai#739). Add both to the HARD GATE invocation rather than the continue-on-error one, so a failure actually reds the job. Both are deterministic, CPU-pure and sub-second under the torch shim, matching that step's criteria.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7ac5a2bcac
ℹ️ 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".
| else None | ||
| ) | ||
| batches = create_preference_batches( | ||
| dataset=self.train_dataset, |
There was a problem hiding this comment.
When train_dataset is a sized text dataset, MLXTrainer.__init__ has already replaced self.train_dataset with _MLXTokenizedDatasetView, whose iterator runs the SFT text/chat renderer to synthesize input_ids before yielding each row. For ORPO/DPO rows that already have valid prompt/chosen/rejected columns but also carry an unused messages or chat-style extra column, this can raise on a missing/mismatched chat template before the preference collator ever reads the explicit preference fields. Pass the train_dataset variable resolved from _train_dataset_for_batches() here so extra columns are not preprocessed by the SFT view.
Useful? React with 👍 / 👎.
|
@BardiaKoopah thanks for the work on this. The core ORPO/DPO implementation looks solid, but I think these three remaining issues need to be addressed before merging:
Please include focused regression tests for these cases. While touching this path, it would also be good to separate the pair-count loss weight from actual trained-token telemetry. |
5 participants
Add ORPO and DPO preference tuning for text models to MLXTrainer
Adds two preference-tuning methods — ORPO and DPO — to the MLX trainer, bringing preference optimization to the Apple-Silicon path. On CUDA, Unsloth gets these by patching TRL's trainers (
rl_replacements.py); since there is no TRL on MLX, these are native implementations built on the existing MLX trainer, sharing a common preference data collator (mirroring TRL's singleDataCollatorForPreference).Enable via
loss_type="orpo"orloss_type="dpo". Default ("sft") is unchanged.ORPO
Combines SFT with a preference signal in one loss, no reference model:
L = L_SFT + beta * L_OR, whereL_OR = -log(sigmoid(log_odds_chosen - log_odds_rejected))using the fullp/(1-p)odds (paper-faithful, not a simplified log-prob ratio), with the1 - exp(logp)term numerically stabilized.Config:
orpo_beta: float = 0.1(TRL default).DPO
Compares the policy against a frozen reference:
L = -log(sigmoid(beta * [(logp_chosen^pol - logp_chosen^ref) - (logp_rejected^pol - logp_rejected^ref)])).Reference via live LoRA-disable (matching TRL's
is_peft_modeldefault — no second model copy). The base model is the reference: we temporarily zero the LoRA scales (mlx_lm'sLoRALinear.__call__isy + scale*z, so scale=0 disables the adapter), run the reference forward undermx.stop_gradient, and restore the scales in afinallyblock. Verified the scale restore is bit-exact and survivesvalue_and_gradtracing.Config:
dpo_beta: float = 0.1,reference_free: bool = False(drops the reference term, matching TRL). Precompute of reference log-probs (TRL'sprecompute_ref_log_probs) is a possible follow-up optimization; this PR computes the reference live each step.Shared design
concatenated_forward: chosen and rejected stacked into one batch (chosen block then rejected block), single forward, split in the loss. Verified MLXvalue_and_graddifferentiates cleanly through the concat-and-split.create_preference_batches: orders pairs perdataset_order—"default"(length-sort, and the default, so existing runs are unchanged),"sequential", or"torch_randperm", withpreserve_dataset_order=Trueforcing sequential — then dynamic-pads to a multiple of 32 (Apple-Silicon padding) and produces(batch, lengths, None)consumed by the existing text step path. Both ORPO and DPO use it.make_baseline_loss_fn's(model, batch, lengths, labels)signature; the batch path short-circuits in_prepare_data.Scope
Text models only. VLMs raise a clear not-supported error; eval during preference tuning prints a skip notice and trains normally. Both are follow-ups.
Preference collation accepts both row formats: plain-string
prompt/chosen/rejected, and conversational (message-list) rows, which are rendered through the tokenizer chat template. Conversational rows also handle OpenAI-style content parts (content=[{"type": "text", ...}]), per-rowtools, and per-rowchat_template_kwargs.The returned batch list exposes
cycle_length, the micro-batch count of one dataset pass, so the trainer can force the epoch-final optimizer step on a ragged tail instead of folding it into the next pass.Not yet supported, both follow-ups:
max_prompt_length/max_completion_lengthtruncation knobs. Over-long rows are truncated prompt-first, keeping the end of the prompt so the completion survives with a non-empty loss span, but the separate named length budgets do not exist.chosen/rejectedcolumns. An explicitpromptcolumn is required.Testing (Qwen2.5-0.5B, small preference sets)
0.6931 = -log(0.5)exactly (untrained adapter → policy == reference → zero logits), the theoretically-correct starting value, then moves as the policy diverges. DPO loss math verified directionally in a controlled test (GOOD < NEUTRAL < BAD, NEUTRAL = -log(0.5)). Reference scale-toggle verified bit-exact and gradient-safe across multiple steps (no leak).test_cce_*) that fail identically onmainwithout this change — a torch-shim issue unrelated to this work.Smoke tests confirm the objectives run correctly and the losses behave as expected; they do not measure downstream preference-alignment quality (needs a full eval, out of scope). Tested on small datasets.