Fix hardcoded grain_worker_count in HF data processing pipelines to prevent OOM while running Multimodal SFT via pathway

src/maxtext/configs/types.py

@@ -3499,6 +3499,8 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         self.hf_eval_split = "train"
       if self.eval_interval > 0 and not self.hf_eval_split:
         raise ValueError("Please specify hf_eval_split or set eval_interval to <=0.")
+      if self.grain_worker_count > 1:
+        raise ValueError("Only supports <= 1 for now, more workers results in duplicated data")
     elif self.dataset_type == DatasetType.GRAIN:
       if not self.grain_train_files and not self.grain_train_mixture_config_path:
         raise ValueError("When dataset_type=grain, please set grain_train_files or grain_train_mixture_config_path")

src/maxtext/input_pipeline/hf_data_processing.py

@@ -181,7 +181,7 @@ def vision_sft_preprocessing_pipeline(
       data_source=dataset,
       operations=operations,
       sampler=dummy_index_sampler,
-      worker_count=1,  # only supports <=1 for now, more workers results in duplicated data
+      worker_count=config.grain_worker_count,
       worker_buffer_size=1,
       read_options=grain.ReadOptions(num_threads=1, prefetch_buffer_size=batch_size * 4),
   )
@@ -446,6 +446,7 @@ def make_hf_train_iterator(
         chat_template=config.chat_template,
         formatting_func_path=config.formatting_func_path,
         formatting_func_kwargs=config.formatting_func_kwargs,
+        grain_worker_count=config.grain_worker_count,
     )
   return train_iter
 
@@ -506,5 +507,6 @@ def make_hf_eval_iterator(
         chat_template=config.chat_template,
         formatting_func_path=config.formatting_func_path,
         formatting_func_kwargs=config.formatting_func_kwargs,
+        grain_worker_count=config.grain_worker_count,
     )
   return eval_iter