Skip to content

v1.9.0

Latest
Latest

Choose a tag to compare

View all tags
@qgallouedec qgallouedec released this 21 Jul 20:36
v1.9.0
35def3e
This commit was created on GitHub.com and signed with GitHub’s verified signature.
GPG key ID: B5690EEEBB952194
Verified
Learn about vigilant mode.

Features

Iterable / streaming datasets in GRPO and RLOO

Long-standing request finally landed. GRPO and RLOO rely on RepeatSampler to repeat each prompt num_generations times and group them across processes — but samplers can't attach to iterable datasets (no length, no indexing), and the old behavior silently bypassed the sampler, quietly corrupting advantage computation.

A new repeat_iterable_dataset generator reproduces RepeatSampler's exact ordering by transforming the stream itself instead of reordering indices. Trainers wrap iterable datasets via IterableDataset.from_generator, force dispatch_batches=False, and shard the stream per-process to preserve prompt groups after cross-process gathering.

from datasets import load_dataset
from trl import GRPOConfig, GRPOTrainer

dataset = load_dataset("trl-lib/DeepMath-103K", split="train", streaming=True)

trainer = GRPOTrainer(
    model="Qwen/Qwen3-4B",
    args=GRPOConfig(max_steps=1000),  # required for iterable datasets
    train_dataset=dataset,
    reward_funcs=accuracy_reward,
)

max_steps is required for iterable datasets (no length). A clear error also fires if dispatch_batches=True is combined with an iterable dataset.

by @albertvillanova in #6351

Environment-owned datasets

When an environment_factory is provided, the environment can now own the data. train_dataset becomes optional: reset() returns the prompt. No more fabricating a dummy dataset just to drive the loop.

class WordleEnv:
    def reset(self, **kwargs) -> str:      # returns the prompt
        self._target = sample(words)
        return f"Guess a {len(self._target)}-letter word."
    def get_reward(self) -> float:
        return 1.0 if self._solved else 0.0
    def guess(self, word: str) -> str:     # exposed tool
        self._solved = word == self._target
        return _feedback(word, self._target)

trainer = GRPOTrainer(
    model="Qwen/Qwen3-4B",
    args=GRPOConfig(max_steps=1000),
    environment_factory=WordleEnv,          # no train_dataset needed
)

Applies to both GRPOTrainer and AsyncGRPOTrainer. A provided train_dataset still works exactly as before; multi-environment routing datasets can now be environment-only (no prompt column required).

by @qgallouedec in #6349

Message-level rollouts in AsyncGRPO

An opt-in way to build training rows from a multi-turn conversation. Message mode keeps the conversation as messages and re-tokenizes the whole thing each turn, then checks whether the fresh tokens still start with the tokens held so far. If yes → append the new part (same as token mode). If no → a rewrite happened → close the row and open a new one that matches what the model actually read.

AsyncGRPOConfig(
    rollout_protocol="message",   # "token" (default) | "message"
    fork_threshold_tokens=1024,   # message-mode only
)

Per-turn drift is classified as CLEAN (append), REALIGN (last-answer tail wobble → overwrite same row), or FORK (real rewrite → new row). One advantage per conversation, stamped on every row it produced; under token-mean loss a fork is invisible. Sets up the plumbing for future tree-trajectory support — scoring is already branch-agnostic (groups by rollout_id).

by @AmineDiro in #6250

AsyncGRPO: VLLMClient and decoupled weight sync

A follow-up to v1.8's rollout worker split. A small VLLMClient becomes the single place that speaks HTTP to the vLLM server (named methods for wait_for_server_ready, get_max_model_len, pause/resume, weight-update endpoints — stateless, so picklable). Generation deliberately stays on its own async aiohttp session in the rollout child.

Weight sync is now built from config independent of rollout_worker (previously a custom rollout worker force-set it to None), and is injectable via a new WeightTransferProtocol. token_budget now waits for /health before defaulting to get_max_model_len(), so a slow-starting vLLM waits instead of failing the run.

by @AmineDiro in #6269

DistillationTrainer refactor

