Skip to content

v1.8.0

Latest

Choose a tag to compare

@qgallouedec qgallouedec released this 09 Jul 18:41
95809b9

Features

πŸŽ“ KTO is now a stable trainer

After many cycles of KTOTrainer ↔ DPOTrainer alignment work, KTO graduates from trl.experimental.kto to the top-level trl package. Same API as DPO/GRPO/SFT β€” imports move from experimental, tests move to the main test tree, docs no longer flag it as experimental. The experimental path still works and emits a FutureWarning (removal in v2.0.0).

# Before
from trl.experimental.kto import KTOConfig, KTOTrainer

# Now
from trl import KTOConfig, KTOTrainer

Per our telemetry, KTO is the 4th most used trainer in TRL β€” this graduation was overdue.

by @albertvillanova in #6175, #6287 and #6345

Environment-owned rewards & multi-environment support

Three interrelated changes make agentic RL training with environments substantially more ergonomic.

Environment-owned reward. If your environment_factory env defines a reserved get_reward() method (no args β†’ float), it's called once per completed rollout and treated as a reward source. reward_funcs becomes optional β€” no more leaking env state back out to trainer-owned reward funcs.

class WordleEnv:
    def reset(self, **kwargs):
        self._target = sample(words); self._solved = False

    def get_reward(self) -> float:       # optional, reserved (not a tool)
        return 1.0 if self._solved else 0.0

    def guess(self, word: str) -> str:   # exposed as a tool
        self._solved = word == self._target; ...

trainer = GRPOTrainer(
    model=model,
    train_dataset=dataset,
    environment_factory=WordleEnv,       # no reward_funcs needed
)

Multi-environment support. environment_factory now accepts dict[str, factory] in addition to a single callable. Each dataset row selects its environment via an environment column, and only that env's tools are exposed in that row's prompt β€” so a coding task and a game can train together in one run without leaking each other's tool schemas. Single-callable usage is unchanged.

Same wiring lands in GRPO, AsyncGRPO, DPPO, and GRPO-with-replay-buffer.

Env-owned reward by @qgallouedec in #6238; multi-env in #6001 and #6002

Entropy regularization for GRPO

GRPOConfig now supports both static and adaptive entropy regularization (Skywork-OR1). The bonus encourages exploration and helps prevent premature policy collapse.

GRPOConfig(
    entropy_coef=0.01,          # static
    # or:
    use_adaptive_entropy=True,  # adjust to target entropy
    entropy_target=1.5,
    entropy_coef_delta=0.01,
    entropy_coef_min=0.0,
    entropy_coef_max=0.1,
)

Adaptive mode adjusts the coefficient once per optimizer step from window-aggregated entropy (gradient accumulation-aware) and persists entropy_coef in the checkpoint for resume. Not compatible with the Liger kernel.

by @albertvillanova in #6140

quantization_config trainer argument (streamlined QLoRA)

QLoRA no longer requires reaching into model_init_kwargs or pre-loading the model manually.

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

Added to SFTTrainer, DPOTrainer, GRPOTrainer, RLOOTrainer, RewardTrainer, and KTOTrainer. 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. Also drops the redundant get_kbit_device_map() line β€” QLoRA trains identically without it across all tested configurations.

by @qgallouedec in #6157 and #6276

MoE aux loss extends to DPO and KTO

v1.7 added the router load-balancing auxiliary loss to GRPOTrainer / RLOOTrainer / AsyncGRPOTrainer. It's now on DPOTrainer and KTOTrainer too, so post-training MoE models with preference data keeps experts balanced.

by @qgallouedec in #6208 and #6275

Neuron-friendly chunked_nll via static-shape token packing

The chunked NLL path had data-dependent indexing that broke XLA/Neuron compilation. This PR reworks it to pack valid tokens to the front via a stable argsort on the validity mask, iterate over ceil(n_valid / chunk_size) * chunk_size whole chunks (a tensor, no Python-int sync), and use ignore_index=-100 inside each chunk. GPU behavior is unchanged; Neuron now works too β€” same tiny memory footprint (~1.33 GiB vs the naive 5.94 GiB alternative) with no per-mask recompilation.

by @michaelbenayoun in #6314

Packing-aware dynamic batching in AsyncGRPO

AsyncGRPO micro-batching becomes packing-aware and token-bounded on top of the padding-free path from v1.7:

  • Ξ£ Lα΅’Β² row balancing β€” a greedy longest-first assignment gives every DP row the same Ξ£ Lα΅’Β² instead of the same token count (attention is O(LΒ²), FFN is O(L), so equal token counts don't equalize wall-clock). Cross-rank stragglers vanish; +19% MFU at 4B in the benchmark.
  • Token-budget packing (opt-in) β€” cap each row at token_budget tokens with a variable sample count, decoupling peak memory from per_device_train_batch_size. Useful in memory-bound regimes (no gradient checkpointing, long context, very large models).

Both ride HF Trainer's existing gradient accumulation β€” no training-loop surgery, FSDP/EP collectives stay in lockstep.

by @AmineDiro in #6092

VLM support in GOLDTrainer

GOLDTrainer now supports vision-language models end-to-end.

by @Strongich in #5969

Support tool calling in KTO

by @qgallouedec in #6259

Support PEFT + Liger in DPO and KTO

by @albertvillanova in #6159 and #6277

Per-dataset fraction in dataset mixtures

You can now weight a DatasetMixtureConfig by fraction instead of only by explicit counts.

by @qgallouedec in #6199

SFT: truncate during dataset preparation

Continuing the label-refactor from v1.7 (#6037), truncation moves out of the collator and into dataset preparation. SFTTrainer and DPOTrainer now truncate consistently at the same phase.

by @qgallouedec in #6155

Simplify tokenization [1-5/N]

A big refactor pass merging DPO / SFT / Reward / KTO tokenization into one shared implementation:

  • Remove redundant is_vlm parameter β€” #6298
  • Make _tokenize a module-level function β€” #6301
  • Factor _tokenize into a single shared function β€” #6302
  • Bundle apply_chat_template kwargs β€” #6305
  • Pass chat_template as apply_chat_template_kwargs β€” #6315

All by @albertvillanova. Related: align collators across DPO / SFT / Reward / KTO by @qgallouedec in #6178.

Other

Deprecations and removals

Fixes

Documentation and Examples

  • Docs: fix Jobs requirement β€” positive credit balance, not a Pro/Team/Enterprise plan by @davanstrien in #6306
  • Docs: clarify dtype defaults between transformers v5 and TRL by @casinca in #5457
  • Document that max_steps is required for iterable train datasets by @albertvillanova in #6333
  • Align KTO doc with DPO and fix Logged-metrics wording by @qgallouedec in #6258
  • Align epsilon help/docstring wording by @qgallouedec in #6014 (v1.7 window)
  • Fix KTO max_length docstring (truncation direction) by @qgallouedec in #6254
  • Fix KTO compute_metrics type hint (EvalLoopOutput β†’ EvalPrediction) by @qgallouedec in #6253

CI

New Contributors

What's Changed

Full Changelog: v1.7.0...v1.8.0