Adding basic elastic training

Added guards to only use fast-resume if the proxy backend is used.
Added the changes to the jobset for elastic training
Temporary changes to the configuration to decrease batch size
Adding a stop_trace to cancel any ongoing traces
Changing the batch size to match the chip count and the checkpoint step interval to avoid any checkpoints for testing

Author: lukebaumann

axlearn/cloud/gcp/pathways_utils.py

@@ -320,6 +320,8 @@ def _build_pathways_head_sidecar_containers(self) -> list[Nested[Any]]:
             f"--resource_manager_address=localhost:{_PATHWAYS_RESOURCE_MANAGER_PORT}",
             f"--server_port={_PATHWAYS_PROXY_PORT}",
             f"--gcs_scratch_location={staging_location}",
+            # This should be made configurable
+            f"--num_elastic_slices={cfg.accelerator.num_replicas}"
         ]
         cmd_args.extend(xla_flags_from_options(self._xla_options).split())
 
@@ -586,14 +588,19 @@ def _build_pathways_worker_job(
         annotations.update(
             {"alpha.jobset.sigs.k8s.io/exclusive-topology": "cloud.google.com/gke-nodepool"}
         )
+        # Default value for suspend and resume.
+        # References:
+        # https://github.com/google/pathways-job/blob/4417de7aa23d3c2316e400a3a327512834374475/internal/controller/pathwaysjob_controller.go#L651
+        # backoffLimit = system.vms_per_slice * 4
+
+        # This backoffLimit is just for verifying elastic fast-resume
+        large_number = 1000
+        backoffLimit = system.vms_per_slice * 4 * large_number
 
         spec = dict(
             parallelism=system.vms_per_slice,
             completions=system.vms_per_slice,
-            # Default value for suspend and resume.
-            # References:
-            # https://github.com/google/pathways-job/blob/4417de7aa23d3c2316e400a3a327512834374475/internal/controller/pathwaysjob_controller.go#L651
-            backoffLimit=system.vms_per_slice * 4,
+            backoffLimit=backoffLimit,
             template=self._build_pathways_worker_pod(pathways_worker_replicated_job_index),
         )
         worker_job = dict(

axlearn/common/launch_trainer.py

@@ -147,8 +147,34 @@ def run_trainer(trainer_config: SpmdTrainer.Config) -> Any:
                 f,
             )
 
-    trainer: SpmdTrainer = trainer_config.instantiate(parent=None)
-    prng_key = jax.random.PRNGKey(seed=FLAGS.trainer_prng_seed)
-    output = trainer.run(prng_key)
+    if FLAGS.jax_backend == "proxy":
+        # pylint: disable-next=import-error,import-outside-toplevel
+        from pathwaysutils.elastic import manager
+        elastic_manager = manager.Manager()
+        while True:
+            try:
+                trainer: SpmdTrainer = trainer_config.instantiate(parent=None)
+                prng_key = jax.random.PRNGKey(seed=FLAGS.trainer_prng_seed)
+                output = trainer.run(prng_key)
+                break
+            except jax.errors.JaxRuntimeError as error:
+                if not elastic_manager.is_error_due_to_slice_down(error):
+                    raise
+                try:
+                    logging.info("Trying to clean up ongoing traces")
+                    jax.stop_trace()
+                    logging.info("Successfully cleaned up ongoing traces")
+                except ValueError as e:
+                  logging.info("No ongoing traces to clean up", exc_info=True)
+                except Exception as e:
+                  logging.exception("Error trying to clean up ongoing traces")
+                  raise
+                ten_minutes = 10 * 60
+                elastic_manager.wait_for_slices(timeout=ten_minutes)
+    else:
+        trainer: SpmdTrainer = trainer_config.instantiate(parent=None)
+        prng_key = jax.random.PRNGKey(seed=FLAGS.trainer_prng_seed)
+        output = trainer.run(prng_key)
+
     measurement.record_event(measurement.Event.END_JOB)
     return output

axlearn/experiments/text/gpt/fuji.py

@@ -249,6 +249,7 @@ def get_trainer_kwargs(
     max_step = TOTAL_TOKENS[version][model_size] // tokens_per_batch
     max_sequence_length = MAX_SEQUENCE_LENGTH[version]
     train_batch_size = tokens_per_batch // max_sequence_length
+    train_batch_size = 32
 
     # Whether to use grouped query attention.
     num_kv_heads = None
@@ -380,6 +381,7 @@ def get_trainer_kwargs(
             max_sequence_length=max_sequence_length,
             train_batch_size=train_batch_size,
             max_step=max_step,
+            save_every_n_steps=1000000,
             mesh_shape=mesh_shape_from_axes(data=-1, fsdp=8),
             mesh_rules=(
                 # Step time: