Skip to content

Internal change #211

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 #211

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
177 changes: 166 additions & 11 deletions recml/core/utils/keras_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import jax
import keras
import orbax.checkpoint as ocp
from recml.core.utils import sparsecore_checkpoint_utils
import tensorflow as tf


Expand All @@ -38,6 +39,9 @@
NON_TRAINABLE_PATHS_KEY = "non_trainable_paths"
OPTIMIZER_PATHS_KEY = "optimizer_paths"
SHARED_VARIABLE_MAP_KEY = "shared_variable_map"
# How the SparseCore stacked tables were sharded, only saved for models with
# SparseCore tables (see `sparsecore_checkpoint_utils.sparsecore_topology`).
SPARSECORE_TOPOLOGY_KEY = sparsecore_checkpoint_utils.SPARSECORE_TOPOLOGY_KEY
ORBAX_CHECKPOINT_DEFAULT_KEY = "default"


Expand Down Expand Up @@ -241,6 +245,7 @@ def save_model_variables(
args=ocp.args.Composite(**{
STATE_CHECKPOINT_KEY: ocp.args.StandardSave(state),
CONFIG_CHECKPOINT_KEY: ocp.args.JsonSave(config),
**sparsecore_checkpoint_utils.sparsecore_topology_save_args(model),
}),
metrics=logs,
)
Expand All @@ -263,6 +268,11 @@ def restore_model_variables(self, model: keras.Model, epoch: int):
# TODO(zixiangzhou): Update variables to use a nested dictionary and index
# map instead of flattened list.

checkpoint_path, _ = resolve_orbax_checkpoint_path(
str(self.directory), epoch
)
sparsecore_restores = _prepare_sparsecore_restore(model, checkpoint_path)

# Construct abstract variables to ensure the checkpoint is restored with
# the same sharding as the current variables. This is so we can delete the
# variables from device memory to reduce peak memory usage.
Expand Down Expand Up @@ -291,6 +301,8 @@ def restore_model_variables(self, model: keras.Model, epoch: int):
):
var._value = restored_var # pylint: disable=protected-access

sparsecore_checkpoint_utils.finish_sparsecore_restore(sparsecore_restores)


class KerasOrbaxCheckpointManagerV3(ocp.CheckpointManager):
"""An Orbax checkpoint manager for Keras 3 with dictionary state.
Expand Down Expand Up @@ -335,6 +347,7 @@ def __init__(
NON_TRAINABLE_PATHS_KEY,
OPTIMIZER_PATHS_KEY,
SHARED_VARIABLE_MAP_KEY,
SPARSECORE_TOPOLOGY_KEY,
),
options=ocp.CheckpointManagerOptions(
save_interval_steps=save_interval_epochs,
Expand Down Expand Up @@ -403,6 +416,7 @@ def save_model_variables(
NON_TRAINABLE_PATHS_KEY: ocp.args.JsonSave(non_trainable_paths),
OPTIMIZER_PATHS_KEY: ocp.args.JsonSave(optimizer_paths),
SHARED_VARIABLE_MAP_KEY: ocp.args.JsonSave(shared_variable_map),
**sparsecore_checkpoint_utils.sparsecore_topology_save_args(model),
}),
metrics=logs,
)
Expand Down Expand Up @@ -432,6 +446,13 @@ def restore_model_variables(self, model: keras.Model, epoch: int):
OPTIMIZER_VARIABLES_KEY: optimizer_variables,
}

step_path = os.path.join(self.directory, str(epoch))
# Read once for the SparseCore restore and `_prepare_v3_restore`.
saved_state_metadata = _validate_v3_checkpoint(step_path)
sparsecore_restores = _prepare_sparsecore_restore(
model, step_path, saved_state_metadata
)

# 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)
Expand All @@ -441,12 +462,12 @@ def restore_model_variables(self, model: keras.Model, epoch: int):

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,
saved_state_metadata=saved_state_metadata,
)

restored_items = self.restore(
Expand Down Expand Up @@ -479,6 +500,8 @@ def restore_model_variables(self, model: keras.Model, epoch: int):
for path, var in var_dict.items():
var._value = restored_var_dict[path] # pylint: disable=protected-access

sparsecore_checkpoint_utils.finish_sparsecore_restore(sparsecore_restores)


def resolve_orbax_checkpoint_path(
checkpoint_dir: str, epoch: int | None = None
Expand Down Expand Up @@ -595,6 +618,25 @@ def is_v3_checkpoint(checkpoint_dir: str, epoch: int | None = None) -> bool:
return _detect_checkpoint_version(checkpoint_path) == CheckpointVersion.V3


def _load_shared_variable_map(checkpoint_path: str) -> dict[str, str]:
"""Returns the shared variable map of a V3 checkpoint.

