From 9267700a76f1ef156312c488f4aba4a9cde27053 Mon Sep 17 00:00:00 2001 From: RecML authors Date: Sat, 24 Oct 2026 18:35:40 -0700 Subject: [PATCH] Internal change PiperOrigin-RevId: 1006505061 --- recml/core/utils/keras_utils.py | 606 ++++++++++++++++++++++++++- recml/core/utils/keras_utils_test.py | 32 ++ 2 files changed, 637 insertions(+), 1 deletion(-) diff --git a/recml/core/utils/keras_utils.py b/recml/core/utils/keras_utils.py index ab6af89..c07d3d1 100644 --- a/recml/core/utils/keras_utils.py +++ b/recml/core/utils/keras_utils.py @@ -14,17 +14,23 @@ """Utilities for training Keras models on Jax backend.""" from collections.abc import Container, Mapping, Sequence +import copy import dataclasses import datetime import enum +import math import os import re -from typing import Any +from typing import Any, cast from absl import logging from etils import epath import jax +from jax_tpu_embedding.sparsecore.lib.nn import embedding as jte_embedding +from jax_tpu_embedding.sparsecore.lib.nn import embedding_spec +from jax_tpu_embedding.sparsecore.lib.nn import table_stacking as jte_table_stacking import keras +import numpy as np import orbax.checkpoint as ocp import tensorflow as tf @@ -38,6 +44,8 @@ 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 (see `_sparsecore_topology`). +SPARSECORE_TOPOLOGY_KEY = "sparsecore_topology" ORBAX_CHECKPOINT_DEFAULT_KEY = "default" @@ -241,6 +249,7 @@ def save_model_variables( args=ocp.args.Composite(**{ STATE_CHECKPOINT_KEY: ocp.args.StandardSave(state), CONFIG_CHECKPOINT_KEY: ocp.args.JsonSave(config), + **_sparsecore_topology_save_args(model), }), metrics=logs, ) @@ -263,6 +272,12 @@ 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. + sparsecore_restores = ( + _prepare_sparsecore_restore(model, str(self.directory), epoch) + if _has_sparsecore_tables(model) + else [] + ) + # 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. @@ -291,6 +306,9 @@ def restore_model_variables(self, model: keras.Model, epoch: int): ): var._value = restored_var # pylint: disable=protected-access + if sparsecore_restores: + _finish_sparsecore_restore(sparsecore_restores) + class KerasOrbaxCheckpointManagerV3(ocp.CheckpointManager): """An Orbax checkpoint manager for Keras 3 with dictionary state. @@ -335,6 +353,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, @@ -403,6 +422,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_topology_save_args(model), }), metrics=logs, ) @@ -432,6 +452,12 @@ def restore_model_variables(self, model: keras.Model, epoch: int): OPTIMIZER_VARIABLES_KEY: optimizer_variables, } + sparsecore_restores = ( + _prepare_sparsecore_restore(model, str(self.directory), epoch) + if _has_sparsecore_tables(model) + else [] + ) + # 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) @@ -479,6 +505,9 @@ 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 + if sparsecore_restores: + _finish_sparsecore_restore(sparsecore_restores) + def resolve_orbax_checkpoint_path( checkpoint_dir: str, epoch: int | None = None @@ -1199,6 +1228,12 @@ def restore_keras_checkpoint( # TODO(zixiangzhou): Update variables to use a nested dictionary and index map # instead of flattened list. + sparsecore_restores = ( + _prepare_sparsecore_restore(model, checkpoint_dir, epoch) + if _has_sparsecore_tables(model) + else [] + ) + # 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) @@ -1239,6 +1274,9 @@ def restore_keras_checkpoint( ): var._value = restored_var # pylint: disable=protected-access + if sparsecore_restores: + _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: @@ -1247,6 +1285,572 @@ def restore_keras_checkpoint( return model +class _RestoreValuePlaceholder: + """Stands in for `keras.Variable._value` to restore a differently shaped value. + + Used by `_prepare_sparsecore_restore` for SparseCore variables saved for + another topology. The restore paths of this module build the Orbax restore + target from `var.value.shape`, `.dtype` and `.sharding` (see + `_to_shape_dtype_struct`), so these attributes make Orbax restore the saved + shape. They then call `var.value.delete()` to free each variable before + restoring; the placeholder holds no buffer, so `delete` is a no-op. + """ + + __slots__ = ("shape", "dtype", "sharding") + + def __init__( + self, + shape: tuple[int, ...], + dtype: Any, + sharding: jax.sharding.Sharding, + ): + self.shape = shape + self.dtype = dtype + self.sharding = sharding + + def delete(self) -> None: + pass + + +def _settings_in_stack( + table_stack: Sequence[embedding_spec.TableSpec], + num_shards: int, + num_sc_per_device: int, +) -> Mapping[str, embedding_spec.TableSettingInStack]: + """Returns each table's `TableSettingInStack` for a SparseCore topology. + + Recomputes the row offsets and shard rotations that + `jte_table_stacking.stack_tables` assigns to the tables of a stack for + `num_shards` SparseCores, keeping the table order of the existing stack. + `keras_rs.layers.DistributedEmbedding` stacks tables with the default + `rotation=None`, which `stack_tables` resolves to `num_sc_per_device`. + + Args: + table_stack: The stacked `TableSpec`s of one stack (for any topology). + num_shards: The total number of SparseCores of the topology. + num_sc_per_device: The number of SparseCores per chip of the topology. + + Returns: + Table name to its setting in the stack for the topology. + """ + ordered = sorted( + table_stack, key=lambda spec: spec.setting_in_stack.row_offset_in_shard + ) + first_setting = ordered[0].setting_in_stack + _, table_to_padded_vocab_size = ( + jte_table_stacking.round_up_dim_and_vocab_size( + {spec.name: spec for spec in ordered}, num_shards + ) + ) + return jte_table_stacking.compute_table_to_setting_in_stack( + stack_name=first_setting.stack_name, + table_names=[spec.name for spec in ordered], + padded_embedding_dim=first_setting.padded_embedding_dim, + table_to_padded_vocab_size=table_to_padded_vocab_size, + num_shards=num_shards, + rotation=num_sc_per_device, + ) + + +def _reshard_stacked_array( + stacked_array: jax.Array, + table_stack: Sequence[embedding_spec.TableSpec], + *, + src_num_shards: int, + src_num_sc_per_device: int, + dst_num_shards: int, + dst_num_sc_per_device: int, + out_sharding: jax.sharding.Sharding, +) -> jax.Array: + """Reshards a stacked SparseCore table or slot tensor across topologies. + + Unstacks `stacked_array` (laid out for the source topology) into its logical + tables with `jte_table_stacking.unshard_and_unstack_tables`, then restacks + them for the destination topology of `table_stack` with + `jte_table_stacking.stack_and_shard_tables`, all in one `jax.jit` + computation. XLA shards the intermediate tables (whose vocab sizes need not + be divisible by the device count), so they stay distributed across the + devices of all hosts, and every process must call this in lockstep. + + `stacked_array` is donated to reduce the peak memory usage. + + Args: + stacked_array: The `[stack_vocab_size, stack_embedding_dim]` stacked tensor + as saved for the source topology, in SparseCore shard order. + table_stack: The stacked `TableSpec`s of the stack for the destination + topology. + src_num_shards: The total number of SparseCores of the source topology. + src_num_sc_per_device: The number of SparseCores per chip of the source + topology. + dst_num_shards: The total number of SparseCores of the destination topology. + dst_num_sc_per_device: The number of SparseCores per chip of the destination + topology. + out_sharding: The sharding of the result. + + Returns: + The stacked tensor laid out for the destination topology. + + Raises: + ValueError: If the stacking layout of `table_stack` differs from the one + recomputed for the destination topology, e.g. because the + jax_tpu_embedding stacking scheme changed. + """ + # TODO(b/577354074): Replace with keras_rs's cross-topology resharding once + # supported. + dst_settings = _settings_in_stack( + table_stack, dst_num_shards, dst_num_sc_per_device + ) + for spec in table_stack: + if spec.setting_in_stack != dst_settings[spec.name]: + raise ValueError( + f"Unexpected SparseCore stacking layout for table '{spec.name}':" + f" {spec.setting_in_stack} != {dst_settings[spec.name]}. Cannot" + " reshard the checkpoint across TPU topologies." + ) + src_settings = _settings_in_stack( + table_stack, src_num_shards, src_num_sc_per_device + ) + src_table_stack = [] + for spec in table_stack: + src_spec = copy.copy(spec) + src_spec.setting_in_stack = src_settings[spec.name] + src_table_stack.append(src_spec) + stack_name = table_stack[0].stacked_table_spec.stack_name + + def _reshard(stacked: jax.Array) -> jax.Array: + tables = jte_table_stacking.unshard_and_unstack_tables( + src_table_stack, {stack_name: stacked}, src_num_shards + ) + restacked = jte_table_stacking.stack_and_shard_tables( + list(table_stack), tables, dst_num_shards, pad_value=0.0 + )[stack_name] + restacked = cast(jax.Array, restacked) + return restacked.reshape(-1, stacked.shape[-1]).astype(stacked.dtype) + + return jax.jit(_reshard, out_shardings=out_sharding, donate_argnums=0)( + stacked_array + ) + + +def _sparsecore_topology(model: keras.Model) -> dict[str, dict[str, int]]: + """Returns the SparseCore topology of each stacked table of `model`. + + Saved with the checkpoint so that a restore on another TPU topology knows + how the stacked tables (and their optimizer slots, which share the layout of + their table) were sharded. + + Args: + model: The model being saved. + + Returns: + Stacked table variable path to `{"num_shards": ..., "num_sc_per_device": + ...}`. Empty for models without SparseCore tables. + """ + if keras.backend.backend() != "jax": + return {} + topology = {} + for config, _, table_vars in _sparsecore_layer_states(model): + num_sc_per_device = int(config.num_sc_per_device) + num_shards = int(config.mesh.devices.size) * num_sc_per_device + for emb_vars in table_vars.values(): + topology[emb_vars.table.path] = { + "num_shards": num_shards, + "num_sc_per_device": num_sc_per_device, + } + return topology + + +def _sparsecore_topology_save_args(model: keras.Model) -> dict[str, Any]: + """Returns the Orbax save args of `_sparsecore_topology`, if any. + + Only models with SparseCore tables get the item, so the checkpoints of other + models are unchanged. + + Args: + model: The model being saved. + """ + topology = _sparsecore_topology(model) + if not topology: + return {} + return { + SPARSECORE_TOPOLOGY_KEY: ocp.args.JsonSave( + {SPARSECORE_TOPOLOGY_KEY: topology} + ) + } + + +def _saved_sparsecore_topology( + checkpoint_path: str, +) -> dict[str, dict[str, int]] | None: + """Returns the saved `_sparsecore_topology`, or None if it wasn't saved. + + Args: + checkpoint_path: Path to the checkpoint directory of one step. + """ + topology_path = epath.Path(checkpoint_path) / SPARSECORE_TOPOLOGY_KEY + if not topology_path.exists(): + return None + checkpointer = ocp.Checkpointer(ocp.handlers.JsonCheckpointHandler()) + try: + return checkpointer.restore( + os.fspath(topology_path), args=ocp.args.JsonRestore() + )[SPARSECORE_TOPOLOGY_KEY] + finally: + checkpointer.close() + + +def _sparsecore_layer_states( + model: keras.Model, +) -> list[tuple[Any, Any, Mapping[str, Any]]]: + """Returns the SparseCore state of the SparseCore layers of `model`. + + Matches the JAX `keras_rs.layers.DistributedEmbedding` layers with tables on + SparseCore, including the one wrapped by `SparseCoreEmbedding`. + + Args: + model: The model being restored. + + Returns: + For each layer, its SparseCore lookup config, SparseCore `TensorLayout` and + stacked table and slot variables by stack name. + """ + states = [] + for layer in model._flatten_layers(include_self=True): # pylint: disable=protected-access + if not getattr(layer, "built", False) or not getattr( + layer, "_sparsecore_built", False + ): + continue + state = ( + getattr(layer, "_config", None), + getattr(layer, "_sparsecore_layout", None), + getattr(layer, "_table_and_slot_variables", None), + ) + if all(x is not None for x in state): + states.append(state) + return states + + +def _has_sparsecore_tables(model: keras.Model) -> bool: + """Returns whether `model` has embedding tables on SparseCore.""" + return keras.backend.backend() == "jax" and bool( + _sparsecore_layer_states(model) + ) + + +@dataclasses.dataclass(frozen=True) +class _StackRestore: + """How to restore one stacked SparseCore table and its optimizer slots. + + Attributes: + variables: The stacked table variable followed by its slot variables. + saved_shapes: The shape of each variable in the checkpoint. + table_stack: The (destination) `TableSpec`s of the tables in the stack. + layout: The SparseCore `TensorLayout` of the variables. + src_num_shards: The number of SparseCores the checkpoint was saved with. + src_num_sc_per_device: The SparseCores per chip it was saved with. + dst_num_shards: The number of SparseCores of the current topology. + dst_num_sc_per_device: The SparseCores per chip of the current topology. + """ + + variables: Sequence[keras.Variable] + saved_shapes: Sequence[tuple[int, ...]] + table_stack: Sequence[embedding_spec.TableSpec] + layout: Any + src_num_shards: int + src_num_sc_per_device: int + dst_num_shards: int + dst_num_sc_per_device: int + + @property + def needs_reshard(self) -> bool: + """Whether the checkpoint was saved for another SparseCore topology.""" + return (self.src_num_shards, self.src_num_sc_per_device) != ( + self.dst_num_shards, + self.dst_num_sc_per_device, + ) + + +def _saved_trainable_metadata( + model: keras.Model, checkpoint_dir: str, epoch: int | None +) -> dict[str, Any]: + """Returns the checkpoint metadata of the model's trainable variables by path. + + Args: + model: The model being restored. + checkpoint_dir: Directory containing the Orbax checkpoint(s). + epoch: Optional step/epoch being restored. + + Returns: + Variable path to its Orbax `ArrayMetadata`, or an empty dict if the + checkpoint has no Orbax state item. + """ + try: + checkpoint_path, _ = resolve_orbax_checkpoint_path(checkpoint_dir, epoch) + state_dir = epath.Path(checkpoint_path) / STATE_CHECKPOINT_KEY + if not state_dir.exists(): + return {} + saved_state_metadata = ocp.handlers.PyTreeCheckpointHandler().metadata( + state_dir + ) + 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 (epoch %s) to" + " restore the SparseCore tables across TPU topologies.", + checkpoint_dir, + epoch, + ) + raise + + saved_trainable_meta = saved_state_metadata.get(TRAINABLE_VARIABLES_KEY) + if isinstance(saved_trainable_meta, Mapping): # V3: keyed by variable path. + return { + var.path: saved_trainable_meta[var.path] + for var in model.trainable_variables + if var.path in saved_trainable_meta + } + 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 _plan_sparsecore_restore( + model: keras.Model, checkpoint_dir: str, epoch: int | None +) -> list[_StackRestore]: + """Returns the SparseCore stacks of `model` that must be resharded. + + Stacks saved for the current topology are left out: Orbax restores them as + is. Stacks of checkpoints saved without their topology are restored as is if + their shapes match. + + Args: + model: The model being restored. + checkpoint_dir: Directory containing the Orbax checkpoint(s). + epoch: Optional step/epoch being restored. + + Raises: + ValueError: If a stack saved with another shape has no saved topology. + """ + sc_layer_states = _sparsecore_layer_states(model) + if not sc_layer_states: + return [] + if not model._jax_state_synced: # pylint: disable=protected-access + model.jax_state_sync() + meta_by_path = _saved_trainable_metadata(model, checkpoint_dir, epoch) + if not meta_by_path: + return [] + checkpoint_path, _ = resolve_orbax_checkpoint_path(checkpoint_dir, epoch) + saved_topology = _saved_sparsecore_topology(checkpoint_path) or {} + + plans = [] + for config, sc_layout, table_vars in sc_layer_states: + dst_num_sc_per_device = config.num_sc_per_device + dst_num_shards = config.mesh.devices.size * dst_num_sc_per_device + table_stacks = jte_table_stacking.get_table_stacks( + jte_embedding.get_table_specs(config.feature_specs) + ) + for stack_name, table_stack in table_stacks.items(): + emb_vars = table_vars[stack_name] + stored_meta = meta_by_path.get(emb_vars.table.path) + if stored_meta is None: + continue + variables = [emb_vars.table, *keras.tree.flatten(emb_vars.slot)] + # Slots have the shape of their table. + saved_shapes = [ + tuple(meta_by_path.get(var.path, stored_meta).shape) + for var in variables + ] + shapes_match = all( + shape == tuple(var.shape) + for var, shape in zip(variables, saved_shapes) + ) + src_topology = saved_topology.get(emb_vars.table.path) + if src_topology is None: + # Checkpoints saved before the topology was persisted. + if not shapes_match: + raise ValueError( + f"The SparseCore stack '{stack_name}' was saved with shapes" + f" {saved_shapes} but without its SparseCore topology, so it" + " can't be resharded for the current shapes" + f" {[tuple(var.shape) for var in variables]}. Restore it on the" + " TPU topology it was saved on, and save it again with" + " `KerasOrbaxCheckpointManagerV2/V3`." + ) + logging.warning( + "The SparseCore stack '%s' was saved without its SparseCore" + " topology. Restoring it as is, assuming it was saved on the" + " current one (%d SparseCores, %d per chip); if not, the restored" + " SparseCore tables are wrong.", + stack_name, + dst_num_shards, + dst_num_sc_per_device, + ) + continue + plan = _StackRestore( + variables=variables, + saved_shapes=saved_shapes, + table_stack=table_stack, + layout=sc_layout, + src_num_shards=int(src_topology["num_shards"]), + src_num_sc_per_device=int(src_topology["num_sc_per_device"]), + dst_num_shards=dst_num_shards, + dst_num_sc_per_device=dst_num_sc_per_device, + ) + if plan.needs_reshard or not shapes_match: + plans.append(plan) + return plans + + +def _row_sharding( + sharding: jax.sharding.NamedSharding, num_rows: int +) -> jax.sharding.NamedSharding: + """Returns a sharding of the rows of an array over the devices of `sharding`. + + `sharding` must shard the rows over all its devices. If `num_rows` is not + divisible by the device count, the rows are split into + `gcd(num_rows, device_count)` shards, each replicated on consecutive devices, + instead of replicating the whole array on every device. + + Args: + sharding: The sharding of the rows over all its devices. + num_rows: The number of rows of the array. + """ + devices = sharding.mesh.devices.reshape(-1) + num_row_shards = math.gcd(num_rows, devices.size) + if num_row_shards == devices.size: + return sharding + mesh = jax.sharding.Mesh( + devices.reshape(num_row_shards, -1), ("rows", "replicas") + ) + return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec("rows")) + + +# TODO(b/577354074): Remove this cross-topology restore (and the +# jax_tpu_embedding dependency of this module) once keras_rs +# `DistributedEmbedding` supports resharding SparseCore tables and optimizer +# slots across TPU topologies. Until then, consider moving it behind a +# layer-provided checkpoint restore adapter next to `SparseCoreEmbedding`. +def _prepare_sparsecore_restore( + model: keras.Model, checkpoint_dir: str, epoch: int | None +) -> list[_StackRestore]: + """Prepares the SparseCore tables of `model` for a cross-topology restore. + + `jax_tpu_embedding` stacks, pads and mod-shards the SparseCore tables (and + their optimizer slots) for the total number of SparseCores of the topology + (`global_device_count * num_sc_per_device`), so both the shape and the row + layout of a stacked variable change across topologies (e.g. `glp=4x4` train + vs. `glp=2x2` eval). The restore paths of this module restore each variable + into `jax.ShapeDtypeStruct(var.value.shape, ...)`, so restoring such a + checkpoint as is would fail with a shape mismatch, or load rows sharded for + the wrong topology when the shapes happen to match. + + For models with SparseCore tables (see `_has_sparsecore_tables`), call this + before building the Orbax restore target and `_finish_sparsecore_restore` + after the restore. This reads the source topology saved with the checkpoint + (see `_sparsecore_topology`) and, for each stacked variable saved with + another shape, replaces its value with a `_RestoreValuePlaceholder` of the + saved shape, so that Orbax restores the raw saved tensor. Stacks saved for + the current topology are not touched. + + If the restore fails in between, the model must be rebuilt: like the restore + paths themselves, this frees the current buffers of the restored variables. + + Args: + model: The model being restored. + checkpoint_dir: Directory containing the Orbax checkpoint(s). + epoch: Optional step/epoch being restored. + + Returns: + The stacks to reshard, to pass to `_finish_sparsecore_restore`. + + Raises: + ValueError: If a stack saved with another shape has no saved topology. + Exception: If the checkpoint metadata can't be read. + """ + plans = _plan_sparsecore_restore(model, checkpoint_dir, epoch) + for plan in plans: + sc_sharding = plan.layout.backend_layout.sharding + for var, saved_shape in zip(plan.variables, plan.saved_shapes): + # Restore the raw saved value instead of laying it out for the current + # topology. + var._layout = None # pylint: disable=protected-access + if saved_shape != tuple(var.shape): + dtype = var.value.dtype + sharding = _row_sharding(sc_sharding, saved_shape[0]) + num_devices = sc_sharding.mesh.devices.size + num_row_shards = sharding.mesh.devices.shape[0] + if num_row_shards < num_devices: + # TODO(b/577354074): Shard the saved rows unevenly instead, so that + # every device holds 1/num_devices of the stack. + logging.warning( + "Restoring SparseCore variable %s with saved shape %s in %d row" + " shards, each replicated on %d of the %d devices (%.2f GiB per" + " device), since its rows are not divisible by the device count.", + var.path, + saved_shape, + num_row_shards, + num_devices // num_row_shards, + num_devices, + math.prod(saved_shape) + * np.dtype(dtype).itemsize + / num_row_shards + / 2**30, + ) + var.value.delete() + var._value = _RestoreValuePlaceholder( # pylint: disable=protected-access + shape=saved_shape, + dtype=dtype, + sharding=sharding, + ) + return plans + + +def _finish_sparsecore_restore(plans: Sequence[_StackRestore]) -> None: + """Reshards the restored SparseCore stacks for the current topology. + + Args: + plans: The result of `_prepare_sparsecore_restore`, after the restore. + """ + for plan in plans: + if plan.needs_reshard: + # TODO(b/577354074): Validate the cross-topology restore on multi-host + # TPU topologies (e.g. 16x16 train -> 2x4 eval and back). + logging.info( + "Resharding SparseCore stack %s (saved shape %s) from %d SparseCores" + " (%d per chip) to %d (%d per chip) across %d processes.", + plan.table_stack[0].stacked_table_spec.stack_name, + plan.saved_shapes[0], + plan.src_num_shards, + plan.src_num_sc_per_device, + plan.dst_num_shards, + plan.dst_num_sc_per_device, + jax.process_count(), + ) + for var in plan.variables: + var._layout = plan.layout # pylint: disable=protected-access + restored = var._value # pylint: disable=protected-access + if plan.needs_reshard or tuple(restored.shape) != tuple(var.shape): + restored = _reshard_stacked_array( + restored, + plan.table_stack, + src_num_shards=plan.src_num_shards, + src_num_sc_per_device=plan.src_num_sc_per_device, + dst_num_shards=plan.dst_num_shards, + dst_num_sc_per_device=plan.dst_num_sc_per_device, + out_sharding=plan.layout.backend_layout.sharding, + ) + var.assign(restored) + + def restore_partial_checkpoint( checkpoint_dir: str, partial_variables: Mapping[str, Any], diff --git a/recml/core/utils/keras_utils_test.py b/recml/core/utils/keras_utils_test.py index d130f55..2b24c18 100644 --- a/recml/core/utils/keras_utils_test.py +++ b/recml/core/utils/keras_utils_test.py @@ -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) @@ -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) @@ -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) @@ -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)