separate out train and eval dataset streaming

src/axolotl/utils/data/sft.py

@@ -44,6 +44,48 @@ from axolotl.utils.trainer import (
 LOG = get_logger(__name__)
 
 
+def _is_streaming_enabled_for_split(
+    cfg: DictDefault, split: Literal["train", "test"]
+) -> bool:
+    """Check if streaming is enabled for a specific split."""
+    if split == "test":
+        # For eval datasets, check eval_streaming first, then fall back to streaming
+        eval_streaming = cfg.get("eval_streaming")
+        if eval_streaming is not None:
+            return eval_streaming
+
+    # Fall back to main streaming setting
+    streaming = cfg.get("streaming")
+    if streaming is True:
+        return True
+
+    # Check if pretraining dataset exists (defaults to streaming)
+    has_pretraining = cfg.get("pretraining_dataset") is not None
+    streaming_default_for_pretraining = has_pretraining and streaming is None
+
+    return streaming_default_for_pretraining
+
+
+def _get_streaming_config_for_split(
+    cfg: DictDefault, split: Literal["train", "test"]
+) -> DictDefault:
+    """Get a modified config object with split-specific streaming settings."""
+    if split != "test":
+        return cfg
+
+    # Override with eval-specific configs if they exist
+    streaming_cfg = DictDefault(cfg)
+    eval_strategy = cfg.get("eval_streaming_dataset_mixing_strategy")
+    eval_weights = cfg.get("eval_streaming_mixing_weights")
+
+    if eval_strategy is not None:
+        streaming_cfg["streaming_dataset_mixing_strategy"] = eval_strategy
+    if eval_weights is not None:
+        streaming_cfg["streaming_mixing_weights"] = eval_weights
+
+    return streaming_cfg
+
+
 @retry_on_request_exceptions(max_retries=3, delay=5)
 def prepare_datasets(
     cfg: DictDefault,
@@ -267,10 +309,14 @@ def _load_tokenized_prepared_datasets(
     datasets_configs = cfg.datasets if split == "train" else cfg.test_datasets
     prompters: list[Prompter | None] = []
 
-    # For streaming datasets, skip caching and load raw datasets directly
+    # Check if streaming is enabled for this split
-    if cfg.streaming:
+    use_streaming = _is_streaming_enabled_for_split(cfg, split)
+
+    if use_streaming:
+        # For streaming datasets, skip caching and load raw datasets directly
+        streaming_cfg = _get_streaming_config_for_split(cfg, split)
         dataset, prompters = _load_raw_datasets(
-            cfg,
+            streaming_cfg,
             datasets_configs,
             tokenizer,
             split,

src/axolotl/utils/data/shared.py

@@ -550,10 +550,7 @@ def merge_datasets(
             return ds
         return ds.shuffle(seed=cfg.seed)
 
-    # Check if we have any IterableDatasets
+    if any(isinstance(ds, IterableDataset) for ds in datasets):
-    has_iterable = any(isinstance(ds, IterableDataset) for ds in datasets)
-
-    if has_iterable:
         LOG.info("Merging streaming datasets...")
         merged_dataset = _merge_streaming_datasets(datasets, cfg)
     else:

src/axolotl/utils/schemas/config.py

@@ -935,7 +935,13 @@ class AxolotlInputConfig(
     streaming: bool | None = Field(
         default=None,
         json_schema_extra={
-            "description": "Whether to use streaming datasets (IterableDataset) for processing large datasets that don't fit in memory. When True, data is loaded on-demand during training without upfront preprocessing. Requires max_steps to be set. Pre-training datasets default to streaming unless explicitly set to False."
+            "description": "Whether to use streaming datasets (IterableDataset) for training datasets. When True, data is loaded on-demand during training without upfront preprocessing. Requires max_steps to be set. Pre-training datasets default to streaming unless explicitly set to False."
+        },
+    )
+    eval_streaming: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to use streaming datasets for evaluation datasets. If not set, falls back to the 'streaming' setting. Useful for streaming large training data while keeping smaller eval datasets in memory."
         },
     )
     streaming_dataset_mixing_strategy: str | None = Field(
@@ -950,6 +956,18 @@ class AxolotlInputConfig(
             "description": "Weights for weighted mixing strategy when using multiple streaming datasets. Must sum to 1.0 and have same length as datasets list. Only used when streaming_dataset_mixing_strategy='weighted'."
         },
     )
+    eval_streaming_dataset_mixing_strategy: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Strategy for mixing multiple streaming evaluation datasets. If not set, falls back to streaming_dataset_mixing_strategy. Options: 'round_robin', 'weighted', 'random'."
+        },
+    )
+    eval_streaming_mixing_weights: list[float] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Weights for weighted mixing strategy for evaluation datasets. Must sum to 1.0 and have same length as evaluation datasets list."
+        },
+    )
 
     # INTERNALS - document for now, generally not set externally
     is_preprocess: bool | None = None

