Add: Working Replica Resize with fuji-test-v2.

Some changes to fuji-test-v2 hidden_dim=8 --> hidden_dim=32 fixed mesh shape error.
Change poll time pathways elastic manager 10 --> 30s

Author: camiloCienet

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, live_devices
+from axlearn.common.utils import Nested, Tensor, live_devices, live_slice_indices
+
 
 
 class ElasticSpmdInputDispatcher(BaseInputDispatcher):
@@ -79,6 +80,10 @@ def is_in_elastic_mode(self) -> bool:
         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:
@@ -95,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.")
 
@@ -120,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(
@@ -128,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(
@@ -178,6 +186,7 @@ def fid2pids(feed_id):
 
         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())
             )

axlearn/common/launch_trainer.py

@@ -149,7 +149,15 @@ def get_trainer_config(
         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(live_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

@@ -67,7 +67,7 @@ def pre_callback():
 
             train = utils.elastic_manager.replica_resize(
                 max_resizes=10,  # Handle up to 10 slice up or slice down transitions
-                poll_interval=10,  # Monitor thread checks inactive slice health every 10 seconds
+                poll_interval=30,  # Monitor thread checks inactive slice health every 30 seconds
                 pre_callback=pre_callback,
             )(train)
 

axlearn/common/utils.py

@@ -98,6 +98,9 @@ def live_devices():
     device_list = jax.devices()
 
     if pathwaysutils.is_pathways_backend_used() and elastic_manager is not None:
+        # print(f"Live devices: {device_list}, active slices: {elastic_manager.active_slice_indices}")
+        # return_cnt = [d for d in device_list if d.slice_index in elastic_manager.active_slice_indices]
+        # print(f"Live Devices by camilo return value : {return_cnt}")
         return [d for d in device_list if d.slice_index in elastic_manager.active_slice_indices]
     else:
         return device_list
@@ -1820,11 +1823,14 @@ def create_device_mesh(
     # Check if the devices are part of a multi-granule configuration.
     # <https://github.com/google/jax/blob/b81b79c1b0d2ec/jax/experimental/mesh_utils.py#L313>
     device_platform = devices[0].platform
+
+    print(f"Devices before granule calculation: {devices}")
     device_attr = "process_index" if device_platform != "tpu" else "slice_index"
     is_multi_granule_env = hasattr(devices[0], device_attr)
     if not all(el.platform == device_platform for el in devices):
         raise NotImplementedError(f"Not all devices had platform: {device_platform}.")
-
+    print("IS MULTI GRANULE by camilo: {is_multi_granule_env}")
+    print(f"Granule devices: {(getattr(el, device_attr) for el in devices.flatten())} + 1, by camilo")
     num_granules = (
         max(getattr(el, device_attr) for el in devices.flatten()) + 1 if is_multi_granule_env else 1
     )
@@ -1833,6 +1839,9 @@ def create_device_mesh(
         num_devices % num_granules == 0
     ), "Number of devices must be divisible by number of granules."
     num_devices_per_granule = num_devices // num_granules
+    print(f"Num Granules by camilo: {num_granules}")
+    print(f"Num Devices by camilo: {num_devices}")
+    print(f"Num Devices per Granule by camilo: {num_devices_per_granule}")
 
     # Fallback to a standard mesh if on GPU with incompatible multi-granule mesh.
     if (
@@ -1898,6 +1907,8 @@ def create_device_mesh(
     if num_granules == 1:
         return build_standard_mesh(mesh_shape.ici_mesh_shape, devices=devices)
 
+    print(f"Mesh shape CREATE DEVICE MESH by camilo: {devices}")
+
     return create_hybrid_device_mesh(
         mesh_shape,
         devices=devices,

axlearn/experiments/text/gpt/fuji.py

@@ -295,11 +295,12 @@ def get_trainer_kwargs(
 
     # dict() is more readable here.
     # pylint: disable=use-dict-literal
+    # Camilo: Changes made to fuji-test to make replica resize work.
     if model_size == "test":
         trainer_kwargs = dict(
             model_kwargs=dict(
                 num_layers=4,
-                hidden_dim=8,
+                hidden_dim=32,
                 ffn_dim=scaled_hidden_dim(scale=8 / 3, round_up_to_multiples_of=16),
                 num_heads=4,
                 num_kv_heads=2,
@@ -310,12 +311,12 @@ def get_trainer_kwargs(
             ),
             learner_kwargs=dict(peak_lr=6e-4, weight_decay=0.01),
             max_sequence_length=64,
-            train_batch_size=32,
-            eval_batch_size=32,
-            max_step=3000,
-            eval_every_n_steps=1500,
-            save_every_n_steps=500,
-            mesh_shape=mesh_shape_from_axes(data=-1),
+            train_batch_size=train_batch_size,
+            eval_batch_size=train_batch_size,
+            max_step=10000,
+            eval_every_n_steps=50000,
+            save_every_n_steps=1000,
+            mesh_shape=mesh_shape_from_axes(data=-1, fsdp=8),
         )
     elif model_size == "1B":
         trainer_kwargs = dict(