Use active elastic devices for Pathways MTC init

Author: SujeethJinesh

src/maxtext/utils/elastic_utils.py

@@ -16,6 +16,7 @@
 
 import functools
 from collections import Counter
+from types import SimpleNamespace
 
 import jax
 from maxtext.utils import gcs_utils
@@ -226,3 +227,34 @@ def maybe_elastic_scale_up(config, checkpoint_manager):
       checkpoint_manager.wait_until_finished()
     max_logging.log("Checkpoint save completed. Interrupting")
     raise manager.ScaleUpSignalError()
+
+
+def single_controller_mtc_init_kwargs(raw_keys):
+  """Returns topology kwargs for single-controller MTC initialization."""
+  kwargs = {
+      "data_parallelism": raw_keys["mtc_data_parallelism"],
+      "num_slices": raw_keys["num_slices"],
+  }
+  if not raw_keys.get("elastic_enabled", False):
+    return kwargs
+
+  config = SimpleNamespace(**raw_keys)
+  if not should_use_elastic(config):
+    return kwargs
+
+  active_devices = tuple(live_devices(config))
+  active_slice_indices = {getattr(device, "slice_index", 0) for device in active_devices if device is not None}
+  if not active_devices or not active_slice_indices:
+    raise ValueError("Elastic single-controller MTC initialization found no active devices.")
+
+  kwargs["devices"] = active_devices
+  kwargs["num_slices"] = len(active_slice_indices)
+  if not kwargs["data_parallelism"]:
+    kwargs["data_parallelism"] = kwargs["num_slices"]
+  max_logging.log(
+      "Using active elastic devices for single-controller MTC initialization: "
+      f"active_num_slices={kwargs['num_slices']}, "
+      f"active_device_count={len(active_devices)}, "
+      f"configured_num_slices={raw_keys['num_slices']}."
+  )
+  return kwargs

src/maxtext/utils/max_utils.py

@@ -248,14 +248,14 @@ def maybe_initialize_jax_distributed_system(raw_keys):
     max_logging.log("Skipping jax distributed system since its not needed for single controller.")
     if raw_keys["enable_multi_tier_checkpointing"]:
       max_logging.log("Initializing multi-tier checkpointing for single controller...")
+      mtc_init_kwargs = elastic_utils.single_controller_mtc_init_kwargs(raw_keys)
       initialize_multi_tier_checkpointing(
           local_checkpoint_directory=raw_keys["local_checkpoint_directory"],
           backup_interval_minutes=raw_keys["multi_tier_checkpointing_backup_interval_minutes"],
           run_name=raw_keys["run_name"],
           jax_initialization_timeout_seconds=raw_keys["jax_distributed_initialization_timeout"],
-          data_parallelism=raw_keys["mtc_data_parallelism"],
-          num_slices=raw_keys["num_slices"],
           use_colocated_python=True,
+          **mtc_init_kwargs,
       )
     return
   if jax.distributed.is_initialized():

tests/unit/elastic_utils_test.py

@@ -245,6 +245,45 @@ def test_live_slice_indices(self):
     indices = elastic_utils.live_slice_indices(config)
     self.assertEqual(indices, {0, 1})
 
+  def _base_mtc_keys(self, **overrides):
+    keys = {
+        "elastic_enabled": True,
+        "mtc_data_parallelism": 1,
+        "num_slices": 2,
+    }
+    keys.update(overrides)
+    return keys
+
+  def test_single_controller_mtc_init_kwargs_uses_active_elastic_devices(self):
+    self.fake_pathwaysutils.is_pathways_backend_used.return_value = True
+    device0 = FakeDevice(slice_index=0)
+    device1 = FakeDevice(slice_index=1)
+    self.fake_jax.devices.return_value = [device0, device1]
+    self.fake_manager.active_slice_indices = {0}
+
+    kwargs = elastic_utils.single_controller_mtc_init_kwargs(self._base_mtc_keys(mtc_data_parallelism=0))
+
+    self.assertEqual(kwargs["devices"], (device0,))
+    self.assertEqual(kwargs["num_slices"], 1)
+    self.assertEqual(kwargs["data_parallelism"], 1)
+
+  def test_single_controller_mtc_init_kwargs_raises_if_empty(self):
+    self.fake_pathwaysutils.is_pathways_backend_used.return_value = True
+    device0 = FakeDevice(slice_index=0)
+    self.fake_jax.devices.return_value = [device0]
+    self.fake_manager.active_slice_indices = {1}
+
+    with self.assertRaisesRegex(ValueError, "Elastic single-controller MTC initialization found no active devices."):
+      elastic_utils.single_controller_mtc_init_kwargs(self._base_mtc_keys())
+
+  def test_single_controller_mtc_init_kwargs_non_elastic(self):
+    kwargs = elastic_utils.single_controller_mtc_init_kwargs(
+        self._base_mtc_keys(elastic_enabled=False, mtc_data_parallelism=3, num_slices=4)
+    )
+
+    self.assertEqual(kwargs, {"data_parallelism": 3, "num_slices": 4})
+    self.assertIsNone(elastic_utils.elastic_manager)
+
   def test_get_devices_per_host(self):
     device0 = FakeDevice(slice_index=0, process_index=0, task_id=0)
     device1 = FakeDevice(slice_index=0, process_index=0, task_id=0)

tests/unit/max_utils_test.py

@@ -492,6 +492,44 @@ def test_single_controller_multi_tier_checkpointing_uses_colocated_python(self,
         use_colocated_python=True,
     )
 
+  @mock.patch("maxtext.utils.max_utils.elastic_utils.single_controller_mtc_init_kwargs")
+  @mock.patch("maxtext.utils.max_utils.initialize_multi_tier_checkpointing")
+  @mock.patch("jax.distributed.initialize")
+  def test_single_controller_multi_tier_checkpointing_uses_elastic_utils_kwargs(
+      self, mock_init, mock_mtc, mock_mtc_init_kwargs
+  ):
+    active_devices = (
+        mock.Mock(slice_index=0),
+        mock.Mock(slice_index=0),
+    )
+    mock_mtc_init_kwargs.return_value = {
+        "data_parallelism": 1,
+        "num_slices": 1,
+        "devices": active_devices,
+    }
+    raw_keys = self._base_keys(
+        enable_single_controller=True,
+        enable_multi_tier_checkpointing=True,
+        elastic_enabled=True,
+        mtc_data_parallelism=0,
+        num_slices=2,
+    )
+
+    max_utils.maybe_initialize_jax_distributed_system(raw_keys)
+
+    mock_init.assert_not_called()
+    mock_mtc_init_kwargs.assert_called_once_with(raw_keys)
+    mock_mtc.assert_called_once_with(
+        local_checkpoint_directory=self._base_keys()["local_checkpoint_directory"],
+        backup_interval_minutes=self._base_keys()["multi_tier_checkpointing_backup_interval_minutes"],
+        run_name=self._base_keys()["run_name"],
+        jax_initialization_timeout_seconds=self._base_keys()["jax_distributed_initialization_timeout"],
+        data_parallelism=1,
+        num_slices=1,
+        use_colocated_python=True,
+        devices=active_devices,
+    )
+
   @mock.patch("jax.distributed.initialize")
   def test_tpu_checkpointing_no_emergency_calls_jax_init(self, mock_init):
     raw_keys = self._base_keys(enable_checkpointing=True, compile_topology_num_slices=-1)