diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index edfa492488..cf51083597 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -15,46 +15,45 @@ """Create an Orbax CheckpointManager with specified (Async or not) Checkpointer.""" +import datetime import time from typing import Any -from absl import flags -import datetime from etils import epath from flax import nnx from flax.training import train_state +from grain.experimental import ElasticIterator import jax -from maxtext.utils.globals import DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE +from maxtext.checkpoint_conversion.utils.load_dynamic import load_safetensors_dynamic_state +from maxtext.common import emergency_checkpointing +from maxtext.common import grain_utility +from maxtext.common import train_state_nnx from maxtext.input_pipeline.multihost_dataloading import MultiHostDataLoadIterator from maxtext.input_pipeline.multihost_dataloading import RemoteIteratorWrapper from maxtext.input_pipeline.synthetic_data_processing import PlaceHolderDataIterator -from maxtext.common import grain_utility -from maxtext.common import train_state_nnx +from maxtext.utils import elastic_utils from maxtext.utils import exceptions -from maxtext.utils import max_logging from maxtext.utils import gcs_utils -from maxtext.utils import elastic_utils -from maxtext.checkpoint_conversion.utils.load_dynamic import load_safetensors_dynamic_state - +from maxtext.utils import max_logging +from maxtext.utils.globals import DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE 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 -from grain.experimental import ElasticIterator -CheckpointManager = ocp.CheckpointManager CheckpointManagerOptions = ocp.CheckpointManagerOptions Composite = ocp.args.Composite PyTreeCheckpointHandler = ocp.PyTreeCheckpointHandler -EmergencyCheckpointManager = emergency_checkpoint_manager.CheckpointManager -LocalCheckpointOptions = emergency_checkpoint_manager.LocalCheckpointOptions -PersistentCheckpointOptions = emergency_checkpoint_manager.PersistentCheckpointOptions -EmergencyReplicatorCheckpointManager = emergency_replicator_checkpoint_manager.ReplicatorCheckpointManager +# Backward compatibility aliases for v0 emergency managers. +EmergencyCheckpointManager = emergency_checkpointing.CheckpointManager +EmergencyReplicatorCheckpointManager = emergency_checkpointing.ReplicatorCheckpointManager +create_orbax_emergency_checkpoint_manager = emergency_checkpointing.create_emergency_checkpoint_manager +create_orbax_emergency_replicator_checkpoint_manager = emergency_checkpointing.create_replicator_checkpoint_manager + +# Union of CheckpointManager / the emergency factories return; used in type hints. +CheckpointManager = ocp.CheckpointManager | EmergencyCheckpointManager | EmergencyReplicatorCheckpointManager def _weight_mismatches(want, have, path=()): @@ -336,7 +335,7 @@ def create_orbax_checkpoint_manager( async_options = ocp.AsyncOptions( timeout_secs=int(datetime.timedelta(minutes=60).total_seconds()), ) - manager = CheckpointManager( + manager = ocp.CheckpointManager( p, item_names=item_names, item_handlers=item_handlers, @@ -356,115 +355,6 @@ def create_orbax_checkpoint_manager( return manager -def create_orbax_emergency_checkpoint_manager( - local_checkpoint_dir: str, - persistent_checkpoint_dir: str, - global_mesh: jax.sharding.Mesh, - abstract_state: Any, - local_save_interval_steps: int, - persistent_save_interval_steps: int, - orbax_logger: Any = None, # pytype: disable=attribute-error -): - """Returns an emergency checkpoint manager.""" - flags.FLAGS.experimental_orbax_use_distributed_process_id = True - max_logging.log("Creating emergency checkpoint manager...") - - # Only create local directories if running on GPUs as the previous directory structure might be assumed by TPUs. - if global_mesh.devices.flatten()[0].platform == "gpu": - # pylint: disable=protected-access - local_checkpoint_dir = f"{local_checkpoint_dir}/{jax._src.distributed.global_state.process_id}" - local_p = epath.Path(local_checkpoint_dir) - local_p.mkdir(exist_ok=True, parents=True) - - persistent_p = gcs_utils.mkdir_and_check_permissions(persistent_checkpoint_dir) - - # pure_nnx saves via to_checkpoint_dict (Linen params/opt_state/step plus an nnx_aux - # subtree), but the emergency manager restores against the abstract it is built with. - # Convert it the same way so it matches what is on disk; restore reshapes back to NNX. - if isinstance(abstract_state, nnx.State): - abstract_state = train_state_nnx.to_checkpoint_dict(abstract_state) - - manager = EmergencyCheckpointManager( - local_checkpoint_dir, - persistent_p, - 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), - ), - logger=orbax_logger, - ) - - max_logging.log("Emergency checkpoint manager created!") - return manager - - -def create_orbax_emergency_replicator_checkpoint_manager( - local_checkpoint_dir: str, - save_interval_steps: int, - global_mesh: jax.sharding.Mesh, - colocated_python_checkpointing: bool = False, -): - """Returns an emergency replicator checkpoint manager.""" - flags.FLAGS.experimental_orbax_use_distributed_process_id = True - max_logging.log("Creating emergency replicator checkpoint manager...") - - manager = EmergencyReplicatorCheckpointManager( - epath.Path(local_checkpoint_dir), - options=emergency_replicator_checkpoint_manager.ReplicatorCheckpointManagerOptions( - save_interval_steps=save_interval_steps, - use_colocated_python=colocated_python_checkpointing, - ), - global_mesh=global_mesh, - ) - - max_logging.log("Emergency replicator checkpoint manager created!") - return manager - - -def replicator_error_handler(config: Any): - """Replicator error handler to handle errors in replicator service.""" - if config.enable_multi_tier_checkpointing: - local_dir = config.local_checkpoint_directory - replicator_errors_file = f"{local_dir}/replicator.errors" - replicator_failed_file = f"{local_dir}/replicator.failed" - process_replicator_error_file(replicator_errors_file) - - # if the replicator.failed file exists, then we have a fatal error - is_fatal = process_replicator_error_file(replicator_failed_file) - if is_fatal: - raise ValueError("Replicator fatal error found in replicator.failed file.") - - -def process_replicator_error_file(error_file: str) -> bool: - """Handles replicator errors by reading, logging, cleaning the error file.""" - error_file_path_exists = epath.Path(error_file).exists() - if error_file_path_exists: - max_logging.log(f"replicator_error_handler: file found: {error_file}.") - read_replicator_error_file(error_file) - cleanup_replicator_error_file(error_file) - - return error_file_path_exists - - -def read_replicator_error_file(error_file: str): - """Read replicator errors file.""" - try: - error_data = epath.Path(error_file).read_text() - max_logging.log(f"Contents of replicator error file:\n{error_data}") - except (OSError, ValueError) as e: - max_logging.log("replicator_error_handler: Failed to read contents of failed" f" file: {e}") - - -def cleanup_replicator_error_file(error_file: str): - """Clean up replicator errors file.""" - try: - epath.Path(error_file).unlink() - except (OSError, ValueError) as e: - max_logging.log("replicator_error_handler: Failed to remove replicator errors file:" f" {e}") - - def print_save_message(step, async_checkpointing): if async_checkpointing: max_logging.log(f"Started an asynchronous checkpoint save for step {step}") @@ -944,7 +834,7 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator= case (checkpoint_manager, _, _) if isinstance( checkpoint_manager, (EmergencyCheckpointManager, EmergencyReplicatorCheckpointManager) ): - replicator_error_handler(config) + emergency_checkpointing.replicator_error_handler(config) return checkpoint_manager.save(step, args=Composite(state=checkpoint_args), force=force) case _: return checkpoint_manager.save( diff --git a/src/maxtext/common/emergency_checkpointing.py b/src/maxtext/common/emergency_checkpointing.py new file mode 100644 index 0000000000..fba2f29546 --- /dev/null +++ b/src/maxtext/common/emergency_checkpointing.py @@ -0,0 +1,178 @@ +# Copyright 2023–2026 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. + +"""Orbax v0 emergency / multi-tier checkpoint managers (no v1 equivalent). + +MaxText's standard checkpointing runs on Orbax v1 (see ``common/checkpointing.py``). +The emergency and multi-tier replicator managers have no v1 counterpart, so they +stay on Orbax v0 — isolated here so the standard path stays v0-free. +""" + +from typing import Any + +from absl import flags +from etils import epath +from flax import nnx +import jax +from maxtext.common import train_state_nnx +from maxtext.utils import gcs_utils +from maxtext.utils import globals as maxtext_globals +from maxtext.utils import max_logging +import orbax.checkpoint as ocp +import orbax.checkpoint.experimental.emergency.checkpoint_manager as emergency_checkpoint_manager +import orbax.checkpoint.experimental.emergency.replicator_checkpoint_manager as emergency_replicator_checkpoint_manager + + +DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE = maxtext_globals.DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE +CheckpointManager = emergency_checkpoint_manager.CheckpointManager +LocalCheckpointOptions = emergency_checkpoint_manager.LocalCheckpointOptions +PersistentCheckpointOptions = emergency_checkpoint_manager.PersistentCheckpointOptions +ReplicatorCheckpointManager = emergency_replicator_checkpoint_manager.ReplicatorCheckpointManager + +_MANAGERS = (CheckpointManager, ReplicatorCheckpointManager) + + +def create_emergency_checkpoint_manager( + local_checkpoint_dir: str, + persistent_checkpoint_dir: str, + global_mesh: jax.sharding.Mesh, + abstract_state: Any, + local_save_interval_steps: int, + persistent_save_interval_steps: int, + orbax_logger: Any = None, # pytype: disable=attribute-error +): + """Returns an emergency checkpoint manager.""" + flags.FLAGS.experimental_orbax_use_distributed_process_id = True + max_logging.log("Creating emergency checkpoint manager...") + + # Only create local directories if running on GPUs as the previous directory structure might be assumed by TPUs. + if global_mesh.devices.flatten()[0].platform == "gpu": + # pylint: disable=protected-access + local_checkpoint_dir = f"{local_checkpoint_dir}/{jax._src.distributed.global_state.process_id}" + local_p = epath.Path(local_checkpoint_dir) + local_p.mkdir(exist_ok=True, parents=True) + + persistent_p = gcs_utils.mkdir_and_check_permissions(persistent_checkpoint_dir) + + # pure_nnx saves via to_checkpoint_dict (Linen params/opt_state/step plus an nnx_aux + # subtree), but the emergency manager restores against the abstract it is built with. + # Convert it the same way so it matches what is on disk; restore reshapes back to NNX. + if isinstance(abstract_state, nnx.State): + abstract_state = train_state_nnx.to_checkpoint_dict(abstract_state) + + manager = CheckpointManager( + local_checkpoint_dir, + persistent_p, + 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), + ), + logger=orbax_logger, + ) + + max_logging.log("Emergency checkpoint manager created!") + return manager + + +def create_replicator_checkpoint_manager( + local_checkpoint_dir: str, + save_interval_steps: int, + global_mesh: jax.sharding.Mesh, + colocated_python_checkpointing: bool = False, +): + """Returns an emergency replicator checkpoint manager.""" + flags.FLAGS.experimental_orbax_use_distributed_process_id = True + max_logging.log("Creating emergency replicator checkpoint manager...") + + manager = ReplicatorCheckpointManager( + epath.Path(local_checkpoint_dir), + options=emergency_replicator_checkpoint_manager.ReplicatorCheckpointManagerOptions( + save_interval_steps=save_interval_steps, + use_colocated_python=colocated_python_checkpointing, + ), + global_mesh=global_mesh, + ) + + max_logging.log("Emergency replicator checkpoint manager created!") + return manager + + +def save(checkpoint_manager, step, state, config, force): + """v0 emergency save: writes the state under the ``state`` checkpointable.""" + chunk_byte_size = ( + config.checkpoint_storage_target_data_file_size_bytes if config else DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE + ) + checkpoint_args = ocp.args.PyTreeSave( + item=state, + save_args=jax.tree.map(lambda _: ocp.SaveArgs(chunk_byte_size=chunk_byte_size), state), + ocdbt_target_data_file_size=chunk_byte_size, + ) + replicator_error_handler(config) + return checkpoint_manager.save(step, args=ocp.args.Composite(state=checkpoint_args), force=force) + + +def restore(checkpoint_manager, step, abstract_unboxed_pre_state): + """v0 emergency restore: returns the restored state pytree.""" + restore_target = abstract_unboxed_pre_state + if isinstance(abstract_unboxed_pre_state, nnx.State): + restore_target = abstract_unboxed_pre_state.to_pure_dict() + restore_args = jax.tree_util.tree_map( + lambda data: ocp.type_handlers.ArrayRestoreArgs(sharding=data.sharding), restore_target + ) + checkpoint_args = ocp.args.PyTreeRestore(item=restore_target, restore_args=restore_args, partial_restore=True) + return checkpoint_manager.restore(step, args=ocp.args.Composite(state=checkpoint_args)).state + + +def replicator_error_handler(config: Any): + """Replicator error handler to handle errors in replicator service.""" + if config.enable_multi_tier_checkpointing: + local_dir = config.local_checkpoint_directory + replicator_errors_file = f"{local_dir}/replicator.errors" + replicator_failed_file = f"{local_dir}/replicator.failed" + process_replicator_error_file(replicator_errors_file) + + # if the replicator.failed file exists, then we have a fatal error + is_fatal = process_replicator_error_file(replicator_failed_file) + if is_fatal: + raise ValueError("Replicator fatal error found in replicator.failed file.") + + +def process_replicator_error_file(error_file: str) -> bool: + """Handles replicator errors by reading, logging, cleaning the error file.""" + error_file_path_exists = epath.Path(error_file).exists() + if error_file_path_exists: + max_logging.log(f"replicator_error_handler: file found: {error_file}.") + read_replicator_error_file(error_file) + cleanup_replicator_error_file(error_file) + + return error_file_path_exists + + +def read_replicator_error_file(error_file: str): + """Read replicator errors file.""" + try: + error_data = epath.Path(error_file).read_text() + max_logging.log(f"Contents of replicator error file:\n{error_data}") + except (OSError, ValueError) as e: + max_logging.log("replicator_error_handler: Failed to read contents of failed" f" file: {e}") + + +def cleanup_replicator_error_file(error_file: str): + """Clean up replicator errors file.""" + try: + epath.Path(error_file).unlink() + except (OSError, ValueError) as e: + max_logging.log("replicator_error_handler: Failed to remove replicator errors file:" f" {e}") diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py index 515e906f5f..bfc21de4fe 100644 --- a/src/maxtext/utils/train_utils.py +++ b/src/maxtext/utils/train_utils.py @@ -25,6 +25,7 @@ from flax.linen import partitioning as nn_partitioning from maxtext.common import checkpointing +from maxtext.common import emergency_checkpointing from maxtext.common import train_state_nnx from maxtext.common.common_types import ReorderStrategy from maxtext.common.data_loader import create_dataloader @@ -52,7 +53,7 @@ def create_checkpoint_manager(config, mesh, init_state_fn): # pass in model for muon logger = checkpointing.setup_checkpoint_logger(config) if config.enable_multi_tier_checkpointing: - checkpoint_manager = checkpointing.create_orbax_emergency_replicator_checkpoint_manager( + checkpoint_manager = emergency_checkpointing.create_replicator_checkpoint_manager( config.local_checkpoint_directory, config.local_checkpoint_period, mesh, @@ -60,7 +61,7 @@ def create_checkpoint_manager(config, mesh, init_state_fn): ) elif config.enable_emergency_checkpoint: abstract_state, _, _ = maxtext_utils.get_abstract_state(config, mesh, init_state_fn, is_training=True) - checkpoint_manager = checkpointing.create_orbax_emergency_checkpoint_manager( + checkpoint_manager = emergency_checkpointing.create_emergency_checkpoint_manager( config.local_checkpoint_directory, config.checkpoint_dir, mesh, diff --git a/tests/unit/checkpointing_emergency_nnx_test.py b/tests/unit/checkpointing_emergency_nnx_test.py index d8013d26d1..2547e84fe4 100644 --- a/tests/unit/checkpointing_emergency_nnx_test.py +++ b/tests/unit/checkpointing_emergency_nnx_test.py @@ -148,8 +148,10 @@ def _abstract_handed_to_manager(self, abstract_state): """Calls the constructor with Orbax mocked, and returns the abstract it was handed.""" mesh = jax.sharding.Mesh(jax.devices(), ("x",)) with ( - mock.patch.object(checkpointing, "EmergencyCheckpointManager") as manager_cls, - mock.patch.object(checkpointing.gcs_utils, "mkdir_and_check_permissions", side_effect=epath.Path), + mock.patch.object(checkpointing.emergency_checkpointing, "CheckpointManager") as manager_cls, + mock.patch.object( + checkpointing.emergency_checkpointing.gcs_utils, "mkdir_and_check_permissions", side_effect=epath.Path + ), tempfile.TemporaryDirectory() as d, ): checkpointing.create_orbax_emergency_checkpoint_manager( diff --git a/tests/unit/train_state_nnx_checkpoint_test.py b/tests/unit/train_state_nnx_checkpoint_test.py index b94cf51382..ef19ac427e 100644 --- a/tests/unit/train_state_nnx_checkpoint_test.py +++ b/tests/unit/train_state_nnx_checkpoint_test.py @@ -73,8 +73,8 @@ def test_colocated_python_option_is_forwarded(self): mesh = object() with mock.patch.object( - checkpointing, - "EmergencyReplicatorCheckpointManager", + checkpointing.emergency_checkpointing, + "ReplicatorCheckpointManager", return_value=checkpoint_manager, ) as manager_cls: result = checkpointing.create_orbax_emergency_replicator_checkpoint_manager( diff --git a/tests/unit/train_utils_test.py b/tests/unit/train_utils_test.py index e6436877e1..d7857c24b2 100644 --- a/tests/unit/train_utils_test.py +++ b/tests/unit/train_utils_test.py @@ -217,8 +217,8 @@ def test_single_controller_mtc_registers_colocated_python_handlers(self): with ( mock.patch.object( - train_utils.checkpointing, - "create_orbax_emergency_replicator_checkpoint_manager", + train_utils.emergency_checkpointing, + "create_replicator_checkpoint_manager", return_value=checkpoint_manager, ) as create_manager, mock.patch.object(