Args:
checkpoint_path: The resolved checkpoint directory path of one step.

Returns:
Map from the alias path of each shared variable to its canonical path.
"""
map_path = epath.Path(checkpoint_path) / SHARED_VARIABLE_MAP_KEY
checkpointer = ocp.Checkpointer(ocp.handlers.JsonCheckpointHandler())
try:
return checkpointer.restore(
os.fspath(map_path), args=ocp.args.JsonRestore()
)[SHARED_VARIABLE_MAP_KEY]
finally:
checkpointer.close()


def _validate_v3_checkpoint(
checkpoint_path: str,
saved_state_metadata: Any | None = None,
Expand Down Expand Up @@ -783,6 +825,7 @@ def _prepare_v3_restore(
restore_optimizer_vars: bool = False,
transforms: Mapping[str, Any] | None = None,
prefix_map: Mapping[str, str] | None = None,
saved_state_metadata: Any | None = None,
) -> tuple[Mapping[str, Any], Mapping[str, Any]]:
"""Prepares the abstract state and constructs transforms for V3 restore.

Expand All @@ -804,23 +847,20 @@ def _prepare_v3_restore(
transforms: An optional mapping of custom transforms for variable mapping.
prefix_map: An optional map from target path prefix to source path prefix
for trainable variables. See `restore_partial_checkpoint`.
saved_state_metadata: Optional state metadata already read from the
checkpoint (see `_validate_v3_checkpoint`), so it isn't read again.

Returns:
A tuple of (filtered_abstract_state, state_transforms) to be passed to
the Orbax restore call.
"""
# Validate that underlying checkpoint is healthy and retrieve state metadata.
saved_state_metadata = _validate_v3_checkpoint(checkpoint_path)
saved_state_metadata = _validate_v3_checkpoint(
checkpoint_path, saved_state_metadata
)

# The shared variable map is a required V3 item, checked above.
map_path = epath.Path(checkpoint_path) / SHARED_VARIABLE_MAP_KEY
checkpointer = ocp.Checkpointer(ocp.handlers.JsonCheckpointHandler())
try:
shared_var_map = checkpointer.restore(
os.fspath(map_path), args=ocp.args.JsonRestore()
)[SHARED_VARIABLE_MAP_KEY]
finally:
checkpointer.close()
shared_var_map = _load_shared_variable_map(checkpoint_path)

# Load index path metadata for non-trainable and optimizer variables if
# needed.
Expand Down Expand Up @@ -1199,14 +1239,26 @@ def restore_keras_checkpoint(
# TODO(zixiangzhou): Update variables to use a nested dictionary and index map
# instead of flattened list.

# Read once for the SparseCore restore and `_prepare_v3_restore`.
saved_state_metadata = (
_validate_v3_checkpoint(checkpoint_path) if is_v3 else None
)
sparsecore_restores = _prepare_sparsecore_restore(
model, checkpoint_path, saved_state_metadata
)

# Construct abstract variables to ensure the checkpoint is restored with
# 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
checkpoint_path,
abstract_state,
model,
restore_optimizer_vars,
saved_state_metadata=saved_state_metadata,
)

# Delete the variables from device memory to reduce peak memory usage.
Expand Down Expand Up @@ -1239,6 +1291,8 @@ def restore_keras_checkpoint(
):
var._value = restored_var # pylint: disable=protected-access

sparsecore_checkpoint_utils.finish_sparsecore_restore(sparsecore_restores)

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:
Expand All @@ -1247,6 +1301,107 @@ def restore_keras_checkpoint(
return model


def _saved_trainable_metadata(
model: keras.Model,
checkpoint_path: str,
saved_state_metadata: Any | None = None,
) -> dict[str, Any]:
"""Returns the checkpoint metadata of the model's trainable variables by path.

Matches the saved trainable variables like the restore does: by path for V3
checkpoints (following the aliases of shared variables, see
`_find_stored_path`), and by order for V2 checkpoints.

Args:
model: The model being restored.
checkpoint_path: The resolved checkpoint directory path of one step.
saved_state_metadata: Optional state metadata already read from the
checkpoint, so it isn't read again.

Returns:
Variable path to its Orbax `ArrayMetadata`, or an empty dict if the
checkpoint has no Orbax state item.
"""
try:
if saved_state_metadata is None:
state_dir = epath.Path(checkpoint_path) / STATE_CHECKPOINT_KEY
if not state_dir.exists():
return {}
saved_state_metadata = ocp.handlers.PyTreeCheckpointHandler().metadata(
state_dir
)
saved_trainable_meta = saved_state_metadata.get(TRAINABLE_VARIABLES_KEY)
shared_var_map = {}
if (
isinstance(saved_trainable_meta, Mapping)
and (epath.Path(checkpoint_path) / SHARED_VARIABLE_MAP_KEY).exists()
):
shared_var_map = _load_shared_variable_map(checkpoint_path)
except Exception:
# Without the metadata, the SparseCore tables can't be resharded, and
# restoring them as is would fail with a shape mismatch or silently load
# rows sharded for another topology.
logging.exception(
"Failed to read the metadata of the checkpoint in %s to restore the"
" SparseCore tables across TPU topologies.",
checkpoint_path,
)
raise

if isinstance(saved_trainable_meta, Mapping): # V3: keyed by variable path.
meta_by_path = {}
for var in model.trainable_variables:
stored_path = _find_stored_path(
var.path, saved_trainable_meta, shared_var_map
)
if stored_path is not None:
meta_by_path[var.path] = saved_trainable_meta[stored_path]
return meta_by_path
if isinstance(saved_trainable_meta, Sequence) and len(
saved_trainable_meta
) == len(
model.trainable_variables
): # V2: in variable order.
return {
var.path: meta
for var, meta in zip(model.trainable_variables, saved_trainable_meta)
}
return {}


def _prepare_sparsecore_restore(
model: keras.Model,
checkpoint_path: str,
saved_state_metadata: Any | None = None,
) -> list[sparsecore_checkpoint_utils.StackRestore]:
"""Prepares the SparseCore tables of `model`, if any, for the restore.

Call this before building the Orbax restore target and
`sparsecore_checkpoint_utils.finish_sparsecore_restore` after the restore,
so that SparseCore tables saved on another TPU topology are resharded (see
`sparsecore_checkpoint_utils.prepare_sparsecore_restore`).

Args:
model: The model being restored.
checkpoint_path: The resolved checkpoint directory path of one step.
saved_state_metadata: Optional state metadata already read from the
checkpoint, so it isn't read again.

Returns:
The SparseCore stacks to reshard after the restore.
"""
if not sparsecore_checkpoint_utils.has_sparsecore_tables(model):
return []
saved_trainable_metadata = _saved_trainable_metadata(
model, checkpoint_path, saved_state_metadata
)
if not saved_trainable_metadata:
return []
return sparsecore_checkpoint_utils.prepare_sparsecore_restore(
model, checkpoint_path, saved_trainable_metadata
)


def restore_partial_checkpoint(
checkpoint_dir: str,
partial_variables: Mapping[str, Any],
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 @@ -163,6 +163,14 @@ def test_keras_orbax_checkpointer_v2(
)
checkpoint_manager.save_model_variables(bert_pretrainer, 0)
checkpoint_manager.wait_until_finished()
# Only models with SparseCore tables save their SparseCore topology.
self.assertFalse(
os.path.exists(
os.path.join(
checkpoint_dir, "0", keras_utils.SPARSECORE_TOPOLOGY_KEY
)
)
)

preds = bert_pretrainer(dummy_inputs)

Expand Down Expand Up @@ -242,6 +250,14 @@ def test_keras_orbax_checkpointer_v3(
)
checkpoint_manager.save_model_variables(bert_pretrainer, 0)
checkpoint_manager.wait_until_finished()
# Only models with SparseCore tables save their SparseCore topology.
self.assertFalse(
os.path.exists(
os.path.join(
checkpoint_dir, "0", keras_utils.SPARSECORE_TOPOLOGY_KEY
)
)
)

preds = bert_pretrainer(dummy_inputs)

Expand Down Expand Up @@ -287,6 +303,14 @@ def test_restore_keras_checkpoint(self):
)
checkpoint_manager.save_model_variables(bert_pretrainer, epoch=1)
checkpoint_manager.close()
# Only models with SparseCore tables save their SparseCore topology.
self.assertFalse(
os.path.exists(
os.path.join(
checkpoint_dir, "1", keras_utils.SPARSECORE_TOPOLOGY_KEY
)
)
)

restored_model = keras_utils.restore_keras_checkpoint(checkpoint_dir)
preds_after_restoration = restored_model(dummy_inputs)
Expand Down Expand Up @@ -318,6 +342,14 @@ def test_restore_keras_checkpoint_v3(self):
)
checkpoint_manager.save_model_variables(bert_pretrainer, epoch=1)
checkpoint_manager.close()
# Only models with SparseCore tables save their SparseCore topology.
self.assertFalse(
os.path.exists(
os.path.join(
checkpoint_dir, "1", keras_utils.SPARSECORE_TOPOLOGY_KEY
)
)
)

restored_model = keras_utils.restore_keras_checkpoint(checkpoint_dir)
preds_after_restoration = restored_model(dummy_inputs)
Expand Down
Loading
Loading