Add continuous checkpoint option in MaxText.

PiperOrigin-RevId: 845008539

Author: lyglst

src/MaxText/checkpointing.py

@@ -30,6 +30,8 @@
 import orbax.checkpoint as ocp
 from orbax.checkpoint import v1 as ocp_v1
 from orbax.checkpoint._src.arrays import sharding as sharding_utils
+from orbax.checkpoint._src.checkpoint_managers import preservation_policy as preservation_policy_lib
+from orbax.checkpoint._src.checkpoint_managers import save_decision_policy as save_decision_policy_lib
 import orbax.checkpoint.experimental.emergency.checkpoint_manager as emergency_checkpoint_manager
 import orbax.checkpoint.experimental.emergency.replicator_checkpoint_manager as emergency_replicator_checkpoint_manager
 # pylint: disable=too-many-positional-arguments
@@ -194,6 +196,8 @@ def create_orbax_checkpoint_manager(
     orbax_logger: Any = None,  # pytype: disable=attribute-error
     use_ocdbt: bool = True,
     use_zarr3: bool = True,
+    enable_continuous_checkpointing: bool = False,
+    max_num_checkpoints_to_keep: int = 10,
 ):
   """Returns specified Orbax (async or not) CheckpointManager or None if checkpointing is disabled."""
   if not enable_checkpointing:
@@ -214,15 +218,28 @@ def create_orbax_checkpoint_manager(
   # local storage checkpoint needs parent directory created
   p = epath.Path(checkpoint_dir)
   p.mkdir(exist_ok=True, parents=True)
+  if enable_continuous_checkpointing:
+    save_decision_policy = save_decision_policy_lib.ContinuousCheckpointingPolicy()
+    preservation_policy = preservation_policy_lib.LatestN(
+        max_num_checkpoints_to_keep
+    )
+  else:
+    save_decision_policy = save_decision_policy_lib.FixedIntervalPolicy(
+        interval=save_interval_steps
+    )
+    preservation_policy = preservation_policy_lib.LatestN(
+        max_num_checkpoints_to_keep
+    )
   manager = CheckpointManager(
       p,
       item_names=item_names,
       item_handlers=item_handlers,
       options=CheckpointManagerOptions(
           create=True,
-          save_interval_steps=save_interval_steps,
           enable_async_checkpointing=use_async,
-      ),
+          save_decision_policy=save_decision_policy,
+          preservation_policy=preservation_policy,
+          ),
       logger=orbax_logger,
   )
 
@@ -259,8 +276,12 @@ def create_orbax_emergency_checkpoint_manager(
       global_mesh=global_mesh,
       abstract_state=abstract_state,
       options=emergency_checkpoint_manager.CheckpointManagerOptions(
-          local=LocalCheckpointOptions(save_interval_steps=local_save_interval_steps),
-          persistent=PersistentCheckpointOptions(save_interval_steps=persistent_save_interval_steps),
+          local=LocalCheckpointOptions(
+              save_interval_steps=local_save_interval_steps
+          ),
+          persistent=PersistentCheckpointOptions(
+              save_interval_steps=persistent_save_interval_steps
+          ),
       ),
       logger=orbax_logger,
   )

src/MaxText/configs/base.yml

@@ -52,6 +52,7 @@ save_checkpoint_on_completion: True
 async_checkpointing: True
 checkpoint_period: 10_000
 max_num_checkpoints_to_keep: None
+enable_continuous_checkpointing: False
 # enables one replica to read the ckpt then broadcast to the rest
 enable_single_replica_ckpt_restoring: False
 

src/MaxText/configs/types.py

@@ -279,6 +279,9 @@ class Checkpointing(BaseModel):
   save_checkpoint_on_completion: bool = Field(
       True, description="If True, saves a final checkpoint upon training completion."
   )
+  enable_continuous_checkpointing: bool = Field(
+      False, description="If True, enables continuous checkpointing."
+  )
 
 
 class OrbaxStorage(BaseModel):

src/MaxText/train_utils.py

@@ -73,6 +73,8 @@ def create_training_tools(config, model, mesh):
         logger,
         use_ocdbt,
         use_zarr3,
+        config.enable_continuous_checkpointing,
+        config.max_num_checkpoints_to_keep,
     )
 
   return init_rng, checkpoint_manager, learning_rate_schedule, tx