migrate remaining tests off tf

Author: aireenmei

src/maxtext/common/checkpointing.py

@@ -236,7 +236,7 @@ def create_orbax_checkpoint_manager(
     enable_checkpointing: bool,
     use_async: bool,
     save_interval_steps: int,
-    dataset_type: None | str = "tfds",
+    dataset_type: None | str = None,
     orbax_logger: Any = None,  # pytype: disable=attribute-error
     use_ocdbt: bool = True,
     use_zarr3: bool = True,
@@ -269,7 +269,7 @@ def create_orbax_checkpoint_manager(
       )
   }
 
-  if dataset_type == "grain":
+  if dataset_type is not None and dataset_type == "grain":
     item_names += ("iter",)
     item_handlers["iter"] = GrainCheckpointHandler()
 

src/maxtext/input_pipeline/input_pipeline_interface.py

@@ -25,10 +25,6 @@
 from maxtext.input_pipeline.hf_data_processing import make_hf_eval_iterator
 from maxtext.input_pipeline.olmo_grain_data_processing import make_olmo_grain_train_iterator
 from maxtext.input_pipeline.olmo_grain_data_processing import make_olmo_grain_eval_iterator
-from maxtext.input_pipeline.tfds_data_processing import make_tfds_train_iterator
-from maxtext.input_pipeline.tfds_data_processing import make_tfds_eval_iterator
-from maxtext.input_pipeline.tfds_data_processing_c4_mlperf import make_c4_mlperf_train_iterator
-from maxtext.input_pipeline.tfds_data_processing_c4_mlperf import make_c4_mlperf_eval_iterator
 from maxtext.input_pipeline.synthetic_data_processing import SyntheticDataIterator
 from maxtext.input_pipeline.synthetic_data_processing import PlaceHolderDataIterator
 from maxtext.utils import max_logging
@@ -71,12 +67,15 @@ def create_data_iterator(config: pyconfig.HyperParameters, mesh):
     eval_iterator = SyntheticDataIterator(config, mesh) if config.eval_interval > 0 else None
     return SyntheticDataIterator(config, mesh), eval_iterator
   dataset_type_to_train_eval_iterator = {
-      "tfds": (make_tfds_train_iterator, make_tfds_eval_iterator),
       "grain": (make_grain_train_iterator, make_grain_eval_iterator),
       "hf": (make_hf_train_iterator, make_hf_eval_iterator),
-      "c4_mlperf": (make_c4_mlperf_train_iterator, make_c4_mlperf_eval_iterator),
       "olmo_grain": (make_olmo_grain_train_iterator, make_olmo_grain_eval_iterator),
   }
+  if config.dataset_type in ("tfds", "c4_mlperf"):
+    from maxtext.input_pipeline.tfds_data_processing import make_tfds_train_iterator, make_tfds_eval_iterator  # pylint: disable=import-outside-toplevel
+    from maxtext.input_pipeline.tfds_data_processing_c4_mlperf import make_c4_mlperf_train_iterator, make_c4_mlperf_eval_iterator  # pylint: disable=import-outside-toplevel
+    dataset_type_to_train_eval_iterator["tfds"] = (make_tfds_train_iterator, make_tfds_eval_iterator)
+    dataset_type_to_train_eval_iterator["c4_mlperf"] = (make_c4_mlperf_train_iterator, make_c4_mlperf_eval_iterator)
 
   # Collect train and eval iterators
   if config.dataset_type in ["tfds", "grain", "hf", "c4_mlperf", "olmo_grain"]:

src/maxtext/input_pipeline/multihost_dataloading.py

@@ -25,7 +25,12 @@
 import json
 
 from etils import epath
-import tensorflow as tf  # pylint: disable=g-import-not-at-top
+try:
+  import tensorflow as tf
+  _TF_RETRYABLE_ERRORS = (tf.errors.FailedPreconditionError,)
+except ImportError:
+  tf = None  # type: ignore[assignment]
+  _TF_RETRYABLE_ERRORS = ()
 
 import numpy as np
 
@@ -74,14 +79,14 @@ class MultiHostDataLoadIterator:
 
   def __init__(
       self,
-      dataloader: tf.data.Dataset | Iterable,
+      dataloader: Iterable,
       global_mesh: Mesh,
       generate_padding_batch: bool = False,
       expansion_loading_factor_for_grain: int = -1,
   ):
     self.global_mesh = global_mesh
     self.dataloader = dataloader
-    if isinstance(self.dataloader, tf.data.Dataset):
+    if hasattr(self.dataloader, "as_numpy_iterator"):
       self.local_iterator = self.dataloader.as_numpy_iterator()
     elif isinstance(self.dataloader, Iterable):
       self.local_iterator = iter(self.dataloader)
