From 62df8dd007a85a3b665ddce1cf65f74b708e9150 Mon Sep 17 00:00:00 2001 From: Muzi Li Date: Fri, 31 Jul 2026 21:38:24 +0000 Subject: [PATCH] Add standalone checkpointer benchmark configuration flags and loop features - Add per_step_interval, drop_page_cache_before_restore, standalone_checkpointer_enable_restore_in_loop, and standalone_checkpointer_start_from_checkpoint configuration flags to types.py and base.yml for standalone checkpointer storage benchmarking. - Support optional checkpoint loading at startup via load_state_if_possible when standalone_checkpointer_start_from_checkpoint is enabled. - Update standalone_checkpointer.py loop to execute save, optional interval wait, optional Linux OS page cache eviction, and restore on each step iteration. - Remove deprecated legacy GCS CSV metric uploads and redundant startup log from standalone_checkpointer.py (metrics are collected by independent downstream evaluators). --- src/maxtext/configs/base.yml | 9 +++ src/maxtext/configs/types.py | 25 +++++++ src/maxtext/utils/standalone_checkpointer.py | 78 +++++++++++++++++--- 3 files changed, 102 insertions(+), 10 deletions(-) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 0f2960effb..ef558ee931 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -63,6 +63,15 @@ checkpoint_todelete_subdir: None # Full path to move checkpoints to before deletion. checkpoint_todelete_full_path: None +# Interval in seconds between iterations in standalone checkpointer benchmark loop +per_step_interval: 0.0 +# Whether to execute sudo sh -c 'echo 3 > /proc/sys/vm/drop_caches' before restoring a checkpoint in standalone_checkpointer loop +drop_page_cache_before_restore: false +# Whether to also restore checkpoint after saving in each step of standalone_checkpointer loop +standalone_checkpointer_enable_restore_in_loop: true +# Whether to start by attempting to load an existing checkpoint in standalone_checkpointer +standalone_checkpointer_start_from_checkpoint: false + force_unroll: false # during generate_param_only_checkpoint should we unroll the loop? # checkpointing using orbax has two important parameters: array driver diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 29d870f7ab..5badae3927 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -344,6 +344,31 @@ class Checkpointing(BaseModel): description="Subdirectory to move checkpoints to before deletion. (Ignored if directory is prefixed with gs://)", ) checkpoint_todelete_full_path: str | None = Field(None, description="Full path to move checkpoints to before deletion.") + per_step_interval: float = Field( + 0.0, + description="Interval in seconds between iterations in standalone checkpointer benchmark loop.", + ) + drop_page_cache_before_restore: bool = Field( + False, + description=( + "Whether to execute sudo sh -c 'echo 3 > /proc/sys/vm/drop_caches' before restoring a checkpoint in" + " standalone_checkpointer loop (use for storage benchmarking only)." + ), + ) + standalone_checkpointer_enable_restore_in_loop: bool = Field( + True, + description=( + "In standalone_checkpointer loop, whether to restore checkpoint after saving in each step (defaults to True" + " for bidirectional storage read/write benchmarking)." + ), + ) + standalone_checkpointer_start_from_checkpoint: bool = Field( + False, + description=( + "In standalone_checkpointer, whether to start by attempting to load an existing checkpoint before setting" + " up training state (for checkpoint restore benchmarking)." + ), + ) force_unroll: bool = Field( False, description="During param-only checkpoint generation, whether to unroll the loop.", diff --git a/src/maxtext/utils/standalone_checkpointer.py b/src/maxtext/utils/standalone_checkpointer.py index db9c4f0193..adb85ec9df 100644 --- a/src/maxtext/utils/standalone_checkpointer.py +++ b/src/maxtext/utils/standalone_checkpointer.py @@ -21,10 +21,12 @@ import datetime from functools import partial import os +import time from typing import Sequence from absl import app from flax import nnx +from flax.linen import partitioning as nn_partitioning import jax from jax import numpy as jnp from maxtext.configs import pyconfig @@ -79,18 +81,43 @@ def init_state_fn(): # A barrier to sync all hosts before starting to restore checkpoint jax.experimental.multihost_utils.sync_global_devices("Barrier before load") - checkpoint_load_start = datetime.datetime.now() - # Delegate checkpoint restoration or state initialization to setup_training_state - state, _, _, _, was_restored = maxtext_utils.setup_training_state(None, config, mesh, checkpoint_manager, init_state_fn) + was_restored = False + state = None + + if config.standalone_checkpointer_start_from_checkpoint: + unboxed_abstract_state, _, _ = maxtext_utils.get_abstract_state(config, mesh, init_state_fn, is_training=True) + with nn_partitioning.axis_rules(config.logical_axis_rules): + loaded_state, _ = checkpointing.load_state_if_possible( + checkpoint_manager, + None, + config.load_parameters_path, + config.load_full_state_path, + config.checkpoint_storage_concurrent_gb, + unboxed_abstract_state, + config.enable_single_replica_ckpt_restoring, + config.dataset_type, + use_ocdbt=config.checkpoint_storage_use_ocdbt, + use_zarr3=config.checkpoint_storage_use_zarr3, + enable_orbax_v1=config.enable_orbax_v1, + checkpoint_conversion_fn=config.checkpoint_conversion_fn, + source_checkpoint_layout=config.source_checkpoint_layout, + expansion_factor_real_data=config.expansion_factor_real_data, + maxtext_config=config, + ) + if loaded_state: + state = loaded_state.get("items", loaded_state) + was_restored = True + + if state is None: + # Delegate checkpoint restoration or state initialization to setup_training_state + state, _, _, _, was_restored = maxtext_utils.setup_training_state( + model, config, mesh, checkpoint_manager, init_state_fn + ) + jax.block_until_ready(state) - checkpoint_load_end = datetime.datetime.now() - if was_restored: - if jax.process_index() == 0: - max_logging.log( - "STANDALONE CHECKPOINTER : Checkpoint restored in :" f" {checkpoint_load_end - checkpoint_load_start}" - ) - else: # Checkpoint was unavailable, fresh state needs to be perturbed with entropy + if not was_restored: + # Checkpoint was unavailable, fresh state needs to be perturbed with entropy state = add_entropy_to_checkpoint(state) start_step = get_first_step(model, state) # this is the start_step for training @@ -107,6 +134,37 @@ def init_state_fn(): max_logging.log( "STANDALONE CHECKPOINTER : Checkpoint saved in" f" {end_time - start_time} ,step {step}, on host 0" ) + elapsed_time = datetime.datetime.now() - start_time + time_to_wait = config.per_step_interval - elapsed_time.total_seconds() + if time_to_wait > 0: + time.sleep(time_to_wait) + jax.experimental.multihost_utils.sync_global_devices("Barrier after step") + + if config.standalone_checkpointer_enable_restore_in_loop: + # Optional OS Page Cache Eviction (for Checkpointing Benchmarks): + # When saving a checkpoint to storage and immediately restoring it on the same host, + # the Linux kernel OS page cache holds the newly written blocks in RAM. + # Without dropping the cache, the restore operation will read from host RAM rather + # than actual backing storage (e.g., GCS / Lustre / persistent disk), which artificially + # inflates restore speeds and distorts storage benchmark metrics. + # + # NOTE: Executing this command requires `sudo` privileges on Linux (`sudo sh -c 'echo 3 > /proc/sys/vm/drop_caches'`). + # It defaults to False for compatibility with standard non-sudo MaxText environments, + # and should only be enabled in dedicated benchmarking environments. + if jax.process_index() == 0 and config.drop_page_cache_before_restore: + max_logging.log("STANDALONE CHECKPOINTER : Dropping OS page cache before restore...") + os.system("sudo sh -c 'echo 3 > /proc/sys/vm/drop_caches'") + + restore_start = datetime.datetime.now() + restored_state = checkpoint_manager.restore(int(step)) + if restored_state: + restored_state = restored_state.get("items", restored_state) + jax.block_until_ready(restored_state) + restore_end = datetime.datetime.now() + if jax.process_index() == 0: + max_logging.log( + f"STANDALONE CHECKPOINTER : Checkpoint restored in {restore_end - restore_start} ,step {step}, on host 0" + ) return state