Merge pull request #2785 from AI-Hypercomputer:aireen/fix_interleave

PiperOrigin-RevId: 841811659

Author: Google-ML-Automation

src/MaxText/configs/base.yml

@@ -599,8 +599,10 @@ grain_train_mixture_config_path: '' # Path to a JSON file specifying the mixture
 grain_file_type: 'arrayrecord' # arrayrecord or parquet
 grain_worker_count: 1 # Set to -1 to enable auto-tuning: automatically determines optimal worker count. See https://google-grain.readthedocs.io/en/latest/_autosummary/grain.experimental.pick_performance_config.html
 grain_per_worker_buffer_size: 1
-# num_threads and prefetch_buffer_size are per-worker per-dataset. Used in ReadOptions (https://google-grain.readthedocs.io/en/latest/tutorials/data_loader_tutorial.html#per-worker-readoptions)
+# num_threads and prefetch_buffer_size are per-worker per-dataset. 
+# When using array_records, they are used in ReadOptions (https://google-grain.readthedocs.io/en/latest/tutorials/data_loader_tutorial.html#per-worker-readoptions)
 # The default value matches that in the Grain package. If mixing multiple data sources, consider lowering these values to reduce memory usage.
+# When using parquet, grain_num_threads is the number of files to read and interleave in parallel
 grain_num_threads: 16
 grain_prefetch_buffer_size: 500
 grain_worker_count_eval: 1

src/MaxText/input_pipeline/_grain_data_processing.py

@@ -47,6 +47,30 @@ def find_data_files(data_file_pattern):
   return data_files
 
 
+def _apply_mapdataset_transforms(
+    dataset,
+    shuffle,
+    shuffle_seed,
+    num_epoch,
+    dataloading_host_index,
+    dataloading_host_count,
+    grain_num_threads,
+    grain_prefetch_buffer_size,
+):
+  """Apply standard shuffle, repeat, shard, and iter conversion transforms."""
+  if shuffle:
+    dataset = dataset.shuffle(seed=shuffle_seed)
+  dataset = dataset.repeat(num_epoch)
+  dataset = dataset[dataloading_host_index::dataloading_host_count]  # sharding
+  dataset = dataset.to_iter_dataset(
+      read_options=grain.ReadOptions(
+          num_threads=grain_num_threads,
+          prefetch_buffer_size=grain_prefetch_buffer_size,
+      )
+  )
+  return dataset
+
+
 def get_datasets(
     data_file_pattern,
     data_file_type,
@@ -84,12 +108,16 @@ def create_dataset_from_pattern(pattern):
       datasets_dict = dict(zip(mixture_config.keys(), dataset_list))
 
       for name, ds in datasets_dict.items():
-        if shuffle:
-          ds = ds.shuffle(seed=shuffle_seed)
-        ds = ds.repeat(num_epoch)
-        ds = ds[dataloading_host_index::dataloading_host_count]  # sharding
-        ds = ds.to_iter_dataset()
-        datasets_dict[name] = ds
+        datasets_dict[name] = _apply_mapdataset_transforms(
+            ds,
+            shuffle,
+            shuffle_seed,
+            num_epoch,
+            dataloading_host_index,
+            dataloading_host_count,
+            grain_num_threads,
+            grain_prefetch_buffer_size,
+        )
 
       # Normalize weights
       total_weight = sum(weights)
@@ -111,15 +139,15 @@ def create_dataset_from_pattern(pattern):
 
       # Apply shuffle, repeat, sharding, and conversion to IterDataset to each dataset before mixing
       for d, _ in enumerate(dataset_list):
-        if shuffle:
-          dataset_list[d] = dataset_list[d].shuffle(seed=shuffle_seed)
-        dataset_list[d] = dataset_list[d].repeat(num_epoch)
-        dataset_list[d] = dataset_list[d][dataloading_host_index::dataloading_host_count]  # sharding
-        dataset_list[d] = dataset_list[d].to_iter_dataset(
-            read_options=grain.ReadOptions(
-                num_threads=grain_num_threads,
-                prefetch_buffer_size=grain_prefetch_buffer_size,
-            )
+        dataset_list[d] = _apply_mapdataset_transforms(
+            dataset_list[d],
+            shuffle,
+            shuffle_seed,
+            num_epoch,
+            dataloading_host_index,
+            dataloading_host_count,
+            grain_num_threads,
+            grain_prefetch_buffer_size,
         )
       # Use IterDataset.mix instead of MapDataset.mix in order to have per-mixture component checkpoints
       # for supporting changing the mixture after checkpointing
@@ -128,15 +156,15 @@ def create_dataset_from_pattern(pattern):
     else:
       # Single pattern case - no need for parallelization
       dataset = create_dataset_from_pattern(data_file_pattern)
-      if shuffle:
-        dataset = dataset.shuffle(seed=shuffle_seed)
-      dataset = dataset.repeat(num_epoch)
-      dataset = dataset[dataloading_host_index::dataloading_host_count]  # sharding
-      dataset = dataset.to_iter_dataset(
-          read_options=grain.ReadOptions(
-              num_threads=grain_num_threads,
-              prefetch_buffer_size=grain_prefetch_buffer_size,
-          )
+      dataset = _apply_mapdataset_transforms(
+          dataset,
+          shuffle,
+          shuffle_seed,
+          num_epoch,
+          dataloading_host_index,
+          dataloading_host_count,
+          grain_num_threads,
+          grain_prefetch_buffer_size,
       )
       return dataset
   elif data_file_type == "parquet":
@@ -152,8 +180,10 @@ def create_dataset_from_pattern(pattern):
         f"Please lower grain_worker_count or increase file shard count."
     )
     dataset = dataset.map(grain.experimental.ParquetIterDataset)
-    dataset = grain.experimental.InterleaveIterDataset(dataset, cycle_length=len(dataset))
-    dataset = grain.experimental.WindowShuffleIterDataset(dataset, window_size=100, seed=shuffle_seed)
+    cycle_length = min(len(dataset) // num_epoch, grain_num_threads)
+    dataset = grain.experimental.InterleaveIterDataset(dataset, cycle_length=cycle_length)
+    if shuffle:
+      dataset = grain.experimental.WindowShuffleIterDataset(dataset, window_size=100, seed=shuffle_seed)
     return dataset
   else:
     raise ValueError(f"grain pipeline supports (arrayrecord, parquet) as grain_file_type, but got {data_file_type}")