Merge pull request #3121 from bghira/issue/3099

Make DDP dataset selection deterministic

Author: bghira

simpletuner/helpers/data_backend/runtime/batch_fetcher.py

@@ -9,6 +9,7 @@
 
 from simpletuner.helpers.training.multi_process import should_log
 
+from .dataloader_iterator import _set_epoch_step_for_training_step
 from .dataloader_iterator import random_dataloader_iterator as _runtime_random_iterator
 
 DEFAULT_KEEP_RUNNING_PROPERTY = None
@@ -49,6 +50,7 @@ def _resolve_random_iterator():
 
 
 class BatchFetcher:
+    """Prefetch batches using the last consumed training step as the first schedule hint."""
 
     def __init__(self, step: int, max_size: int = 10, datasets: Optional[Dict[str, Any]] = None) -> None:
         if DEFAULT_KEEP_RUNNING_PROPERTY is not None:
@@ -57,6 +59,7 @@ def __init__(self, step: int, max_size: int = 10, datasets: Optional[Dict[str, A
         self.datasets = datasets or {}
         self._keep_running = True
         self.step = step
+        self._next_fetch_step = step
 
     def start_fetching(self) -> threading.Thread:
         thread = threading.Thread(target=self.fetch_responses)
@@ -72,7 +75,14 @@ def fetch_responses(self) -> None:
                 if self.queue.qsize() < self.queue.maxsize:
                     prefetch_log_debug(f"Queue size: {self.queue.qsize()}. Fetching more data.")
                     try:
-                        item = iterator_fn(self.step, self.datasets)
+                        if iterator_fn is DEFAULT_RANDOM_ITERATOR:
+                            item = iterator_fn(
+                                self._next_fetch_step,
+                                self.datasets,
+                                update_epoch_step=False,
+                            )
+                        else:
+                            item = iterator_fn(self._next_fetch_step, self.datasets)
                     except ValueError:
                         prefetch_log_debug("No datasets available during prefetch; stopping fetch thread.")
                         self._keep_running = False
@@ -84,6 +94,7 @@ def fetch_responses(self) -> None:
                         logger.debug(f"BatchFetcher encountered exception: {exc}")
                         break
                     self.queue.put(item)
+                    self._next_fetch_step += 1
                     if self.queue.qsize() >= self.queue.maxsize:
                         prefetch_log_debug("Completed fetching data. Queue is full.")
                         if threading.current_thread() is threading.main_thread():
@@ -107,7 +118,9 @@ def next_response(self, step: int) -> Any:
         while self.queue.empty():
             continue
         prefetch_log_debug("Queue has data. Yielding next item.")
-        return self.queue.get()
+        item = self.queue.get()
+        _set_epoch_step_for_training_step(step)
+        return item
 
     def stop_fetching(self) -> None:
         self._keep_running = False

simpletuner/helpers/data_backend/runtime/dataloader_iterator.py

@@ -24,6 +24,7 @@
     prefetch_log.setLevel(logging.ERROR)
 
 _SCALED_SAMPLERS: dict[int, "ScaledDatasetSampler"] = {}
+_MAX_TORCH_SEED = 2**63 - 1
 
 
 def prefetch_log_debug(message: str) -> None:
@@ -32,14 +33,33 @@ def prefetch_log_debug(message: str) -> None:
     prefetch_log.debug(f"{rank_info()} {message}")
 
 
+def _set_epoch_step_for_training_step(step: int) -> None:
+    args = StateTracker.get_args()
+    grad_steps = getattr(args, "gradient_accumulation_steps", 1) if args is not None else 1
+    if isinstance(grad_steps, (int, float)):
+        gradient_accumulation_steps = max(1, int(grad_steps))
+    else:
+        gradient_accumulation_steps = 1
+    StateTracker.set_epoch_step(int(step / gradient_accumulation_steps))
+
+
+def _selection_generator(step: int) -> torch.Generator:
+    args = StateTracker.get_args()
+    base_seed = getattr(args, "seed", 0) if args is not None else 0
+    generator = torch.Generator(device="cpu")
+    generator.manual_seed((int(base_seed or 0) + int(step)) % _MAX_TORCH_SEED)
+    return generator
+
+
 def select_dataloader_index(step: int, backends: Dict[str, Any]) -> Optional[str]:
     if not backends:
         raise ValueError("No data backends available for selection.")
 
     weights = []
     backend_ids = []
     inactive_ids = []
-    for backend_id, backend in backends.items():
+    for backend_id in sorted(backends):
+        backend = backends[backend_id]
         weight = get_backend_weight(backend_id, backend, step)
         if weight <= 0:
             inactive_ids.append(backend_id)
@@ -69,7 +89,7 @@ def select_dataloader_index(step: int, backends: Dict[str, Any]) -> Optional[str
         raise ValueError("All backend sampling weights are zero.")
     weights_tensor /= total
 
-    chosen_index = torch.multinomial(weights_tensor, 1).item()
+    chosen_index = torch.multinomial(weights_tensor, 1, generator=_selection_generator(step)).item()
     chosen_backend_id = backend_ids[chosen_index]
 
     return chosen_backend_id
@@ -171,7 +191,8 @@ def __init__(self, slider_key: str = "slider_strength") -> None:
 
     def rebuild(self, backends: Dict[str, Any]) -> None:
         self._groups = {"positive": [], "negative": [], "neutral": []}
-        for backend_id, backend in backends.items():
+        for backend_id in sorted(backends):
+            backend = backends[backend_id]
             strength = _read_slider_strength(backend_id, backend, self.slider_key)
             if strength is None:
                 self._groups["neutral"].append(backend_id)
@@ -226,7 +247,7 @@ def _choose_from_group(self, group: str, step: int, backends: Dict[str, Any]) ->
         if total <= 0:
             return None
         weights_tensor /= total
-        chosen_index = torch.multinomial(weights_tensor, 1).item()
+        chosen_index = torch.multinomial(weights_tensor, 1, generator=_selection_generator(step)).item()
         return filtered_ids[chosen_index]
 
     def next_backend_id(self, step: int, backends: Dict[str, Any]) -> Optional[str]:
@@ -264,21 +285,20 @@ def _get_scaled_sampler(backends: Dict[str, Any]) -> Optional[ScaledDatasetSampl
     return sampler
 
 
-def random_dataloader_iterator(step: int, backends: Dict[str, Any]) -> Union[Any, bool]:
+def random_dataloader_iterator(
+    step: int,
+    backends: Dict[str, Any],
+    *,
+    update_epoch_step: bool = True,
+) -> Union[Any, bool]:
     if not backends:
         raise ValueError("No data backends provided to iterator.")
 
     prefetch_log_debug("Random dataloader iterator launched.")
-    args = StateTracker.get_args()
-    grad_steps = getattr(args, "gradient_accumulation_steps", 1) if args is not None else 1
-    if isinstance(grad_steps, (int, float)):
-        gradient_accumulation_steps = max(1, int(grad_steps))
-    else:
-        gradient_accumulation_steps = 1
     logger.debug(f"Backends to select from {backends}")
     while backends:
-        epoch_step = int(step / gradient_accumulation_steps)
-        StateTracker.set_epoch_step(epoch_step)
+        if update_epoch_step:
+            _set_epoch_step_for_training_step(step)
 
         sampler = _get_scaled_sampler(backends)
         chosen_backend_id = None

tests/test_runtime_components.py

@@ -11,12 +11,14 @@
 import threading
 import time
 import unittest
+from types import SimpleNamespace
 from typing import Any, Dict
 from unittest.mock import MagicMock, Mock, patch
 
 import torch
 
 from simpletuner.helpers.data_backend.runtime import BatchFetcher
+from simpletuner.helpers.data_backend.runtime import batch_fetcher as batch_fetcher_module
 from simpletuner.helpers.data_backend.runtime import dataloader_iterator as dataloader_module
 from simpletuner.helpers.data_backend.runtime import get_backend_weight, random_dataloader_iterator, select_dataloader_index
 
@@ -110,6 +112,29 @@ def check_iteration():
         # Verify iterator was called multiple times
         self.assertGreater(mock_iterator.call_count, 1)
 
+    @patch("simpletuner.helpers.data_backend.runtime.batch_fetcher.random_dataloader_iterator")
+    def test_prefetch_starts_from_last_consumed_schedule_hint(self, mock_iterator):
+        mock_iterator.side_effect = lambda step, datasets: {"step": step}
+        fetcher = BatchFetcher(step=100, max_size=3, datasets=self.datasets)
+
+        fetcher.fetch_responses()
+
+        self.assertEqual(
+            [call.args[0] for call in mock_iterator.call_args_list],
+            [100, 101, 102],
+        )
+
+    def test_runtime_iterator_does_not_advance_epoch_state_during_prefetch(self):
+        mock_iterator = Mock(return_value={"batch": "data"})
+        with (
+            patch.object(batch_fetcher_module, "DEFAULT_RANDOM_ITERATOR", mock_iterator),
+            patch.object(batch_fetcher_module, "random_dataloader_iterator", mock_iterator),
+        ):
+            fetcher = BatchFetcher(step=100, max_size=1, datasets=self.datasets)
+            fetcher.fetch_responses()
+
+        mock_iterator.assert_called_once_with(100, self.datasets, update_epoch_step=False)
+
     @patch("simpletuner.helpers.data_backend.runtime.batch_fetcher.random_dataloader_iterator")
     def test_fetch_responses_queue_full_behavior(self, mock_iterator):
         """Test fetch_responses behavior when queue is full"""
@@ -134,7 +159,8 @@ def limited_keep_running():
         # Queue should be at maximum capacity
         self.assertEqual(fetcher.queue.qsize(), 1)
 
-    def test_next_response_with_data_available(self):
+    @patch("simpletuner.helpers.data_backend.runtime.batch_fetcher._set_epoch_step_for_training_step")
+    def test_next_response_with_data_available(self, mock_set_epoch_step):
         """Test next_response method when data is available"""
         fetcher = BatchFetcher(step=100, max_size=5, datasets=self.datasets)
 
@@ -150,8 +176,10 @@ def test_next_response_with_data_available(self):
 
         # Verify correct data was returned
         self.assertEqual(result, test_data)
+        mock_set_epoch_step.assert_called_once_with(101)
 
-    def test_next_response_with_empty_queue(self):
+    @patch("simpletuner.helpers.data_backend.runtime.batch_fetcher._set_epoch_step_for_training_step")
+    def test_next_response_with_empty_queue(self, mock_set_epoch_step):
         """Test next_response method blocks when queue is empty"""
         fetcher = BatchFetcher(step=100, max_size=5, datasets=self.datasets)
 
@@ -179,6 +207,7 @@ def add_data_later():
 
         # Verify step was updated
         self.assertEqual(fetcher.step, 101)
+        mock_set_epoch_step.assert_called_once_with(101)
 
     def test_stop_fetching(self):
         """Test stop_fetching method"""
@@ -326,6 +355,19 @@ def test_select_dataloader_index_multiple_backends(self):
         self.assertTrue(len(results) >= 1)  # At least one backend selected
         self.assertTrue(all(r in ["backend1", "backend2"] for r in results))
 
+    @patch("simpletuner.helpers.data_backend.runtime.dataloader_iterator.StateTracker.get_args")
+    def test_select_dataloader_index_is_independent_of_global_rng(self, mock_get_args):
+        mock_get_args.return_value = SimpleNamespace(seed=1234, data_backend_sampling="uniform")
+        backends = {"backend1": self.mock_backend1, "backend2": self.mock_backend2}
+
+        torch.manual_seed(1)
+        first_sequence = [select_dataloader_index(step=step, backends=backends) for step in range(32)]
+        torch.manual_seed(987654)
+        torch.rand(127)
+        second_sequence = [select_dataloader_index(step=step, backends=backends) for step in range(32)]
+
+        self.assertEqual(first_sequence, second_sequence)
+
     def test_select_dataloader_index_empty_backends(self):
         """Test select_dataloader_index with empty backends"""
         backends = {}
@@ -394,6 +436,17 @@ def test_random_dataloader_iterator_single_backend(self, mock_select):
         # Verify correct dataloader was returned
         self.assertEqual(result, "dataloader1")
 
+    @patch("simpletuner.helpers.data_backend.runtime.dataloader_iterator.StateTracker")
+    @patch("simpletuner.helpers.data_backend.runtime.dataloader_iterator.select_dataloader_index")
+    def test_random_dataloader_iterator_can_defer_epoch_state_update(self, mock_select, mock_state_tracker):
+        mock_select.return_value = "backend1"
+        backends = {"backend1": self.mock_backend1}
+
+        result = random_dataloader_iterator(step=100, backends=backends, update_epoch_step=False)
+
+        self.assertEqual(result, "dataloader1")
+        mock_state_tracker.set_epoch_step.assert_not_called()
+
     @patch("simpletuner.helpers.data_backend.runtime.dataloader_iterator.select_dataloader_index")
     def test_random_dataloader_iterator_multiple_backends(self, mock_select):
         """Test random_dataloader_iterator with multiple backends"""