Small update to keep elastic training working

Adds input_data_shardings to elastic_handler

Author: lukebaumann

MaxText/configs/base.yml

@@ -745,4 +745,4 @@ projector_output_dim_for_vit: 4096
 rope_theta_for_vit: 10000
 vision_output_dim_for_vit: 4096
 pixel_shuffle_ratio_for_vit: 0.5
-projector_dropout_for_vit: 0.0
+projector_dropout_for_vit: 0.0

MaxText/elastic_train.py

@@ -122,6 +122,7 @@ def elastic_handler(
     )
     with mesh:
       data_iterator, _ = create_data_iterator(config, mesh)
+      input_data_shardings = maxtext_utils.get_input_data_sharding(config, mesh)
 
       step, snapshot_jax_arrays, _ = elastic_manager.get_resharded_snapshot(mesh)
 
@@ -183,6 +184,7 @@ def elastic_handler(
       example_batch,
       learning_rate_schedule,
       metric_logger,
+      input_data_shardings,
   )
 
 
@@ -301,6 +303,7 @@ def train_loop(config, elastic_manager, recorder, state=None):
             example_batch,
             learning_rate_schedule,
             metric_logger,
+            input_data_shardings,
         ) = ret
 
       if step == start_step:
@@ -335,6 +338,7 @@ def train_loop(config, elastic_manager, recorder, state=None):
             example_batch,
             learning_rate_schedule,
             metric_logger,
+            input_data_shardings,
         ) = ret
     except exceptions.StopTraining as error:
       max_logging.log(f"Training stopped: {str(error)}")