From 4eec16b3473ef7d8141bf81861b2ade6f78e2eb9 Mon Sep 17 00:00:00 2001 From: CJstate <1507965754@qq.com> Date: Wed, 9 Sep 2026 15:46:58 +0800 Subject: [PATCH 1/2] fix: validate ngram_size in get_repetition_penalty_reward Zero or negative ngram_size values were accepted at construction time but caused a ZeroDivisionError (or silently empty n-grams) during reward computation. Add construction-time validation consistent with the existing max_penalty check. Fixes #7015. --- tests/test_rewards.py | 5 +++++ trl/rewards/other_rewards.py | 2 ++ 2 files changed, 7 insertions(+) diff --git a/tests/test_rewards.py b/tests/test_rewards.py index f5fd582f68d..25beec5b573 100644 --- a/tests/test_rewards.py +++ b/tests/test_rewards.py @@ -161,6 +161,11 @@ def test_positive_max_penalty_raises(self): with pytest.raises(ValueError): get_repetition_penalty_reward(ngram_size=2, max_penalty=0.5) + @pytest.mark.parametrize("ngram_size", [0, -1]) + def test_non_positive_ngram_size_raises(self, ngram_size): + with pytest.raises(ValueError): + get_repetition_penalty_reward(ngram_size=ngram_size, max_penalty=-1.0) + def test_extra_kwargs_are_ignored(self): """Trainers pass prompts/completions/etc. as kwargs; the reward must accept and ignore them.""" reward_fn = get_repetition_penalty_reward(ngram_size=2, max_penalty=-1.0) diff --git a/trl/rewards/other_rewards.py b/trl/rewards/other_rewards.py index 4848d10d183..dd30114593f 100644 --- a/trl/rewards/other_rewards.py +++ b/trl/rewards/other_rewards.py @@ -56,6 +56,8 @@ def get_repetition_penalty_reward(ngram_size: int = 3, max_penalty: float = -1.0 """ if max_penalty > 0: raise ValueError(f"max_penalty {max_penalty} should not be positive") + if ngram_size <= 0: + raise ValueError(f"ngram_size {ngram_size} should be greater than 0") return _RepetitionPenalty(ngram_size, max_penalty) From 8668b23e7a315b0aa0283c22584e5178a5a8ac51 Mon Sep 17 00:00:00 2001 From: CJstate <1507965754@qq.com> Date: Thu, 10 Sep 2026 12:54:41 +0800 Subject: [PATCH 2/2] fix: validate fractions in get_dataset to prevent ZeroDivisionError - Add construction-time validation for negative fractions and zero-sum fractions - Reject negative values with clear ValueError - Require sum of fractions to be positive - Fixes #6981 Co-authored-by: CJstate --- tests/test_cli_utils.py | 20 ++++++++++++++++++++ trl/scripts/utils.py | 12 ++++++++++-- 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/tests/test_cli_utils.py b/tests/test_cli_utils.py index c3eeff391fa..c890a7f5498 100644 --- a/tests/test_cli_utils.py +++ b/tests/test_cli_utils.py @@ -393,6 +393,26 @@ def test_dataset_fraction_streaming_raises_error(self): with pytest.raises(ValueError, match="not supported with streaming datasets"): get_dataset(mixture_config) + def test_dataset_fraction_negative_raises_error(self): + mixture_config = DatasetMixtureConfig( + datasets=[ + DatasetConfig(path="trl-internal-testing/zen", name="standard_language_modeling", fraction=0.5), + DatasetConfig(path="trl-internal-testing/zen", name="standard_language_modeling", fraction=-0.5), + ] + ) + with pytest.raises(ValueError, match="All `fraction` values must be non-negative"): + get_dataset(mixture_config) + + def test_dataset_fraction_zero_sum_raises_error(self): + mixture_config = DatasetMixtureConfig( + datasets=[ + DatasetConfig(path="trl-internal-testing/zen", name="standard_language_modeling", fraction=0.0), + DatasetConfig(path="trl-internal-testing/zen", name="standard_language_modeling", fraction=0.0), + ] + ) + with pytest.raises(ValueError, match="Sum of `fraction` values must be positive"): + get_dataset(mixture_config) + def test_dataset_mixture_with_test_split(self): mixture_config = DatasetMixtureConfig( datasets=[DatasetConfig(path="trl-internal-testing/zen", name="standard_language_modeling")], diff --git a/trl/scripts/utils.py b/trl/scripts/utils.py index 89919961e20..815387179d7 100644 --- a/trl/scripts/utils.py +++ b/trl/scripts/utils.py @@ -441,6 +441,16 @@ def get_dataset(mixture_config: DatasetMixtureConfig) -> "DatasetDict": """ import datasets + # Validate fractions before loading datasets + fractions = [dataset_config.fraction for dataset_config in mixture_config.datasets] + if any(fraction is not None for fraction in fractions): + if any(fraction is None for fraction in fractions): + raise ValueError("`fraction` must be set for either all datasets in the mixture or none of them.") + if any(fraction < 0 for fraction in fractions): + raise ValueError(f"All `fraction` values must be non-negative, got {fractions}") + if sum(fractions) <= 0: + raise ValueError(f"Sum of `fraction` values must be positive, got {fractions} (sum={sum(fractions)})") + logger.info(f"Creating dataset mixture with {len(mixture_config.datasets)} datasets") datasets_list = [] for dataset_config in mixture_config.datasets: @@ -462,8 +472,6 @@ def get_dataset(mixture_config: DatasetMixtureConfig) -> "DatasetDict": # each dataset, where `total` is the largest mixture size such that no dataset contributes more rows than it has. fractions = [dataset_config.fraction for dataset_config in mixture_config.datasets] if any(fraction is not None for fraction in fractions): - if any(fraction is None for fraction in fractions): - raise ValueError("`fraction` must be set for either all datasets in the mixture or none of them.") if mixture_config.streaming: raise ValueError("Using a dataset `fraction` is not supported with streaming datasets.") weights = [fraction / sum(fractions) for fraction in fractions]