Align KTO with DPO: quantization_config trainer argument

[qgallouedec]

Ports the quantization_config trainer argument added to DPOTrainer in #6157 to KTOTrainer, to streamline QLoRA (pass a BitsAndBytesConfig directly instead of stuffing it into model_init_kwargs).

Mirrors DPO exactly. No behavior change unless quantization_config is passed.

---

Note

Low Risk
Additive API and script refactor aligned with existing DPO behavior; default training paths are unchanged unless quantization_config is passed.

Overview
KTOTrainer now accepts an optional quantization_config (BitsAndBytesConfig) when loading from a model id, matching DPOTrainer: it is merged into model_init_kwargs for the policy and auto-created reference model, with a ValueError if both the argument and args.model_init_kwargs set quantization, and a warning if the model is already instantiated.

The trl/scripts/kto.py entry point no longer loads model, ref model, and tokenizer up front; it passes the model path to the trainer, sets training_args.model_init_kwargs from ModelConfig, and wires get_quantization_config(model_args) and get_peft_config(model_args) for QLoRA-style runs.

^{Reviewed by Cursor Bugbot for commit 6fd39d8. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 6a7e848. Configure here.}

[cursor]

Ref load ignores quantization warning

Medium Severity

When model is already instantiated, the trainer logs that quantization_config will be ignored, but auto-created reference models still inject that config into ref_model_init_kwargs. Full fine-tuning without PEFT can then pair an unquantized policy with a separately loaded quantized reference, producing inconsistent reference log-probs.

Additional Locations (1)

• trl/experimental/kto/kto_trainer.py#L770-L772

 

^{Reviewed by Cursor Bugbot for commit 6a7e848. Configure here.}

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 6a7e848653

ℹ️ 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".

[chatgpt-codex-connector]

Avoid quantizing only the implicit reference model

When callers pass an already-instantiated non-PEFT model together with quantization_config, the trainer warns above that the config is ignored because the model is already instantiated, but this branch still injects it while auto-creating the reference model. In that scenario KTO compares an unquantized policy against a 4/8-bit reference instead of the initial policy, changing the loss despite the argument being reported as ignored; only propagate this config here when it was also used to load the policy model, or require an explicit matching ref_model.

Useful? React with 👍 / 👎.