@@ -93,7 +98,7 @@ def __init__(
     self.expansion_loading_factor_for_grain = expansion_loading_factor_for_grain
 
   def reset(self):
-    if isinstance(self.dataloader, tf.data.Dataset):
+    if hasattr(self.dataloader, "as_numpy_iterator"):
       self.local_iterator = self.dataloader.as_numpy_iterator()
     elif isinstance(self.dataloader, Iterable):
       self.local_iterator = iter(self.dataloader)
@@ -132,7 +137,7 @@ def _get_next_batch_sharded(self) -> jax.Array:
               local_data_list.append(next_batch)
             local_data = jtu.tree_map(lambda *xs: np.concatenate(xs, axis=0), *local_data_list)
           break  # exit the loop on success
-        except tf.errors.FailedPreconditionError as e:
+        except _TF_RETRYABLE_ERRORS as e:
           max_logging.log(f"Failed to get next data batch due to {e}, retrying")
           time.sleep(SLEEP_TIME)
         except StopIteration as e:
@@ -188,7 +193,7 @@ def __init__(self, get_ds_fn, preprocessing_fn, global_shape, checkpoint_path, e
   def reset(self):
     ds = self.get_ds_fn(dataloading_host_index=jax.process_index(), dataloading_host_count=jax.process_count())
     dataloader = self.preprocessing_fn(dataset=ds)
-    if isinstance(dataloader, tf.data.Dataset):
+    if hasattr(dataloader, "as_numpy_iterator"):
       self.iterator = dataloader.as_numpy_iterator()
     elif isinstance(dataloader, Iterable):
       self.iterator = iter(dataloader)

tests/end_to_end/tpu/test_convergence_1b_params.sh

@@ -18,7 +18,7 @@ export LOSS_THRESHOLD=100.0 # Set to large value so test is guaranteed to pass.
 export STEPS=20400 # Run for 20B tokens for a 1B sized mode for "chinchilla" scaling https://arxiv.org/abs/2203.15556
 export EVAL_STEPS=160
 export EVAL_INTERVAL=100
-export DATASET_TYPE=tfds
+export DATASET_TYPE=grain
 export MTP_NUM_LAYERS=0 # Disable MTP by default
 export PER_DEVICE_BATCH_SIZE=8.0 # With the default learning rate (3e-4) this should have global batch of 512, with 2k sequence length (1M global batch in tokens)
 

tests/integration/checkpoint_compatibility_test.py

@@ -26,10 +26,16 @@
 """
 
 from datetime import datetime
+import importlib.util
 import json
 import os
 import pytest
 from maxtext.trainers.pre_train.train import main as train_main
+
+pytestmark = pytest.mark.skipif(
+    importlib.util.find_spec("tensorflow") is None,
+    reason="tensorflow not installed; skip testing checkpoint compatibility between tfds and grain",
+)
 from maxtext.utils.globals import MAXTEXT_REPO_ROOT
 from tests.integration.checkpointing_test import get_checkpointing_command
 

tests/unit/attention_test.py

@@ -309,7 +309,7 @@ def fake_to_nnx(*args, **kwargs):  # pylint: disable=unused-argument
         context_parallel_strategy="ring",
         context_parallel_load_balance=False,
         packing=True,
-        dataset_type="tfds",
+        dataset_type="grain",
         max_segments_per_seq=4,
         head_dim=2,
         attention_kernel="cudnn_flash_te",

tests/unit/multihost_dataloading_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 # pylint: disable=missing-module-docstring, missing-function-docstring
+import itertools
 import sys
 import unittest
 
@@ -23,9 +24,6 @@
 import jax
 from jax.sharding import Mesh
 from jax.experimental import mesh_utils
-from jax.sharding import PartitionSpec
-
-import tensorflow as tf
 
 from maxtext.configs import pyconfig
 from maxtext.input_pipeline import multihost_dataloading
@@ -51,16 +49,16 @@ def setUp(self):
         data_sharding=["data"],
         enable_checkpointing=False,
     )
-    global_data_shape = PartitionSpec(batch_size, config.max_target_length)
     mesh_shape_1d = (len(jax.devices()),)
     self.mesh = Mesh(mesh_utils.create_device_mesh(mesh_shape_1d), config.mesh_axes)
-    # creating 2 batches of data
-    global_data = np.arange(np.prod(global_data_shape) * 2).reshape((batch_size * 2, config.max_target_length))
-
-    dataset = tf.data.Dataset.from_tensor_slices(global_data)
-    dataset = dataset.repeat()
-    dataset = dataset.batch(batch_size)
-    self.multihost_gen = multihost_dataloading.MultiHostDataLoadIterator(dataset, self.mesh)
+    # Create 2 distinct batches and cycle through them infinitely.
+    global_data = np.arange(batch_size * 2 * config.max_target_length, dtype=np.int32).reshape(
+        (batch_size * 2, config.max_target_length)
+    )
+    data_batches = [global_data[:batch_size], global_data[batch_size:]]
+    self.multihost_gen = multihost_dataloading.MultiHostDataLoadIterator(
+        itertools.cycle(data_batches), self.mesh
+    )
 
   @pytest.mark.tpu_only
   def test_batch_sharded_data_pipeline(self):

tests/unit/tfds_data_processing_test.py

@@ -17,12 +17,14 @@
 import sys
 import unittest
 
+import pytest
+
 import jax
 from jax.sharding import Mesh
 from jax.experimental import mesh_utils
 
-import tensorflow as tf
-import tensorflow_datasets as tfds
+tf = pytest.importorskip("tensorflow")
+tfds = pytest.importorskip("tensorflow_datasets")
 
 from maxtext.configs import pyconfig
 from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT

tests/unit/train_utils_test.py

@@ -32,7 +32,7 @@ class MockConfig:
   quantization: str = ""
   gradient_accumulation_steps: int = 1
   packing: bool = False
-  dataset_type: str = "tfds"
+  dataset_type: str = "synthetic"
 
   # Fields needed for create_training_optimizer
   opt_type: str = "adamw"