diff --git a/recml/core/training/keras_trainer.py b/recml/core/training/keras_trainer.py index c8e121e..0b020bb 100644 --- a/recml/core/training/keras_trainer.py +++ b/recml/core/training/keras_trainer.py @@ -122,9 +122,32 @@ def __init__( max_checkpoints_to_keep: int = 5, checkpoint_save_interval_epochs: int = 1, rng_seed: int = core.DEFAULT_RNG_SEED, - legacy_checkpoint_format: bool = True, + checkpoint_version: keras_utils.CheckpointVersion | str = "v2", + legacy_checkpoint_format: bool | None = None, ): - """Initializes the instance.""" + """Initializes the instance. + + Args: + distribution: The distribution strategy to use. + model_dir: The directory to save checkpoints and logs. + train_steps: Total number of training steps. + steps_per_eval: Number of steps between evaluations. + continuous_eval_timeout: Timeout for continuous evaluation. + steps_per_loop: Number of steps per training loop. + max_checkpoints_to_keep: Maximum number of checkpoints to keep. + checkpoint_save_interval_epochs: Interval in epochs to save checkpoints. + rng_seed: Random seed. + checkpoint_version: The checkpoint version to use. Supported versions: + "v1" (legacy V1), "v2" (V2, default), "v3" (V3). + legacy_checkpoint_format: Deprecated. Use checkpoint_version instead. If + set, True maps to V1, and False maps to V2. TODO(b/542602169): Remove + this in v2. + resume_training_launcher: Launcher to resume training. + enable_xmanager_measurements: Whether to enable XManager measurements. + enable_autoxprof: Whether to enable AutoXprof. + autoxprof_settings: Settings for AutoXprof. + s2_logging_settings: Settings for S2 logging. + """ keras.utils.set_random_seed(rng_seed) @@ -153,25 +176,42 @@ def __init__( self._checkpoint_dir = os.path.join(model_dir, core.CHECKPOINT_DIR) self._max_checkpoints_to_keep = max_checkpoints_to_keep self._checkpoint_save_interval_epochs = checkpoint_save_interval_epochs - self._legacy_checkpoint_format = legacy_checkpoint_format + if legacy_checkpoint_format is not None: + logging.warning( + "legacy_checkpoint_format is deprecated, use checkpoint_version" + " instead." + ) + self._checkpoint_version = ( + keras_utils.CheckpointVersion.V1 + if legacy_checkpoint_format + else keras_utils.CheckpointVersion.V2 + ) + else: + self._checkpoint_version = keras_utils.CheckpointVersion( + checkpoint_version + ) @functools.cached_property def train_callbacks(self) -> list[keras.callbacks.Callback]: """Returns the training callbacks.""" if keras.backend.backend() == "jax": - if self._legacy_checkpoint_format: - checkpoint_manager = keras_utils.KerasOrbaxCheckpointManager( - checkpoint_dir=self._checkpoint_dir, - max_to_keep=self._max_checkpoints_to_keep, - save_interval_epochs=self._checkpoint_save_interval_epochs, - ) + if self._checkpoint_version == keras_utils.CheckpointVersion.V1: + manager_cls = keras_utils.KerasOrbaxCheckpointManager + elif self._checkpoint_version == keras_utils.CheckpointVersion.V2: + manager_cls = keras_utils.KerasOrbaxCheckpointManagerV2 + elif self._checkpoint_version == keras_utils.CheckpointVersion.V3: + manager_cls = keras_utils.KerasOrbaxCheckpointManagerV3 else: - checkpoint_manager = keras_utils.KerasOrbaxCheckpointManagerV2( - checkpoint_dir=self._checkpoint_dir, - max_to_keep=self._max_checkpoints_to_keep, - save_interval_epochs=self._checkpoint_save_interval_epochs, + raise ValueError( + f"Unsupported checkpoint version: {self._checkpoint_version}" ) + checkpoint_manager = manager_cls( + checkpoint_dir=self._checkpoint_dir, + max_to_keep=self._max_checkpoints_to_keep, + save_interval_epochs=self._checkpoint_save_interval_epochs, + ) + callbacks = [ keras_utils.EpochSummaryCallback( log_dir=os.path.join(self._model_dir, core.LOG_DIR), @@ -379,7 +419,7 @@ def timeout_fn() -> bool: else: steps_msg = "running complete evaluation..." - use_legacy_checkpoint_format = self._legacy_checkpoint_format + checkpoint_version = self._checkpoint_version class _RestoreCallback(keras.callbacks.Callback): """Callback for restoring the model from the latest checkpoint.""" @@ -393,14 +433,21 @@ def __init__( self._epoch = epoch def on_test_begin(self, logs: Mapping[str, Any] | None = None): - if use_legacy_checkpoint_format: + if checkpoint_version == keras_utils.CheckpointVersion.V1: keras_utils.restore_keras_model( model, self._checkpoint_dir, step=self._epoch ) - else: + elif checkpoint_version in ( + keras_utils.CheckpointVersion.V2, + keras_utils.CheckpointVersion.V3, + ): keras_utils.restore_keras_checkpoint( self._checkpoint_dir, model=model, epoch=self._epoch ) + else: + raise ValueError( + f"Unsupported checkpoint version: {checkpoint_version}" + ) history = None for epoch in ocp.checkpoint_utils.checkpoints_iterator( diff --git a/recml/core/training/keras_trainer_test.py b/recml/core/training/keras_trainer_test.py index 7a63f6f..552cc42 100644 --- a/recml/core/training/keras_trainer_test.py +++ b/recml/core/training/keras_trainer_test.py @@ -55,6 +55,7 @@ class KerasTrainerTest(parameterized.TestCase): def setUp(self): super().setUp() + keras.backend.clear_session() # Workaround to make `create_tempdir` work with pytest. if not flags.FLAGS.is_parsed(): flags.FLAGS.mark_as_parsed() @@ -63,26 +64,38 @@ def setUp(self): {"testcase_name": "train", "mode": core.Trainer.Mode.TRAIN}, {"testcase_name": "eval", "mode": core.Trainer.Mode.EVAL}, { - "testcase_name": "train_and_eval", + "testcase_name": "train_and_eval_v3", "mode": core.Trainer.Mode.TRAIN_AND_EVAL, + "checkpoint_version": "v3", }, { - "testcase_name": "continuous_eval_", - "mode": core.Trainer.Mode.CONTINUOUS_EVAL, + "testcase_name": "train_and_eval_v2", + "mode": core.Trainer.Mode.TRAIN_AND_EVAL, + "checkpoint_version": "v2", }, { - "testcase_name": "train_and_eval_legacy_checkpoint_format", + "testcase_name": "train_and_eval_v1", "mode": core.Trainer.Mode.TRAIN_AND_EVAL, - "legacy_checkpoint_format": True, + "checkpoint_version": "v1", + }, + { + "testcase_name": "continuous_eval_v3", + "mode": core.Trainer.Mode.CONTINUOUS_EVAL, + "checkpoint_version": "v3", + }, + { + "testcase_name": "continuous_eval_v2", + "mode": core.Trainer.Mode.CONTINUOUS_EVAL, + "checkpoint_version": "v2", }, { - "testcase_name": "continuous_eval_legacy_checkpoint_format", + "testcase_name": "continuous_eval_v1", "mode": core.Trainer.Mode.CONTINUOUS_EVAL, - "legacy_checkpoint_format": True, + "checkpoint_version": "v1", }, ) def test_keras_task_and_trainer( - self, mode: str, legacy_checkpoint_format: bool = False + self, mode: str, checkpoint_version: str = "v3" ): if keras.backend.backend() == "jax": distribution = keras.distribution.DataParallel() @@ -98,13 +111,14 @@ def test_keras_task_and_trainer( steps_per_loop=2, model_dir=self.create_tempdir().full_path, continuous_eval_timeout=5, - legacy_checkpoint_format=legacy_checkpoint_format, + checkpoint_version=checkpoint_version, ) experiment = core.Experiment(_KerasTask(), trainer) if mode == core.Trainer.Mode.CONTINUOUS_EVAL: # Produce one checkpoint so there is something to evaluate. core.run_experiment(experiment, core.Trainer.Mode.TRAIN) + keras.backend.clear_session() history = core.run_experiment(experiment, mode) diff --git a/recml/core/utils/keras_utils.py b/recml/core/utils/keras_utils.py index 31eb566..dd5ae40 100644 --- a/recml/core/utils/keras_utils.py +++ b/recml/core/utils/keras_utils.py @@ -13,15 +13,17 @@ # limitations under the License. """Utilities for training Keras models on Jax backend.""" -from collections.abc import Mapping +from collections.abc import Mapping, Sequence import dataclasses import datetime +import enum import os import re from typing import Any from absl import logging +from etils import epath import jax import keras import orbax.checkpoint as ocp @@ -33,9 +35,18 @@ NON_TRAINABLE_VARIABLES_KEY = "non_trainable_variables" OPTIMIZER_VARIABLES_KEY = "optimizer_variables" CONFIG_CHECKPOINT_KEY = "config" +FORMAT_VERSION_KEY = "format_version" +NON_TRAINABLE_PATHS_KEY = "non_trainable_paths" +OPTIMIZER_PATHS_KEY = "optimizer_paths" ORBAX_CHECKPOINT_DEFAULT_KEY = "default" +class CheckpointVersion(enum.StrEnum): + V1 = "v1" + V2 = "v2" + V3 = "v3" + + def _assert_variables_built(model: keras.Model): if not model.built or not model.optimizer.built: raise ValueError( @@ -56,6 +67,27 @@ def _assert_all_layers_built(model: keras.Model): ) +def _variables_to_path_dict( + variables: Sequence[keras.Variable], + collection_name: str, +) -> dict[str, keras.Variable]: + """Converts a sequence of variables to a dict mapped by path, checking for duplicates.""" + var_dict = {} + duplicates = [] + for v in variables: + if v.path in var_dict: + duplicates.append(v.path) + else: + var_dict[v.path] = v + if duplicates: + raise ValueError( + f"Duplicate variable paths detected in {collection_name}. Ensure" + " unique layer names (e.g. name_layers=True). Duplicates: " + f"{duplicates}" + ) + return var_dict + + def _to_shape_dtype_struct(x: keras.Variable) -> jax.ShapeDtypeStruct: if not isinstance(x, keras.Variable): raise ValueError(f"Expected a `keras.Variable`, got {type(x)}.") @@ -173,7 +205,183 @@ def restore_model_variables(self, model: keras.Model, epoch: int): var._value = restored_var # pylint: disable=protected-access -def _resolve_orbax_checkpoint_path( +class KerasOrbaxCheckpointManagerV3(ocp.CheckpointManager): + """An Orbax checkpoint manager for Keras 3 with dictionary state. + + This manager saves the full training state (trainable, non-trainable, and + optimizer variables). For training resume and preemption recovery, the full + state is restored via `restore_keras_checkpoint`. + + For selective weight transfer (warm-starting from a checkpoint of a model + with a different architecture), use `restore_partial_checkpoint`. Note that + partial restoration is restricted to trainable variables (weights). + Non-trainable and optimizer variables are specific to the training run and + are not supported for partial transfer. + """ + + def __init__( + self, + checkpoint_dir: str, + max_to_keep: int = 5, + save_interval_epochs: int = 1, + ): + """Initializes a KerasOrbaxCheckpointManagerV3. + + Args: + checkpoint_dir: The directory to save checkpoints to. + max_to_keep: The maximum number of checkpoints to keep. + save_interval_epochs: The interval (in epochs) to save checkpoints. + """ + if keras.backend.backend() != "jax": + raise ValueError( + "`KerasOrbaxCheckpointManagerV3` is only supported on a `jax`" + " backend." + ) + super().__init__( + directory=checkpoint_dir, + item_names=( + STATE_CHECKPOINT_KEY, + CONFIG_CHECKPOINT_KEY, + FORMAT_VERSION_KEY, + NON_TRAINABLE_PATHS_KEY, + OPTIMIZER_PATHS_KEY, + ), + options=ocp.CheckpointManagerOptions( + save_interval_steps=save_interval_epochs, + max_to_keep=max_to_keep, + ), + ) + + def save_model_variables( + self, + model: keras.Model, + epoch: int, + logs: Mapping[str, Any] | None = None, + ): + """Saves the model variables and optimizer variables to a checkpoint.""" + _assert_variables_built(model) + _assert_all_layers_built(model) + + if not model._jax_state_synced: # pylint: disable=protected-access + model.jax_state_sync() + + trainable_variables = _variables_to_path_dict( + model.trainable_variables, TRAINABLE_VARIABLES_KEY + ) + non_trainable_variables = _variables_to_path_dict( + model.non_trainable_variables, NON_TRAINABLE_VARIABLES_KEY + ) + optimizer_variables = _variables_to_path_dict( + model.optimizer.variables, OPTIMIZER_VARIABLES_KEY + ) + + # Extract values from keras.Variable instances + state = { + TRAINABLE_VARIABLES_KEY: { + k: v.value for k, v in trainable_variables.items() + }, + NON_TRAINABLE_VARIABLES_KEY: { + k: v.value for k, v in non_trainable_variables.items() + }, + OPTIMIZER_VARIABLES_KEY: { + k: v.value for k, v in optimizer_variables.items() + }, + } + config = keras.utils.serialize_keras_object(model) + non_trainable_paths = { + "paths": [v.path for v in model.non_trainable_variables] + } + optimizer_paths = {"paths": [v.path for v in model.optimizer.variables]} + logging.info("SAVED non_trainable_paths: %s", non_trainable_paths) + logging.info("SAVED optimizer_paths: %s", optimizer_paths) + + logging.info("Saving checkpoint for epoch %s...", epoch) + self.save( + step=epoch, + args=ocp.args.Composite(**{ + STATE_CHECKPOINT_KEY: ocp.args.PyTreeSave(state), + CONFIG_CHECKPOINT_KEY: ocp.args.JsonSave(config), + FORMAT_VERSION_KEY: ocp.args.JsonSave({"version": 3}), + NON_TRAINABLE_PATHS_KEY: ocp.args.JsonSave(non_trainable_paths), + OPTIMIZER_PATHS_KEY: ocp.args.JsonSave(optimizer_paths), + }), + metrics=logs, + ) + + def restore_model_variables(self, model: keras.Model, epoch: int): + """Restores the model variables and optimizer variables during training.""" + + _assert_variables_built(model) + _assert_all_layers_built(model) + + if not model._jax_state_synced: # pylint: disable=protected-access + model.jax_state_sync() + + trainable_variables = _variables_to_path_dict( + model.trainable_variables, TRAINABLE_VARIABLES_KEY + ) + non_trainable_variables = _variables_to_path_dict( + model.non_trainable_variables, NON_TRAINABLE_VARIABLES_KEY + ) + optimizer_variables = _variables_to_path_dict( + model.optimizer.variables, OPTIMIZER_VARIABLES_KEY + ) + + variables = { + TRAINABLE_VARIABLES_KEY: trainable_variables, + NON_TRAINABLE_VARIABLES_KEY: non_trainable_variables, + OPTIMIZER_VARIABLES_KEY: optimizer_variables, + } + + # Construct abstract variables to ensure the checkpoint is restored with + # the same sharding as the current variables. + abstract_variables = jax.tree.map(_to_shape_dtype_struct, variables) + for var in jax.tree.flatten(variables)[0]: + var.value.delete() + var._value = None # pylint: disable=protected-access + + logging.info("Restoring checkpoint for epoch %s...", epoch) + + step_path = os.path.join(self.directory, str(epoch)) + abstract_variables, state_transforms = _prepare_v3_restore( + step_path, + abstract_variables, + model, + restore_optimizer_vars=True, + ) + + restored_items = self.restore( + step=epoch, + args=ocp.args.Composite(**{ + STATE_CHECKPOINT_KEY: ocp.args.PyTreeRestore( + abstract_variables, + transforms=state_transforms, + restore_args=ocp.checkpoint_utils.construct_restore_args( + abstract_variables + ), + ) + }), + ) + restored_variables = restored_items[STATE_CHECKPOINT_KEY] + + logging.info("Restored checkpoint for epoch %s.", epoch) + + model._initial_epoch = epoch + 1 # pylint: disable=protected-access + + keras.tree.assert_same_structure(variables, restored_variables) + + for key in [ + TRAINABLE_VARIABLES_KEY, + NON_TRAINABLE_VARIABLES_KEY, + OPTIMIZER_VARIABLES_KEY, + ]: + var_dict = variables[key] + restored_var_dict = restored_variables[key] + for path, var in var_dict.items(): + var._value = restored_var_dict[path] # pylint: disable=protected-access + + +def resolve_orbax_checkpoint_path( checkpoint_dir: str, epoch: int | None = None ) -> tuple[str, int | None]: """Resolves the checkpoint path and epoch for an Orbax checkpoint. @@ -227,6 +435,237 @@ def _resolve_orbax_checkpoint_path( return os.fspath(checkpoint_path), epoch +def _is_v1_checkpoint_path(checkpoint_path: str) -> bool: + """Checks if a resolved checkpoint path is in V1 format.""" + return gfile.Exists( + os.path.join(checkpoint_path, ORBAX_CHECKPOINT_DEFAULT_KEY) + ) + + +def _is_v3_checkpoint_path(checkpoint_path: str) -> bool: + """Checks if a resolved checkpoint path is in V3 format.""" + if not gfile.Exists(os.path.join(checkpoint_path, FORMAT_VERSION_KEY)): + return False + + version_checkpointer = ocp.Checkpointer( + ocp.CompositeCheckpointHandler( + **{FORMAT_VERSION_KEY: ocp.handlers.JsonCheckpointHandler()} # pyrefly: ignore[bad-argument-type] + ) + ) + is_v3 = False + try: + version_info = version_checkpointer.restore( + checkpoint_path, + args=ocp.args.Composite(**{FORMAT_VERSION_KEY: ocp.args.JsonRestore()}), + )[FORMAT_VERSION_KEY] + if version_info.get("version") == 3: + is_v3 = True + except (ValueError, KeyError, OSError) as e: + logging.warning( + "Failed to read format version from %s: %s", checkpoint_path, e + ) + finally: + version_checkpointer.close() + return is_v3 + + +def is_v3_checkpoint(checkpoint_dir: str, epoch: int | None = None) -> bool: + """Checks if a checkpoint is in V3 format.""" + try: + checkpoint_path, _ = resolve_orbax_checkpoint_path(checkpoint_dir, epoch) + return _is_v3_checkpoint_path(checkpoint_path) + except (FileNotFoundError, ValueError): + return False + + +def is_v1_checkpoint(checkpoint_dir: str, epoch: int | None = None) -> bool: + """Checks if a checkpoint is in V1 format.""" + try: + checkpoint_path, _ = resolve_orbax_checkpoint_path(checkpoint_dir, epoch) + return _is_v1_checkpoint_path(checkpoint_path) + except (FileNotFoundError, ValueError): + return False + + +def _prepare_v3_restore( + checkpoint_path: str, + abstract_state: Mapping[str, Any], + model: keras.Model | None = None, + restore_optimizer_vars: bool = False, +) -> tuple[Mapping[str, Any], Mapping[str, Any]]: + """Prepares abstract state and state transforms for V3 checkpoint restore.""" + metadata_handler = ocp.handlers.PyTreeCheckpointHandler() + state_checkpoint_path = epath.Path(checkpoint_path) / STATE_CHECKPOINT_KEY + saved_state_metadata = metadata_handler.metadata(state_checkpoint_path) + + has_paths = gfile.Exists( + os.path.join(checkpoint_path, NON_TRAINABLE_PATHS_KEY) + ) + non_trainable_paths = None + optimizer_paths = None + if has_paths and model is not None: + paths_checkpointer = ocp.Checkpointer( + ocp.CompositeCheckpointHandler(**{ # pyrefly: ignore[bad-argument-type] + NON_TRAINABLE_PATHS_KEY: ocp.handlers.JsonCheckpointHandler(), + OPTIMIZER_PATHS_KEY: ocp.handlers.JsonCheckpointHandler(), + }) + ) + restored_paths = paths_checkpointer.restore( + checkpoint_path, + args=ocp.args.Composite(**{ + NON_TRAINABLE_PATHS_KEY: ocp.args.JsonRestore(), + OPTIMIZER_PATHS_KEY: ocp.args.JsonRestore(), + }), + ) + non_trainable_paths = restored_paths[NON_TRAINABLE_PATHS_KEY]["paths"] + optimizer_paths = restored_paths[OPTIMIZER_PATHS_KEY]["paths"] + paths_checkpointer.close() + + filtered_abstract_state = {} + state_transforms = {} + keys = [TRAINABLE_VARIABLES_KEY, NON_TRAINABLE_VARIABLES_KEY] + if restore_optimizer_vars: + keys.append(OPTIMIZER_VARIABLES_KEY) + + for key in keys: + if key in abstract_state and key in saved_state_metadata: + filtered_abstract_state[key] = {} + state_transforms[key] = {} + + if key in [NON_TRAINABLE_VARIABLES_KEY, OPTIMIZER_VARIABLES_KEY]: + if model is None: + raise ValueError(f"Model must be provided to restore key {key}.") + target_paths = ( + [v.path for v in model.non_trainable_variables] + if key == NON_TRAINABLE_VARIABLES_KEY + else [v.path for v in model.optimizer.variables] + ) + source_paths = ( + non_trainable_paths + if key == NON_TRAINABLE_VARIABLES_KEY + else optimizer_paths + ) + + if has_paths: + assert source_paths is not None + # Map by index using the saved paths ordering + for i, target_path in enumerate(target_paths): + struct = abstract_state[key][target_path] + if i < len(source_paths): + source_path = source_paths[i] + filtered_abstract_state[key][target_path] = struct + if target_path != source_path: + state_transforms[key][target_path] = ( + ocp.transform_utils.Transform( + original_key=f"{key}/{source_path}" + ) + ) + logging.info( + "Mapping target path %s to source path %s by index %d", + target_path, + source_path, + i, + ) + else: + logging.warning( + "No source path at index %d for target path %s", + i, + target_path, + ) + else: + raise ValueError( + "Index information (non_trainable_paths / optimizer_paths) " + f"not found in checkpoint for key {key}. " + "Unable to perform index-based mapping." + ) + else: + # Strict matching for trainable variables + missing_paths = [] + for target_path, struct in abstract_state[key].items(): + if target_path in saved_state_metadata[key]: + filtered_abstract_state[key][target_path] = struct + else: + missing_paths.append(target_path) + + if missing_paths: + raise ValueError( + f"Failed to restore variables for key {key}. " + f"Missing paths in checkpoint: {missing_paths}" + ) + elif key in abstract_state: + logging.warning("Key %s not found in checkpoint metadata", key) + + return filtered_abstract_state, state_transforms + + +def _assign_v3_restored_values( + variables: Mapping[str, Any], + restored_state: Mapping[str, Any], + restore_optimizer_vars: bool, +): + """Assigns restored V3 values back to Keras variables.""" + for key in [ + TRAINABLE_VARIABLES_KEY, + NON_TRAINABLE_VARIABLES_KEY, + ]: + if key in variables: + var_dict = variables[key] + restored_var_dict = restored_state[key] + for path, var in var_dict.items(): + if path in restored_var_dict: + logging.info("Restoring variable %s for key %s", path, key) + var._value = restored_var_dict[path] # pylint: disable=protected-access + else: + logging.warning( + "Path %s not found in restored state for key %s", path, key + ) + if restore_optimizer_vars: + key = OPTIMIZER_VARIABLES_KEY + var_dict = variables[key] + restored_var_dict = restored_state[key] + + # Try exact match first + missing_paths = [] + for path, var in var_dict.items(): + if path in restored_var_dict: + var._value = restored_var_dict[path] # pylint: disable=protected-access + else: + missing_paths.append((path, var)) + + if missing_paths: + logging.warning( + "Some optimizer paths did not match exactly. Trying heuristic" + " matching." + ) + # Heuristic match: compare suffixes or ignore optimizer name prefix + # e.g. 'adam_3/var_name' vs 'adam_1/var_name' + # Let's try matching by the part after the first '/' + restored_by_suffix = {} + for k, v in restored_var_dict.items(): + parts = k.split("/", 1) + if len(parts) > 1: + restored_by_suffix[parts[1]] = v + else: + restored_by_suffix[k] = v + + still_missing = [] + for path, var in missing_paths: + parts = path.split("/", 1) + suffix = parts[1] if len(parts) > 1 else path + if suffix in restored_by_suffix: + var._value = restored_by_suffix[suffix] # pylint: disable=protected-access + logging.info( + "Matched optimizer variable %s by suffix %s", path, suffix + ) + else: + still_missing.append(path) + + if still_missing: + raise ValueError( + f"Failed to restore optimizer variables for paths: {still_missing}" + ) + + def restore_keras_checkpoint( checkpoint_dir: str, *, @@ -282,7 +721,15 @@ def restore_keras_checkpoint( " True, a model must be provided." ) - checkpoint_path, epoch = _resolve_orbax_checkpoint_path(checkpoint_dir, epoch) + checkpoint_path, epoch = resolve_orbax_checkpoint_path(checkpoint_dir, epoch) + + if _is_v1_checkpoint_path(checkpoint_path): + raise ValueError( + f"The checkpoint in {checkpoint_dir} is in V1 format (list-based)" + f" at step {epoch}." + " `restore_keras_checkpoint` is only compatible with V2/V3 checkpoints." + " Please use `restore_keras_model` instead." + ) if model is None: cfg = {**load_keras_model_config(checkpoint_dir, epoch=epoch)} @@ -305,18 +752,40 @@ def restore_keras_checkpoint( _assert_all_layers_built(model) - variables = { - TRAINABLE_VARIABLES_KEY: model.trainable_variables, - NON_TRAINABLE_VARIABLES_KEY: model.non_trainable_variables, - } - if restore_optimizer_vars: - if not model.optimizer.built: - raise ValueError( - "To use `restore_keras_checkpoint` on an existing model with" - " `restore_optimizer_vars` set to True, the optimizer must be" - " built." + is_v3 = _is_v3_checkpoint_path(checkpoint_path) + + if is_v3: + variables = { + TRAINABLE_VARIABLES_KEY: _variables_to_path_dict( + model.trainable_variables, TRAINABLE_VARIABLES_KEY + ), + NON_TRAINABLE_VARIABLES_KEY: _variables_to_path_dict( + model.non_trainable_variables, NON_TRAINABLE_VARIABLES_KEY + ), + } + if restore_optimizer_vars: + if not model.optimizer.built: + raise ValueError( + "To use `restore_keras_checkpoint` on an existing model with" + " `restore_optimizer_vars` set to True, the optimizer must be" + " built." + ) + variables[OPTIMIZER_VARIABLES_KEY] = _variables_to_path_dict( + model.optimizer.variables, OPTIMIZER_VARIABLES_KEY ) - variables[OPTIMIZER_VARIABLES_KEY] = model.optimizer.variables + else: + variables = { + TRAINABLE_VARIABLES_KEY: model.trainable_variables, + NON_TRAINABLE_VARIABLES_KEY: model.non_trainable_variables, + } + if restore_optimizer_vars: + if not model.optimizer.built: + raise ValueError( + "To use `restore_keras_checkpoint` on an existing model with" + " `restore_optimizer_vars` set to True, the optimizer must be" + " built." + ) + variables[OPTIMIZER_VARIABLES_KEY] = model.optimizer.variables # TODO(zixiangzhou): Update variables to use a nested dictionary and index map # instead of flattened list. @@ -325,13 +794,124 @@ def restore_keras_checkpoint( # the same sharding as the current variables. abstract_state = jax.tree.map(_to_shape_dtype_struct, variables) + state_transforms = {} + if is_v3: + abstract_state, state_transforms = _prepare_v3_restore( + checkpoint_path, abstract_state, model, restore_optimizer_vars + ) + # Delete the variables from device memory to reduce peak memory usage. - for var in jax.tree.flatten(variables)[0]: - var.value.delete() - var._value = None # pylint: disable=protected-access + # Only delete variables that we are actually trying to restore. + if is_v3: + for key, path_dict in abstract_state.items(): + for path in path_dict.keys(): + var = variables[key][path] + var.value.delete() + var._value = None # pylint: disable=protected-access + else: + for var in jax.tree.flatten(variables)[0]: + var.value.delete() + var._value = None # pylint: disable=protected-access + + # Always use PyTreeCheckpointHandler for restoring, as it is more flexible + # (supports transforms) + state_handler = ocp.handlers.PyTreeCheckpointHandler( + restore_concurrent_gb=96, + ) + checkpointer = ocp.Checkpointer( + ocp.CompositeCheckpointHandler(**{ # pyrefly: ignore[bad-argument-type] + STATE_CHECKPOINT_KEY: state_handler, + }) + ) + + restore_args = ocp.args.Composite(**{ + STATE_CHECKPOINT_KEY: ocp.args.PyTreeRestore( + abstract_state, + transforms=state_transforms if is_v3 else {}, + restore_args=ocp.checkpoint_utils.construct_restore_args( + abstract_state + ), + ), + }) + + restored_state = checkpointer.restore( + checkpoint_path, + args=restore_args, + )[STATE_CHECKPOINT_KEY] + checkpointer.close() + + if is_v3: + _assign_v3_restored_values( + variables, restored_state, restore_optimizer_vars + ) + else: + keras.tree.assert_same_structure(variables, restored_state) + for var, restored_var in zip( + jax.tree.flatten(variables)[0], jax.tree.flatten(restored_state)[0] + ): + var._value = restored_var # pylint: disable=protected-access + + if restore_model_epoch: + model._initial_epoch = epoch + 1 # pylint: disable=protected-access # pyrefly: ignore[unsupported-operation] + if restore_optimizer_vars and not restore_iterations: + model.optimizer.iterations.assign(0) + + return model + + +def restore_partial_checkpoint( + checkpoint_dir: str, + partial_variables: Mapping[str, Any], + epoch: int | None = None, +) -> Mapping[str, Any]: + """Restores partial variables from an Orbax checkpoint. + + Args: + checkpoint_dir: The directory containing the Orbax checkpoint(s). + partial_variables: A dictionary mapping keys (e.g. + TRAINABLE_VARIABLES_KEY) to dictionaries mapping variable paths to + keras.Variable instances. + epoch: The epoch to restore. If None, latest is used. + + Returns: + The restored state dictionary (containing Jax Arrays). + """ + checkpoint_path, _ = resolve_orbax_checkpoint_path(checkpoint_dir, epoch) + + is_v3 = _is_v3_checkpoint_path(checkpoint_path) + if not is_v3: + raise ValueError( + "restore_partial_checkpoint only supports V3 (dictionary-based)" + " checkpoints." + ) + + # Partial restoration is restricted to trainable variables because they are + # the primary targets for selective weight transfer (e.g. sequence encoder). + # Non-trainable and optimizer variables are training-run specific and their + # mapping by index is fragile, so they are not supported for partial restore. + if ( + NON_TRAINABLE_VARIABLES_KEY in partial_variables + and partial_variables[NON_TRAINABLE_VARIABLES_KEY] + ) or ( + OPTIMIZER_VARIABLES_KEY in partial_variables + and partial_variables[OPTIMIZER_VARIABLES_KEY] + ): + raise ValueError( + "Partial restoration is only supported for trainable variables." + ) + + abstract_state = jax.tree.map(_to_shape_dtype_struct, partial_variables) + + # Delete variables from device memory to reduce peak memory usage. + for var in jax.tree.flatten(partial_variables)[0]: + if var._value is not None: # pylint: disable=protected-access + var.value.delete() + var._value = None # pylint: disable=protected-access + + abstract_state, state_transforms = _prepare_v3_restore( + checkpoint_path, abstract_state, model=None, restore_optimizer_vars=False + ) - # TODO(aahil): Look into converging the logic here with the checkpointing - # logic in KerasOrbaxCheckpointManagerV2. checkpointer = ocp.Checkpointer( ocp.CompositeCheckpointHandler(**{ # pyrefly: ignore[bad-argument-type] STATE_CHECKPOINT_KEY: ocp.handlers.PyTreeCheckpointHandler( @@ -344,28 +924,26 @@ def restore_keras_checkpoint( args=ocp.args.Composite(**{ STATE_CHECKPOINT_KEY: ocp.args.PyTreeRestore( abstract_state, - transforms={}, + transforms=state_transforms, restore_args=ocp.checkpoint_utils.construct_restore_args( abstract_state ), - ), + ) }), )[STATE_CHECKPOINT_KEY] - checkpointer.close() - - # TODO(zixiangzhou): Unflatten the variables based on index here. - keras.tree.assert_same_structure(variables, restored_state) - for var, restored_var in zip( - jax.tree.flatten(variables)[0], jax.tree.flatten(restored_state)[0] - ): - var._value = restored_var # pylint: disable=protected-access - - if restore_model_epoch: - model._initial_epoch = epoch + 1 # pylint: disable=protected-access # pyrefly: ignore[unsupported-operation] - if restore_optimizer_vars and not restore_iterations: - model.optimizer.iterations.assign(0) + # Assign restored values back to partial_variables in-place + key = TRAINABLE_VARIABLES_KEY + if key in partial_variables and key in restored_state: + var_dict = partial_variables[key] + restored_var_dict = restored_state[key] + for path, var in var_dict.items(): + if path in restored_var_dict: + var._value = restored_var_dict[path] # pylint: disable=protected-access + else: + logging.warning("Path %s was NOT restored for key %s", path, key) - return model + checkpointer.close() + return restored_state def load_keras_model_config( @@ -377,7 +955,7 @@ def load_keras_model_config( "This function only supports loading a Keras 3 Jax backend model." ) - checkpoint_path, _ = _resolve_orbax_checkpoint_path(checkpoint_dir, epoch) + checkpoint_path, _ = resolve_orbax_checkpoint_path(checkpoint_dir, epoch) json_checkpointer = ocp.Checkpointer( ocp.CompositeCheckpointHandler( @@ -531,7 +1109,9 @@ class EpochOrbaxCheckpointAndRestoreCallback(keras.callbacks.Callback): def __init__( self, checkpoint_manager: ( - KerasOrbaxCheckpointManager | KerasOrbaxCheckpointManagerV2 + KerasOrbaxCheckpointManager + | KerasOrbaxCheckpointManagerV2 + | KerasOrbaxCheckpointManagerV3 ), marker_path: str | None = None, ): @@ -583,7 +1163,8 @@ def restore_keras_model( """Restores a Keras 3 Jax backend model from an Orbax checkpoint. This is only compatible with `KerasOrbaxCheckpointManager`. If you are using - `KerasOrbaxCheckpointManagerV2`, use `restore_keras_checkpoint` instead. + `KerasOrbaxCheckpointManagerV2` or `KerasOrbaxCheckpointManagerV3`, use + `restore_keras_checkpoint` instead. Args: model: The Keras model to restore. @@ -633,6 +1214,18 @@ def restore_keras_model( f"{ocp.path.step.checkpoint_steps(checkpoint_dir)}" ) + checkpoint_path = ocp.path.step.build_step_path( + checkpoint_dir, ocp.path.step.standard_name_format(), step + ) + + if gfile.Exists(os.path.join(checkpoint_path, STATE_CHECKPOINT_KEY)): + raise ValueError( + f"The checkpoint in {checkpoint_dir} is in V2/V3 format" + f" (dictionary-based) at step {step}." + " `restore_keras_model` is only compatible with legacy V1 checkpoints." + " Please use `restore_keras_checkpoint` instead." + ) + checkpointer = ocp.Checkpointer( ocp.CompositeCheckpointHandler(**{ # pyrefly: ignore[bad-argument-type] ORBAX_CHECKPOINT_DEFAULT_KEY: ocp.handlers.PyTreeCheckpointHandler() @@ -651,9 +1244,7 @@ def restore_keras_model( jax.tree.map( lambda x: x.delete() if isinstance(x, jax.Array) else None, state ) - checkpoint_path = ocp.path.step.build_step_path( - checkpoint_dir, ocp.path.step.standard_name_format(), step - ) + # TODO(zixiangzhou): 'transforms' is a walkaround to avoid the error of # loading a checkpoint that has a different number of variables than the # current state because we don't want to load metrics_variables. But this diff --git a/recml/core/utils/keras_utils_test.py b/recml/core/utils/keras_utils_test.py index 421f257..4402baf 100644 --- a/recml/core/utils/keras_utils_test.py +++ b/recml/core/utils/keras_utils_test.py @@ -71,6 +71,45 @@ def _create_model(input_shapes: Sequence[int]) -> keras.Model: return model +@keras.saving.register_keras_serializable(package="Recml") +class MyNonTrainableLayer(keras.layers.Layer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.non_trainable_weight = self.add_weight( + shape=(10,), initializer="ones", trainable=False, name="weight" + ) + + def call(self, x): + return x + + +@keras.saving.register_keras_serializable(package="Recml") +class MyTestModel(keras.Model): + + def __init__(self, layer_name, **kwargs): + super().__init__(**kwargs) + self.layer_name = layer_name + self.my_layer = MyNonTrainableLayer(name=layer_name) + self.dense = keras.layers.Dense(5, name="my_trainable_dense") + self.direct_non_trainable = self.add_weight( + shape=(5,), + initializer="ones", + trainable=False, + name="direct_non_trainable", + ) + + def call(self, x): + x = self.my_layer(x) + x = self.dense(x) + return x + + def get_config(self): + config = super().get_config() + config.update({"layer_name": self.layer_name}) + return config + + class KerasUtilsTest(parameterized.TestCase): def setUp(self): @@ -157,6 +196,85 @@ def test_keras_orbax_checkpointer_v2( # Ensures predictions are identical. np.testing.assert_allclose(preds, preds_after_restoration) + @parameterized.named_parameters( + { + "testcase_name": "single_core", + "data_parallel": False, + "restore_with_checkpointer": True, + }, + { + "testcase_name": "data_parallel", + "data_parallel": True, + "restore_with_checkpointer": True, + }, + { + "testcase_name": "restore_without_checkpointer_single_core", + "data_parallel": False, + "restore_with_checkpointer": False, + }, + { + "testcase_name": "restore_without_checkpointer_data_parallel", + "data_parallel": True, + "restore_with_checkpointer": False, + }, + ) + def test_keras_orbax_checkpointer_v3( + self, data_parallel: bool, restore_with_checkpointer: bool + ): + if data_parallel: + keras.distribution.set_distribution(keras.distribution.DataParallel()) + else: + keras.distribution.set_distribution(None) + + checkpoint_dir = self.create_tempdir().full_path + checkpoint_manager = keras_utils.KerasOrbaxCheckpointManagerV3( + checkpoint_dir, max_to_keep=5 + ) + dummy_inputs = _create_dummy_inputs() + + bert_pretrainer = _create_model(jax.tree.map(jnp.shape, dummy_inputs)) + + state = ( + [v.value for v in bert_pretrainer.trainable_variables], + [v.value for v in bert_pretrainer.non_trainable_variables], + [v.value for v in bert_pretrainer.optimizer.variables], + ) + checkpoint_manager.save_model_variables(bert_pretrainer, 0) + checkpoint_manager.wait_until_finished() + + preds = bert_pretrainer(dummy_inputs) + + bert_pretrainer = _create_model(jax.tree.map(jnp.shape, dummy_inputs)) + if restore_with_checkpointer: + checkpoint_manager.restore_model_variables(bert_pretrainer, 0) + else: + keras_utils.restore_keras_checkpoint( + checkpoint_dir, model=bert_pretrainer, restore_optimizer_vars=True + ) + + checkpoint_manager.close() + + restored_state = ( + [v.value for v in bert_pretrainer.trainable_variables], + [v.value for v in bert_pretrainer.non_trainable_variables], + [v.value for v in bert_pretrainer.optimizer.variables], + ) + preds_after_restoration = bert_pretrainer(dummy_inputs) + + keras.tree.assert_same_structure(state, restored_state) + for expected, observed in zip( + jax.tree.flatten(state)[0], jax.tree.flatten(restored_state)[0] + ): + # Ensures the objects are different but the values are the same. + self.assertNotEqual(id(expected), id(observed)) + self.assertEqual(expected.shape, observed.shape) + self.assertEqual(expected.dtype, observed.dtype) + self.assertEqual(expected.sharding, observed.sharding) + np.testing.assert_allclose(observed, expected) + + # Ensures predictions are identical. + np.testing.assert_allclose(preds, preds_after_restoration) + def test_restore_keras_checkpoint(self): dummy_inputs = _create_dummy_inputs() bert_pretrainer = _create_model(jax.tree.map(jnp.shape, dummy_inputs)) @@ -188,6 +306,119 @@ def test_restore_keras_checkpoint(self): ) np.testing.assert_allclose(preds, preds_after_restoration) + def test_restore_keras_checkpoint_v3(self): + dummy_inputs = _create_dummy_inputs() + bert_pretrainer = _create_model(jax.tree.map(jnp.shape, dummy_inputs)) + preds = bert_pretrainer(dummy_inputs) + + checkpoint_dir = self.create_tempdir().full_path + checkpoint_manager = keras_utils.KerasOrbaxCheckpointManagerV3( + checkpoint_dir + ) + checkpoint_manager.save_model_variables(bert_pretrainer, epoch=1) + checkpoint_manager.close() + + restored_model = keras_utils.restore_keras_checkpoint(checkpoint_dir) + preds_after_restoration = restored_model(dummy_inputs) + + for expected, observed in zip( + [v.value for v in bert_pretrainer.variables], + [v.value for v in restored_model.variables], + ): + self.assertNotEqual(id(expected), id(observed)) + self.assertEqual(expected.shape, observed.shape) + self.assertEqual(expected.dtype, observed.dtype) + self.assertEqual(expected.sharding, observed.sharding) + np.testing.assert_allclose(observed, expected) + + self.assertDictEqual( + bert_pretrainer.get_config(), restored_model.get_config() + ) + np.testing.assert_allclose(preds, preds_after_restoration) + + def test_restore_shape_mismatch_fails(self): + dummy_inputs = _create_dummy_inputs() + bert_pretrainer = _create_model(jax.tree.map(jnp.shape, dummy_inputs)) + + checkpoint_dir = self.create_tempdir().full_path + checkpoint_manager = keras_utils.KerasOrbaxCheckpointManagerV3( + checkpoint_dir, max_to_keep=5 + ) + checkpoint_manager.save_model_variables(bert_pretrainer, 0) + checkpoint_manager.wait_until_finished() + + # Create a model with a different intermediate_dim, causing shape mismatches + different_pretrainer = keras_hub.models.BertMaskedLM( + backbone=keras_hub.models.BertBackbone( + vocabulary_size=2048, + num_layers=4, + num_heads=8, + hidden_dim=32, + intermediate_dim=128, # Different shape! + max_sequence_length=128, + num_segments=8, + dropout=0.1, + ) + ) + optimizer = keras.optimizers.Adam() + different_pretrainer.compile(optimizer=optimizer) + different_pretrainer.build(jax.tree.map(jnp.shape, dummy_inputs)) + different_pretrainer.optimizer.build( + different_pretrainer.trainable_variables + ) + + with self.assertRaises(ValueError): + checkpoint_manager.restore_model_variables(different_pretrainer, 0) + + checkpoint_manager.close() + + def test_restore_renamed_non_trainable_and_optimizer_variables(self): + checkpoint_dir = self.create_tempdir().full_path + checkpoint_manager = keras_utils.KerasOrbaxCheckpointManagerV3( + checkpoint_dir, max_to_keep=5 + ) + + # Save side: model with my_layer, optimizer named adam_save + model_save = MyTestModel(layer_name="my_layer", name="my_test_model") + model_save(np.zeros((1, 10))) + optimizer_save = keras.optimizers.Adam(name="adam_save") + model_save.compile(optimizer=optimizer_save) + optimizer_save.build(model_save.trainable_variables) + + # Initialize non-trainable and optimizer slot values to distinct values + for i, var in enumerate(model_save.non_trainable_variables): + var.assign(np.ones(var.shape) * (i + 7.0)) + for i, var in enumerate(model_save.optimizer.variables): + var.assign(np.ones(var.shape) * (i + 9.0)) + + checkpoint_manager.save_model_variables(model_save, 0) + checkpoint_manager.wait_until_finished() + + # Restore side: model with my_layer, optimizer named adam_restore + model_restore = MyTestModel(layer_name="my_layer", name="my_test_model") + model_restore(np.zeros((1, 10))) + optimizer_restore = keras.optimizers.Adam(name="adam_restore") + model_restore.compile(optimizer=optimizer_restore) + optimizer_restore.build(model_restore.trainable_variables) + + # Initialize to zeros + for var in model_restore.non_trainable_variables: + var.assign(np.zeros(var.shape)) + for var in model_restore.optimizer.variables: + var.assign(np.zeros(var.shape)) + + # Restore should succeed and map optimizer variables by index ordering + checkpoint_manager.restore_model_variables(model_restore, 0) + checkpoint_manager.close() + + # Verify non-trainable variables were restored + for i, var in enumerate(model_restore.non_trainable_variables): + np.testing.assert_allclose(var.value, np.ones(var.shape) * (i + 7.0)) + + # Verify optimizer variables were restored + for i, var in enumerate(model_restore.optimizer.variables): + np.testing.assert_allclose(var.value, np.ones(var.shape) * (i + 9.0)) + def test_load_keras_model_config(self): dummy_inputs = _create_dummy_inputs() bert_pretrainer = _create_model(jax.tree.map(jnp.shape, dummy_inputs)) @@ -275,7 +506,7 @@ def test_resolve_orbax_checkpoint_path_success( checkpoint_manager.close() test_path = os.path.join(checkpoint_dir, test_path_suffix) - resolved_path, resolved_epoch = keras_utils._resolve_orbax_checkpoint_path( + resolved_path, resolved_epoch = keras_utils.resolve_orbax_checkpoint_path( test_path, epoch=input_epoch ) @@ -291,13 +522,13 @@ def test_resolve_orbax_checkpoint_path_missing_step(self): checkpoint_manager.close() with self.assertRaisesRegex(ValueError, "Step 99 not found"): - keras_utils._resolve_orbax_checkpoint_path(test_dir, epoch=99) + keras_utils.resolve_orbax_checkpoint_path(test_dir, epoch=99) def test_resolve_orbax_checkpoint_path_empty_dir(self): test_dir = self.create_tempdir().full_path with self.assertRaisesRegex(FileNotFoundError, "No checkpoints found"): - keras_utils._resolve_orbax_checkpoint_path(test_dir, epoch=None) + keras_utils.resolve_orbax_checkpoint_path(test_dir, epoch=None) @parameterized.named_parameters( { @@ -384,6 +615,63 @@ def test_restore_keras_model_error_cases(self): with self.assertRaises(FileNotFoundError): keras_utils.restore_keras_model(bert_pretrainer, "not_found_dir") + def test_restore_keras_model_fails_on_v3_checkpoint(self): + dummy_inputs = _create_dummy_inputs() + bert_pretrainer = _create_model(jax.tree.map(jnp.shape, dummy_inputs)) + + checkpoint_dir = self.create_tempdir().full_path + checkpoint_manager = keras_utils.KerasOrbaxCheckpointManagerV3( + checkpoint_dir + ) + checkpoint_manager.save_model_variables(bert_pretrainer, epoch=1) + checkpoint_manager.wait_until_finished() + + with self.assertRaisesRegex( + ValueError, "is in V2/V3 format.*restore_keras_checkpoint" + ): + keras_utils.restore_keras_model(bert_pretrainer, checkpoint_dir, step=1) + + checkpoint_manager.close() + + def test_restore_keras_model_fails_on_v2_checkpoint(self): + dummy_inputs = _create_dummy_inputs() + bert_pretrainer = _create_model(jax.tree.map(jnp.shape, dummy_inputs)) + + checkpoint_dir = self.create_tempdir().full_path + checkpoint_manager = keras_utils.KerasOrbaxCheckpointManagerV2( + checkpoint_dir + ) + checkpoint_manager.save_model_variables(bert_pretrainer, epoch=1) + checkpoint_manager.wait_until_finished() + + with self.assertRaisesRegex( + ValueError, "is in V2/V3 format.*restore_keras_checkpoint" + ): + keras_utils.restore_keras_model(bert_pretrainer, checkpoint_dir, step=1) + + checkpoint_manager.close() + + def test_restore_keras_checkpoint_fails_on_v1_checkpoint(self): + dummy_inputs = _create_dummy_inputs() + bert_pretrainer = _create_model(jax.tree.map(jnp.shape, dummy_inputs)) + + checkpoint_dir = self.create_tempdir().full_path + checkpoint_manager = keras_utils.KerasOrbaxCheckpointManager(checkpoint_dir) + checkpoint_manager.save_model_variables(bert_pretrainer, epoch=1) + checkpoint_manager.wait_until_finished() + + # We expect it to fail because restore_keras_checkpoint expects V2/V3 + # structure. We want it to fail loudly and suggest using + # restore_keras_model. + with self.assertRaisesRegex( + ValueError, "is in V1 format.*restore_keras_model" + ): + keras_utils.restore_keras_checkpoint( + checkpoint_dir, model=bert_pretrainer + ) + + checkpoint_manager.close() + @parameterized.named_parameters( { "testcase_name": "restore_with_checkpointer", @@ -602,5 +890,250 @@ def test_epoch_orbax_checkpoint_and_restore_callback_saves_at_epoch_0( self.assertTrue(os.path.exists(os.path.join(checkpoint_dir, "0"))) +class KerasOrbaxCheckpointUtilsTest(absltest.TestCase): + + def setUp(self): + super().setUp() + self.checkpoint_dir = self.create_tempdir().full_path + + def test_is_v3_checkpoint_true(self): + manager = keras_utils.KerasOrbaxCheckpointManagerV3( + checkpoint_dir=self.checkpoint_dir, + max_to_keep=1, + save_interval_epochs=1, + ) + model = keras.Sequential([keras.layers.Dense(1)]) + model.compile(optimizer="adam") + model.build((1, 1)) + model.optimizer.build(model.trainable_variables) + + manager.save_model_variables(model, epoch=1) + manager.wait_until_finished() + + self.assertTrue(keras_utils.is_v3_checkpoint(self.checkpoint_dir)) + self.assertTrue(keras_utils.is_v3_checkpoint(self.checkpoint_dir, epoch=1)) + + def test_is_v3_checkpoint_false_for_v1(self): + manager = keras_utils.KerasOrbaxCheckpointManager( + checkpoint_dir=self.checkpoint_dir, + max_to_keep=1, + save_interval_epochs=1, + ) + model = keras.Sequential([keras.layers.Dense(1)]) + model.compile(optimizer="adam") + model.build((1, 1)) + model.optimizer.build(model.trainable_variables) + + manager.save_model_variables(model, epoch=1) + manager.wait_until_finished() + + self.assertFalse(keras_utils.is_v3_checkpoint(self.checkpoint_dir)) + + def test_is_v3_checkpoint_false_for_missing(self): + self.assertFalse(keras_utils.is_v3_checkpoint(self.checkpoint_dir)) + self.assertFalse(keras_utils.is_v3_checkpoint(self.checkpoint_dir, epoch=1)) + + def test_restore_partial_checkpoint_success(self): + manager = keras_utils.KerasOrbaxCheckpointManagerV3( + checkpoint_dir=self.checkpoint_dir, + max_to_keep=1, + save_interval_epochs=1, + ) + model = keras.Sequential( + [ + keras.layers.Dense(2, name="dense1"), + keras.layers.Dense(1, name="dense2"), + ], + name="my_model", + ) + model.compile(optimizer="adam") + model.build((1, 1)) + model.optimizer.build(model.trainable_variables) + + model.layers[0].kernel.assign([[1.0, 2.0]]) + model.layers[1].kernel.assign([[3.0], [4.0]]) + + manager.save_model_variables(model, epoch=1) + manager.wait_until_finished() + + new_model = keras.Sequential( + [ + keras.layers.Dense(2, name="dense1"), + keras.layers.Dense(1, name="dense2"), + ], + name="my_model", + ) + new_model.build((1, 1)) + + new_model.layers[0].kernel.assign([[0.0, 0.0]]) + new_model.layers[1].kernel.assign([[0.0], [0.0]]) + + partial_vars = { + keras_utils.TRAINABLE_VARIABLES_KEY: { + new_model.layers[0].kernel.path: new_model.layers[0].kernel, + new_model.layers[0].bias.path: new_model.layers[0].bias, + } + } + + restored_state = keras_utils.restore_partial_checkpoint( + self.checkpoint_dir, partial_vars, epoch=1 + ) + + for key, var_dict in partial_vars.items(): + for path, var in var_dict.items(): + var.assign(restored_state[key][path]) + + np.testing.assert_allclose(new_model.layers[0].kernel.value, [[1.0, 2.0]]) + np.testing.assert_allclose(new_model.layers[1].kernel.value, [[0.0], [0.0]]) + + def test_restore_partial_checkpoint_non_overlapping_architectures(self): + manager = keras_utils.KerasOrbaxCheckpointManagerV3( + checkpoint_dir=self.checkpoint_dir, + max_to_keep=1, + save_interval_epochs=1, + ) + # Model in checkpoint has dense1 and dense2 + model = keras.Sequential( + [ + keras.layers.Dense(2, name="dense1"), + keras.layers.Dense(1, name="dense2"), + ], + name="my_model", + ) + model.compile(optimizer="adam") + model.build((1, 1)) + model.optimizer.build(model.trainable_variables) + + model.layers[0].kernel.assign([[1.0, 2.0]]) + model.layers[1].kernel.assign([[3.0], [4.0]]) + + manager.save_model_variables(model, epoch=1) + manager.wait_until_finished() + + # New model has dense1 and dense3 (dense2 is missing, dense3 is new) + new_model = keras.Sequential( + [ + keras.layers.Dense(2, name="dense1"), + keras.layers.Dense(3, name="dense3"), + ], + name="my_model", + ) + new_model.build((1, 1)) + + new_model.layers[0].kernel.assign([[0.0, 0.0]]) + new_model.layers[1].kernel.assign([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]) + + # We only restore dense1 from the checkpoint + partial_vars = { + keras_utils.TRAINABLE_VARIABLES_KEY: { + new_model.layers[0].kernel.path: new_model.layers[0].kernel, + new_model.layers[0].bias.path: new_model.layers[0].bias, + } + } + + restored_state = keras_utils.restore_partial_checkpoint( + self.checkpoint_dir, partial_vars, epoch=1 + ) + + for key, var_dict in partial_vars.items(): + for path, var in var_dict.items(): + var.assign(restored_state[key][path]) + + # dense1 should be restored + np.testing.assert_allclose(new_model.layers[0].kernel.value, [[1.0, 2.0]]) + # dense3 should remain unchanged (zeros) + np.testing.assert_allclose( + new_model.layers[1].kernel.value, [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]] + ) + + def test_restore_partial_checkpoint_nested_layer_mismatch(self): + manager = keras_utils.KerasOrbaxCheckpointManagerV3( + checkpoint_dir=self.checkpoint_dir, + max_to_keep=1, + save_interval_epochs=1, + ) + + class ParentLayer(keras.layers.Layer): + + def __init__(self, nested_name, **kwargs): + super().__init__(**kwargs) + self.nested = keras.layers.Dense(2, name=nested_name) + + def call(self, x): + return self.nested(x) + + # Model A has nested layer named "dense_a" under "parent" + model_a = keras.Sequential( + [ParentLayer(nested_name="dense_a", name="parent")], name="model" + ) + model_a.compile(optimizer="adam") + model_a.build((1, 1)) + model_a.optimizer.build(model_a.trainable_variables) + manager.save_model_variables(model_a, epoch=1) + manager.wait_until_finished() + + # Model B has nested layer named "dense_b" under "parent" + model_b = keras.Sequential( + [ParentLayer(nested_name="dense_b", name="parent")], name="model" + ) + model_b.build((1, 1)) + + # Try to restore model_b variables (which have paths like + # 'parent/dense_b/kernel') + partial_vars = { + keras_utils.TRAINABLE_VARIABLES_KEY: { + model_b.layers[0].nested.kernel.path: ( + model_b.layers[0].nested.kernel + ), + } + } + + # This should fail because 'parent/dense_b/kernel' is not in the checkpoint + with self.assertRaisesRegex(ValueError, "Missing paths in checkpoint"): + keras_utils.restore_partial_checkpoint( + self.checkpoint_dir, partial_vars, epoch=1 + ) + + def test_restore_partial_checkpoint_invalid_keys(self): + manager = keras_utils.KerasOrbaxCheckpointManagerV3( + checkpoint_dir=self.checkpoint_dir, + max_to_keep=1, + save_interval_epochs=1, + ) + model = keras.Sequential([keras.layers.Dense(1)]) + model.compile(optimizer="adam") + model.build((1, 1)) + model.optimizer.build(model.trainable_variables) + + manager.save_model_variables(model, epoch=1) + manager.wait_until_finished() + + partial_vars = { + keras_utils.NON_TRAINABLE_VARIABLES_KEY: { + "some_path": model.trainable_variables[0] + } + } + with self.assertRaisesRegex( + ValueError, + "Partial restoration is only supported for trainable variables", + ): + keras_utils.restore_partial_checkpoint( + self.checkpoint_dir, partial_vars, epoch=1 + ) + + partial_vars = { + keras_utils.OPTIMIZER_VARIABLES_KEY: { + "some_path": model.trainable_variables[0] + } + } + with self.assertRaisesRegex( + ValueError, + "Partial restoration is only supported for trainable variables", + ): + keras_utils.restore_partial_checkpoint( + self.checkpoint_dir, partial_vars, epoch=1 + ) + + if __name__ == "__main__": absltest.main()