From 14d87868f56d153c5537be930c9b80ebb43278c1 Mon Sep 17 00:00:00 2001 From: RecML authors Date: Mon, 2 Nov 2026 20:52:34 -0800 Subject: [PATCH] Internal change PiperOrigin-RevId: 1012671438 --- recml/core/utils/keras_utils.py | 19 ++ recml/core/utils/keras_utils_test.py | 32 +++ .../core/utils/sparsecore_checkpoint_utils.py | 191 ++++++++++++- .../utils/sparsecore_checkpoint_utils_test.py | 255 +++++++++++++++++- 4 files changed, 487 insertions(+), 10 deletions(-) diff --git a/recml/core/utils/keras_utils.py b/recml/core/utils/keras_utils.py index e271edc..e7315cf 100644 --- a/recml/core/utils/keras_utils.py +++ b/recml/core/utils/keras_utils.py @@ -27,6 +27,7 @@ import jax import keras import orbax.checkpoint as ocp +from recml.core.utils import sparsecore_checkpoint_utils import tensorflow as tf @@ -39,6 +40,14 @@ NON_TRAINABLE_PATHS_KEY = "non_trainable_paths" OPTIMIZER_PATHS_KEY = "optimizer_paths" SHARED_VARIABLE_MAP_KEY = "shared_variable_map" +# Optional checkpoint item recording how the SparseCore stacked tables were +# sharded (`num_shards` and `num_sc_per_device`), saved only for models with +# SparseCore tables (see `sparsecore_checkpoint_utils.sparsecore_topology`). +# Used during restore to unshard/unstack and reshard SparseCore tables when the +# restoring TPU topology differs from the saving topology. +# TODO(b/576848327): Remove `SPARSECORE_TOPOLOGY_KEY` once `keras_rs` +# `DistributedEmbedding` natively supports cross-topology resharding. +SPARSECORE_TOPOLOGY_KEY = sparsecore_checkpoint_utils.SPARSECORE_TOPOLOGY_KEY ORBAX_CHECKPOINT_DEFAULT_KEY = "default" @@ -242,6 +251,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, ) @@ -336,6 +346,14 @@ def __init__( NON_TRAINABLE_PATHS_KEY, OPTIMIZER_PATHS_KEY, SHARED_VARIABLE_MAP_KEY, + # Optional item saved only when the model has SparseCore tables, + # used to reshard stacked tables across TPU topologies on restore. + # TODO(b/576848327): Remove once `keras_rs` supports cross-topology + # resharding natively. + # TODO(b/578977772): Verify that checkpoints with and without this + # item stay restorable with the backward compatibility test once it + # is ready. + SPARSECORE_TOPOLOGY_KEY, ), options=ocp.CheckpointManagerOptions( save_interval_steps=save_interval_epochs, @@ -404,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_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 88ae6bb..182e096 100644 --- a/recml/core/utils/keras_utils_test.py +++ b/recml/core/utils/keras_utils_test.py @@ -164,6 +164,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) @@ -243,6 +251,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) @@ -288,6 +304,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) @@ -319,6 +343,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 index 0252dc9..d43bdc0 100644 --- a/recml/core/utils/sparsecore_checkpoint_utils.py +++ b/recml/core/utils/sparsecore_checkpoint_utils.py @@ -18,8 +18,9 @@ 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 -reshards such variables across topologies, so that checkpoints can be restored -on another topology than the one they were saved on. +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/576848327): Remove this module (and its `jax_tpu_embedding` dependency) once keras_rs `DistributedEmbedding` supports resharding SparseCore tables and @@ -29,12 +30,16 @@ from collections.abc import Mapping, Sequence import copy import math -from typing import cast +import os +from typing import Any, NamedTuple, cast +from etils import epath import jax from jax.experimental import layout as jax_layout 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( @@ -209,3 +214,183 @@ def row_sharding( devices.reshape(num_row_shards, -1), ("rows", "replicas") ) return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec("rows")) + + +# Optional checkpoint item key storing the source SparseCore topology +# (`num_shards` and `num_sc_per_device` per stacked table; see +# `sparsecore_topology`). Saved only for models with SparseCore tables and read +# on restore to unshard/unstack and reshard stacked tables when restoring on a +# different TPU topology. +# TODO(b/576848327): Remove once `keras_rs` `DistributedEmbedding` natively +# supports cross-topology resharding. +SPARSECORE_TOPOLOGY_KEY = "sparsecore_topology" + + +class SparseCoreLayerState(NamedTuple): + """The SparseCore state of a `keras_rs.layers.DistributedEmbedding` layer. + + Attributes: + config: The SparseCore lookup config of the layer (with `feature_specs`, + `mesh` and `num_sc_per_device`). + layout: The SparseCore `TensorLayout` of the stacked variables. + table_and_slot_variables: The stacked table and slot variables by stack + name. + """ + + config: Any + layout: keras.distribution.TensorLayout + table_and_slot_variables: Mapping[str, Any] + + +def sparsecore_layer_states(model: keras.Model) -> list[SparseCoreLayerState]: + """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`. Only these + layers set `_sparsecore_built`, so models on other backends have none. + + Args: + model: The model being saved or restored. Must be built, along with all of + its layers. + + Returns: + The SparseCore state of each built SparseCore layer. + + Raises: + ValueError: If `model` or any of its layers is not built, or if a built + SparseCore layer is missing its SparseCore state. + """ + states = [] + # pylint: disable-next=protected-access + for layer in model._flatten_layers(include_self=True): + if not getattr(layer, "built", False) or not getattr( + layer, "_sparsecore_built", True + ): + raise ValueError( + "All layers of `model` must be built before inspecting SparseCore" + f" state, but layer '{layer.name}' ({type(layer).__name__}) is not" + " built." + ) + if not getattr(layer, "_sparsecore_built", False): + continue + config = getattr(layer, "_config", None) + layout = getattr(layer, "_sparsecore_layout", None) + table_and_slot_variables = getattr(layer, "_table_and_slot_variables", None) + if config is None or layout is None or table_and_slot_variables is None: + missing = [ + name + for name, value in ( + ("_config", config), + ("_sparsecore_layout", layout), + ("_table_and_slot_variables", table_and_slot_variables), + ) + if value is None + ] + raise ValueError( + f"SparseCore layer '{layer.name}' has `_sparsecore_built=True` but" + f" is missing {missing}." + ) + states.append( + SparseCoreLayerState(config, layout, table_and_slot_variables) + ) + return states + + +def has_sparsecore_tables(model: keras.Model) -> bool: + """Returns whether `model` has embedding tables on SparseCore. + + Args: + model: The model being saved or restored. + + Raises: + ValueError: See `sparsecore_layer_states`. + """ + return 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. + + Raises: + ValueError: If two stacked tables have the same variable path (e.g. two + SparseCore layers with the same name and stack names), or see + `sparsecore_layer_states`. + """ + topology = {} + for state in sparsecore_layer_states(model): + num_sc_per_device = int(state.config.num_sc_per_device) + num_shards = int(state.config.mesh.devices.size) * num_sc_per_device + for emb_vars in state.table_and_slot_variables.values(): + # V2 checkpoints store the variables in a list, so only this check + # catches duplicate paths, whose topologies would overwrite each other. + if emb_vars.table.path in topology: + raise ValueError( + "Duplicate SparseCore table variable path" + f" '{emb_vars.table.path}'. Give the SparseCore layers unique" + " names." + ) + 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. + + Returns: + `{SPARSECORE_TOPOLOGY_KEY: ocp.args.JsonSave(...)}` to merge into the Orbax + `Composite` save args, or an empty dict for models without SparseCore + tables. + """ + 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. + + Returns: + The `sparsecore_topology` saved with the checkpoint, or None if the + checkpoint has no `SPARSECORE_TOPOLOGY_KEY` item (e.g. it was saved for a + model without SparseCore tables). + """ + 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 index 7f9ff9e..9c7b83c 100644 --- a/recml/core/utils/sparsecore_checkpoint_utils_test.py +++ b/recml/core/utils/sparsecore_checkpoint_utils_test.py @@ -14,12 +14,15 @@ """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, so the tests run on CPU without -SparseCore. `sparsecore_checkpoint_utils_multidevice_test` runs them on 8 -virtual CPU devices. +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 @@ -28,7 +31,9 @@ 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 @@ -36,12 +41,12 @@ _VOCABS_AND_DIMS = {"t1": (19, 6), "t2": (35, 8), "t3": (10, 4)} -def _table_stack( +def _stacked_feature_specs( num_sc_per_device: int, rotation: int | None = None, global_device_count: int | None = None, -) -> Sequence[embedding_spec.TableSpec]: - """Returns the tables stacked like `keras_rs` (for the current devices).""" +) -> list[embedding_spec.FeatureSpec]: + """Returns features of tables stacked like `keras_rs` (for the devices).""" feature_specs = [] for name, (vocab, dim) in _VOCABS_AND_DIMS.items(): table_spec = embedding_spec.TableSpec( @@ -67,8 +72,21 @@ def _table_stack( 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, + global_device_count: int | None = None, +) -> Sequence[embedding_spec.TableSpec]: + """Returns the tables stacked like `keras_rs` (for the current devices).""" (table_stack,) = jte_table_stacking.get_table_stacks( - jte_embedding.get_table_specs(feature_specs) + jte_embedding.get_table_specs( + _stacked_feature_specs( + num_sc_per_device, rotation, global_device_count + ) + ) ).values() return table_stack @@ -96,6 +114,140 @@ def _sparsecore_sharding() -> jax.sharding.NamedSharding: ) +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 + self._sparsecore_built = False + 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). + # pylint: disable-next=protected-access + layout._backend_layout = jax_layout.Format(None, sharding) + # Like `keras_rs`, the table and its slot are trainable but trained on + # SparseCore, so the Keras optimizer has no slots for them. + table = self.add_weight(shape=shape, initializer="zeros", name="table") + slot = self.add_weight(shape=shape, initializer="zeros", name="slot") + for var in (table, slot): + var.overwrite_with_gradient = True + 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 + + +def _save_checkpoint( + manager_cls, + ckpt_dir: str, + model: _TestModel, +) -> None: + manager = manager_cls(ckpt_dir, choose_store_cell=False) + manager.save_model_variables(model, epoch=5) + manager.wait_until_finished() + manager.close() + + class SparseCoreCheckpointUtilsTest(parameterized.TestCase): @parameterized.parameters(1, 2, 4) @@ -270,6 +422,95 @@ def test_row_sharding_splits_uneven_rows_into_groups(self): {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.layout, 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_layers_raise(self): + unbuilt_model = _TestModel(2) + with self.assertRaisesRegex(ValueError, "must be built"): + sparsecore_checkpoint_utils.sparsecore_layer_states(unbuilt_model) + + model = _build_model(2) + model.sc_layer._sparsecore_built = False + with self.assertRaisesRegex(ValueError, "must be built"): + sparsecore_checkpoint_utils.sparsecore_layer_states(model) + + @parameterized.parameters( + "_config", "_sparsecore_layout", "_table_and_slot_variables" + ) + def test_incomplete_sparsecore_layer_state_raises(self, attr): + model = _build_model(2) + setattr(model.sc_layer, attr, None) + with self.assertRaisesRegex(ValueError, f"is missing .*'{attr}'"): + sparsecore_checkpoint_utils.sparsecore_layer_states(model) + + def test_duplicate_sparsecore_table_paths_raise(self): + model = _build_model(2) + emb_vars = model.sc_layer.stack_variables + # Like two SparseCore layers with the same names. + model.sc_layer._table_and_slot_variables = { + "stack_a": emb_vars, + "stack_b": emb_vars, + } + with self.assertRaisesRegex(ValueError, "Duplicate SparseCore table"): + sparsecore_checkpoint_utils.sparsecore_topology(model) + + @parameterized.named_parameters( + ("v2", keras_utils.KerasOrbaxCheckpointManagerV2), + ("v3", keras_utils.KerasOrbaxCheckpointManagerV3), + ) + def test_sparsecore_topology_is_saved_with_the_checkpoint(self, manager_cls): + model = _build_model(4) + ckpt_dir = self.create_tempdir().full_path + _save_checkpoint(manager_cls, ckpt_dir, model) + self.assertEqual( + sparsecore_checkpoint_utils.saved_sparsecore_topology( + os.path.join(ckpt_dir, "5") + ), + { + model.sc_layer.stack_variables.table.path: { + "num_shards": jax.device_count() * 4, + "num_sc_per_device": 4, + } + }, + ) + + @parameterized.named_parameters( + ("v2", keras_utils.KerasOrbaxCheckpointManagerV2), + ("v3", keras_utils.KerasOrbaxCheckpointManagerV3), + ) + def test_sparsecore_topology_is_not_saved_without_sparsecore_tables( + self, manager_cls + ): + ckpt_dir = self.create_tempdir().full_path + _save_checkpoint(manager_cls, ckpt_dir, _build_model(None)) + self.assertIsNone( + sparsecore_checkpoint_utils.saved_sparsecore_topology( + os.path.join(ckpt_dir, "5") + ) + ) + if __name__ == "__main__": absltest.main()