fix

Author: chiajunglien

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("grain_worker_count must be <= 1 when dataset_type=hf to prevent OOM")
     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")