Updating elastic training to work with pathwaysutils 0.1.1 and pinning the version so that we can control updates better in the future

Author: lukebaumann

MaxText/elastic_train.py

@@ -118,7 +118,7 @@ def elastic_handler(
     with mesh:
       data_iterator, _ = create_data_iterator(config, mesh)
 
-      step, snapshot = elastic_manager.get_resharded_snapshot(mesh)
+      step, snapshot_jax_arrays, _ = elastic_manager.get_resharded_snapshot(mesh)
 
       # We do not want to restore from the previous checkpoint but instead
       # restore from the host offloaded snapshot.
@@ -143,7 +143,7 @@ def elastic_handler(
           checkpoint_manager=None,
       )
 
-      state = state.replace(**snapshot)
+      state = state.replace(**snapshot_jax_arrays)
       state = state.replace(step=state.step.at[None].set(step))
 
       (
@@ -259,7 +259,7 @@ def train_loop(config, elastic_manager, state=None):
 
   elastic_manager.maybe_snapshot(
       step,
-      snapshot={
+      snapshot_jax_arrays={
           "params": state.params,
           "opt_state": state.opt_state,
       },
@@ -314,7 +314,7 @@ def train_loop(config, elastic_manager, state=None):
 
       elastic_manager.maybe_snapshot(
           step=step,
-          snapshot={
+          snapshot_jax_arrays={
               "params": state.params,
               "opt_state": state.opt_state,
           },
@@ -323,7 +323,7 @@ def train_loop(config, elastic_manager, state=None):
 
       ret = elastic_manager.maybe_reshard_up(
           step=step,
-          snapshot={
+          snapshot_jax_arrays={
               "params": state.params,
               "opt_state": state.opt_state,
           },

requirements.txt

@@ -38,5 +38,5 @@ transformers
 mlperf-logging@git+https://github.com/mlperf/logging.git
 google-jetstream@git+https://github.com/AI-Hypercomputer/JetStream.git
 jsonlines
-pathwaysutils@git+https://github.com/AI-Hypercomputer/pathways-utils.git
-omegaconf
+pathwaysutils==0.1.1
+omegaconf

requirements_with_jax_stable_stack.txt

@@ -18,7 +18,7 @@ transformers
 mlperf-logging@git+https://github.com/mlperf/logging.git
 google-jetstream@git+https://github.com/AI-Hypercomputer/JetStream.git
 jsonlines
-pathwaysutils@git+https://github.com/AI-Hypercomputer/pathways-utils.git
+pathwaysutils==0.1.1
 google-api-python-client
 omegaconf
-jaxtyping
+jaxtyping

requirements_with_jax_stable_stack_0_5_2_pipreqs.txt

@@ -27,7 +27,7 @@ omegaconf==2.3.0
 optax==0.2.4
 orbax==0.1.9
 pandas==2.2.3
-pathwaysutils==0.1.0
+pathwaysutils==0.1.1
 # Removing due to conflicts during build
 # protobuf==3.20.3
 protobuf
@@ -53,4 +53,4 @@ tqdm==4.67.1
 transformer_engine==2.1.0
 transformers==4.51.3
 trl==0.16.1
-urllib3==2.4.0
+urllib3==2.4.0