Merge pull request Borklet-Labs#11 from Borklet-Labs/cienet/replica_resize

Add: Elastic Training functionality for Pause & Resums and Replica Resize

Author: camiloCienet

axlearn/cloud/gcp/pathways_utils.py

@@ -556,6 +556,7 @@ def _build_pathways_head_sidecar_containers(self) -> list[Nested[Any]]:
         system = USER_FACING_NAME_TO_SYSTEM_CHARACTERISTICS[self._tpu_type]
         staging_location = f"{cfg.output_dir}/pathways-staging"
         pathways_tpu_version = get_pathways_tpu_version(system.gce_machine_type)
+        num_elastic_slices = 1
 
         # If multi-head, every pathways-head will only
         # be connected to one pathways instance (a pathways-worker replicated job).
@@ -565,6 +566,7 @@ 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}",
+            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,8 @@
 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, live_slice_indices
+
 
 
 class ElasticSpmdInputDispatcher(BaseInputDispatcher):
@@ -72,13 +73,20 @@ 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:
+            print(
+              f"Live_slice_count by lkolluru: {live_slice_indices()}, slices_cnt: {len(live_slice_indices())}",
+              live_slice_indices(),len(live_slice_indices())
+            )
             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.
@@ -92,6 +100,7 @@ def __init__(self, cfg: Config, *, parent: Optional[Module]):
         cfg: ElasticSpmdInputDispatcher.Config = self.config
 
         mesh = thread_resources.env.physical_mesh
+        print("physical mesh by lkolluru: ", mesh)
         if mesh.empty:
             raise ValueError("Expected to be initialized within the context of a mesh.")
 
@@ -117,6 +126,7 @@ def __init__(self, cfg: Config, *, parent: Optional[Module]):
         )
         if self.is_in_elastic_mode:
             num_partitions = num_partitions // slice_count() * cfg.num_max_slices
+            print(f"num_partitions by lkolluru: {num_partitions} and live_devices: {slice_count()} ")
 
         if cfg.global_logical_batch_size % num_partitions != 0:
             raise ValueError(
@@ -125,6 +135,7 @@ def __init__(self, cfg: Config, *, parent: Optional[Module]):
             )
 
         self._device_physical_batch_size = cfg.global_logical_batch_size // num_partitions
+        print("device_physical_batch_size on init by lkolluru: ", self._device_physical_batch_size)
 
         # Infer the physical feeds and feed index along dim=0.
         _, _, pid2fid = get_process_index_and_count_and_mapping(
@@ -156,14 +167,33 @@ 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
+        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
 
-        assert cfg.global_logical_batch_size % self.feed_count == 0
-        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")
+            # I think this should be len(live_slice_indices())....
             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 +432,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 +448,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 +540,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 +556,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 +573,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 +585,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
+    print("pid2fid by lkolluru: ", pid2fid)
     return feed_index, feed_count, pid2fid

axlearn/common/flash_attention/neuron_attention.py

@@ -17,12 +17,13 @@
 from axlearn.common.attention_bias import BaseAttentionBias, CausalAttentionBias, split
 from axlearn.common.flash_attention.common import BaseFlashAttention, repeat_kv_heads
 from axlearn.common.kv_cache.base_kv_cache import BaseKVCache
-from axlearn.common.utils import Nested
+from axlearn.common.utils import Nested, live_devices
 
 # pytype: enable=import-error
 
 Tensor = jax.Array
-lnc = 2 if jax.devices()[0].device_kind == "NC_v3d" else 1
+# lnc = 2 if jax.devices()[0].device_kind == "NC_v3d" else 1
+lnc = 2 if live_devices()[0].device_kind == "NC_v3d" else 1
 
 
 # TODO(apoorvtintin): Add segment IDs as an argument when the kernel supports it.

axlearn/common/launch.py

@@ -45,7 +45,7 @@
 from absl import flags, logging
 
 from axlearn.common.status_server import StatusHTTPServer
-from axlearn.common.utils import get_data_dir
+from axlearn.common.utils import get_data_dir, live_devices
 from axlearn.common.utils_spmd import setup as setup_spmd
 
 # pylint: enable=wrong-import-position
@@ -142,7 +142,8 @@ def setup():
         status_server = StatusHTTPServer(FLAGS.status_port)
         status_server.start()
 
-    devices = jax.devices()
+    # devices = jax.devices()  ##### may be update here ####
+    devices = live_devices()
     logging.info("Devices: %s", devices)
     local_devices = jax.local_devices()
     logging.info("Local Devices: %s", local_devices)

axlearn/common/launch_trainer.py

@@ -13,7 +13,7 @@
 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.utils import MeshShape, get_data_dir, infer_mesh_shape, live_devices
 
 # Trainer-specific flags.
 flags.DEFINE_string(
@@ -145,10 +145,21 @@ def get_trainer_config(
         )
     trainer_config: SpmdTrainer.Config = trainer_config_fn()
     trainer_config.dir = trainer_config.dir or flag_values.trainer_dir
+
+    print(f"Trainer Config Dir: {trainer_config.dir} by Camilo")
     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)
+
+    print("Live devices ", live_devices())
+    if len(live_devices()) == 32:
+      trainer_config.mesh_shape = trainer_config.mesh_shape or (len(live_devices()), 1)
+    if len(live_devices()) == 16:
+      trainer_config.mesh_shape = (1, 1, 1, len(live_devices()), 1, 1)
+
+    print("trainer_config.mesh_shape by camilo: ", trainer_config.mesh_shape)
+
     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]

