Add quantization_config trainer argument (streamline QLoRA)

[qgallouedec]

Adds a quantization_config argument to SFTTrainer, DPOTrainer, GRPOTrainer, RLOOTrainer, and RewardTrainer, so QLoRA no longer requires reaching into model_init_kwargs (or worse, manual model loading)

After:

SFTTrainer(
    model="meta-llama/Llama-2-7b-hf",
    quantization_config=BitsAndBytesConfig(load_in_4bit=True),
    peft_config=LoraConfig(),
    train_dataset=dataset,
)

Compare with before (many ressources are written like this!):

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=BitsAndBytesConfig(load_in_4bit=True),
    device_map="auto",
)
SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=LoraConfig(),
)

Before (the "right" way, but not very popular):

SFTTrainer(
    model="meta-llama/Llama-2-7b-hf",
    args=SFTConfig(model_init_kwargs={"quantization_config": BitsAndBytesConfig(load_in_4bit=True)}),
    peft_config=LoraConfig(),
    train_dataset=dataset,
)

It sits next to peft_config (the other non-serializable QLoRA ingredient), flows into from_pretrained, and raises if also set in args.model_init_kwargs.

Changes

• New quantization_config arg on the five trainers above (+ docstrings).
• The trl/scripts/{sft,dpo,grpo,rloo,reward}.py CLIs now pass it directly instead of injecting into model_init_kwargs.
• This drops the redundant model_init_kwargs["device_map"] = get_kbit_device_map() line: verified on 8×H100 that QLoRA trains identically with and without it, across transformers 4.56.2 (min supported) and 5.13; distributed runs override device_map to None anyway, and single-process runs auto-place quantized weights on the current CUDA device. See Remove redundant get_kbit_device_map() #6158
• Updated the QLoRA example in docs/source/peft_integration.md.

---

Note

Medium Risk
Touches model loading for all major TRL trainers and reference-model paths; behavior change for QLoRA users but scoped to optional loading kwargs with explicit conflict checks.

Overview
Adds a quantization_config trainer argument (alongside peft_config) on SFTTrainer, DPOTrainer, GRPOTrainer, RLOOTrainer, and RewardTrainer, so QLoRA can pass a model id string and let the trainer load/quantize via from_pretrained instead of pre-loading with AutoModelForCausalLM.

When the model is loaded from a string, the trainer merges quantization_config into model_init_kwargs (and the same for reference models where applicable), errors if it is also set in args.model_init_kwargs, and warns if a pre-instantiated model is passed. CLI entrypoints and example scripts now pass get_quantization_config(model_args) directly to the trainer and no longer inject quantization_config / get_kbit_device_map() into model_init_kwargs.

Docs and Colab notebooks are updated to the new pattern (model id + quantization_config on the trainer, model_init_kwargs for attn/dtype where needed).

^{Reviewed by Cursor Bugbot for commit 7be97c4. Bugbot is set up for automated code reviews on this repo. 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: 45d6a2decd

ℹ️ 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]

Preserve positional peft_config compatibility

Adding quantization_config before the existing peft_config parameter shifts any current positional peft_config argument into quantization_config because this public constructor is not keyword-only. In existing calls that pass peft_config positionally, a model id will forward a PeftConfig object to from_pretrained(..., quantization_config=...) and fail, while an already-instantiated model will ignore it and train without the adapter; the same signature insertion appears in the other updated trainers. Put the new argument after peft_config or otherwise preserve the old positional layout.

Useful? React with 👍 / 👎.

[qgallouedec]

although not specifically disallowed, it would be very surprising that peft_config is used as positional arg

[cursor]

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

There are 3 total unresolved issues (including 2 from previous reviews).

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

^{Reviewed by Cursor Bugbot for commit 0bb426c. Configure here.}

[sergiopaniego]

the example scripts/notebooks in the examples/ folder should also be reviewed and updated

[qgallouedec]

right @sergiopaniego , updated!