diff --git a/recml/core/utils/keras_utils.py b/recml/core/utils/keras_utils.py index ab6af89..17017b5 100644 --- a/recml/core/utils/keras_utils.py +++ b/recml/core/utils/keras_utils.py @@ -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 @@ -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" @@ -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, ) @@ -335,6 +340,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 +409,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, ) 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) diff --git a/recml/core/utils/sparsecore_checkpoint_utils.py b/recml/core/utils/sparsecore_checkpoint_utils.py new file mode 100644 index 0000000..47fa784 --- /dev/null +++ b/recml/core/utils/sparsecore_checkpoint_utils.py @@ -0,0 +1,294 @@ +# Copyright 2024 RecML authors . +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Checkpoint utilities for the SparseCore embedding tables of Keras models. + +`jax_tpu_embedding` stacks, pads and mod-shards the SparseCore embedding tables +of `keras_rs.layers.DistributedEmbedding` (and their optimizer slots) for the +total number of SparseCores of the TPU 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). This module +saves the topology of the stacked variables with the checkpoint, and reshards +them across topologies, so that checkpoints can be restored on another topology +than the one they were saved on. + +TODO(b/577354074): Remove this module (and its `jax_tpu_embedding` dependency) +once keras_rs `DistributedEmbedding` supports resharding SparseCore tables and +optimizer slots across TPU topologies. +""" + +from collections.abc import Mapping, Sequence +import copy +import math +import os +from typing import Any, cast + +from etils import epath +import jax +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 orbax.checkpoint as ocp + + +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 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")) + + +# How the SparseCore stacked tables were sharded (see `sparsecore_topology`). +SPARSECORE_TOPOLOGY_KEY = "sparsecore_topology" + + +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 `SparseCoreEmbeddings`. + + Args: + model: The model being saved or 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) + ) + + +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() diff --git a/recml/core/utils/sparsecore_checkpoint_utils_test.py b/recml/core/utils/sparsecore_checkpoint_utils_test.py new file mode 100644 index 0000000..452c400 --- /dev/null +++ b/recml/core/utils/sparsecore_checkpoint_utils_test.py @@ -0,0 +1,383 @@ +# Copyright 2024 RecML authors . +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for `sparsecore_checkpoint_utils`. + +The SparseCore tables are described by `jax_tpu_embedding` `TableSpec`s stacked +for a given number of SparseCores per chip, and models hold them in a layer with +the SparseCore state of `keras_rs.layers.DistributedEmbedding`, so the tests run +on CPU without SparseCore. `sparsecore_checkpoint_utils_multidevice_test` runs +them on 8 virtual CPU devices. +""" + +from collections.abc import Mapping, Sequence +import os +import types + +from absl.testing import absltest +from absl.testing import parameterized +import jax +from jax.experimental import layout as jax_layout +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 +from recml.core.utils import keras_utils +from recml.core.utils import sparsecore_checkpoint_utils + +# Three tables in one stack (all dims pad to 8) with vocab sizes that pad +# differently for different numbers of SparseCores. +_VOCABS_AND_DIMS = {"t1": (19, 6), "t2": (35, 8), "t3": (10, 4)} + + +def _stacked_feature_specs( + num_sc_per_device: int, + rotation: int | None = None, +) -> list[embedding_spec.FeatureSpec]: + """Returns features of tables stacked for the devices, like `keras_rs`.""" + feature_specs = [] + for name, (vocab, dim) in _VOCABS_AND_DIMS.items(): + table_spec = embedding_spec.TableSpec( + name=name, + vocabulary_size=vocab, + embedding_dim=dim, + initializer=jax.nn.initializers.zeros, + optimizer=embedding_spec.SGDOptimizerSpec(), + combiner="sum", + ) + feature_specs.append( + embedding_spec.FeatureSpec( + name=f"f_{name}", + table_spec=table_spec, + input_shape=(8, 1), + output_shape=(8, dim), + ) + ) + jte_table_stacking.stack_tables( + feature_specs, + list(_VOCABS_AND_DIMS), + global_device_count=jax.device_count(), + num_sc_per_device=num_sc_per_device, + rotation=rotation, + ) + return feature_specs + + +def _table_stack( + num_sc_per_device: int, rotation: int | None = None +) -> Sequence[embedding_spec.TableSpec]: + """Returns the tables stacked for the current devices, like `keras_rs`.""" + (table_stack,) = jte_table_stacking.get_table_stacks( + jte_embedding.get_table_specs( + _stacked_feature_specs(num_sc_per_device, rotation) + ) + ).values() + return table_stack + + +def _stack( + table_stack: Sequence[embedding_spec.TableSpec], + tables: Mapping[str, np.ndarray], + num_shards: int, +) -> np.ndarray: + """Returns `tables` stacked and sharded like `table_stack`.""" + stack_name = table_stack[0].stacked_table_spec.stack_name + stacked = jte_table_stacking.stack_and_shard_tables( + list(table_stack), + [tables[spec.name] for spec in table_stack], + num_shards, + pad_value=0.0, + )[stack_name] + return np.asarray(stacked).reshape(-1, stacked.shape[-1]) + + +def _sparsecore_sharding() -> jax.sharding.NamedSharding: + mesh = jax.sharding.Mesh(np.array(jax.devices()), ("sparsecore",)) + return jax.sharding.NamedSharding( + mesh, jax.sharding.PartitionSpec("sparsecore") + ) + + +class _FakeSparseCoreLayer(keras.layers.Layer): + """Holds the SparseCore state of a `keras_rs.layers.DistributedEmbedding`. + + `sparsecore_checkpoint_utils` reads the SparseCore lookup config, layout and + stacked table and slot variables of the layers. This layer has one stack of + `_VOCABS_AND_DIMS`, stacked like `keras_rs` for `num_sc_per_device` + SparseCores per chip, and sharded across the devices like the SparseCore + tables, so that checkpoints can be tested without SparseCore. + """ + + def __init__(self, num_sc_per_device: int, **kwargs): + super().__init__(**kwargs) + self._num_sc_per_device = num_sc_per_device + feature_specs = _stacked_feature_specs(num_sc_per_device) + (table_stack,) = jte_table_stacking.get_table_stacks( + jte_embedding.get_table_specs(feature_specs) + ).values() + # In a namespace, so that Keras doesn't wrap the lists for tracking. + self._specs = types.SimpleNamespace( + feature_specs=feature_specs, table_stack=table_stack + ) + + @property + def table_stack(self) -> Sequence[embedding_spec.TableSpec]: + return self._specs.table_stack + + @property + def num_shards(self) -> int: + return jax.device_count() * self._num_sc_per_device + + def build(self, input_shape=None): + del input_shape + stacked_spec = self.table_stack[0].stacked_table_spec + shape = (stacked_spec.stack_vocab_size, stacked_spec.stack_embedding_dim) + sharding = _sparsecore_sharding() + layout = keras.distribution.TensorLayout( + ("sparsecore", None), + keras.distribution.DeviceMesh( + (jax.device_count(),), ("sparsecore",), jax.devices() + ), + ) + # Like `keras_rs`, a JAX `Format` (here without the TPU tiling). + layout._backend_layout = jax_layout.Format(None, sharding) # pylint: disable=protected-access + table = self.add_weight(shape=shape, initializer="zeros", name="table") + # Trained on SparseCore, so the Keras optimizer has no slots for it. + table.overwrite_with_gradient = True + slot = self.add_weight( + shape=shape, initializer="zeros", name="slot", trainable=False + ) + for var in (table, slot): + var._layout = layout # pylint: disable=protected-access + var.assign(var.value) + self._config = types.SimpleNamespace( + num_sc_per_device=self._num_sc_per_device, + mesh=sharding.mesh, + feature_specs=self._specs.feature_specs, + ) + self._sparsecore_layout = layout + self._table_and_slot_variables = { + stacked_spec.stack_name: types.SimpleNamespace( + table=table, slot=(slot,) + ) + } + self._sparsecore_built = True + self.built = True + + @property + def stack_variables(self) -> types.SimpleNamespace: + (emb_vars,) = self._table_and_slot_variables.values() + return emb_vars + + def set_tables( + self, tables: Mapping[str, np.ndarray], slots: Mapping[str, np.ndarray] + ) -> None: + self.stack_variables.table.assign( + _stack(self.table_stack, tables, self.num_shards) + ) + (slot,) = self.stack_variables.slot + slot.assign(_stack(self.table_stack, slots, self.num_shards)) + + def get_tables(self, var: keras.Variable) -> dict[str, np.ndarray]: + stack_name = self.table_stack[0].stacked_table_spec.stack_name + tables = jte_table_stacking.unshard_and_unstack_tables( + list(self.table_stack), {stack_name: var.value}, self.num_shards + ) + return { + spec.name: np.asarray(table) + for spec, table in zip(self.table_stack, tables) + } + + +class _TestModel(keras.Model): + + def __init__(self, num_sc_per_device: int | None): + super().__init__(name="test_model") + self._num_sc_per_device = num_sc_per_device + self.sc_layer = ( + _FakeSparseCoreLayer(num_sc_per_device, name="sc_emb") + if num_sc_per_device + else None + ) + self.dense = keras.layers.Dense(2, name="head") + + def build(self, input_shape=None): + if self.sc_layer is not None: + self.sc_layer.build() + self.dense.build((None, 8)) + self.built = True + + def get_config(self): + # Saved with the checkpoints. + return {"num_sc_per_device": self._num_sc_per_device} + + +def _build_model(num_sc_per_device: int | None) -> _TestModel: + model = _TestModel(num_sc_per_device) + model.build() + model.compile(optimizer=keras.optimizers.Adam(1e-3)) + model.optimizer.build(model.trainable_variables) + return model + + +class SparseCoreCheckpointUtilsTest(parameterized.TestCase): + + @parameterized.parameters(1, 2, 4) + def test_settings_in_stack_match_stack_tables(self, num_sc_per_device): + # Recomputed from a stack for another topology. + table_stack = _table_stack(num_sc_per_device=8) + expected = { + spec.name: spec.setting_in_stack + for spec in _table_stack(num_sc_per_device) + } + self.assertEqual( + sparsecore_checkpoint_utils.settings_in_stack( + table_stack, + jax.device_count() * num_sc_per_device, + num_sc_per_device, + ), + expected, + ) + + @parameterized.parameters((1, 4), (2, 4), (4, 2), (4, 4)) + def test_reshard_stacked_array_matches_direct_stacking( + self, src_num_sc_per_device, dst_num_sc_per_device + ): + rng = np.random.default_rng(0) + tables = { + name: rng.normal(size=vocab_and_dim).astype(np.float32) + for name, vocab_and_dim in _VOCABS_AND_DIMS.items() + } + src_num_shards = jax.device_count() * src_num_sc_per_device + dst_num_shards = jax.device_count() * dst_num_sc_per_device + src_stacked = _stack( + _table_stack(src_num_sc_per_device), tables, src_num_shards + ) + dst_table_stack = _table_stack(dst_num_sc_per_device) + expected = _stack(dst_table_stack, tables, dst_num_shards) + sharding = _sparsecore_sharding() + + actual = sparsecore_checkpoint_utils.reshard_stacked_array( + jax.device_put( + src_stacked, + sparsecore_checkpoint_utils.row_sharding( + sharding, src_stacked.shape[0] + ), + ), + dst_table_stack, + src_num_shards=src_num_shards, + src_num_sc_per_device=src_num_sc_per_device, + dst_num_shards=dst_num_shards, + dst_num_sc_per_device=dst_num_sc_per_device, + out_sharding=sharding, + ) + + self.assertTrue(actual.sharding.is_equivalent_to(sharding, actual.ndim)) + np.testing.assert_array_equal(np.asarray(actual), expected) + + def test_reshard_stacked_array_rejects_unexpected_layout(self): + # Stacked with another shard rotation than `keras_rs` uses. + table_stack = _table_stack(num_sc_per_device=4, rotation=1) + num_shards = jax.device_count() * 4 + stacked = np.zeros( + (table_stack[0].stacked_table_spec.stack_vocab_size, 8), np.float32 + ) + with self.assertRaisesRegex( + ValueError, "Unexpected SparseCore stacking layout" + ): + sparsecore_checkpoint_utils.reshard_stacked_array( + jax.numpy.asarray(stacked), + table_stack, + src_num_shards=num_shards, + src_num_sc_per_device=4, + dst_num_shards=num_shards, + dst_num_sc_per_device=4, + out_sharding=_sparsecore_sharding(), + ) + + def test_row_sharding_keeps_divisible_rows(self): + sharding = _sparsecore_sharding() + self.assertIs( + sparsecore_checkpoint_utils.row_sharding( + sharding, 8 * jax.device_count() + ), + sharding, + ) + + def test_row_sharding_splits_uneven_rows_into_groups(self): + num_devices = jax.device_count() + if num_devices < 4: + self.skipTest("Requires at least 4 devices.") + # E.g. 8 devices and 4 * 3 rows: 4 row shards, each on 2 devices. + num_rows = (num_devices // 2) * 3 + sharding = sparsecore_checkpoint_utils.row_sharding( + _sparsecore_sharding(), num_rows + ) + self.assertEqual(sharding.mesh.devices.shape, (num_devices // 2, 2)) + x = jax.device_put(np.arange(num_rows * 2.0).reshape(num_rows, 2), sharding) + self.assertEqual( + {shard.data.shape for shard in x.addressable_shards}, {(3, 2)} + ) + + @parameterized.parameters(2, 4) + def test_sparsecore_topology(self, num_sc_per_device): + model = _build_model(num_sc_per_device) + self.assertTrue(sparsecore_checkpoint_utils.has_sparsecore_tables(model)) + (state,) = sparsecore_checkpoint_utils.sparsecore_layer_states(model) + self.assertIs(state[1], model.sc_layer._sparsecore_layout) + self.assertEqual( + sparsecore_checkpoint_utils.sparsecore_topology(model), + { + model.sc_layer.stack_variables.table.path: { + "num_shards": jax.device_count() * num_sc_per_device, + "num_sc_per_device": num_sc_per_device, + } + }, + ) + + def test_models_without_sparsecore_tables(self): + model = _build_model(None) + self.assertFalse(sparsecore_checkpoint_utils.has_sparsecore_tables(model)) + self.assertEqual(sparsecore_checkpoint_utils.sparsecore_topology(model), {}) + self.assertEqual( + sparsecore_checkpoint_utils.sparsecore_topology_save_args(model), {} + ) + + def test_unbuilt_sparsecore_layers_are_skipped(self): + model = keras.Sequential([_FakeSparseCoreLayer(2)]) + self.assertEmpty(sparsecore_checkpoint_utils.sparsecore_layer_states(model)) + + @parameterized.named_parameters( + ("v2", keras_utils.KerasOrbaxCheckpointManagerV2), + ("v3", keras_utils.KerasOrbaxCheckpointManagerV3), + ) + def test_sparsecore_topology_is_saved_with_the_checkpoint(self, manager_cls): + for num_sc_per_device in (None, 4): + with self.subTest(num_sc_per_device=num_sc_per_device): + model = _build_model(num_sc_per_device) + ckpt_dir = self.create_tempdir().full_path + manager = manager_cls(ckpt_dir, choose_store_cell=False) + manager.save_model_variables(model, epoch=5) + manager.wait_until_finished() + manager.close() + self.assertEqual( + sparsecore_checkpoint_utils.saved_sparsecore_topology( + os.path.join(ckpt_dir, "5") + ), + sparsecore_checkpoint_utils.sparsecore_topology(model) or None, + ) + + +if __name__ == "__main__": + absltest.main()