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,19 @@ 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 +761,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

@@ -332,6 +332,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