individual save and restore works

axlearn/cloud/gcp/pathways_utils.py

@@ -604,10 +604,12 @@ def _build_pathways_head_sidecar_containers(self) -> list[Nested[Any]]:
         # If multi-head, every pathways-head will only
         # be connected to one pathways instance (a pathways-worker replicated job).
         pathways_instance_count = cfg.accelerator.num_replicas if self._is_single_head else 1
+        num_elastic_slices = 1
 
         cmd_args = [
             f"--resource_manager_address=localhost:{_PATHWAYS_RESOURCE_MANAGER_PORT}",
             f"--server_port={_PATHWAYS_PROXY_PORT}",
+            f"--num_elastic_slices={num_elastic_slices}",
         ]
         if self._colocated_python.is_colocated_python_enabled:
             cmd_args.append("--sidecar_name=external")

axlearn/common/elastic_input.py

@@ -49,7 +49,7 @@
 from axlearn.common.config import REQUIRED, Required, config_class, maybe_set_config
 from axlearn.common.input_dispatch import BaseInputDispatcher, _validate_logical_feed_shapes
 from axlearn.common.module import Module
-from axlearn.common.utils import Nested, Tensor
+from axlearn.common.utils import Nested, Tensor, live_devices
 
 
 class ElasticSpmdInputDispatcher(BaseInputDispatcher):
@@ -72,13 +72,16 @@ class Config(BaseInputDispatcher.Config):
 
     @property
     def is_in_elastic_mode(self) -> bool:
+        print("In is_in_elastic_mode by lkolluru")
         cfg = self.config
+        print("cfg.num_max_slices by lkolluru: ", cfg.num_max_slices)
+        print("slice_count by lkolluru: ", slice_count())
         if cfg.num_max_slices is None:
             return False
         else:
             if slice_count() < cfg.num_max_slices:
                 return True
-            elif slice_count() == cfg.num_max_slices:
+            elif slice_count() >= cfg.num_max_slices:
                 return False
             else:
                 # TODO (jtian22): consider supporting scaling up in the future.
@@ -156,14 +159,32 @@ def fid2pids(feed_id):
 
         self.feed_count = len(set(pid2fid.values())) // slice_count() * cfg.num_max_slices
         self.feed_index = pid2fid[jax.process_index()]
+        print("feed_count by lkolluru: ", self.feed_count)
+        print("global_logical_batch_size by lkolluru: ", cfg.global_logical_batch_size)
 
-        assert cfg.global_logical_batch_size % self.feed_count == 0
-        self._feed_logical_batch_size = cfg.global_logical_batch_size // self.feed_count
+        # assert cfg.global_logical_batch_size % self.feed_count == 0
+        if self.feed_count == 0:
+            self._feed_logical_batch_size = cfg.global_logical_batch_size
+        else:
+            self._feed_logical_batch_size = cfg.global_logical_batch_size // self.feed_count
+
+        adjusted_device_physical_batch_size = math.ceil(
+            self._device_physical_batch_size * (cfg.num_max_slices / slice_count())
+        )
+        print(
+            "adjusted_device_physical_batch_size outside elastic by lkolluru: ",
+            adjusted_device_physical_batch_size,
+        )
 
         if self.is_in_elastic_mode:
+            print(" In elastic mode lkolluru")
             adjusted_device_physical_batch_size = math.ceil(
                 self._device_physical_batch_size * (cfg.num_max_slices / slice_count())
             )
+            print(
+                "adjusted_device_physical_batch_size inside elastic by lkolluru: ",
+                adjusted_device_physical_batch_size,
+            )
             padding_per_device = (
                 adjusted_device_physical_batch_size - self._device_physical_batch_size
             )
@@ -402,7 +423,14 @@ def _padded_select(path, x, y):
 
 def slice_count() -> int:
     """Returns the number of slices."""
-    return len(set(d.slice_index for d in jax.devices() if hasattr(d, "slice_index"))) or 1
+    # slice_cnt_val=len(set(d.slice_index for d in jax.devices() if hasattr(d, "slice_index"))) or 1
+    # print("slice_count by lkolluru: ", slice_cnt_val)
+    # return len(set(d.slice_index for d in jax.devices() if hasattr(d, "slice_index"))) or 1
+    slice_cnt_val = (
+        len(set(d.slice_index for d in live_devices() if hasattr(d, "slice_index"))) or 1
+    )
+    print("slice_count by lkolluru: ", slice_cnt_val)
+    return len(set(d.slice_index for d in live_devices() if hasattr(d, "slice_index"))) or 1
 
 
 def process_count_per_slice() -> int:
