fix: validate dataset mixture fractions

[cananoo]

What does this PR do?

Fixes #6981.

This change validates dataset mixture fractions before loading any datasets. It rejects negative fractions and zero-sum fractions with actionable ValueError messages, while preserving mixtures that include individual zero fractions when the total is positive.

The regression tests also verify that invalid configurations fail before datasets.load_dataset is called.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? The existing fraction documentation remains accurate, so no documentation change is needed.
• Did you write any new necessary tests?

AI writing disclosure

• No AI usage: the PR was written entirely by a human.
• AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.
• AI-generated: the PR was mostly or fully generated by an AI tool.

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag members/contributors who may be interested in your PR.

---

Note

Low Risk
Input validation and ordering change in dataset loading; behavior for valid mixtures is unchanged aside from earlier errors for invalid configs.

Overview
get_dataset now checks mixture fraction settings before any datasets.load_dataset calls, so bad configs fail fast with clear ValueError messages instead of after expensive I/O.

New rules: each fraction must be non-negative, and the sum must be > 0 (all-zero sums are rejected; individual zeros remain OK when the total is positive). Existing checks—fraction on all-or-none datasets and no fractions with streaming—are part of the same upfront block. DatasetConfig docs note the non-negative / positive-sum requirement.

Tests add a parametrized case that mocks load_dataset and asserts it is never called when fractions are invalid (zero sum or negatives).

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

[behroozazarkhalili]

Verdict: the fix is correct, with two small changes worth making before merge.

What I checked

Interpreter /scratch/ermia/venvs/hf_trl/bin/python, worktree at the PR head, base 02dcb656.

New and neighbouring tests on the PR head:

$ python -m pytest tests/test_cli_utils.py::TestGetDataset::test_dataset_fraction_invalid_raises_error_before_loading \
    tests/test_cli_utils.py::TestGetDataset::test_dataset_fraction_streaming_raises_error \
    tests/test_cli_utils.py::TestGetDataset::test_dataset_fraction_partial_raises_error \
    tests/test_cli_utils.py::TestGetDataset::test_dataset_fraction -p no:randomly -q
5 passed, 1 warning in 14.33s

Same test on a copy of the head worktree with trl/scripts/utils.py reverted to the base commit (new tests kept):

FAILED ...test_dataset_fraction_invalid_raises_error_before_loading[fractions0-must be greater than zero]
FAILED ...test_dataset_fraction_invalid_raises_error_before_loading[fractions1-must be non-negative]
2 failed, 1 warning in 11.21s

So the test is a real oracle: red before, green after.

The minimal case from issue #6981, with the real trl-internal-testing/zen dataset:

--- HEAD ---
[0.0, 0.0]: ValueError: The sum of dataset `fraction` values must be greater than zero.
[-0.5, 1.5]: ValueError: Dataset `fraction` values must be non-negative.
[0.0, 0.5]: OK rows=17
--- BASE ---
[0.0, 0.0]: ZeroDivisionError: float division by zero
[-0.5, 1.5]: OK rows=17
[0.0, 0.5]: OK rows=17

The reported crash is gone, the silently accepted negative fraction is now rejected, and a single zero share alongside a positive one still works.

get_dataset is the only copy of this logic in the tree, so there is no sibling to keep in sync (grep -rn "sum(fractions)" trl/ returns only trl/scripts/utils.py:453). ruff check and ruff format --check on both changed files pass. total_fraction is assigned and read under the same any(fraction is not None ...) guard over an unchanged fractions list, so there is no path that reads it unbound.

Findings

1. tests/test_cli_utils.py:386-402: the negative case only puts the negative value in position 0, so the test does not pin down that every entry is checked. I replaced any(fraction < 0 for fraction in fractions) with fractions[0] < 0 in a copy of the head worktree and the new test still passed (3 passed, 1 warning in 14.75s). Adding a ([1.5, -0.5], "must be non-negative") parameter case closes that gap; on the real code that ordering already raises correctly.

2. trl/scripts/utils.py:63-68: the fraction docstring already spells out the other two constraints (set for all datasets or none, unsupported with streaming) but not the two this PR adds. The PR body says the existing documentation stays accurate; it stays accurate but is now incomplete. One sentence, such as "Values must be non-negative and their sum must be greater than zero", keeps the docstring the single place a user learns the rules.

3. trl/scripts/utils.py:451-455: a NaN fraction slips through both guards, since nan < 0 and nan <= 0 are both false, and then fails later with ValueError: min() arg is an empty sequence (observed on the head for the mixture [nan, 1.0]). This is pre-existing rather than a regression, and it is fine to leave out of scope, but writing the guard as if not fraction >= 0 would cover it at no extra cost.

Not reproduced

Nothing. Every claim above was executed.

Unevaluable

Nothing. All cases fit the tiny trl-internal-testing/zen dataset and ran locally.

[cananoo]

Thanks for the detailed review. I pushed 7d7784c with both requested follow-ups: (1) the regression matrix now also covers [1.5, -0.5], proving a negative fraction is rejected regardless of position; (2) DatasetConfig.fraction documents the non-negative and positive-sum constraints. ruff check and Python compilation pass locally. The full targeted pytest requires the project environment datasets and torch dependencies, which are not installed in this Windows worktree.

[cananoo]

The follow-up head 7d7784c is pushed and Cursor Bugbot is passing. The Tests, documentation, and secret-scan workflows are currently action_required and did not execute; a maintainer approval or rerun is needed for the updated head.

[cananoo]

Updated the PR branch onto the latest huggingface/trl main (56f6675) as requested. The change itself is unchanged; the new head is ecb3079. Local ruff, Python compilation, and diff check pass. Cursor Bugbot has started running; the other workflows may still require maintainer approval. Please re-review when convenient.

[qgallouedec]

Closing: we don't review fully AI-generated PRs from first-time contributors, see the AI usage policy.

[behroozazarkhalili]

@cananoo after talking with @qgallouedec I reopened this work as #7142 from my fork so it can be reviewed. Your commits are cherry-picked there with your authorship intact and the PR body credits you; if you would rather carry it yourself, say so there and I will close mine.