Add autocheckpoint feature to MaxText.

PiperOrigin-RevId: 892064416

Author: abhinavclemson

src/maxtext/common/checkpointing.py

@@ -221,6 +221,7 @@ def create_orbax_checkpoint_manager(
     enable_single_controller: bool = False,
     colocated_python_checkpointing: bool = False,
     enable_single_replica_ckpt_restoring: bool = False,
+    enable_autocheckpoint: bool = False,
 ):
   """Returns specified Orbax (async or not) CheckpointManager or None if checkpointing is disabled."""
   if not enable_checkpointing:
@@ -248,11 +249,21 @@ def create_orbax_checkpoint_manager(
   # local storage checkpoint needs parent directory created
   p = gcs_utils.mkdir_and_check_permissions(checkpoint_dir)
   if enable_continuous_checkpointing:
+    max_logging.log("Enabling policy for continuous checkpointing.")
     save_decision_policy = save_decision_policy_lib.ContinuousCheckpointingPolicy()
-    preservation_policy = preservation_policy_lib.LatestN(max_num_checkpoints_to_keep)
+  elif enable_autocheckpoint:
+    max_logging.log("Enabling policy for autocheckpoint.")
+    save_decision_policy = save_decision_policy_lib.AnySavePolicy(
+        [
+            save_decision_policy_lib.PreemptionCheckpointingPolicy(),
+            save_decision_policy_lib.FixedIntervalPolicy(save_interval_steps),
+        ]
+    )
   else:
+    max_logging.log("Enabling policy for fixed interval checkpointing.")
     save_decision_policy = save_decision_policy_lib.FixedIntervalPolicy(interval=save_interval_steps)
-    preservation_policy = preservation_policy_lib.LatestN(max_num_checkpoints_to_keep)
+  preservation_policy = preservation_policy_lib.LatestN(max_num_checkpoints_to_keep)
+
   async_options = None
   if enable_continuous_checkpointing:
     async_options = ocp.AsyncOptions(
@@ -752,6 +763,7 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=
         or (step % config.checkpoint_period == 0 and not config.enable_continuous_checkpointing)
         or (step % config.checkpoint_period == 0)
         or (config.enable_emergency_checkpoint and step % config.local_checkpoint_period == 0)
+        or (config.enable_autocheckpoint and checkpoint_manager.reached_preemption(step))
     ):
       blocking_until_ready_start = time.time()
       max_logging.log(f"Waiting for step {step} to finish before checkpoint...")

src/maxtext/configs/base.yml

@@ -83,6 +83,9 @@ source_checkpoint_layout: "orbax"
 
 # Only applicable to Single Controller/Pathways on Cloud. Experimental feature, under testing 
 colocated_python_checkpointing: False
+
+# enables autocheckpoint, which saves a checkpoint at the preemption step.
+enable_autocheckpoint: False
 ############################### end checkpointing ##################################
 
 

src/maxtext/configs/types.py

@@ -333,6 +333,9 @@ class Checkpointing(BaseModel):
       False,
       description="If True, enables checkpointing from remote TPU VMs instead of head node on pathways.",
   )
+  enable_autocheckpoint: bool = Field(
+      False, description="If True, enables autocheckpoint or preemption induced checkpointing."
+  )
 
 
 class OrbaxStorage(BaseModel):

src/maxtext/utils/train_utils.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 # pylint: disable=bare-except, consider-using-generator
-""" Utils that are only interesting for training in MaxText. """
+"""Utils that are only interesting for training in MaxText."""
 
 import os
 import jax
@@ -82,6 +82,7 @@ def create_training_tools(config, model, mesh):
         config.enable_single_controller,
         config.colocated_python_checkpointing,
         config.enable_single_replica_ckpt_restoring,
+        config.enable_autocheckpoint,
     )
 
   return init_rng, checkpoint_manager, learning_rate_schedule, tx

tests/autocheckpoint_test.py

@@ -0,0 +1,90 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for autocheckpoint feature."""
+
+from unittest import mock
+from absl.testing import absltest
+from maxtext.common import checkpointing
+from maxtext.utils import exceptions
+import orbax.checkpoint as ocp
+from orbax.checkpoint._src.checkpoint_managers import save_decision_policy as save_decision_policy_lib
+
+
+class AutocheckpointTest(absltest.TestCase):
+
+  @mock.patch("maxtext.src.maxtext.common.checkpointing.PyTreeCheckpointHandler")
+  @mock.patch("maxtext.src.maxtext.common.checkpointing.CheckpointManager")
+  @mock.patch("maxtext.src.maxtext.utils.gcs_utils.mkdir_and_check_permissions")
+  def test_create_checkpoint_manager_with_autocheckpoint(self, mock_mkdir, mock_manager_cls, mock_handler_cls):
+    mock_mkdir.return_value = "/tmp/checkpoint"
+
+    manager = checkpointing.create_orbax_checkpoint_manager(
+        checkpoint_dir="/tmp/checkpoint",
+        enable_checkpointing=True,
+        use_async=False,
+        save_interval_steps=100,
+        enable_autocheckpoint=True,
+    )
+
+    self.assertIsNotNone(manager)
+    mock_manager_cls.assert_called_once()
+    _, kwargs = mock_manager_cls.call_args
+    options = kwargs.get("options")
+    self.assertIsInstance(options.save_decision_policy, save_decision_policy_lib.AnySavePolicy)
+    # AnySavePolicy internally has policies
+    policies = options.save_decision_policy.policies
+    self.assertTrue(any(isinstance(p, save_decision_policy_lib.PreemptionCheckpointingPolicy) for p in policies))
+    self.assertTrue(any(isinstance(p, save_decision_policy_lib.FixedIntervalPolicy) for p in policies))
+
+  @mock.patch("jax.block_until_ready")
+  @mock.patch("maxtext.src.maxtext.common.checkpointing.ocp.args.PyTreeSave")
+  def test_save_checkpoint_triggers_on_preemption(self, mock_save_args, mock_block):
+    mock_manager = mock.MagicMock(spec=ocp.CheckpointManager)
+    mock_manager.reached_preemption.return_value = True
+
+    config = mock.MagicMock()
+    config.enable_checkpointing = True
+    config.checkpoint_period = 1000
+    config.enable_autocheckpoint = True
+    config.checkpoint_storage_target_data_file_size_bytes = 1024
+
+    state = mock.MagicMock()
+
+    # Step 5 is not a multiple of checkpoint_period (1000), but reached_preemption is True
+    checkpointing.save_checkpoint(mock_manager, 5, state, config=config)
+
+    mock_manager.save.assert_called_once()
+    mock_manager.reached_preemption.assert_called_with(5)
+
+  def test_maybe_save_checkpoint_handles_preemption(self):
+    mock_manager = mock.MagicMock(spec=ocp.CheckpointManager)
+    mock_manager.reached_preemption.return_value = True
+
+    config = mock.MagicMock()
+    config.checkpoint_period = 1000
+    config.enable_autocheckpoint = True
+
+    state = mock.MagicMock()
+    state.step = 6
+
+    with self.assertRaisesRegex(exceptions.StopTraining, "Job is preempted."):
+      # step=None means it will use state.step - 1 = 5
+      checkpointing.maybe_save_checkpoint(mock_manager, state, config, data_iterator=None, step=None)
+
+    mock_manager.wait_until_finished.assert_called_once()
+
+
+if __name__ == "__main__":
+  absltest.main()