@@ -411,7 +439,8 @@ def process_count_per_slice() -> int:
         len(
             set(
                 d.process_index
-                for d in jax.devices()
+                # for d in jax.devices()
+                for d in live_devices()
                 if hasattr(d, "slice_index") and d.slice_index == 0
             )
         )
@@ -502,6 +531,8 @@ def get_process_index_and_count_and_mapping(
     # compatible with any mesh with num_devices.
     device_map = tensor_sharding.devices_indices_map((tensor_sharding.num_devices,) * ndims)
 
+    print("device_map by lkolluru: ", device_map)
+
     # Get the slices for 'dim' for all devices.
     global_slice = {k: v[dim] for k, v in device_map.items()}
 
@@ -516,11 +547,15 @@ def get_process_index_and_count_and_mapping(
         process_to_slice[d.process_index].add(key)
         all_slices.add(key)
 
+    print("process_to_slice by lkolluru: ", process_to_slice)
+
     # Get the set of slices for the current process which we will use to compute
     # the index of the current process.
     current_pid = next(iter(tensor_sharding.addressable_devices)).process_index
     addressable_slices = frozenset(process_to_slice[current_pid])
 
+    print("addressable_slices by lkolluru: ", addressable_slices)
+
     # Verify that all processes have the same number of slices.
     slices_per_process = len(addressable_slices)
     if any(len(x) != slices_per_process for x in process_to_slice.values()):
@@ -529,6 +564,7 @@ def get_process_index_and_count_and_mapping(
             "different number of slices."
         )
     unique_processes = list({frozenset(x) for x in process_to_slice.values()})
+    print("unique_processes by lkolluru: ", unique_processes)
 
     # After removing duplicate processes all unique slices should
     # cover the dimension exactly once. If they don't it means that
@@ -540,6 +576,8 @@ def get_process_index_and_count_and_mapping(
     # !!! patch begin
     pid2fid = {}
     for pid, _ in process_to_slice.items():
+        print("pid by lkolluru: ", pid)
         pid2fid[pid] = unique_processes.index(frozenset(process_to_slice[pid]))
     # !!! patch end
-    return feed_index, feed_count, pid2fid
+    print("pid2fid by lkolluru: ", pid2fid)
+    return feed_index, feed_count, pid2fid

axlearn/common/launch_trainer.py

@@ -4,6 +4,8 @@
 
 import json
 import os
+import time
+import gc
 from typing import Any, Optional
 
 import jax
@@ -12,8 +14,11 @@
 from axlearn.common import file_system as fs
 from axlearn.common import measurement
 from axlearn.common.config import TrainerConfigFn, get_named_trainer_config
-from axlearn.common.trainer import SpmdTrainer, select_mesh_config
-from axlearn.common.utils import MeshShape, get_data_dir, infer_mesh_shape
+from axlearn.common.trainer import SpmdTrainer, select_mesh_config, sync_restore_class_vars, sync_store_class_vars
+from axlearn.common.utils import MeshShape, get_data_dir, infer_mesh_shape, live_devices, set_elastic_manager
+from pathwaysutils.elastic import manager, elastic
+from pathwaysutils.debug import watchdog
+
 
 # Trainer-specific flags.
 flags.DEFINE_string(
@@ -116,6 +121,8 @@
 
 FLAGS = flags.FLAGS
 
+elastic_snapshotting_enabled = True
+
 
 def get_trainer_config(
     trainer_config_fn: Optional[TrainerConfigFn] = None,
@@ -148,7 +155,8 @@ def get_trainer_config(
     if flag_values.mesh_selector is not None:
         select_mesh_config(trainer_config, mesh_selector=flag_values.mesh_selector)
     trainer_config.mesh_axis_names = trainer_config.mesh_axis_names or ("data", "model")
-    trainer_config.mesh_shape = trainer_config.mesh_shape or (len(jax.devices()), 1)
+    #trainer_config.mesh_shape = trainer_config.mesh_shape or (len(jax.devices()), 1)
+    trainer_config.mesh_shape = trainer_config.mesh_shape or (len(live_devices()), 1)
     if isinstance(trainer_config.mesh_shape, MeshShape):
         trainer_config.mesh_shape = infer_mesh_shape(trainer_config.mesh_shape)
     trainer_config.start_trace_steps = [int(el) for el in flag_values.trace_at_steps]
@@ -189,6 +197,16 @@ def get_trainer_config(
     return trainer_config
 
 
+def is_retryable_error(e: Exception) -> bool:
+    if isinstance(e, jax.errors.JaxRuntimeError):
+        err_str = str(e)
+        if elastic.is_error_due_to_slice_down(e):
+            return True
+        if "UNAVAILABLE" in err_str or "RESOURCE_EXHAUSTED" in err_str:
+            return True
+    return False
+
+
 def run_trainer(trainer_config: SpmdTrainer.Config) -> Any:
     measurement.record_event(measurement.Event.START_JOB)
     trainer_config_debug_string = trainer_config.debug_string()
@@ -207,9 +225,74 @@ def run_trainer(trainer_config: SpmdTrainer.Config) -> Any:
                 },
                 f,
             )
+
+    elastic_manager = None
+    elastic_manager_initialized = False
+
+    output = None
+    jax_device_state = {}
+    python_vars = {}
+    immutable_data = {}
+    trainer = None
+    logging.info("[ELASTIC] Starting elastic training loop cycle.")
+    while True:
+        try:
+            if not elastic_manager_initialized:
+                if elastic_snapshotting_enabled:
+                    logging.info("[ELASTIC] Initializing elastic manager...")
+                    elastic_manager = manager.Manager()
+                    set_elastic_manager(elastic_manager)
+                    logging.info("[ELASTIC] Elastic manager initialized.")
+                else:
+                    logging.info("[ELASTIC] Elastic snapshotting disabled or not supported (no slice_index).")
+                elastic_manager_initialized = True
+
+            clean_trainer: SpmdTrainer = trainer_config.instantiate(parent=None)
+            logging.info("[ELASTIC] Instantiated clean trainer.")
+
+            if elastic_manager and elastic_manager.new_slice_event.is_set():
+                logging.info("[ELASTIC] New slice event is set. Restoring from snapshot...")
+                elastic_manager.new_slice_event.clear()
+                trainer, prng_key = sync_restore_class_vars(clean_trainer, jax_device_state, python_vars, immutable_data)
+                logging.info("[ELASTIC] Restored trainer state from class vars.")
+            else:
+                logging.info("[ELASTIC] Starting fresh trainer (no elastic recovery triggered).")
+                trainer = clean_trainer
+                prng_key = jax.random.PRNGKey(seed=FLAGS.trainer_prng_seed)
+
+            logging.info("[ELASTIC] Starting trainer.run().")
+            output = trainer.run(prng_key)
+            logging.info("[ELASTIC] trainer.run() completed.")
+            measurement.record_event(measurement.Event.END_JOB)
+            break
+
+        except jax.errors.JaxRuntimeError as e:
+            if is_retryable_error(e):
+                logging.warning("[ELASTIC] Caught retryable error: %s. Retrying...", e)
+                if trainer is not None:
+                    jax_device_state = getattr(trainer, "_jax_device_state", {})
+                    python_vars = getattr(trainer, "_python_vars", {})
+                    immutable_data = getattr(trainer, "_immutable_data", {})
 
-    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
+                #     jax_device_state.pop("_mesh", None)
+                #     # Free massive XLA executables and module caches from device memory
+                #     jax_device_state.pop("_compiled_train_step", None)
+                #     jax_device_state.pop("_jit_train_step", None)
+                #     jax_device_state.pop("model", None)
+                #     jax_device_state.pop("learner", None)
+
+                # # Clear old trainer objects and JAX caches to release TPU memory.
+                # # We keep the extracted state dicts above to restore onto the new mesh.
+                # del trainer
+                # del clean_trainer
+                # jax.clear_caches()
+                # gc.collect()
+
+                if elastic_manager:
+                    elastic_manager.new_slice_event.set()
+                time.sleep(10)
+                continue
+            else:
+                logging.error("[ELASTIC] Caught non-retryable error: %s", e)
+                raise e
+    return output