working replica-resize

Author: lkolluru05

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,7 +72,10 @@ 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:
@@ -165,11 +168,23 @@ def fid2pids(feed_id):
         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
             )
@@ -408,9 +423,14 @@ def _padded_select(path, x, y):
 
 def slice_count() -> int:
     """Returns the number of slices."""
-    slice_cnt_val = 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 jax.devices() if hasattr(d, "slice_index"))) or 1
+    return len(set(d.slice_index for d in live_devices() if hasattr(d, "slice_index"))) or 1
 
 
 def process_count_per_slice() -> int:
@@ -419,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
             )
         )

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(
@@ -148,7 +148,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]

axlearn/common/launch_trainer_main.py

@@ -17,17 +17,24 @@
 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()
+    # trainer_config = launch_trainer.get_trainer_config()
+    # trainer_config.set(recorder=config_for_function(lambda: measurement.global_recorder))
+    # measurement.start_monitoring()
 
     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(
@@ -66,6 +73,11 @@ def pre_callback():
 
         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)
 
 

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.

axlearn/common/trainer.py

@@ -18,6 +18,7 @@
 from jax import numpy as jnp
 from jax.experimental import multihost_utils
 from jax.experimental.pjit import pjit
+from pathwaysutils.elastic import manager
 
 from axlearn.common import file_system as fs
 from axlearn.common import measurement, utils
@@ -59,6 +60,7 @@
     count_model_params,
     flatten_items,
     host_to_global_specs,
+    live_devices,
     match_regex_rules,
     thread_stack_traces,
 )
@@ -407,7 +409,8 @@ def model_params_for_eval(self):
 
     def _step_log(self, msg, *args, **kwargs):
         logging.info(
-            "%s process % 3d step % 8d] " + msg,
+            "live_devices: %d ; %s process % 3d step % 8d ] " + msg,
+            len(live_devices()),
             self.path(),
             jax.process_index(),
             -1 if self.step is None else self.step,
@@ -646,7 +649,7 @@ def run(
                         )
                         self.vlog(3, "Done step %s", self.step)
                         num_steps += 1
-                        if num_steps % 100 == 0:
+                        if num_steps % 5 == 0:
                             now = time.perf_counter()
                             average_step_time = (now - start_time) / num_steps
                             self._step_log("Average step time: %s seconds", average_step_time)
@@ -1018,6 +1021,10 @@ def save_checkpoint(self, evaler_summaries: Optional[dict[str, Any]]) -> Optiona
             self.checkpointer.save(
                 step=self.step, state=ckpt_state, evaler_summaries=evaler_summaries
             )
+            print("Checkpoint saved !!!!")
+            if utils.elastic_manager is not None and utils.elastic_manager.new_slice_event.is_set():
+                logging.info("Interrupting as a new elastic slice event is seen")
+                raise manager.ScaleUpSignalError()
 
     def _restore_from_builder(self) -> Optional[TrainerStateBuilder.State]:
         """Restores trainer state by building it with init_state_builder."""
@@ -1069,7 +1076,8 @@ def _get_compiled_train_step_fn(
         cfg: SpmdTrainer.Config = self.config
         # Get device kinds and assert that they are homogenous.
         # TODO(markblee): Get devices from self._mesh.devices.
-        device_kinds = set(d.device_kind for d in jax.devices())
+        # device_kinds = set(d.device_kind for d in jax.devices())  #### may be here #####
+        device_kinds = set(d.device_kind for d in live_devices())
         if len(device_kinds) != 1:
             raise RuntimeError(f"Heterogenous device kinds ({device_kinds}) are not supported.")
         device_kind = device_kinds.pop()

axlearn/common/utils.py

@@ -43,6 +43,7 @@
 import attr
 import jax
 import numpy as np
+import pathwaysutils
 from absl import logging
 from jax import numpy as jnp
 from jax._src.ad_checkpoint import name_p
@@ -93,6 +94,19 @@
 elastic_manager: manager.Manager | None = None
 
 
+def live_devices():
+    device_list = jax.devices()
+
+    if pathwaysutils.is_pathways_backend_used() and elastic_manager is not None:
+        return [d for d in device_list if d.slice_index in elastic_manager.active_slice_indices]
+    else:
+        return device_list
+
+
+def live_slice_indices() -> set[int]:
+    return {d.slice_index for d in live_devices()}
+
+
 @dataclasses.dataclass
 class HybridMeshShape:
     """A mesh shape for hybrid (i.e., ICI and DCN) parallelism.
@@ -1799,7 +1813,8 @@ def create_device_mesh(
         NotImplementedError: If not all devices have the same platform.
     """
     if devices is None:
-        devices = jax.devices()
+        # devices = jax.devices()  ##### may be update here ###
+        devices = live_devices()
     devices = np.asarray(devices)
 
     # Check if the devices are part of a multi-granule configuration.
@@ -1910,7 +1925,8 @@ def infer_mesh_shape(mesh_shape: MeshShape, *, num_devices: Optional[int] = None
     # Handle the case with one -1.
     prod = math.prod(mesh_shape, start=-1)
     if num_devices is None:
-        num_devices = len(jax.devices())
+        # num_devices = len(jax.devices())  ###### may be update here ####
+        num_devices = len(live_devices())
     if num_devices % prod != 0:
         raise ValueError(
             f"Unable to infer -1 in mesh shape {mesh_shape} as num_devices {num_devices} "

axlearn/experiments/text/gpt/common.py

@@ -765,7 +765,7 @@ def config_fn() -> InstantiableConfig:
 
         train_dispatcher_cfg = elastic_input.ElasticSpmdInputDispatcher.default_config().set(
             global_logical_batch_size=train_batch_size,
-            num_max_slices=1,
+            num_max_slices=2,
         )
 
         cfg.input = elastic_input.ElasticInput.default_config().set(
@@ -809,7 +809,7 @@ def config_fn() -> InstantiableConfig:
             cfg.evalers[name] = evaler_cfg
         # Summaries and checkpoints.
         cfg.checkpointer.save_policy = config_for_function(every_n_steps_and_last_policy).set(
-            n=save_every_n_steps or min(eval_every_n_steps, 5_000),
+            n=save_every_n_steps or min(eval_every_n_steps, 50),
             max_step=max_step,
         )
         cfg.checkpointer.keep_every_n_steps = min(max_step, keep_every_n_steps)