Re-enable one step ahead device_put of data from datasets on JAX. (#22353)

The `_prefectch_numpy_iterator` feature in the JAX trainer had been disabled based on the idea that many dataset implementations already support prefectching. However, this came from a misundestanding of what `_prefetch_numpy_iterator` does, probably because of the misleading name.

`_prefetch_numpy_iterator` triggers `device_put` for all the arrays of a batch of data to transfer arrays from CPU to the accelerator (GPU or TPU). `device_put` is asynchronous. By triggering `device_put` one step ahead. we parallelize the transfer of the next batch of data with the running of the model with the current batch of data. Without this, the batch of data is immediately used after called `device_put`, which causes a wait until `device_put` is complete.

Prefetching happens only on CPU and is independent.

- Removed the `builtin_prefetch` property, which now does not serve any purpose anymore.
- Renamed `_prefetch_numpy_iterator` to `_one_batch_ahead_iterator`, which more accurately describes its purpose.

Author: hertschuh

keras/src/backend/jax/trainer.py

@@ -1,4 +1,3 @@
-import collections
 import itertools
 import warnings
 from functools import partial
@@ -1088,10 +1087,8 @@ def _get_iterator(self):
         distribution = distribution_lib.distribution()
         if distribution is not None:
             return self._get_distributed_iterator(distribution)
-        if self.data_adapter.builtin_prefetch:
-            return self.data_adapter.get_jax_iterator()
         else:
-            return self._prefetch_numpy_iterator(
+            return self._one_batch_ahead_iterator(
                 self.data_adapter.get_jax_iterator()
             )
 
@@ -1108,27 +1105,23 @@ def _get_distributed_iterator(self, distribution):
                 )
             yield _distribute_data(data, layouts)
 
-    def _prefetch_numpy_iterator(self, numpy_iterator):
-        """Shard and prefetch batches on device.
+    def _one_batch_ahead_iterator(self, numpy_iterator):
+        """Initiate transfers to the device one batch ahead.
 
-        Most of the implementation has been borrowed from
-        `flax.jax_utils.prefetch_to_device`
-
-        This utility takes an iterator and returns a new iterator which fills an
-        on device prefetch buffer. Eager prefetching can improve the performance
-        of training loops significantly by overlapping compute and data
-        transfer.
+        This utility takes an iterator and returns a new iterator which
+        initiates the transfer to device one step ahead. This can improve the
+        performance of training loops significantly by overlapping compute and
+        data transfer.
         """
-        queue = collections.deque()
-
-        # If you're training on GPUs, 2 is generally the best choice because
-        # this guarantees that you can overlap a training step on GPU with a
-        # data prefetch step on CPU.
-        def enqueue(n=2):
-            for data in itertools.islice(numpy_iterator, n):
-                queue.append(_distribute_data(data))
-
-        enqueue(n=2)  # TODO: should we make `n` configurable?
-        while queue:
-            yield queue.popleft()
-            enqueue(1)
+        next_batch = None
+        for batch in numpy_iterator:
+            batch = _distribute_data(batch)
+            if next_batch is None:
+                next_batch = batch
+            else:
+                current_batch = next_batch
+                next_batch = batch
+                yield current_batch
+
+        if next_batch is not None:
+            yield next_batch

keras/src/trainers/data_adapters/data_adapter.py

@@ -46,21 +46,6 @@ def get_torch_dataloader(self):
         """
         raise NotImplementedError
 
-    @property
-    def builtin_prefetch(self):
-        """Whether the DataAdapter has built-in prefetching capabilities.
-
-        Prefetching is an optimization technique where data is loaded and
-        prepared in advance while the model is processing the current batch,
-        reducing training time by overlapping data loading with computation.
-
-        Returns:
-            bool: True if the DataAdapter implements its own prefetching
-            mechanism and handles data loading asynchronously. False if the
-            caller should implement prefetching externally.
-        """
-        return False
-
     @property
     def num_batches(self):
         """Return the size (number of batches) for the dataset created.

keras/src/trainers/data_adapters/grain_dataset_adapter.py

@@ -201,10 +201,6 @@ def __iter__(self):
             ConverterIterableDataset(self._dataset), batch_size=None
         )
 
-    @property
-    def builtin_prefetch(self):
-        return True
-
     @property
     def num_batches(self):
         return None

keras/src/trainers/data_adapters/grain_dataset_adapter_test.py

@@ -185,11 +185,6 @@ def test_multiple_calling_on_iterators(self):
                 bx, by = batch
                 self.assertEqual(bx.dtype, by.dtype)
 
-    def test_builtin_prefetch(self):
-        dataset = grain.MapDataset.source(Range2DSource(0, 42))
-        adapter = grain_dataset_adapter.GrainDatasetAdapter(dataset)
-        self.assertTrue(adapter.builtin_prefetch)
-
     def test_num_batches(self):
         dataset = grain.MapDataset.source(Range2DSource(0, 42))
         adapter = grain_dataset_adapter.GrainDatasetAdapter(dataset)

keras/src/trainers/data_adapters/tf_dataset_adapter.py

@@ -62,10 +62,6 @@ def get_tf_dataset(self):
     def get_torch_dataloader(self):
         return data_adapter_utils.get_torch_dataloader(self._dataset)
 
-    @property
-    def builtin_prefetch(self):
-        return True
-
     @property
     def num_batches(self):
         cardinality = self._dataset.cardinality

keras/src/trainers/data_adapters/tf_dataset_adapter_test.py

@@ -84,11 +84,6 @@ def test_class_weights_int_targets(self):
     def test_class_weights_categorical_targets(self):
         self._test_class_weights(target_encoding="categorical")
 
-    def test_builtin_prefetch(self):
-        dataset = tf.data.Dataset.range(42)
-        adapter = tf_dataset_adapter.TFDatasetAdapter(dataset)
-        self.assertTrue(adapter.builtin_prefetch)
-
     def test_num_batches(self):
         dataset = tf.data.Dataset.range(42)
         cardinality = int(dataset.cardinality())

keras/src/trainers/data_adapters/torch_data_loader_adapter.py

@@ -65,14 +65,6 @@ def get_tf_dataset(self):
     def get_torch_dataloader(self):
         return self._dataloader
 
-    @property
-    def builtin_prefetch(self):
-        prefetch_factor = self._dataloader.prefetch_factor
-        if prefetch_factor is not None and prefetch_factor > 0:
-            return True
-        else:
-            return False
-
     @property
     def num_batches(self):
         return self._num_batches

keras/src/trainers/data_adapters/torch_data_loader_adapter_test.py

@@ -171,17 +171,3 @@ def test_with_different_shapes(self):
             else:
                 self.assertEqual(bx.shape, (2, 6))
                 self.assertEqual(by.shape, (2, 2))
-
-    @parameterized.named_parameters(named_product(num_workers=[0, 2]))
-    def test_builtin_prefetch(self, num_workers):
-        x = torch.normal(2, 3, size=(34, 4))
-        y = torch.normal(1, 3, size=(34, 2))
-        ds = torch.utils.data.TensorDataset(x, y)
-        dataloader = torch.utils.data.DataLoader(
-            ds, batch_size=16, num_workers=num_workers
-        )
-        adapter = TorchDataLoaderAdapter(dataloader)
-        if num_workers > 0:
-            self.assertTrue(adapter.builtin_prefetch)
-        else:
-            self.assertFalse(adapter.builtin_prefetch)