axlearn/common/launch_trainer_main.py

@@ -2,19 +2,85 @@
 
 """Main function for launching the trainer."""
 
+import pathwaysutils
+import functools
 from absl import app, flags
+from pathwaysutils.elastic import elastic, manager
 
-from axlearn.common import launch, launch_trainer, measurement
+from axlearn.common import launch, launch_trainer, measurement, utils
 from axlearn.common.config import config_for_function
 
+enable_elastic_training = True
+enable_pause_resume = False
+enable_replica_resize = True
+
 
 def main(_):
     measurement.initialize(flags.FLAGS)
     launch.setup()
     trainer_config = launch_trainer.get_trainer_config()
-    trainer_config.set(recorder=config_for_function(lambda: measurement.global_recorder))
-    measurement.start_monitoring()
-    launch_trainer.run_trainer(trainer_config)
+    # trainer_config.set(recorder=config_for_function(lambda: measurement.global_recorder))
+    # measurement.start_monitoring()
+    clean_up_checkpoints = functools.partial(utils.clean_up_checkpoints, checkpoint_dir=trainer_config.dir)
+    if pathwaysutils.is_pathways_backend_used() and enable_elastic_training:
+
+        def train():
+            # measurement.initialize(flags.FLAGS)
+            # launch.setup()
+            trainer_config = launch_trainer.get_trainer_config()
+            trainer_config.set(recorder=config_for_function(lambda: measurement.global_recorder))
+            measurement.start_monitoring()
+            launch_trainer.run_trainer(trainer_config)
+
+        utils.elastic_manager = manager.Manager()
+
+        # utils.elastic_manager.live_devices = live_devices()
+
+        if enable_pause_resume:
+            print("Pathways backend with pause resume being used")
+            train = utils.elastic_manager.pause_resume(
+                max_retries=10,  # Handle up to 10 disruptions before restarting
+                poll_interval=10,  # While paused, checks every 10 seconds for health
+                timeout=300,  # Waits for slices to rejoin for 5 minutes
+                # on_elastic_event_callback=clean_up_checkpoints,
+            )(train)
+
+        if enable_replica_resize:
+
+            def pre_callback():
+                # Wait up to 1 minute before starting if there are any inactive slices
+                if utils.elastic_manager.inactive_slice_indices:
+                    try:
+                        utils.elastic_manager.active_slice_indices = elastic.wait_for_slices(
+                            slice_count=utils.elastic_manager.total_slice_count,
+                            slice_to_devices=utils.elastic_manager.slice_to_devices,
+                            poll_interval=10,
+                            timeout=60,
+                        )
+                    except TimeoutError:
+                        # If there are still inactive slices, we must update the active
+                        # slices with one final check and then proceed.
+                        utils.elastic_manager.active_slice_indices = (
+                            elastic.get_active_slice_indices(
+                                slice_to_devices=utils.elastic_manager.slice_to_devices,
+                            )
+                        )
+
+            train = utils.elastic_manager.replica_resize(
+                max_resizes=10,  # Handle up to 10 slice up or slice down transitions
+                poll_interval=30,  # Monitor thread checks inactive slice health every 30 seconds
+                pre_callback=pre_callback,
+                on_elastic_event_callback=clean_up_checkpoints,
+            )(train)
+
+        train()
+    else:
+        measurement.initialize(flags.FLAGS)
+        launch.setup()
+        trainer_config = launch_trainer.get_trainer_config()
+        trainer_config.set(recorder=config_for_function(lambda: measurement.global_recorder))
+        measurement.start_monitoring()
+        launch_trainer.run_trainer(trainer_config)
 
 
 if __name__ == "__main__":

axlearn/common/state_builder.py

@@ -47,6 +47,7 @@
     flatten_items,
     get_data_dir,
     infer_mesh_shape,
+    live_devices,
     set_data_dir,
 )
 
@@ -479,7 +480,10 @@ def target_to_source(self, target: Builder.State) -> tuple[Builder.State, Any]:
                 cfg.mesh_axis_names or trainer_cfg.mesh_axis_names or ("data", "model")
             )
             trainer_cfg.mesh_shape = infer_mesh_shape(
-                cfg.mesh_shape or trainer_cfg.mesh_shape or (len(jax.devices()), 1)
+                # cfg.mesh_shape or trainer_cfg.mesh_shape or (len(jax.devices()), 1)
+                cfg.mesh_shape
+                or trainer_cfg.mesh_shape
+                or (len(live_devices()), 1)
             )
             # Reset datasets and evalers for the pretrained model config.
             # This input is not used. Set global_batch_size to 0 by default.