Support DatasetDict and IterableDatasetDict as eval_dataset in trainers

[albertvillanova]

This PR adds explicit support for passing a DatasetDict or IterableDatasetDict as eval_dataset when initializing the stable trainers (SFT, DPO, KTO, Reward, GRPO, RLOO).

This PR adds and tests support for passing DatasetDict and IterableDatasetDict objects as evaluation datasets (eval_dataset) to all TRL trainers, not just plain dictionaries or datasets. This makes it easier to use the raw output of load_dataset for evaluation, including both map-style and streaming datasets where supported. The tests are updated to cover these cases, and the trainer docstrings and type annotations are updated accordingly. For trainers that do not support streaming (IterableDatasetDict), explicit error handling is added.

Motivation

load_dataset returns a DatasetDict (or an IterableDatasetDict when streaming) whenever no split is specified, so users commonly hold these types. They already flowed through at runtime because both subclass dict, but the trainers neither declared nor documented them, so the support was implicit and untested.

Solution

The eval_dataset parameter of each stable trainer's __init__ now accepts DatasetDict and IterableDatasetDict explicitly, in its type annotation and docstring. Each split keeps being prepared independently at initialization. GRPO and RLOO still do not support streaming datasets, so an IterableDatasetDict raises the existing informative error.

Changes

• Import DatasetDict and IterableDatasetDict from datasets in the SFT, DPO, KTO, Reward, GRPO and RLOO trainers
• Extend the eval_dataset type annotation and docstring in each trainer's init
• Add tests verifying a DatasetDict and an IterableDatasetDict eval dataset are accepted and prepared per split
• Add tests verifying GRPO and RLOO reject an IterableDatasetDict eval dataset with a clear error

Detailed changes

Trainer support for DatasetDict and IterableDatasetDict as eval datasets:

• All trainer classes (DPOTrainer, KTOTrainer, RewardTrainer, SFTTrainer, GRPOTrainer, RLOOTrainer) now accept eval_dataset as a DatasetDict or IterableDatasetDict, in addition to the previously supported types. The relevant type annotations and docstrings are updated for clarity.

Expanded and improved test coverage:

• New and updated tests in all trainer test files verify initialization and behavior when eval_dataset is a DatasetDict or IterableDatasetDict, including checks that splits are tokenized independently and that keys are handled as expected.

• For trainers that do not support streaming datasets (GRPOTrainer, RLOOTrainer), tests are added to ensure that passing an IterableDatasetDict raises a clear NotImplementedError.

Consistency and minor imports cleanup:

• All affected test files now import DatasetDict and IterableDatasetDict where needed for clarity and consistency.

This change improves the flexibility and robustness of evaluation dataset handling in TRL trainers, especially for users working with Hugging Face Datasets in both map-style and streaming modes.

---

Note

Low Risk
Mostly annotations, docs, and tests; no new eval preparation logic in the diff.

Overview
Stable trainers (SFT, DPO, KTO, Reward, GRPO, RLOO) now document and type eval_dataset as accepting Hugging Face DatasetDict and IterableDatasetDict, not only plain dict or single datasets—matching common load_dataset outputs without a split.

Imports and __init__ annotations/docstrings are updated across those trainers; behavior for multi-split eval (each split prepared at init) is unchanged but now covered by tests (map and streaming where supported). GRPO and RLOO tests assert IterableDatasetDict still raises the existing streaming NotImplementedError.

^{Reviewed by Cursor Bugbot for commit d76596c. 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.

[qgallouedec]

nit: we're just initializing here, so these memory reduction trick could be removed