Add native MLX SimPO preference training with explicit objective parity #1999

Source baseline

Soup origin/main at 3da204fadbd278c01ad5a4b23906b443d15e9dd1.

• Soup SimPO configuration and data contract
• Existing Transformers wrapper
• Current native task refusal

Problem and current evidence

task: simpo can train paired preferences on the Transformers path, but MLX rejects it at config load. The public loader refuses an otherwise valid dense config with task='simpo'; the matching Transformers control loads. SimPO offers preference alignment without a separate reference policy, useful on memory-constrained Apple Silicon.

Proposed initial scope and objective

• Implement native text-only causal SimPO with standard LoRA and quantization: none, initially one documented architecture. Use existing task: simpo, simpo_gamma and cpo_alpha; handle the existing task: preference / preference_loss: simpo dispatch consistently if supported, otherwise refuse that alias explicitly. No Torch/TRL model execution or backend fallback on the native path.
• Match Soup's actual objective rather than silently substituting paper-only SimPO. Soup forwards loss_type="simpo", gamma and the positive NLL weight cpo_alpha to CPOConfig, but does not forward a beta field. At the supported TRL floor v0.29.0, beta defaults to 0.1 and label smoothing to zero: for completion-mean log probabilities a and b, the preference term is mean(softplus(-(beta * (a - b) - gamma))); total loss adds cpo_alpha * chosen_token_NLL. The NLL is a mean over valid chosen tokens, not a mean of per-sequence means. Make these effective settings explicit and persist them. Do not repurpose dpo_beta or orpo_beta, remove the NLL, or enable cpo_alpha: 0 without a separate schema/API decision.
• Use causal shifting, prompt/padding exclusion and completion token counts consistently. Define EOS/chat-template behavior and independent per-answer truncation under the total max_length cap. Refuse an answer with no supervised token after preprocessing before updates; do not copy known TRL truncation defects tracked in SimPO and KTO rows run past max_length on trl 0.29: the two wrappers never call enforce_preference_sequence_limit #1185/Preference prompt cap: should DPO and ORPO cut the prompt only when the row does not fit? #1209 (see active cap-fix PR Fix preference max length enforcement across TRL wrappers #1196) or regress SimPO trains to NaN, behind a logged loss of 0.0, when chosen is longer than max_length #1208. Log preference and NLL components separately with meaningful aggregation.
• Freeze base weights, train actual LoRA tensors, honor the supported optimizer/accumulation settings and save complete native adapter/tokenizer/objective metadata. Refuse unsupported quantization, full tuning, DoRA/LoRA+/FA and other task-specific combinations before costly loading until separately qualified. Do not claim DPO, online RL, VLM or every preference objective as part of this issue.

Acceptance tests

• Public config/train routing accepts the implemented slice; otherwise valid unsupported combinations fail with specific diagnostics before model/output mutation. Include the alias decision and preserve Transformers routing. Use paired prompt/chosen/rejected fixtures, malformed rows and unequal completion lengths.
• An independent numeric oracle verifies the preference term, chosen-token NLL, their weighted total and gradients on nonconstant logits with unequal lengths. Change gamma and cpo_alpha independently; summing instead of averaging completion scores, reversing the pair, omitting the margin or removing NLL must each break the appropriate assertion. Include padded batches, causal shifts and long prompts that formerly erased a completion.
• A tiny native public training run gives finite losses and actual LoRA updates while frozen tensors remain unchanged. Compare accumulated and equivalent effective batches under the stated reduction; check the final partial batch. The measured held-out preference score should move in the intended direction on a controlled learnable fixture, with a base/untrained control.
• Save and reload in a fresh native process, reproducing logits/scores and a nonzero adapter effect within documented tolerances. Missing/mismatched base or adapter metadata, failed writes and incompatible artifacts must give explicit errors without a complete-looking partial result. Separately qualify a pinned pretrained model and report pairwise held-out metrics without promising general alignment gains.

Coordination and alternatives

PR #1846 is native DPO, not SimPO. Reuse suitable routing, token scoring and artifact primitives if available without expanding that PR. #1977 owns exact training-state resume; #1985 owns broader message/history masking. Transformers is the current alternative. Versioned numerical references: TRL v0.29.0 CPO config and trainer. Audit later supported TRL versions before declaring cross-version parity.

Implementation validation and evidence limits

These are acceptance requirements for future implementation, not completed native qualification. Start with a relevant failing regression on the original source, exercise public entry points and real behavior, and verify test sensitivity by restoring the old behavior or removing the relevant condition. Assert results, gradients/updates, artifacts and specific refusal diagnostics; config acceptance, call counts and source-string checks alone do not establish support. Keep default tests offline and deterministic, without model downloads. Native tests must use real MLX operations; mocked or CPU-only tests do not qualify Apple Silicon training/inference.

Run the required non-smoke suite and clean Ruff for an implementation; update feature docs, supported-backend tables, changelog and generated references when affected. Explicitly select new native tests in the existing macOS MLX CI job (or an appropriately scoped evaluation/server job when extra dependencies are necessary); the current job does not discover every new native file automatically. Record source SHA, model/tokenizer revision, data/fixture identity, dependency versions, commands, numerical tolerances and passed/skipped/unexecuted checks. Qualify the declared dependency floor or add an explicit compatibility guard. Keep heavy imports lazy and preserve existing Transformers behavior.

This unclaimed proposal is for maintainer triage. It does not reserve work or start an implementation deadline. Maintainers retain area/difficulty assignment and the maintained-area first-refusal process. Preparation used pinned source/docs, live issue/PR duplicate checks and offline probes where stated; no native training, performance or pretrained-quality result is claimed.

Area

Training (SFT/DPO/GRPO)

Suggested difficulty

Not sure — maintainer triage requested.