Merge pull request AI-Hypercomputer#4407 from AI-Hypercomputer:sujinesh/avoid-mxla-hang-eager-scale-up-upstream

PiperOrigin-RevId: 949675411

Author: Google-ML-Automation

src/maxtext/common/checkpointing.py

@@ -1094,8 +1094,8 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step
     checkpoint_saved = save_checkpoint(checkpoint_manager, actual_step, state, config, data_iterator, force_ckpt_save)
     if checkpoint_saved:
       print_save_message(actual_step, config.async_checkpointing)
-      if config.elastic_enabled:
-        elastic_utils.maybe_elastic_scale_up(config, checkpoint_manager)
+    if config.elastic_enabled:
+      elastic_utils.maybe_elastic_scale_up(config, checkpoint_manager)
   except elastic_utils.manager.ScaleUpSignalError as e:
     if config.elastic_enabled:
       max_logging.log(f"Elastic event detected, letting exception bubble up: {e}")

src/maxtext/trainers/pre_train/train.py

@@ -680,6 +680,8 @@ def training_loop_iteration(
   dump_hlo_upload_all = immutable_data["dump_hlo_upload_all"]
 
   prof.maybe_activate_profiler(step, state)
+  if config.elastic_enabled:
+    elastic_utils.maybe_elastic_scale_up(config, checkpoint_manager)
 
   with jax.profiler.StepTraceAnnotation("train", step_num=step):
     example_batch = data_loader.load_next_batch(rampup_manager=rampup_manager)

tests/unit/train_state_nnx_checkpoint_test.py

@@ -363,6 +363,7 @@ def _config(self, **overrides):
         "enable_multi_tier_checkpointing": False,
         "local_checkpoint_period": 0,
         "enable_autocheckpoint": False,
+        "elastic_enabled": False,
     }
     values.update(overrides)
     return SimpleNamespace(**values)
@@ -598,6 +599,26 @@ def test_maybe_save_checkpoint_allows_mtc_period_with_continuous_policy(
     to_checkpoint_dict_mock.assert_called_once_with(state)
     save_checkpoint_mock.assert_called_once()
 
+  def test_maybe_save_checkpoint_checks_scale_up_after_unsaved_dispatch(self):
+    """Elastic scale-up is checked after save dispatch even when no checkpoint was saved."""
+    state = mock.Mock()
+    config = self._config(checkpoint_period=1, elastic_enabled=True)
+    mgr = mock.MagicMock()
+    mgr.latest_step.return_value = None
+    mgr.reached_preemption.return_value = False
+    save_checkpoint_mock = mock.MagicMock(return_value=False)
+
+    with (
+        mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock),
+        mock.patch.object(train_state_nnx, "to_checkpoint_dict", return_value={}) as to_checkpoint_dict_mock,
+        mock.patch.object(checkpointing.elastic_utils, "maybe_elastic_scale_up") as mock_maybe_scale_up,
+    ):
+      checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5)
+
+    to_checkpoint_dict_mock.assert_called_once_with(state)
+    save_checkpoint_mock.assert_called_once()
+    mock_maybe_scale_up.assert_called_once_with(config, mgr)
+
 
 class TestLinenCheckpointFormatConverters(unittest.TestCase):
   """to_linen_checkpoint_dict / from_linen_checkpoint_dict (NNX <-> Linen on-disk layout)."""