From 410abb8c3aa66055ab93f643b7186d85476e7477 Mon Sep 17 00:00:00 2001 From: RecML authors Date: Tue, 30 Jun 2026 15:38:16 -0700 Subject: [PATCH] Internal change PiperOrigin-RevId: 940694716 --- recml/core/training/keras_trainer.py | 4 +- recml/core/utils/keras_utils.py | 678 +++++++++++++++++++++++++-- recml/core/utils/keras_utils_test.py | 207 +++++++- 3 files changed, 847 insertions(+), 42 deletions(-) diff --git a/recml/core/training/keras_trainer.py b/recml/core/training/keras_trainer.py index c8e121e..4412edc 100644 --- a/recml/core/training/keras_trainer.py +++ b/recml/core/training/keras_trainer.py @@ -122,7 +122,7 @@ 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, + legacy_checkpoint_format: bool = False, ): """Initializes the instance.""" @@ -166,7 +166,7 @@ def train_callbacks(self) -> list[keras.callbacks.Callback]: save_interval_epochs=self._checkpoint_save_interval_epochs, ) else: - checkpoint_manager = keras_utils.KerasOrbaxCheckpointManagerV2( + checkpoint_manager = keras_utils.KerasOrbaxCheckpointManagerV3( checkpoint_dir=self._checkpoint_dir, max_to_keep=self._max_checkpoints_to_keep, save_interval_epochs=self._checkpoint_save_interval_epochs, diff --git a/recml/core/utils/keras_utils.py b/recml/core/utils/keras_utils.py index 1109a8f..51b2c9c 100644 --- a/recml/core/utils/keras_utils.py +++ b/recml/core/utils/keras_utils.py @@ -21,6 +21,7 @@ from typing import Any from absl import logging +from etils import epath import jax import keras import orbax.checkpoint as ocp @@ -172,7 +173,276 @@ 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.""" + + 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", + "non_trainable_paths", + "optimizer_paths", + ), + 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 = {} + duplicates = [] + for v in model.trainable_variables: + if v.path in trainable_variables: + duplicates.append(v.path) + else: + trainable_variables[v.path] = v + if duplicates: + raise ValueError( + "Duplicate variable paths detected in trainable_variables. Ensure" + " unique layer names (e.g. name_layers=True). Duplicates: " + f"{duplicates}" + ) + non_trainable_variables = {v.path: v for v in model.non_trainable_variables} + optimizer_variables = {v.path: v for v in model.optimizer.variables} + + # 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": ocp.args.JsonSave({"version": 3}), + "non_trainable_paths": ocp.args.JsonSave(non_trainable_paths), + "optimizer_paths": 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 = {v.path: v for v in model.trainable_variables} + if len(trainable_variables) != len(model.trainable_variables): + raise ValueError( + "Duplicate variable paths detected in trainable_variables. Ensure" + " unique layer names (e.g. name_layers=True)." + ) + non_trainable_variables = {v.path: v for v in model.non_trainable_variables} + + optimizer_variables = {v.path: v for v in model.optimizer.variables} + + 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) + + # For V3, Orbax StandardRestore is strict. We need to make sure + # abstract_variables only contains keys that actually exist in the + # checkpoint, or use transforms to map them. + state_transforms = {} + metadata_handler = ocp.handlers.PyTreeCheckpointHandler() + step_path = os.path.join(self.directory, str(epoch)) + + state_checkpoint_path = epath.Path(step_path) / STATE_CHECKPOINT_KEY + saved_state_metadata = metadata_handler.metadata(state_checkpoint_path) + + has_paths = gfile.Exists(os.path.join(step_path, "non_trainable_paths")) + non_trainable_paths = None + optimizer_paths = None + if has_paths: + paths_checkpointer = ocp.Checkpointer( + ocp.CompositeCheckpointHandler( + non_trainable_paths=ocp.handlers.JsonCheckpointHandler(), + optimizer_paths=ocp.handlers.JsonCheckpointHandler(), + ) + ) + restored_paths = paths_checkpointer.restore( + step_path, + args=ocp.args.Composite( + non_trainable_paths=ocp.args.JsonRestore(), + optimizer_paths=ocp.args.JsonRestore(), + ), + ) + non_trainable_paths = restored_paths["non_trainable_paths"]["paths"] + optimizer_paths = restored_paths["optimizer_paths"]["paths"] + paths_checkpointer.close() + + filtered_abstract_variables = {} + for key in [ + TRAINABLE_VARIABLES_KEY, + NON_TRAINABLE_VARIABLES_KEY, + OPTIMIZER_VARIABLES_KEY, + ]: + if key in abstract_variables and key in saved_state_metadata: + filtered_abstract_variables[key] = {} + state_transforms[key] = {} + + if key in [NON_TRAINABLE_VARIABLES_KEY, OPTIMIZER_VARIABLES_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 + ) + logging.info("TARGET paths for key %s: %s", key, target_paths) + logging.info("RESTORED paths for key %s: %s", key, source_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_variables[key][target_path] + if i < len(source_paths): + source_path = source_paths[i] + filtered_abstract_variables[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_variables[key].items(): + if target_path in saved_state_metadata[key]: + filtered_abstract_variables[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_variables: + logging.warning("Key %s not found in checkpoint metadata", key) + + abstract_variables = filtered_abstract_variables + + 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. @@ -226,6 +496,36 @@ def _resolve_orbax_checkpoint_path( return os.fspath(checkpoint_path), epoch +def _is_v3_checkpoint_path(checkpoint_path: str) -> bool: + """Checks if a resolved checkpoint path is in V3 format.""" + version_checkpointer = ocp.Checkpointer( + ocp.CompositeCheckpointHandler( + **{"format_version": ocp.handlers.JsonCheckpointHandler()} + ) + ) + is_v3 = False + try: + version_info = version_checkpointer.restore( + checkpoint_path, + args=ocp.args.Composite(**{"format_version": ocp.args.JsonRestore()}), + )["format_version"] + if version_info.get("version") == 3: + is_v3 = True + except Exception: # pylint: disable=broad-exception-caught + pass + 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 Exception: # pylint: disable=broad-exception-caught + return False + + def restore_keras_checkpoint( checkpoint_dir: str, *, @@ -281,7 +581,7 @@ 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 model is None: cfg = {**load_keras_model_config(checkpoint_dir, epoch=epoch)} @@ -304,18 +604,38 @@ 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." - ) - variables[OPTIMIZER_VARIABLES_KEY] = model.optimizer.variables + is_v3 = _is_v3_checkpoint_path(checkpoint_path) + + if is_v3: + variables = { + TRAINABLE_VARIABLES_KEY: {v.path: v for v in model.trainable_variables}, + NON_TRAINABLE_VARIABLES_KEY: { + v.path: v for v in 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] = { + v.path: v for v in 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. @@ -324,13 +644,298 @@ def restore_keras_checkpoint( # the same sharding as the current variables. abstract_state = jax.tree.map(_to_shape_dtype_struct, variables) + state_transforms = {} + # For V3, Orbax StandardRestore is strict. We need to make sure abstract_state + # only contains keys that actually exist in the checkpoint, or use + # transforms to map them. + if is_v3: + # Use PyTreeCheckpointHandler to read the structure/metadata of the + # checkpoint without fully restoring it. + metadata_handler = ocp.handlers.PyTreeCheckpointHandler() + # This returns a PyTree of metadata, often matching the saved structure + 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") + ) + non_trainable_paths = None + optimizer_paths = None + if has_paths: + paths_checkpointer = ocp.Checkpointer( + ocp.CompositeCheckpointHandler( + non_trainable_paths=ocp.handlers.JsonCheckpointHandler(), + optimizer_paths=ocp.handlers.JsonCheckpointHandler(), + ) + ) + restored_paths = paths_checkpointer.restore( + checkpoint_path, + args=ocp.args.Composite( + non_trainable_paths=ocp.args.JsonRestore(), + optimizer_paths=ocp.args.JsonRestore(), + ), + ) + non_trainable_paths = restored_paths["non_trainable_paths"]["paths"] + optimizer_paths = restored_paths["optimizer_paths"]["paths"] + paths_checkpointer.close() + + filtered_abstract_state = {} + for key in [ + TRAINABLE_VARIABLES_KEY, + NON_TRAINABLE_VARIABLES_KEY, + OPTIMIZER_VARIABLES_KEY, + ]: + 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]: + 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) + + abstract_state = filtered_abstract_state + # 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 not is_v3: + 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() + + # TODO(zixiangzhou): Unflatten the variables based on index here. + if is_v3: + 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( + "Failed to restore optimizer variables for paths:" + f" {still_missing}" + ) + 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." + ) + + 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 + + state_transforms = {} + metadata_handler = ocp.handlers.PyTreeCheckpointHandler() + state_checkpoint_path = epath.Path(checkpoint_path) / STATE_CHECKPOINT_KEY + saved_state_metadata = metadata_handler.metadata(state_checkpoint_path) + + filtered_abstract_state = {} + key = TRAINABLE_VARIABLES_KEY + if key in abstract_state and key in saved_state_metadata: + filtered_abstract_state[key] = {} + state_transforms[key] = {} + 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) + + abstract_state = filtered_abstract_state - # 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( @@ -343,28 +948,25 @@ 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 + 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( @@ -376,7 +978,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( @@ -530,7 +1132,9 @@ class EpochOrbaxCheckpointAndRestoreCallback(keras.callbacks.Callback): def __init__( self, checkpoint_manager: ( - KerasOrbaxCheckpointManager | KerasOrbaxCheckpointManagerV2 + KerasOrbaxCheckpointManager + | KerasOrbaxCheckpointManagerV2 + | KerasOrbaxCheckpointManagerV3 ), marker_path: str | None = None, ): diff --git a/recml/core/utils/keras_utils_test.py b/recml/core/utils/keras_utils_test.py index 421f257..b98065c 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,89 @@ def test_restore_keras_checkpoint(self): ) 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 +476,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 +492,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( {