Skip to content

Internal change #206

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Internal change #206

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 14 additions & 10 deletions recml/core/utils/keras_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -309,15 +309,16 @@ class KerasOrbaxCheckpointManagerV3(ocp.CheckpointManager):
def __init__(
self,
checkpoint_dir: str,
max_to_keep: int = 5,
max_to_keep: int | None = 5,
save_interval_epochs: int = 1,
choose_store_cell: bool = True,
):
"""Initializes a KerasOrbaxCheckpointManagerV3.

Args:
checkpoint_dir: The directory to save checkpoints to.
max_to_keep: The maximum number of checkpoints to keep.
max_to_keep: The maximum number of checkpoints to keep, or None to keep
every checkpoint.
save_interval_epochs: The interval (in epochs) to save checkpoints.
choose_store_cell: Whether to dynamically select the CNS2 store cell.
"""
Expand Down Expand Up @@ -561,7 +562,7 @@ def _is_v3_checkpoint_path(checkpoint_path: str) -> bool:
version_checkpointer.close()


def _detect_checkpoint_version(checkpoint_path: str) -> CheckpointVersion:
def detect_checkpoint_version(checkpoint_path: str) -> CheckpointVersion:
"""Detects the version of the checkpoint at the given path.

Detection order and discriminator criteria:
Expand Down Expand Up @@ -592,7 +593,7 @@ def _detect_checkpoint_version(checkpoint_path: str) -> CheckpointVersion:
def is_v3_checkpoint(checkpoint_dir: str, epoch: int | None = None) -> bool:
"""Checks if a checkpoint is in V3 format."""
checkpoint_path, _ = resolve_orbax_checkpoint_path(checkpoint_dir, epoch)
return _detect_checkpoint_version(checkpoint_path) == CheckpointVersion.V3
return detect_checkpoint_version(checkpoint_path) == CheckpointVersion.V3


def _validate_v3_checkpoint(
Expand Down Expand Up @@ -823,9 +824,10 @@ def _prepare_v3_restore(
checkpointer.close()

# Load index path metadata for non-trainable and optimizer variables if
# needed.
# present in `abstract_state` (partial restore omits these keys, whereas full
# restore includes them even when the dict is empty `{}`).
non_trainable_paths = None
if model is not None and abstract_state.get(NON_TRAINABLE_VARIABLES_KEY):
if model is not None and NON_TRAINABLE_VARIABLES_KEY in abstract_state:
nt_path = epath.Path(checkpoint_path) / NON_TRAINABLE_PATHS_KEY
checkpointer = ocp.Checkpointer(ocp.handlers.JsonCheckpointHandler())
try:
Expand All @@ -839,7 +841,7 @@ def _prepare_v3_restore(
if (
model is not None
and restore_optimizer_vars
and abstract_state.get(OPTIMIZER_VARIABLES_KEY)
and OPTIMIZER_VARIABLES_KEY in abstract_state
):
opt_path = epath.Path(checkpoint_path) / OPTIMIZER_PATHS_KEY
checkpointer = ocp.Checkpointer(ocp.handlers.JsonCheckpointHandler())
Expand All @@ -858,7 +860,9 @@ def _prepare_v3_restore(
state_transforms = {}

for key in keys:
if not abstract_state.get(key):
# Keep empty dicts (e.g. `{}` when a model has no non-trainable variables)
# in `filtered_abstract_state` so the restored PyTree structure matches.
if key not in abstract_state:
continue

filtered_abstract_state[key] = {}
Expand Down Expand Up @@ -1132,7 +1136,7 @@ def restore_keras_checkpoint(

checkpoint_path, epoch = resolve_orbax_checkpoint_path(checkpoint_dir, epoch)

version = _detect_checkpoint_version(checkpoint_path)
version = detect_checkpoint_version(checkpoint_path)
if version == CheckpointVersion.V1:
raise ValueError(
f"The checkpoint in {checkpoint_dir} is in V1 format (list-based)"
Expand Down Expand Up @@ -1298,7 +1302,7 @@ def restore_partial_checkpoint(
"""
checkpoint_path, _ = resolve_orbax_checkpoint_path(checkpoint_dir, epoch)

version = _detect_checkpoint_version(checkpoint_path)
version = detect_checkpoint_version(checkpoint_path)
if version != CheckpointVersion.V3:
raise ValueError(
"restore_partial_checkpoint only supports V3 (dictionary-based)"
Expand Down
32 changes: 32 additions & 0 deletions recml/core/utils/keras_utils_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -276,6 +276,38 @@ def test_keras_orbax_checkpointer_v3(
# Ensures predictions are identical.
np.testing.assert_allclose(preds, preds_after_restoration)

def test_keras_orbax_checkpointer_v3_no_non_trainable_variables(self):
model = keras.Sequential(
[keras.layers.Dense(5, name="dense")], name="model"
)
model.build((2, 10))
model.compile(optimizer=keras.optimizers.Adam(0.01), loss="mse")
model.optimizer.build(model.trainable_variables)
self.assertEmpty(model.non_trainable_variables)

checkpoint_dir = self.create_tempdir().full_path
checkpoint_manager = keras_utils.KerasOrbaxCheckpointManagerV3(
checkpoint_dir, max_to_keep=None
)
checkpoint_manager.save_model_variables(model, 0)
checkpoint_manager.wait_until_finished()

new_model = keras.Sequential(
[keras.layers.Dense(5, name="dense")], name="model"
)
new_model.build((2, 10))
new_model.compile(optimizer=keras.optimizers.Adam(0.01), loss="mse")
new_model.optimizer.build(new_model.trainable_variables)

checkpoint_manager.restore_model_variables(new_model, 0)
checkpoint_manager.close()

for expected, observed in zip(
[v.value for v in model.trainable_variables],
[v.value for v in new_model.trainable_variables],
):
np.testing.assert_allclose(observed, expected)

def test_restore_keras_checkpoint(self):
dummy_inputs = _create_dummy_inputs()
bert_pretrainer = _create_model(jax.tree.map(jnp.shape, dummy_inputs))
Expand Down
Loading