fix: validate dataset mixture fractions

tests/test_cli_utils.py

@@ -385,6 +385,24 @@ def test_dataset_fraction_partial_raises_error(self):
         with pytest.raises(ValueError, match="must be set for either all datasets"):
             get_dataset(mixture_config)
 
+    @pytest.mark.parametrize(
+        ("fractions", "error_message"),
+        [
+            ([0.0, 0.0], "must be greater than zero"),
+            ([-0.5, 1.5], "must be non-negative"),
+            ([1.5, -0.5], "must be non-negative"),
+        ],
+    )
+    def test_dataset_fraction_invalid_raises_error_before_loading(self, fractions, error_message):
+        mixture_config = DatasetMixtureConfig(
+            datasets=[DatasetConfig(path="unused", fraction=fraction) for fraction in fractions]
+        )
+
+        with patch("datasets.load_dataset") as load_dataset_mock, pytest.raises(ValueError, match=error_message):
+            get_dataset(mixture_config)
+
+        load_dataset_mock.assert_not_called()
+
     def test_dataset_fraction_streaming_raises_error(self):
         mixture_config = DatasetMixtureConfig(
             datasets=[DatasetConfig(path="trl-internal-testing/zen", name="standard_language_modeling", fraction=0.5)],

trl/scripts/utils.py

@@ -65,7 +65,7 @@ class DatasetConfig:
             datasets, and the mixture size is capped so that no dataset is oversampled: the first `round(fraction * N)`
             rows of each dataset are kept, where `N` is the largest mixture size that avoids oversampling. Must be set
             for either all datasets in the mixture or none of them. Not supported for streaming datasets. When unset,
-            the full datasets are concatenated.
+            the full datasets are concatenated. Values must be non-negative and their sum must be greater than zero.
     """
 
     path: str
@@ -442,6 +442,20 @@ def get_dataset(mixture_config: DatasetMixtureConfig) -> "DatasetDict":
     import datasets
 
     logger.info(f"Creating dataset mixture with {len(mixture_config.datasets)} datasets")
+
+    # Validate fractions before loading any datasets so invalid configurations fail with an actionable error.
+    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("Dataset `fraction` values must be non-negative.")
+        total_fraction = sum(fractions)
+        if total_fraction <= 0:
+            raise ValueError("The sum of dataset `fraction` values must be greater than zero.")
+        if mixture_config.streaming:
+            raise ValueError("Using a dataset `fraction` is not supported with streaming datasets.")
+
     datasets_list = []
     for dataset_config in mixture_config.datasets:
         logger.info(f"Loading dataset for mixture: {dataset_config.path} (config name: {dataset_config.name})")
@@ -460,13 +474,8 @@ def get_dataset(mixture_config: DatasetMixtureConfig) -> "DatasetDict":
     # If `fraction` is set, treat the values as target shares of the final mixture. They are normalized to sum to one,
     # and the mixture size is capped so that no dataset is oversampled: we keep the first `round(weight * total)` rows of
     # 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]
+        weights = [fraction / total_fraction for fraction in fractions]
         total = min(
             len(dataset) / weight for dataset, weight in zip(datasets_list, weights, strict=False) if weight > 0
         )