fix(data): no shuffling of validation data

src/dalle_mini/data.py

@@ -182,15 +182,20 @@ class Dataset:
                 yield batch
 
         def _dataloader_datasets_streaming(
-            dataset: Dataset, batch_size: int, epoch: int
+            dataset: Dataset, split: str, batch_size: int, epoch: int
         ):
             keys = ["input_ids", "attention_mask", "labels", "decoder_input_ids"]
             batch = {k: [] for k in keys}
-            first_loop = True
-            while self.multi_hosts or first_loop:
+            first_loop = True  # stop after one loop in some cases
+            while (self.multi_hosts and split == "train") or first_loop:
                 # in multi-host, we run forever (no epoch) as hosts need to stop
-                # at the same time and we don't know how much data is on each host
-                dataset.set_epoch(epoch)  # reshuffle data at each epoch
+                # at the same time and training data may not be split equally
+                # For validation data we put the entire set on each host as we could lose
+                # too many samples on pods
+                if epoch is not None:
+                    # reshuffle training data at each epoch (not applicable with validation set)
+                    dataset.set_epoch(epoch)
+                    epoch += 1
                 for item in dataset:
                     for k, v in item.items():
                         batch[k].append(v)
@@ -199,7 +204,6 @@ class Dataset:
                         batch = shard(batch)
                         yield batch
                         batch = {k: [] for k in keys}
-                epoch += 1
                 first_loop = False
 
         if split == "train":
@@ -210,7 +214,7 @@ class Dataset:
             raise ValueError(f'split must be "train" or "eval", got {split}')
 
         if self.streaming:
-            return _dataloader_datasets_streaming(ds, batch_size, epoch)
+            return _dataloader_datasets_streaming(ds, split, batch_size, epoch)
         else:
             if split == "train":
                 self.rng_dataset, input_rng = jax.random.split(self.rng_dataset)