src/axolotl/utils/schemas/validation.py

@@ -1388,20 +1388,29 @@ class GRPOVllmValidationMixin:
 class StreamingValidationMixin:
     """Validation methods related to streaming datasets."""
 
-    @model_validator(mode="after")
+    def _is_streaming_enabled(self, context: str = "train") -> bool:
-    def check_streaming_requires_max_steps(self):
+        """Check if streaming is enabled for a given context (train or eval)."""
-        """Ensure max_steps is set when using streaming datasets."""
+        if context == "eval":
-        # Check if streaming is explicitly enabled
+            eval_streaming = getattr(self, "eval_streaming", None)
-        streaming_enabled = getattr(self, "streaming", None) is True
+            if eval_streaming is not None:
+                return eval_streaming
+
+        # Fall back to main streaming setting
+        streaming = getattr(self, "streaming", None)
+        if streaming is True:
+            return True
 
         # Check if pretraining dataset exists (defaults to streaming)
         has_pretraining = getattr(self, "pretraining_dataset", None) is not None
-        streaming_default_for_pretraining = (
+        streaming_default_for_pretraining = has_pretraining and streaming is None
-            has_pretraining and getattr(self, "streaming", None) is None
-        )
 
-        # If streaming is enabled (explicitly or by default for pretraining)
+        return streaming_default_for_pretraining
-        if streaming_enabled or streaming_default_for_pretraining:
+
+    @model_validator(mode="after")
+    def check_streaming_requires_max_steps(self):
+        """Ensure max_steps is set when using streaming datasets."""
+        # Check if streaming is enabled for training datasets
+        if self._is_streaming_enabled("train"):
             max_steps = getattr(self, "max_steps", None)
             if not max_steps:
                 raise ValueError("max_steps must be set when using streaming datasets")
@@ -1411,17 +1420,8 @@ class StreamingValidationMixin:
     @model_validator(mode="after")
     def check_streaming_validation_splits_conflict(self):
         """Ensure validation splits are not used with streaming datasets."""
