migrate to colocated python class

Author: aireenmei

src/maxtext/common/checkpointing.py

@@ -24,7 +24,7 @@
 import jax
 from maxtext.utils.globals import DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE
 from maxtext.input_pipeline.multihost_dataloading import MultiHostDataLoadIterator
-from maxtext.input_pipeline.multihost_dataloading import RemoteIterator
+from maxtext.input_pipeline.multihost_dataloading import RemoteIteratorWrapper
 from maxtext.input_pipeline.synthetic_data_processing import PlaceHolderDataIterator
 from maxtext.utils import exceptions
 from maxtext.utils import max_logging
@@ -132,11 +132,11 @@ class GrainCheckpointRestore(ocp.args.CheckpointArgs):
 
 
 def _is_remote_iterator(data_iterator):
-  """Check if data_iterator is a RemoteIterator or contains RemoteIterator instances."""
-  if isinstance(data_iterator, RemoteIterator):
+  """Check if data_iterator is a RemoteIteratorWrapper or contains RemoteIteratorWrapper instances."""
+  if isinstance(data_iterator, RemoteIteratorWrapper):
     return True
   if isinstance(data_iterator, list):
-    return any(isinstance(item, RemoteIterator) for item in data_iterator)
+    return any(isinstance(item, RemoteIteratorWrapper) for item in data_iterator)
   return False
 
 
@@ -619,7 +619,7 @@ def map_to_pspec(data):
               None,
           )
         # Case 2: Matches if dataset type is "grain" and the data iterator is not a
-        # PlaceHolderDataIterator or RemoteIterator and a specific checkpoint file exists for the iterator
+        # PlaceHolderDataIterator or RemoteIteratorWrapper and a specific checkpoint file exists for the iterator
         case (
             checkpoint_manager,
             dataset_type,

src/maxtext/input_pipeline/grain_data_processing.py

@@ -454,7 +454,7 @@ def make_grain_train_iterator(
       )
     if config.colocated_python_data_input:
       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)
+      return multihost_dataloading.RemoteIteratorWrapper(get_ds_fn, preprocessing_fn, global_mesh, global_shape)
     else:
       # config.expansion_factor_real_data is between 0 and 1
       num_dataloader_to_restore = int(1 / config.expansion_factor_real_data)
@@ -567,4 +567,4 @@ def make_grain_eval_iterator(
           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)
+    return multihost_dataloading.RemoteIteratorWrapper(get_ds_fn, preprocessing_fn, global_mesh, global_shape)

src/maxtext/input_pipeline/multihost_dataloading.py

@@ -205,9 +205,53 @@ def _colocated_cpu_mesh(mesh: Mesh) -> Mesh:
   return colocated_python.colocated_cpu_devices(mesh)
 
 
+@colocated_python.colocated_python_class
 class RemoteIterator:
-  "iterator class for using colocated python, iterator is initiated remotely and stored in the state of colocated python"
+  "iterator class for using colocated python class"
+
+  def __init__(self, get_ds_fn, preprocessing_fn, global_shape):
+    # self.cpu_devices = _colocated_cpu_devices(jax.local_devices())
+    # self.tpu_devices = jax.local_devices()
+    # self.cpu_mesh = _colocated_cpu_mesh(global_mesh)
+    # self.tpu_sharding = jax.sharding.NamedSharding(global_mesh, PartitionSpec(global_mesh.axis_names))
+    # self.cpu_sharding = jax.sharding.NamedSharding(self.cpu_mesh, PartitionSpec(self.cpu_mesh.axis_names))
+    self.global_shape = global_shape
+    ds = get_ds_fn(dataloading_host_index=jax.process_index(), dataloading_host_count=jax.process_count())
+    dataloader = preprocessing_fn(dataset=ds)
+    self.iterator = None
+    if isinstance(dataloader, tf.data.Dataset):
+      self.iterator = dataloader.as_numpy_iterator()
+    elif isinstance(dataloader, Iterable):
+      self.iterator = iter(dataloader)
+    else:
+      raise ValueError("Type error: dataloader should be Iterable.")
 
+  def get_next(self, dummy_array):
+    local_data = next(self.iterator)
+    def form_global_array_colocated_python(path, array, devices, global_shape, sharding):
+      try:
+        device_arrays = np.split(array, len(devices), axis=0)
+      except ValueError as array_split_error:
+        raise ValueError(
+            f"Unable to put to devices shape {array.shape} with "
+            f"local device count {len(devices)} "
+            f"at {jtu.keystr(path)}"
+        ) from array_split_error
+      device_arrays = jax.device_put(device_arrays, devices)
+      return jax.make_array_from_single_device_arrays(shape=global_shape, sharding=sharding, arrays=device_arrays)
+
+    return jtu.tree_map_with_path(
+      partial(
+          form_global_array_colocated_python,
+          devices=list(dummy_array.sharding.addressable_devices),
+          global_shape=self.global_shape,
+          sharding=dummy_array.sharding,
+      ),
+      local_data,
+    )
+    
+
+class RemoteIteratorWrapper:
   def __init__(self, get_ds_fn, preprocessing_fn, global_mesh, global_shape):
     self.cpu_devices = _colocated_cpu_devices(jax.local_devices())
     self.tpu_devices = jax.local_devices()
