diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 1f8659f7e0..2b3ad1b0f0 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -4101,17 +4101,34 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de raise ValueError("Only supports <= 1 for now, more workers results in duplicated data") elif self.dataset_type == DatasetType.GRAIN: use_hf_parquet = self.hf_path and self.grain_file_type == "parquet" + use_tfds_tfrecord_train = ( + self.grain_file_type == "tfrecord" and self.dataset_path and self.dataset_name and self.train_split + ) + use_tfds_tfrecord_eval = ( + self.grain_file_type == "tfrecord" and self.dataset_path and self.eval_dataset_name and self.eval_split + ) - if not self.grain_train_files and not self.grain_train_mixture_config_path and not use_hf_parquet: + if ( + not self.grain_train_files + and not self.grain_train_mixture_config_path + and not use_hf_parquet + and not use_tfds_tfrecord_train + ): raise ValueError( "When dataset_type=grain, set grain_train_files, " - "grain_train_mixture_config_path, or use hf_path with grain_file_type=parquet." + "grain_train_mixture_config_path, use hf_path with grain_file_type=parquet, or use dataset_path, " + "dataset_name, and train_split with grain_file_type=tfrecord." + ) + if self.eval_interval > 0 and not self.grain_eval_files and not use_hf_parquet and not use_tfds_tfrecord_eval: + raise ValueError( + "Please specify grain_eval_files, use hf_path with grain_file_type=parquet, or use dataset_path, " + "eval_dataset_name, and eval_split with grain_file_type=tfrecord; otherwise set eval_interval to <=0." ) - if self.eval_interval > 0 and not self.grain_eval_files and not use_hf_parquet: - raise ValueError("Please specify grain_eval_files (or hf_path with parquet) or set eval_interval to <=0.") elif self.dataset_type == DatasetType.TFDS: logger.warning( - "tfds pipeline is deprecated. Use dataset_type=grain, grain_file_type=tfrecord, and provide grain_train_files." + "tfds pipeline is deprecated. Use dataset_type=grain and grain_file_type=tfrecord. You can keep the same " + "dataset_path, dataset_name, train_split, eval_dataset_name, and eval_split settings to automatically construct " + "the file paths. Alternatively, provide grain_train_files or grain_eval_files for custom file paths." ) if self.use_dpo: raise ValueError( diff --git a/src/maxtext/input_pipeline/grain_data_processing.py b/src/maxtext/input_pipeline/grain_data_processing.py index 382df3fd16..c075e8ef83 100644 --- a/src/maxtext/input_pipeline/grain_data_processing.py +++ b/src/maxtext/input_pipeline/grain_data_processing.py @@ -66,6 +66,14 @@ def construct_hf_dataset_path(hf_path: str, hf_train_files: str | None = None, s return full_path +def construct_tfds_tfrecord_path(dataset_path: str, dataset_name: str, split: str) -> str: + """Constructs a glob for TFRecords in the standard TFDS prepared-data layout.""" + dataset_dir = dataset_name.strip().strip("/").replace(":", "/") + path = f"{dataset_path.strip().rstrip('/')}/{dataset_dir}/*-{split}.tfrecord-*" + max_logging.log(f"Automatically constructed Grain TFRecord path from TFDS configuration: {path}") + return path + + def find_data_files(data_file_pattern, hf_access_token=None): """Find data files matching the pattern.""" if data_file_pattern.startswith("gs://"): @@ -477,8 +485,22 @@ def make_grain_train_iterator( pipeline_fn = _get_pipeline_fn(config) grain_train_files = config.grain_train_files - if not grain_train_files and not config.grain_train_mixture_config_path and config.hf_path: + if ( + not grain_train_files + and not config.grain_train_mixture_config_path + and config.grain_file_type == "parquet" + and config.hf_path + ): grain_train_files = construct_hf_dataset_path(config.hf_path, split="train") + elif ( + not grain_train_files + and not config.grain_train_mixture_config_path + and config.grain_file_type == "tfrecord" + and config.dataset_path + and config.dataset_name + and config.train_split + ): + grain_train_files = construct_tfds_tfrecord_path(config.dataset_path, config.dataset_name, config.train_split) get_ds_fn = functools.partial( get_datasets, @@ -583,9 +605,17 @@ def make_grain_eval_iterator( pipeline_fn = _get_pipeline_fn(config) grain_eval_files = config.grain_eval_files - if not grain_eval_files and getattr(config, "hf_path", None): + if not grain_eval_files and config.grain_file_type == "parquet" and getattr(config, "hf_path", None): split = getattr(config, "hf_eval_split", None) or "validation" grain_eval_files = construct_hf_dataset_path(config.hf_path, split=split) + elif ( + not grain_eval_files + and config.grain_file_type == "tfrecord" + and config.dataset_path + and config.eval_dataset_name + and config.eval_split + ): + grain_eval_files = construct_tfds_tfrecord_path(config.dataset_path, config.eval_dataset_name, config.eval_split) get_ds_fn = functools.partial( get_datasets, diff --git a/src/maxtext/input_pipeline/tfds_data_processing.py b/src/maxtext/input_pipeline/tfds_data_processing.py index db1ed6958d..8bb7024221 100644 --- a/src/maxtext/input_pipeline/tfds_data_processing.py +++ b/src/maxtext/input_pipeline/tfds_data_processing.py @@ -24,7 +24,12 @@ import tensorflow_datasets as tfds except ImportError as error: raise ImportError( - "TensorFlow and tensorflow-datasets are required. Run `pip install tensorflow tensorflow-datasets`" + "The deprecated TFDS pipeline requires TensorFlow and tensorflow-datasets.\n\n" + "Recommended: migrate to Grain by setting dataset_type=grain and grain_file_type=tfrecord. Existing " + "dataset_path, dataset_name, train_split, eval_dataset_name, and eval_split settings will be used to " + "automatically construct file paths when grain_train_files or grain_eval_files are not provided. Set the " + "grain_train_files or grain_eval_files explicitly for a custom path.\n\n" + "To continue using the deprecated TFDS pipeline, run: `pip install tensorflow tensorflow-datasets`." ) from error import jax diff --git a/tests/unit/grain_data_processing_test.py b/tests/unit/grain_data_processing_test.py index f13d1b3d64..211747974b 100644 --- a/tests/unit/grain_data_processing_test.py +++ b/tests/unit/grain_data_processing_test.py @@ -34,6 +34,27 @@ from tests.utils.test_helpers import get_test_base_output_directory, get_test_config_path, get_test_dataset_path +class TestTfdsTfrecordPathFallback: + """Tests TFDS prepared-data path construction without reading a dataset.""" + + def test_construct_tfds_tfrecord_path(self, mocker): + log = mocker.patch.object(grain_data_processing.max_logging, "log") + assert ( + grain_data_processing.construct_tfds_tfrecord_path(" gs://maxtext-dataset/ ", "c4/en:3.0.1", "train") + == "gs://maxtext-dataset/c4/en/3.0.1/*-train.tfrecord-*" + ) + log.assert_called_once_with( + "Automatically constructed Grain TFRecord path from TFDS configuration: " + "gs://maxtext-dataset/c4/en/3.0.1/*-train.tfrecord-*" + ) + + def test_missing_derived_path_reports_pattern(self, tmp_path): + pattern = str(tmp_path / "c4/en/3.0.1/*-train.tfrecord-*") + + with pytest.raises(FileNotFoundError, match=r"No files found matching pattern: .*\*-train\.tfrecord-\*"): + grain_data_processing.find_data_files(pattern) + + class GrainBaseProcessingTest: """Base mixin with test_train_ds for all grain data processing tests. @@ -532,6 +553,16 @@ class GrainTFRecordProcessingTest(_GrainTFRecordSetup, GrainDeterminismMixin, Gr def setUpClass(cls): super().setUpClass() + def test_config_accepts_tfds_tfrecord_fallback(self): + config = self._make_config( + grain_train_files="", + dataset_path="gs://maxtext-dataset", + dataset_name="c4/en:3.0.1", + train_split="train", + eval_interval=0, + ) + self.assertEqual(config.grain_train_files, "") + class GrainTFRecordPreTokenizedProcessingTest(_GrainTFRecordSetup, GrainBaseProcessingTest, unittest.TestCase): """Test grain data processing with a pre-tokenized TFRecord dataset (tokenize_train_data=False).