-        # Check if streaming is explicitly enabled
+        # Check if streaming is enabled for training datasets
-        streaming_enabled = getattr(self, "streaming", None) is True
+        if self._is_streaming_enabled("train"):
-
-        # Check if pretraining dataset exists (defaults to streaming)
-        has_pretraining = getattr(self, "pretraining_dataset", None) is not None
-        streaming_default_for_pretraining = (
-            has_pretraining and getattr(self, "streaming", None) is None
-        )
-
-        # If streaming is enabled (explicitly or by default for pretraining)
-        if streaming_enabled or streaming_default_for_pretraining:
             val_set_size = getattr(self, "val_set_size", 0.0)
             if val_set_size and val_set_size > 0:
                 raise ValueError(
@@ -1433,17 +1433,8 @@ class StreamingValidationMixin:
     @model_validator(mode="after")
     def check_streaming_preprocessing_conflict(self):
         """Ensure preprocessing is not enabled with streaming datasets."""
-        # Check if streaming is explicitly enabled
+        # Check if streaming is enabled for training or eval datasets
-        streaming_enabled = getattr(self, "streaming", None) is True
+        if self._is_streaming_enabled("train") or self._is_streaming_enabled("eval"):
-
-        # Check if pretraining dataset exists (defaults to streaming)
-        has_pretraining = getattr(self, "pretraining_dataset", None) is not None
-        streaming_default_for_pretraining = (
-            has_pretraining and getattr(self, "streaming", None) is None
-        )
-
-        # If streaming is enabled (explicitly or by default for pretraining)
-        if streaming_enabled or streaming_default_for_pretraining:
             if os.environ.get("AXOLOTL_IS_PREPROCESS") == "1":
                 raise ValueError("preprocess is not supported for streaming datasets")
 
@@ -1452,17 +1443,8 @@ class StreamingValidationMixin:
     @model_validator(mode="after")
     def check_streaming_skip_prepare_dataset(self):
         """Ensure skip_prepare_dataset is set for streaming datasets."""
-        # Check if streaming is explicitly enabled
+        # Check if streaming is enabled for training or eval datasets
-        streaming_enabled = getattr(self, "streaming", None) is True
+        if self._is_streaming_enabled("train") or self._is_streaming_enabled("eval"):
-
-        # Check if pretraining dataset exists (defaults to streaming)
-        has_pretraining = getattr(self, "pretraining_dataset", None) is not None
-        streaming_default_for_pretraining = (
-            has_pretraining and getattr(self, "streaming", None) is None
-        )
-
-        # If streaming is enabled (explicitly or by default for pretraining)
-        if streaming_enabled or streaming_default_for_pretraining:
             skip_prepare = getattr(self, "skip_prepare_dataset", None)
             if skip_prepare is False:
                 LOG.warning(
@@ -1476,45 +1458,73 @@ class StreamingValidationMixin:
     @model_validator(mode="after")
     def check_streaming_mixing_weights(self):
         """Validate streaming_mixing_weights configuration."""
+        valid_strategies = ["round_robin", "weighted", "random"]
+
+        # Check main strategy and weights
         strategy = getattr(self, "streaming_dataset_mixing_strategy", "round_robin")
         weights = getattr(self, "streaming_mixing_weights", None)
+        self._validate_streaming_strategy_and_weights(
+            strategy,
+            weights,
+            "streaming_dataset_mixing_strategy",
+            "streaming_mixing_weights",
+            valid_strategies,
+        )
 
-        # Validate strategy values
+        # Check eval-specific strategy and weights
-        valid_strategies = ["round_robin", "weighted", "random"]
+        eval_strategy = getattr(self, "eval_streaming_dataset_mixing_strategy", None)
+        eval_weights = getattr(self, "eval_streaming_mixing_weights", None)
+
+        if eval_strategy is not None:
+            self._validate_streaming_strategy_and_weights(
+                eval_strategy,
+                eval_weights,
+                "eval_streaming_dataset_mixing_strategy",
+                "eval_streaming_mixing_weights",
+                valid_strategies,
+            )
+        elif eval_weights is not None:
+            LOG.warning(
+                "eval_streaming_mixing_weights provided but eval_streaming_dataset_mixing_strategy is not set. "
+                "Weights will be ignored unless eval_streaming_dataset_mixing_strategy='weighted'."
+            )
+
+        return self
+
+    def _validate_streaming_strategy_and_weights(
+        self, strategy, weights, strategy_field, weights_field, valid_strategies
+    ):
+        """Helper method to validate strategy and weights pair."""
         if strategy not in valid_strategies:
             raise ValueError(
-                f"streaming_dataset_mixing_strategy must be one of {valid_strategies}, "
+                f"{strategy_field} must be one of {valid_strategies}, "
                 f"got '{strategy}'"
             )
 
         if strategy == "weighted":
             if weights is None:
                 raise ValueError(
-                    "streaming_mixing_weights must be provided when "
+                    f"{weights_field} must be provided when "
-                    "streaming_dataset_mixing_strategy='weighted'"
+                    f"{strategy_field}='weighted'"
                 )
 
             if not isinstance(weights, list) or not all(
                 isinstance(w, (int, float)) for w in weights
             ):
-                raise ValueError("streaming_mixing_weights must be a list of numbers")
+                raise ValueError(f"{weights_field} must be a list of numbers")
 
             if any(w < 0 for w in weights):
-                raise ValueError("streaming_mixing_weights must be non-negative")
+                raise ValueError(f"{weights_field} must be non-negative")
 
             if abs(sum(weights) - 1.0) > 1e-6:
-                raise ValueError(
+                raise ValueError(f"{weights_field} must sum to 1.0, got {sum(weights)}")
-                    f"streaming_mixing_weights must sum to 1.0, got {sum(weights)}"
-                )
 
         elif weights is not None and strategy != "weighted":
             LOG.warning(
-                f"streaming_mixing_weights provided but strategy is '{strategy}'. "
+                f"{weights_field} provided but {strategy_field} is '{strategy}'. "
                 "Weights will be ignored."
             )
 
-        return self
-
 
 # pylint: disable=too-many-ancestors
 class ValidationMixin(