@@ -216,42 +260,60 @@ def __init__(self, get_ds_fn, preprocessing_fn, global_mesh, global_shape):
     self.cpu_sharding = jax.sharding.NamedSharding(self.cpu_mesh, PartitionSpec(self.cpu_mesh.axis_names))
     self.dummy_array = jnp.zeros((len(self.cpu_devices)))
     self.dummy_array = jax.device_put(self.dummy_array, self.cpu_sharding)
-
-    @colocated_python.colocated_python
-    def init(dummy_array):
-      colocated_python.global_shape = global_shape
-      ds = get_ds_fn(dataloading_host_index=jax.process_index(), dataloading_host_count=jax.process_count())
-      dataloader = preprocessing_fn(dataset=ds)
-      if isinstance(dataloader, tf.data.Dataset):
-        colocated_python.iterator = dataloader.as_numpy_iterator()
-      elif isinstance(dataloader, Iterable):
-        colocated_python.iterator = iter(dataloader)
-      else:
-        raise ValueError("Type error: dataloader should be either tf.data.Dataset or grain.DataLoader.")
-      return dummy_array
-
-    max_logging.log("Initiating RemoteIterator")
-    try:
-      out = jax.device_get(init(self.dummy_array))
-      if out is not None:
-        max_logging.log(f"RemoteIterator initiated. Test output: {out}")
-    except Exception as e:
-      max_logging.log(f"RemoteIterator init FAILED with error: {type(e).__name__}: {e}")
-      raise
-
-  def __iter__(self):
-    return self
+    self.remote_iterator = RemoteIterator(get_ds_fn, preprocessing_fn, global_shape)
 
   def __next__(self):
-    out = _get_next(self.dummy_array)
-
-    def put_to_tpu_devices(path, array, sharding):
-      try:
-        return jax.device_put(array, sharding)
-      except Exception as e:  # pylint: disable=broad-exception-caught
-        max_logging.log(f"Error putting data to TPU device path{path}, exception={e}")
-        raise
-
-    input_gdas = jtu.tree_map_with_path(partial(put_to_tpu_devices, sharding=self.tpu_sharding), out)
-
-    return input_gdas
+    out = self.remote_iterator.get_next(self.dummy_array)
+    # use tree_map is out is a dict
+    return jax.device_put(out, self.tpu_sharding)
+
+# class RemoteIterator:
+#   "iterator class for using colocated python, iterator is initiated remotely and stored in the state of colocated python"
+
+#   def __init__(self, get_ds_fn, preprocessing_fn, global_mesh, global_shape):
+#     self.cpu_devices = _colocated_cpu_devices(jax.local_devices())
+#     self.tpu_devices = jax.local_devices()
+#     self.cpu_mesh = _colocated_cpu_mesh(global_mesh)
+#     self.tpu_sharding = jax.sharding.NamedSharding(global_mesh, PartitionSpec(global_mesh.axis_names))
+#     self.cpu_sharding = jax.sharding.NamedSharding(self.cpu_mesh, PartitionSpec(self.cpu_mesh.axis_names))
+#     self.dummy_array = jnp.zeros((len(self.cpu_devices)))
+#     self.dummy_array = jax.device_put(self.dummy_array, self.cpu_sharding)
+
+#     @colocated_python.colocated_python
+#     def init(dummy_array):
+#       colocated_python.global_shape = global_shape
+#       ds = get_ds_fn(dataloading_host_index=jax.process_index(), dataloading_host_count=jax.process_count())
+#       dataloader = preprocessing_fn(dataset=ds)
+#       if isinstance(dataloader, tf.data.Dataset):
+#         colocated_python.iterator = dataloader.as_numpy_iterator()
+#       elif isinstance(dataloader, Iterable):
+#         colocated_python.iterator = iter(dataloader)
+#       else:
+#         raise ValueError("Type error: dataloader should be either tf.data.Dataset or grain.DataLoader.")
+#       return dummy_array
+
+#     max_logging.log("Initiating RemoteIterator")
+#     try:
+#       out = jax.device_get(init(self.dummy_array))
+#       if out is not None:
+#         max_logging.log(f"RemoteIterator initiated. Test output: {out}")
+#     except Exception as e:
+#       max_logging.log(f"RemoteIterator init FAILED with error: {type(e).__name__}: {e}")
+#       raise
+
+#   def __iter__(self):
+#     return self
+
+#   def __next__(self):
+#     out = _get_next(self.dummy_array)
+
+#     def put_to_tpu_devices(path, array, sharding):
+#       try:
+#         return jax.device_put(array, sharding)
+#       except Exception as e:  # pylint: disable=broad-exception-caught
+#         max_logging.log(f"Error putting data to TPU device path{path}, exception={e}")
+#         raise
+
+#     input_gdas = jtu.tree_map_with_path(partial(put_to_tpu_devices, sharding=self.tpu_sharding), out)
+
+#     return input_gdas

src/maxtext/input_pipeline/tfds_data_processing.py

@@ -232,7 +232,7 @@ def make_tfds_train_iterator(
         hf_access_token=config.hf_access_token,
     )
     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)
+    return multihost_dataloading.RemoteIteratorWrapper(get_ds_fn, preprocessing_fn, global_mesh, global_shape)
 
 
 def make_tfds_eval_iterator(
@@ -296,4 +296,4 @@ def make_tfds_eval_iterator(
         use_dpo=config.use_dpo,
         hf_access_token=config.hf_access_token,
     )
-    return multihost_dataloading.RemoteIterator(get_ds_fn, preprocessing_fn, config, global_mesh)
+    return multihost_dataloading.RemoteIteratorWrapper(get_ds_fn, preprocessing_fn, config, global_mesh)