fix: validate fractions in get_dataset to prevent ZeroDivisionError

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")],

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)

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)
 
 

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]