Support Grain per_worker_buffer_size

Author: bernardhan33

src/MaxText/configs/base.yml

@@ -572,12 +572,14 @@ hf_access_token: ''
 # For multiple patterns, use semicolon (;) to separate and colon (:) to specify weights.
 # Example: "path/to/data1.array_record*:0.3;path/to/data2.array_record*:0.7"
 # Note: When using multiple files (separated by ';'), only ArrayRecord format is supported.
-# For more details, see https://github.com/google/maxtext/blob/main/getting_started/Data_Input_Pipeline.md#grain-input-pipeline
+# For more details, see https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/data_input_pipeline/data_input_grain.md
 grain_train_files: ''
 grain_eval_files: ''
 grain_file_type: 'arrayrecord' # arrayrecord or parquet
 grain_worker_count: 1
+grain_per_worker_buffer_size: 1
 grain_worker_count_eval: 1
+grain_per_worker_buffer_size_eval: 1
 # for using pathways
 colocated_python_data_input: False  # experimental feature, under testing
 

src/MaxText/input_pipeline/_grain_data_processing.py

@@ -88,7 +88,14 @@ def get_datasets(
   return dataset
 
 
-def pretrain_preprocessing_pipeline(dataset, config, data_columns, tokenize, grain_worker_count):
+def pretrain_preprocessing_pipeline(
+    dataset,
+    config,
+    data_columns,
+    tokenize,
+    grain_worker_count,
+    grain_per_worker_buffer_size,
+):
   """Use grain pipeline to pre-process the dataset and return iterators for pretrain"""
   if config.grain_file_type == "arrayrecord":
     dataset = dataset.map(_input_pipeline_utils.ParseFeatures(data_columns, tokenize))
@@ -153,11 +160,23 @@ def pretrain_preprocessing_pipeline(dataset, config, data_columns, tokenize, gra
           axis=1,
       )
   )
-  dataset = dataset.mp_prefetch(grain.MultiprocessingOptions(num_workers=grain_worker_count))
+  dataset = dataset.mp_prefetch(
+      grain.MultiprocessingOptions(
+          num_workers=grain_worker_count,
+          per_worker_buffer_size=grain_per_worker_buffer_size,
+      )
+  )
   return dataset
 
 
-def dpo_preprocessing_pipeline(dataset, config, data_columns, tokenize, grain_worker_count):
+def dpo_preprocessing_pipeline(
+    dataset,
+    config,
+    data_columns,
+    tokenize,
+    grain_worker_count,
+    grain_per_worker_buffer_size,
+):
   """Use grain to pre-process the dataset and return iterators for dpo fine-tuning"""
   if config.grain_file_type == "arrayrecord":
     dataset = dataset.map(_input_pipeline_utils.ParseFeatures(data_columns, tokenize))
@@ -184,7 +203,12 @@ def dpo_preprocessing_pipeline(dataset, config, data_columns, tokenize, grain_wo
   batch_size = config.global_batch_size_to_load // jax.process_count()
   batch_fn = functools.partial(grain.experimental.batch_and_pad, batch_size=batch_size, pad_value=pad_id)
   dataset = dataset.batch(batch_size, batch_fn=batch_fn)
-  dataset = dataset.mp_prefetch(grain.MultiprocessingOptions(num_workers=grain_worker_count))
+  dataset = dataset.mp_prefetch(
+      grain.MultiprocessingOptions(
+          num_workers=grain_worker_count,
+          per_worker_buffer_size=grain_per_worker_buffer_size,
+      )
+  )
   return dataset
 
 
@@ -215,6 +239,7 @@ def make_grain_train_iterator(
           data_columns=config.train_data_columns,
           tokenize=config.tokenize_train_data,
           grain_worker_count=config.grain_worker_count,
+          grain_per_worker_buffer_size=config.grain_per_worker_buffer_size,
       )
     else:
       train_dataloader = pretrain_preprocessing_pipeline(
@@ -223,6 +248,7 @@ def make_grain_train_iterator(
           data_columns=config.train_data_columns,
           tokenize=config.tokenize_train_data,
           grain_worker_count=config.grain_worker_count,
+          grain_per_worker_buffer_size=config.grain_per_worker_buffer_size,
       )
     return multihost_dataloading.MultiHostDataLoadIterator(
         train_dataloader,
@@ -247,6 +273,7 @@ def make_grain_train_iterator(
           data_columns=config.train_data_columns,
           tokenize=config.tokenize_train_data,
           grain_worker_count=config.grain_worker_count,
+          grain_per_worker_buffer_size=config.grain_per_worker_buffer_size,
       )
     else:
       preprocessing_fn = functools.partial(
@@ -255,6 +282,7 @@ def make_grain_train_iterator(
           data_columns=config.train_data_columns,
           tokenize=config.tokenize_train_data,
           grain_worker_count=config.grain_worker_count,
+          grain_per_worker_buffer_size=config.grain_per_worker_buffer_size,
       )
     if config.colocated_python_data_input:
       global_shape = (config.global_batch_size_to_load, config.max_target_length)
@@ -302,6 +330,7 @@ def make_grain_eval_iterator(
           data_columns=config.eval_data_columns,
           tokenize=config.tokenize_eval_data,
           grain_worker_count=config.grain_worker_count_eval,
+          grain_per_worker_buffer_size=config.grain_per_worker_buffer_size_eval,
       )
     else:
       eval_dataloader = pretrain_preprocessing_pipeline(
@@ -310,6 +339,7 @@ def make_grain_eval_iterator(
           data_columns=config.eval_data_columns,
           tokenize=config.tokenize_eval_data,
           grain_worker_count=config.grain_worker_count_eval,
+          grain_per_worker_buffer_size=config.grain_per_worker_buffer_size_eval,
       )
     return multihost_dataloading.MultiHostDataLoadIterator(
         eval_dataloader, global_mesh, config.generate_padding_batch_eval
@@ -331,6 +361,7 @@ def make_grain_eval_iterator(
           data_columns=config.eval_data_columns,
           tokenize=config.tokenize_eval_data,
           grain_worker_count=config.grain_worker_count_eval,
+          grain_per_worker_buffer_size=config.grain_per_worker_buffer_size_eval,
       )
     else:
       preprocessing_fn = functools.partial(
@@ -339,6 +370,7 @@ def make_grain_eval_iterator(
           data_columns=config.eval_data_columns,
           tokenize=config.tokenize_eval_data,
           grain_worker_count=config.grain_worker_count_eval,
+          grain_per_worker_buffer_size=config.grain_per_worker_buffer_size_eval,
       )
     global_shape = (config.global_batch_size_to_load, config.max_target_length)
     return multihost_dataloading.RemoteIterator(get_ds_fn, preprocessing_fn, global_mesh, global_shape)

tests/grain_data_processing_test.py

@@ -165,6 +165,7 @@ def setUp(self):
         grain_file_type="parquet",
         grain_train_files=os.path.join(temp_dir, "gcsfuse", "hf", "c4", "c4-train-00000-of-01637.parquet"),
         grain_worker_count=1,
+        grain_per_worker_buffer_size=1,
         tokenizer_path=os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizer"),
         enable_checkpointing=False,
     )