A ~13-PR sweep that reshapes the experimental DistillationTrainer for future stability:

  • Extract ServerDistillationTrainer for the teacher-server path (#6454)
  • Move IW-OPD to a dedicated experimental trainer (#6453)
  • Remove the local sparse top-1 loss path from the base trainer (#6455)
  • Remove top-k support from the base generalized JSD loss (#6456)
  • Validate the teacher by vocab size instead of tokenizer identity (#6457)
  • Deprecate & remove lmbda on/off-policy mixing and the off-policy training branch (#6458, #6460); GKD paper reproduction now points at GKDTrainer (#6459)
  • Accept a prompt column alongside messages (#6461); deprecate messages-format datasets (#6474)
  • ⚠️ Fix num_items_in_batch to count generated completion tokens — the loss denominator was wrong for on-policy training (counted dataset completions, not generated ones) and NaN'd on prompt-only datasets. This changes loss values for anyone on the old path (#6478)
  • Preparatory pins on the objective and on/off-policy behavior end to end (#6450, #6451, #6452)

All by @qgallouedec.

New IW-OPD distillation objective

Importance-Weighted On-Policy Distillation lands as an optional objective on the experimental DistillationTrainer — detached IW-OPD advantages built from sampled-token teacher and rollout logprobs (with cached vLLM rollout logprobs when use_vllm=True).

DistillationConfig(
    distillation_objective="iw_opd",
    iw_opd_gamma=1.0,
    iw_opd_epsilon=0.2,
    lmbda=1.0, reverse_kl_top_1_mode="sampled",  # required
)

Now split into its own experimental trainer as part of the refactor above.

by @kashif in #6191

log_multimodal for GRPO / RLOO

New config toggle to control whether images from multimodal completions are logged to trackers. Also applied to experimental DPPO and GRPO-with-replay-buffer.

by @apardyl in #5408 and by @albertvillanova in #6352

vLLM version sweep

Other

Deprecations and removals

  • Remove the experimental DPPO trainer by @qgallouedec in #6402
  • Deprecate messages-format datasets in DistillationTrainer — use prompt column instead — by @qgallouedec in #6474
  • Deprecate lmbda on/off-policy mixing in DistillationTrainer — off-policy path fully removed by @qgallouedec in #6458 and #6460

Fixes

  • Fix biased cross-rank aggregation of token-weighted metrics in GRPO / RLOO — mean-of-means was biased when per-rank token counts differ; now reduces numerator + count jointly for a true global mean. Also fixes clip_ratio/low_min / high_max to be per-completion instead of per-rank-batch-mean (DP-layout-independent). Logging-only, no loss change. By @qgallouedec in #6380
  • Fix DPO/KTO use_liger_kernel under DeepSpeed ZeRO-3 by @qgallouedec in #6372
  • Force use_reentrant=True for PEFT + ZeRO-3 + gradient checkpointing across all trainers by @qgallouedec in #6356
  • Fix precompute_ref_log_probs=True with DeepSpeed in DPO and KTO by @qgallouedec in #6403
  • Fix cast_lm_head_to_fp32 under FSDP2 / DeepSpeed ZeRO-3 by @qgallouedec in #6330
  • Fix corrupted reasoning blocks in Qwen3-family training chat templates — the both-tags guard broke replay of self-generated thinking traces (SFT / distillation on model outputs); reverted to the original </think>-only check. By @qgallouedec in #6363
  • Fix EOS being masked when pad_token_id == eos_token_id in GKD and GOLD by @roycho96 in #6357
  • Fix GRPO truncated-completion loss normalization by @Functionhx in #6439
  • Fix apply_chat_template dropping leading tokens from chosen when the prompt common prefix shrinks against rejected by @sohumt123 in #6433
  • Fix crash when reward function is a functools.partial or callable instance by @qgallouedec in #6313
  • Fix QLoRA + FSDP1 mixed-dtype crash in prepare_peft_model by @sergiopaniego in #6343
  • Truncate GOLD on-policy prompts before generation, keeping the prompt end by @roycho96 in #6359
  • Fix FileNotFoundError in DPO/KTO precompute_ref_log_probs by @Functionhx in #6348
  • Fix GeometricMixtureWrapper.generate() crash after transformers' MoE decode optimization change by @albertvillanova in #6413
  • Fix GKD vocab_size check for multimodal model configs by @sergiopaniego in #6340
  • Fix OpenEnv docker-image example bootstrap by @sergiopaniego in #6329
  • Fix error type for IterableDataset in GRPO and RLOO by @albertvillanova in #6324
  • Fix KTO evaluate crash on streaming datasets with unpaired-preference data by @albertvillanova in #6325
  • Fix DPO/KTO evaluate crash with a dict eval_dataset and precompute_ref_log_probs by @albertvillanova in #6370
  • Fix DPO/KTO double-preprocessing eval_dataset dict passed to evaluate() by @albertvillanova in #6434
  • Raise when DPO/KTO cannot precompute correct reference log-probs at eval time by @albertvillanova in #6443
  • Avoid in-place mutation of input prompts in GRPO and DPPO environment loops by @albertvillanova in #6364 and #6366
  • Align KTO with DPO: drop hasattr guards on lm_head.bias in the Liger path by @qgallouedec in #6425
  • Align KTO with DPO: fix logits/chosen and logits/rejected metrics by @qgallouedec in #6424
  • Fix missing eos_token_id in non-vLLM on-policy GOLD test stubs by @kashif in #6394
  • Fix BEMACallback bias_power docstring by @DaoyuanLi2816 in #6390

Documentation and Examples

CI

New Contributors

What's Changed

Full Changelog: v1.8.0...v1.9.0

Contributors

kashif, albertvillanova, and 11 other contributors
Assets 2
Loading