diff --git a/recml/core/training/batch_folding.py b/recml/core/training/batch_folding.py new file mode 100644 index 0000000..ae4115c --- /dev/null +++ b/recml/core/training/batch_folding.py @@ -0,0 +1,239 @@ +# 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. +"""Folds the dense leaves of training batches into a few arrays. + +The Keras JAX training loop converts and copies every leaf of a batch to the +device one by one, in Python. With many leaves per batch, as in packed sequence +training, this holds the GIL for tens of milliseconds per step, and the thread +that starts the training steps needs the GIL too. + +`fold_dataset` concatenates the dense leaves of each batch in `tf.data`, one +array per dtype and leading dimension, so that only a few arrays cross to the +device. `model_unfold_scope` makes the jitted train step of a model split them +back, so the model sees the original batch. + +Terms: + leaf: one tensor of a batch, in the order of `keras.tree.flatten`. + leaf group: the leaves that fold into one array. They share dtype and leading + dimension, and there are 2 or more of them. + folded array: the concatenation of a leaf group, of shape + `[leading_dim, total_width]`, where the width of a leaf is the product of + its dimensions after the first. +""" + +import collections +from collections.abc import Callable, Iterator, Sequence +import contextlib +import dataclasses +import functools +import math +from typing import Any + +from absl import logging +import keras +import tensorflow as tf + + +@dataclasses.dataclass(frozen=True) +class _LeafGroups: + """Which leaves of a batch fold into which array, and how to undo it. + + Attributes: + structure: The element structure of the dataset before folding. + specs: The `tf.TensorSpec` (or other spec) of each leaf, in flatten order. + folded_leaf_groups: For each folded array, the indices of its leaves in the + order of their columns. + passthrough_leaf_indices: The indices of the leaves that pass through. + """ + + structure: Any + specs: Sequence[Any] + folded_leaf_groups: Sequence[Sequence[int]] + passthrough_leaf_indices: Sequence[int] + + def leaf_width(self, leaf_index: int) -> int: + """Returns the number of columns of a leaf in its folded array.""" + return math.prod(self.specs[leaf_index].shape[1:]) + + # All control flow here is over static Python specs, not tensors. + @tf.autograph.experimental.do_not_convert + def fold(self, *element: Any) -> tuple[tuple[Any, ...], tuple[Any, ...]]: + """Returns `(folded_arrays, passthrough_leaves)` of a dataset element.""" + # `Dataset.map` passes the components of a tuple element as arguments. + if not isinstance(self.structure, tuple): + [element] = element + leaves = keras.tree.flatten(element) + folded_arrays = [] + for leaf_group in self.folded_leaf_groups: + columns = [ # Each entry has shape [leading_dim, leaf_width]. + tf.reshape(leaves[i], [-1, self.leaf_width(i)]) for i in leaf_group + ] + folded_arrays.append(tf.concat(columns, axis=1)) + passthrough_leaves = [leaves[i] for i in self.passthrough_leaf_indices] + return tuple(folded_arrays), tuple(passthrough_leaves) + + def unfold(self, element: Any) -> Any: + """Returns the original element of a `(folded_arrays, ...)` element. + + The leading dimension of the arrays may differ from the one in `specs`, as + when the arrays hold the batches of several hosts. + + Args: + element: `(folded_arrays, passthrough_leaves)` as returned by `fold`. + + Returns: + The element in `structure`. + """ + folded_arrays, passthrough_leaves = element + leaves = [None] * len(self.specs) + for folded_array, leaf_group in zip( + folded_arrays, self.folded_leaf_groups, strict=True + ): + leading_dim = folded_array.shape[0] + first_column = 0 + for i in leaf_group: + last_column = first_column + self.leaf_width(i) + columns = folded_array[:, first_column:last_column] + leaf_shape = (leading_dim, *self.specs[i].shape[1:]) + leaves[i] = keras.ops.reshape(columns, leaf_shape) + first_column = last_column + for i, leaf in zip( + self.passthrough_leaf_indices, passthrough_leaves, strict=True + ): + leaves[i] = leaf + return keras.tree.pack_sequence_as(self.structure, leaves) + + +def fold_dataset( + dataset: tf.data.Dataset, +) -> tuple[tf.data.Dataset, Callable[[Any], Any] | None]: + """Folds the dense leaves of the elements of `dataset`. + + Leaves of rank 1 or more whose dimensions after the first are all known are + grouped by dtype and static leading dimension. Each group of 2 or more leaves + becomes one array of shape `[leading_dim, total_width]`. All other leaves pass + through unchanged. Leaves with an unknown leading dimension form their own + group; `tf.concat` fails if their leading dimensions differ at run time. + + For example, a batch of 2 rows (keys flatten in sorted order) + + {"ids": int32 [[1, 2], [3, 4]], # shape [2, 2], width 2 + "mask": int32 [[[1, 0]], [[0, 1]]], # shape [2, 1, 2], width 2 + "present": bool [True, False], # only bool leaf + "weight": float32 [0.5, 2.0]} # only float32 leaf + + folds into one int32 array, with "ids" in columns 0-1 and "mask" in columns + 2-3, while "present" and "weight" have no leaf to fold with: + + ((int32 [[1, 2, 1, 0], [3, 4, 0, 1]],), (present, weight)) + + Args: + dataset: A batched dataset. + + Returns: + `(folded_dataset, unfold)`. The elements of `folded_dataset` are + `(folded_arrays, passthrough_leaves)` tuples. `unfold` maps such an element + back to the structure of the elements of `dataset`; it accepts arrays with + a different leading dimension, as on multiple hosts, and runs under `jit`. + If no leaves can be folded, returns `(dataset, None)`. + + Raises: + TypeError: If `dataset` is not a `tf.data.Dataset`. + """ + if not isinstance(dataset, tf.data.Dataset): + raise TypeError( + f"Folding batches needs a `tf.data.Dataset`, got {type(dataset)}." + ) + structure = dataset.element_spec + specs = keras.tree.flatten(structure) + leaf_indices_by_dtype_and_leading_dim = collections.defaultdict(list) + for i, spec in enumerate(specs): + if ( + isinstance(spec, tf.TensorSpec) + and spec.shape.rank + and spec.shape[1:].is_fully_defined() + ): + key = (spec.dtype, spec.shape[0]) + leaf_indices_by_dtype_and_leading_dim[key].append(i) + folded_leaf_groups = [ + leaf_indices + for leaf_indices in leaf_indices_by_dtype_and_leading_dim.values() + if len(leaf_indices) > 1 + ] + if not folded_leaf_groups: + return dataset, None + folded_leaf_indices = set().union(*folded_leaf_groups) + leaf_groups = _LeafGroups( + structure=structure, + specs=specs, + folded_leaf_groups=folded_leaf_groups, + passthrough_leaf_indices=[ + i for i in range(len(specs)) if i not in folded_leaf_indices + ], + ) + logging.info( + "Folding %d of %d batch leaves into %d arrays.", + len(folded_leaf_indices), + len(specs), + len(folded_leaf_groups), + ) + return dataset.map(leaf_groups.fold), leaf_groups.unfold + + +@contextlib.contextmanager +def model_unfold_scope( + model: keras.Model, unfold: Callable[[Any], Any] | None +) -> Iterator[None]: + """Makes the train step of `model` unfold its batches. + + `model.fit` must run on the folded dataset inside this context. The model, + its compiled loss and metrics, and its optimizer must be built (for example + with `model._symbolic_build(data_batch=unfolded_dataset.element_spec)`): + Keras builds unbuilt parts from the first batch, which is folded. + + Args: + model: A built Keras model on the JAX backend. + unfold: The function returned by `fold_dataset`. `None` leaves `model` + untouched. + + Yields: + None. + + Raises: + ValueError: If `unfold` is set and `model` is not built. + """ + if unfold is None: + yield + return + if not model.built: + raise ValueError("The model must be built to train on folded batches.") + had_instance_train_step = "train_step" in vars(model) + train_step = model.train_step + + # A 2-line closure over `train_step` and `unfold` reads more simply than a + # class with the same two fields. + @functools.wraps(train_step) + def unfolded_train_step(state, data): + return train_step(state, unfold(data)) + + model.train_step = unfolded_train_step + model.train_function = None + try: + yield + finally: + if had_instance_train_step: + model.train_step = train_step + else: + del model.train_step + model.train_function = None diff --git a/recml/core/training/batch_folding_test.py b/recml/core/training/batch_folding_test.py new file mode 100644 index 0000000..f9e4c56 --- /dev/null +++ b/recml/core/training/batch_folding_test.py @@ -0,0 +1,306 @@ +# 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. +from absl.testing import absltest +import keras +import numpy as np +from recml.core.training import batch_folding +import tensorflow as tf + +_BATCH_SIZE = 4 +_NUM_BATCHES = 3 + + +def _features() -> dict[str, np.ndarray]: + rng = np.random.default_rng(0) + n = _BATCH_SIZE * _NUM_BATCHES + return { + "ids": rng.integers(0, 10, [n, 3]).astype(np.int32), + "segments": rng.integers(0, 10, [n, 2, 2]).astype(np.int32), + "values": rng.random([n, 5]).astype(np.float32), + "weight": rng.random([n]).astype(np.float32), + # The only bool leaf: it has no group to fold with. + "present": rng.random([n]) > 0.5, + } + + +def _labels() -> np.ndarray: + n = _BATCH_SIZE * _NUM_BATCHES + return np.random.default_rng(1).random([n, 1]).astype(np.float32) + + +def _dataset() -> tf.data.Dataset: + return tf.data.Dataset.from_tensor_slices((_features(), _labels())).batch( + _BATCH_SIZE, drop_remainder=True + ) + + +def _assert_trees_equal(actual, expected): + keras.tree.assert_same_structure(actual, expected) + for a, e in zip(keras.tree.flatten(actual), keras.tree.flatten(expected)): + a, e = np.asarray(a), np.asarray(e) + np.testing.assert_array_equal(a, e) + np.testing.assert_equal(a.dtype, e.dtype) + + +class _Model(keras.Model): + """Regression on the features; its loss needs every leaf.""" + + def __init__(self): + super().__init__() + self._dense = keras.layers.Dense(1) + + def call(self, inputs): + ids = keras.ops.cast(inputs["ids"], "float32") + segments = keras.ops.cast(inputs["segments"], "float32") + x = keras.ops.concatenate( + [ids, keras.ops.reshape(segments, (-1, 4)), inputs["values"]], axis=1 + ) + present = keras.ops.cast(inputs["present"], "float32") + return self._dense(x) * (inputs["weight"] * present)[:, None] + + def compute_loss( + self, x=None, y=None, y_pred=None, sample_weight=None, training=True + ): + del x, sample_weight, training + return keras.ops.mean(keras.ops.square(y_pred - y)) + + +def _built_model() -> _Model: + keras.utils.set_random_seed(0) + model = _Model() + model.compile(optimizer=keras.optimizers.SGD(0.1)) + model(next(iter(_dataset()))[0]) + return model + + +class FoldDatasetTest(absltest.TestCase): + + def test_docstring_example(self): + dataset = tf.data.Dataset.from_tensors({ + "ids": np.array([[1, 2], [3, 4]], np.int32), + "mask": np.array([[[1, 0]], [[0, 1]]], np.int32), + "present": np.array([True, False]), + "weight": np.array([0.5, 2.0], np.float32), + }) + + folded_dataset, unfold = batch_folding.fold_dataset(dataset) + + folded = next(folded_dataset.as_numpy_iterator()) + expected_folded = ( + (np.array([[1, 2, 1, 0], [3, 4, 0, 1]], np.int32),), + (np.array([True, False]), np.array([0.5, 2.0], np.float32)), + ) + _assert_trees_equal(folded, expected_folded) + _assert_trees_equal(unfold(folded), next(dataset.as_numpy_iterator())) + + def test_folds_leaves_by_dtype_and_leading_dimension(self): + folded_dataset, _ = batch_folding.fold_dataset(_dataset()) + + folded_arrays, other_leaves = folded_dataset.element_spec + # int32: ids (3) + segments (4). float32: values (5) + weight (1) + + # labels (1). + self.assertCountEqual( + [(spec.dtype, tuple(spec.shape)) for spec in folded_arrays], + [(tf.int32, (_BATCH_SIZE, 7)), (tf.float32, (_BATCH_SIZE, 7))], + ) + self.assertEqual( + [(spec.dtype, tuple(spec.shape)) for spec in other_leaves], + [(tf.bool, (_BATCH_SIZE,))], + ) + + def test_unfold_restores_the_elements(self): + dataset = _dataset() + folded_dataset, unfold = batch_folding.fold_dataset(dataset) + + num_elements = 0 + for element, folded in zip( + dataset.as_numpy_iterator(), folded_dataset.as_numpy_iterator() + ): + _assert_trees_equal(unfold(folded), element) + num_elements += 1 + self.assertEqual(num_elements, _NUM_BATCHES) + + def test_unfold_accepts_another_leading_dimension(self): + dataset = _dataset() + folded_dataset, unfold = batch_folding.fold_dataset(dataset) + elements = list(dataset.as_numpy_iterator()) + folded = list(folded_dataset.as_numpy_iterator()) + + # On several hosts, the arrays of a step hold the batches of all hosts. + joined = keras.tree.map_structure( + lambda *leaves: np.concatenate(leaves, axis=0), *folded + ) + + _assert_trees_equal( + unfold(joined), + keras.tree.map_structure( + lambda *leaves: np.concatenate(leaves, axis=0), *elements + ), + ) + + def test_folds_a_dict_element(self): + dataset = tf.data.Dataset.from_tensor_slices(_features()).batch( + _BATCH_SIZE, drop_remainder=True + ) + folded_dataset, unfold = batch_folding.fold_dataset(dataset) + + for element, folded in zip( + dataset.as_numpy_iterator(), + folded_dataset.as_numpy_iterator(), + strict=True, + ): + _assert_trees_equal(unfold(folded), element) + + def test_ragged_leaves_pass_through(self): + features = _features() + dataset = tf.data.Dataset.from_tensor_slices({ + "ids": features["ids"], + "segments": features["segments"], + "ragged": tf.ragged.constant( + [[1, 2], [3]] * (_BATCH_SIZE * _NUM_BATCHES // 2), dtype=tf.int32 + ), + }).batch(_BATCH_SIZE, drop_remainder=True) + + folded_dataset, unfold = batch_folding.fold_dataset(dataset) + + for element, folded in zip(dataset, folded_dataset, strict=True): + folded_arrays, other_leaves = folded + self.assertLen(folded_arrays, 1) + [ragged] = other_leaves + restored = unfold(folded) + self.assertIs(restored["ragged"], ragged) + self.assertEqual(ragged.to_list(), element["ragged"].to_list()) + np.testing.assert_array_equal(restored["ids"], element["ids"]) + np.testing.assert_array_equal(restored["segments"], element["segments"]) + + def test_one_leaf_per_dtype_is_not_folded(self): + features = _features() + dataset = tf.data.Dataset.from_tensor_slices( + {"ids": features["ids"], "values": features["values"]} + ).batch(_BATCH_SIZE, drop_remainder=True) + + folded_dataset, unfold = batch_folding.fold_dataset(dataset) + + self.assertIs(folded_dataset, dataset) + self.assertIsNone(unfold) + + def test_folds_leaves_with_an_unknown_leading_dimension(self): + # 10 rows batched by 4 without `drop_remainder` yield sizes [4, 4, 2]. + features = {k: v[:10] for k, v in _features().items()} + labels = _labels()[:10] + dataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch( + _BATCH_SIZE + ) + + folded_dataset, unfold = batch_folding.fold_dataset(dataset) + + folded_arrays, other_leaves = folded_dataset.element_spec + self.assertCountEqual( + [(spec.dtype, tuple(spec.shape)) for spec in folded_arrays], + [(tf.int32, (None, 7)), (tf.float32, (None, 7))], + ) + self.assertEqual( + [(spec.dtype, tuple(spec.shape)) for spec in other_leaves], + [(tf.bool, (None,))], + ) + num_elements = 0 + for element, folded in zip( + dataset.as_numpy_iterator(), folded_dataset.as_numpy_iterator() + ): + _assert_trees_equal(unfold(folded), element) + num_elements += 1 + self.assertEqual(num_elements, 3) + + def test_unknown_trailing_dimension_passes_through(self): + dataset = tf.data.Dataset.from_generator( + lambda: [ + {"a": np.zeros([2, 3], np.int32), "b": np.zeros([2, 4], np.int32)} + ], + output_signature={ + "a": tf.TensorSpec(shape=(2, None), dtype=tf.int32), + "b": tf.TensorSpec(shape=(2, None), dtype=tf.int32), + }, + ) + + folded_dataset, unfold = batch_folding.fold_dataset(dataset) + + self.assertIs(folded_dataset, dataset) + self.assertIsNone(unfold) + + def test_rejects_other_dataset_types(self): + with self.assertRaisesRegex(TypeError, "tf.data.Dataset"): + batch_folding.fold_dataset([1, 2, 3]) + + +class ModelUnfoldScopeTest(absltest.TestCase): + + def test_training_on_folded_batches_gives_the_same_weights(self): + dataset = _dataset() + expected = _built_model() + expected.fit(dataset, epochs=2, verbose=0) + model = _built_model() + folded_dataset, unfold = batch_folding.fold_dataset(dataset) + + with batch_folding.model_unfold_scope(model, unfold): + model.fit(folded_dataset, epochs=2, verbose=0) + + self.assertNotEmpty(model.trainable_weights) + for actual, want in zip(model.get_weights(), expected.get_weights()): + np.testing.assert_allclose(actual, want, rtol=1e-6) + initial = _built_model() + self.assertFalse( + np.allclose(model.get_weights()[0], initial.get_weights()[0]) + ) + + def test_restores_the_train_step_on_exit(self): + model = _built_model() + _, unfold = batch_folding.fold_dataset(_dataset()) + + with batch_folding.model_unfold_scope(model, unfold): + self.assertIn("train_step", vars(model)) + self.assertEqual(model.train_step.__name__, "train_step") + + self.assertNotIn("train_step", vars(model)) + self.assertIsNone(model.train_function) + # The model trains on plain batches again. + model.fit(_dataset(), epochs=1, verbose=0) + + def test_restores_an_instance_train_step_on_exit(self): + model = _built_model() + custom_step = lambda state, data: _Model.train_step(model, state, data) + model.train_step = custom_step + _, unfold = batch_folding.fold_dataset(_dataset()) + + with batch_folding.model_unfold_scope(model, unfold): + self.assertIsNot(model.train_step, custom_step) + + self.assertIs(model.train_step, custom_step) + + def test_no_unfold_leaves_the_model_untouched(self): + model = _built_model() + + with batch_folding.model_unfold_scope(model, None): + self.assertNotIn("train_step", vars(model)) + + def test_rejects_an_unbuilt_model(self): + model = _Model() + _, unfold = batch_folding.fold_dataset(_dataset()) + + with self.assertRaisesRegex(ValueError, "must be built"): + with batch_folding.model_unfold_scope(model, unfold): + pass + + +if __name__ == "__main__": + absltest.main() diff --git a/recml/core/training/keras_trainer.py b/recml/core/training/keras_trainer.py index be865fc..7844ecb 100644 --- a/recml/core/training/keras_trainer.py +++ b/recml/core/training/keras_trainer.py @@ -31,6 +31,7 @@ import keras import orbax.checkpoint as ocp import psutil +from recml.core.training import batch_folding from recml.core.training import core from recml.core.training import device_prefetch from recml.core.utils import py_utils @@ -130,6 +131,8 @@ def __init__( eval_callbacks: Sequence[keras.callbacks.Callback] = (), device_prefetch_depth: int = 0, skip_batch_logs: bool = False, + fold_batch_leaves: bool = False, + zero_copy_batches: bool = False, ): """Initializes the instance. @@ -165,13 +168,27 @@ def __init__( `train_and_evaluate`. skip_batch_logs: During `train`, always drop `ProgbarLogger` (`verbose=0`, which suppresses `Epoch N/M` and epoch-end stdout progress lines while - keeping TensorBoard summaries and checkpoint logs) and disable - the per-step hooks of exact `keras_utils.EpochSummaryCallback` - instances (custom subclasses are kept as-is), and also skip the - per-step conversion of `logs` to Python when no callback needs per-step - logs (see `batch_callbacks.ZeroSyncCallbackList`). Epoch-end summaries - and checkpoints are unchanged. Not applied in `train_and_evaluate`. - Defaults to False. + keeping TensorBoard summaries and checkpoint logs) and disable the + per-step hooks of exact `keras_utils.EpochSummaryCallback` instances + (custom subclasses are kept as-is), and also skip the per-step + conversion of `logs` to Python when no callback needs per-step logs (see + `batch_callbacks.ZeroSyncCallbackList`). Epoch-end summaries and + checkpoints are unchanged. Not applied in `train_and_evaluate`. Defaults + to False. + fold_batch_leaves: During `train`, concatenate the dense leaves of each + batch into one array per dtype and leading dimension in `tf.data`, and + split them back inside the jitted train step (see `batch_folding`). + Fewer arrays then cross to the device on each step, which reduces the + host time per step for batches with many leaves. The model sees the same + batches. Requires the JAX backend and a `tf.data.Dataset`. Not applied + in `train_and_evaluate`. Defaults to False. + zero_copy_batches: During `train` on the JAX backend, pass read-only + zero-copy NumPy views of dense integer, floating-point and boolean + `tf.Tensor` batch leaves to `model.fit` instead of copying every leaf + with `np.array` (see `zero_copy_batches`). Other leaves, including + bfloat16 and complex tensors, are still converted (and copied) as Keras + converts them. Requires a `tf.data.Dataset`. Not applied in + `train_and_evaluate`. Defaults to False. resume_training_launcher: Launcher to resume training. enable_xmanager_measurements: Whether to enable XManager measurements. enable_autoxprof: Whether to enable AutoXprof. @@ -179,12 +196,18 @@ def __init__( s2_logging_settings: Settings for S2 logging. Raises: - ValueError: If `device_prefetch_depth` is negative. + ValueError: If `device_prefetch_depth` is negative, or if + `fold_batch_leaves` is set on a backend other than JAX. """ if device_prefetch_depth < 0: raise ValueError( f"device_prefetch_depth must be >= 0, got {device_prefetch_depth}." ) + if fold_batch_leaves and keras.backend.backend() != "jax": + raise ValueError( + "fold_batch_leaves needs the JAX backend, got" + f" {keras.backend.backend()!r}." + ) keras.utils.set_random_seed(rng_seed) @@ -209,6 +232,8 @@ def __init__( self._steps_per_loop = steps_per_loop self._device_prefetch_depth = device_prefetch_depth self._skip_batch_logs = skip_batch_logs + self._fold_batch_leaves = fold_batch_leaves + self._zero_copy_batches = zero_copy_batches self._marker_path = os.path.join( model_dir, core.TRAINING_COMPLETE_MARKER_FILE ) @@ -377,11 +402,27 @@ def train(self, task: KerasTask) -> core.Logs: steps=self._steps_per_loop, model=model, ) - with device_prefetch.model_prefetch_scope( - model, depth=self._device_prefetch_depth + unfold = None + if self._fold_batch_leaves: + # `fit` builds the model, the compiled loss and metrics, and the + # optimizer from its first batch, which would be folded. + model._symbolic_build(data_batch=dataset.element_spec) # pylint: disable=protected-access + dataset, unfold = batch_folding.fold_dataset(dataset) + fit_data: Any = dataset + if self._zero_copy_batches and keras.backend.backend() == "jax": + # Imported here: it needs private Keras modules that the Keras version + # pinned in requirements.txt does not have. + from recml.core.training import zero_copy_batches as zero_copy_batches_lib # pylint: disable=g-import-not-at-top + + fit_data = zero_copy_batches_lib.zero_copy_adapter(dataset) + with ( + batch_folding.model_unfold_scope(model, unfold), + device_prefetch.model_prefetch_scope( + model, depth=self._device_prefetch_depth + ), ): history = model.fit( - dataset, + fit_data, epochs=epochs, steps_per_epoch=self._steps_per_loop, callbacks=callbacks, @@ -431,10 +472,16 @@ def evaluate(self, task: KerasTask) -> core.Logs: def train_and_evaluate(self, task: KerasTask) -> core.Logs: """Trains and evaluates a Keras model.""" - if self._device_prefetch_depth or self._skip_batch_logs: + if ( + self._device_prefetch_depth + or self._skip_batch_logs + or self._fold_batch_leaves + or self._zero_copy_batches + ): logging.warning( - "device_prefetch_depth and skip_batch_logs only apply to `train`;" - " ignoring them in `train_and_evaluate`." + "device_prefetch_depth, skip_batch_logs, fold_batch_leaves and" + " zero_copy_batches only apply to `train`; ignoring them in" + " `train_and_evaluate`." ) train_dataset = task.create_dataset( training=True, **self._maybe_get_dataset_kws(task) diff --git a/recml/core/training/keras_trainer_options_test.py b/recml/core/training/keras_trainer_options_test.py index f140a41..42e3eec 100644 --- a/recml/core/training/keras_trainer_options_test.py +++ b/recml/core/training/keras_trainer_options_test.py @@ -25,7 +25,9 @@ from recml.core.training import batch_callbacks from recml.core.training import device_prefetch from recml.core.training import keras_trainer +from recml.core.training import zero_copy_batches from recml.core.utils import keras_utils +import tensorflow as tf def _trainer(tmpdir, **overrides): @@ -58,7 +60,10 @@ def fit(self, dataset, epochs, steps_per_epoch, callbacks, initial_epoch): self.make_train_function() fn = self.train_function self.seen_in_fit.append((fn, callbacks)) - it = iter(dataset) + if isinstance(dataset, zero_copy_batches._ZeroCopyAdapter): + it = dataset.get_jax_iterator() + else: + it = iter(dataset) for _ in range(steps_per_epoch): batch, _ = fn(None, it) self.batches_in_fit.append(batch) @@ -68,6 +73,32 @@ def summary(self, print_fn=None): del print_fn +class _StepModel(_DummyModel): + """A `_DummyModel` whose train function runs `train_step`, like Keras.""" + + built = True + + def __init__(self): + super().__init__() + self.datasets_in_fit = [] + + def train_step(self, state, data): + return data, state + + def _symbolic_build(self, iterator=None, data_batch=None): + del iterator, data_batch + + def make_train_function(self, force=False): + if self.train_function is not None and not force: + return + train_step = self.train_step + self.train_function = lambda state, it: train_step(state, next(it)) + + def fit(self, dataset, **kwargs): + self.datasets_in_fit.append(dataset) + return super().fit(dataset, **kwargs) + + def _task(model, exported): return types.SimpleNamespace( create_dataset=lambda training: list(range(100)), @@ -147,6 +178,158 @@ def test_train_and_evaluate_warns_that_prefetch_is_not_applied(self): trainer.train_and_evaluate(task) self.assertIn("device_prefetch_depth", logs.output[0]) + def test_fold_batch_leaves_folds_the_dataset_and_unfolds_in_the_step(self): + features = { + "a": np.arange(40, dtype=np.float32).reshape([20, 2]), + "b": -np.arange(60, dtype=np.float32).reshape([20, 3]), + } + dataset = tf.data.Dataset.from_tensor_slices(features).batch( + 2, drop_remainder=True + ) + model = _StepModel() + task = types.SimpleNamespace( + create_dataset=lambda training: dataset, + create_model=lambda: model, + export_model=lambda m, model_dir: None, + ) + + _trainer(self._tmpdir, fold_batch_leaves=True).train(task) + + # `fit` got 1 folded array per batch instead of the 2 leaves. + self.assertLen(model.datasets_in_fit, 1) + dataset_in_fit = model.datasets_in_fit[0] + folded_specs, other_specs = dataset_in_fit.element_spec + self.assertLen(folded_specs, 1) + self.assertEmpty(other_specs) + # The train step got the original batches. + self.assertLen(model.batches_in_fit, 5) + for step, batch in enumerate(model.batches_in_fit): + self.assertCountEqual(batch, ["a", "b"]) + for key, value in batch.items(): + np.testing.assert_array_equal( + value, features[key][2 * step : 2 * step + 2] + ) + self.assertNotIn("train_step", vars(model)) + + def test_rejects_fold_batch_leaves_on_other_backends(self): + with mock.patch.object(keras.backend, "backend", return_value="tensorflow"): + with self.assertRaisesRegex(ValueError, "needs the JAX backend"): + _trainer(self._tmpdir, fold_batch_leaves=True) + + def test_fold_batch_leaves_with_compiled_loss_and_metrics(self): + rng = np.random.default_rng(0) + x = rng.random((20, 4)).astype(np.float32) + y = rng.random((20, 1)).astype(np.float32) + + def build_model(): + keras.utils.set_random_seed(0) + inputs = keras.Input(shape=(4,), dtype="float32") + model = keras.Model(inputs=inputs, outputs=keras.layers.Dense(1)(inputs)) + # The compiled loss and metrics stay unbuilt until the first batch. + model.compile(optimizer="sgd", loss="mse", metrics=["mae"]) + return model + + def train(model_dir, **options): + model = build_model() + os.makedirs(model_dir) + history = _trainer( + model_dir, train_steps=10, steps_per_loop=5, **options + ).train( + types.SimpleNamespace( + create_dataset=lambda training: tf.data.Dataset.from_tensor_slices( + (x, y) + ) + .batch(2, drop_remainder=True) + .repeat(), + create_model=lambda: model, + export_model=lambda m, model_dir: None, + ) + ) + return model, history.history + + want_model, want_history = train(os.path.join(self._tmpdir, "plain")) + model, history = train( + os.path.join(self._tmpdir, "folded"), fold_batch_leaves=True + ) + + for actual, want in zip( + model.get_weights(), want_model.get_weights(), strict=True + ): + np.testing.assert_allclose(actual, want, rtol=1e-6) + self.assertCountEqual(history, want_history) + for key, values in want_history.items(): + np.testing.assert_allclose(history[key], values, rtol=1e-6, err_msg=key) + + def test_train_and_evaluate_warns_that_folding_is_not_applied(self): + trainer = _trainer(self._tmpdir, fold_batch_leaves=True) + task = mock.Mock() + task.create_dataset.side_effect = RuntimeError("stop") + with self.assertLogs(level="WARNING") as logs: + with self.assertRaisesRegex(RuntimeError, "stop"): + trainer.train_and_evaluate(task) + self.assertIn("fold_batch_leaves", logs.output[0]) + + def test_zero_copy_batches_wraps_the_dataset_in_a_zero_copy_adapter(self): + features = np.arange(20, dtype=np.float32).reshape([10, 2]) + dataset = tf.data.Dataset.from_tensor_slices(features).batch( + 2, drop_remainder=True + ) + model = _StepModel() + task = types.SimpleNamespace( + create_dataset=lambda training: dataset, + create_model=lambda: model, + export_model=lambda m, model_dir: None, + ) + + _trainer(self._tmpdir, zero_copy_batches=True).train(task) + + self.assertLen(model.datasets_in_fit, 1) + self.assertIsInstance( + model.datasets_in_fit[0], zero_copy_batches._ZeroCopyAdapter + ) + self.assertLen(model.batches_in_fit, 5) + self.assertFalse(model.batches_in_fit[0].flags.owndata) + + def test_zero_copy_batches_with_fold_batch_leaves(self): + features = { + "a": np.arange(40, dtype=np.float32).reshape([20, 2]), + "b": -np.arange(60, dtype=np.float32).reshape([20, 3]), + } + dataset = tf.data.Dataset.from_tensor_slices(features).batch( + 2, drop_remainder=True + ) + model = _StepModel() + task = types.SimpleNamespace( + create_dataset=lambda training: dataset, + create_model=lambda: model, + export_model=lambda m, model_dir: None, + ) + + _trainer( + self._tmpdir, fold_batch_leaves=True, zero_copy_batches=True + ).train(task) + + self.assertLen(model.datasets_in_fit, 1) + self.assertIsInstance( + model.datasets_in_fit[0], zero_copy_batches._ZeroCopyAdapter + ) + self.assertLen(model.batches_in_fit, 5) + for step, batch in enumerate(model.batches_in_fit): + self.assertCountEqual(batch, ["a", "b"]) + for key, value in batch.items(): + np.testing.assert_array_equal( + value, features[key][2 * step : 2 * step + 2] + ) + + def test_train_and_evaluate_warns_that_zero_copy_batches_is_not_applied(self): + trainer = _trainer(self._tmpdir, zero_copy_batches=True) + task = mock.Mock() + task.create_dataset.side_effect = RuntimeError("stop") + with self.assertLogs(level="WARNING") as logs: + with self.assertRaisesRegex(RuntimeError, "stop"): + trainer.train_and_evaluate(task) + self.assertIn("zero_copy_batches", logs.output[0]) + def test_train_and_evaluate_warns_that_skip_batch_logs_is_not_applied(self): trainer = _trainer(self._tmpdir, skip_batch_logs=True) task = mock.Mock() diff --git a/recml/core/training/state_folding.py b/recml/core/training/state_folding.py new file mode 100644 index 0000000..873e3bf --- /dev/null +++ b/recml/core/training/state_folding.py @@ -0,0 +1,390 @@ +# 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. +"""Folds small replicated leaves of the Keras JAX training state into vectors. + +The Keras JAX training loop passes every variable of the model, the optimizer, +and the metrics to the jitted train step as its own array, and receives one +array back for each, plus one array per scalar log. The host pays per-array +overhead on every step: Python argument parsing, per-chip buffer handles and +output buffer allocation, and Python array wrapper construction/destruction. +Models with hundreds of metrics pass thousands of small arrays this way. + +`model_state_folding_scope` concatenates the small fully-replicated arrays of +each of the 4 state lists into one 1-D vector per dtype. Inside the jitted +train step, the vectors are split and reshaped back before calling the model's +`train_step`, and the updated state is folded back on return. Large or sharded +arrays pass through unchanged. + +Terms: + leaf: one array of a state list. + small leaf: a fully replicated leaf with at most `max_leaf_size` elements. + vector: the 1-D concatenation of the raveled small leaves of one dtype. + passthrough leaf: a leaf that is not in a vector (large, sharded, or the + only small leaf of its dtype). + folded list: the passthrough leaves in their order, then the vectors. +""" + +from collections.abc import Iterator, Sequence +import contextlib +from typing import Any + +from absl import logging +import jax +from jax import flatten_util +import keras +import numpy as np + +_STATE_KEYS = ( + "trainable_variables", + "non_trainable_variables", + "optimizer_variables", + "metrics_variables", +) + + +@contextlib.contextmanager +def _override_attribute(obj: Any, name: str, value: Any) -> Iterator[None]: + """Sets `obj.` to `value` and restores the instance attribute on exit. + + If `obj` had no instance attribute `name` before (only a class attribute or + method), the instance attribute is deleted on exit. + + Args: + obj: The object to change. + name: The attribute name. + value: The value to set while the context is active. + + Yields: + None. + """ + had_instance_attribute = name in vars(obj) + saved_value = vars(obj).get(name) + setattr(obj, name, value) + try: + yield + finally: + if had_instance_attribute: + setattr(obj, name, saved_value) + else: + delattr(obj, name) + + +def _replicated_vector_sharding( + leaf_sharding: jax.sharding.Sharding, +) -> jax.sharding.Sharding: + """Returns a replicated sharding on the devices of a small leaf.""" + if isinstance(leaf_sharding, jax.sharding.NamedSharding): + return jax.sharding.NamedSharding( + leaf_sharding.mesh, jax.sharding.PartitionSpec() + ) + return leaf_sharding + + +class _ListLayout: + """How one state list is folded. + + For example, with `max_leaf_size=2` the list + + [w: float32[2, 2], b: float32[2], c: float32[], n: int32[]] + + has small leaves b, c (float32) and n (int32). w has 4 elements, so it + passes through. n is the only small int32 leaf, so it passes through too. + The folded list is + + [w, n, float32[3] vector = concat(ravel(b), ravel(c))] + + and `passthrough_leaf_indices == [0, 3]`, `vector_leaf_indices == [[1, 2]]`. + """ + + def __init__(self, leaves: Sequence[jax.Array], max_leaf_size: int): + self.num_leaves = len(leaves) + self.leaf_shardings = [leaf.sharding for leaf in leaves] + small_leaf_indices_by_dtype: dict[np.dtype[Any], list[int]] = {} + for i, leaf in enumerate(leaves): + if leaf.size <= max_leaf_size and leaf.sharding.is_fully_replicated: + dtype = np.dtype(leaf.dtype) + small_leaf_indices_by_dtype.setdefault(dtype, []).append(i) + # For each vector, the indices of its leaves in the order of their entries. + self.vector_leaf_indices = [ + leaf_indices + for leaf_indices in small_leaf_indices_by_dtype.values() + if len(leaf_indices) > 1 + ] + leaf_indices_in_vectors = set().union(*self.vector_leaf_indices) + self.passthrough_leaf_indices = [ + i for i in range(self.num_leaves) if i not in leaf_indices_in_vectors + ] + self.vector_shardings = [ + _replicated_vector_sharding(self.leaf_shardings[leaf_indices[0]]) + for leaf_indices in self.vector_leaf_indices + ] + # `unravel` splits a vector back into its leaves with their shapes. + self._unravel_vectors = [ + flatten_util.ravel_pytree(vector_leaves)[1] + for vector_leaves in self.leaves_by_vector(leaves) + ] + + @property + def folded_list_length(self) -> int: + """Returns the length of the folded list.""" + return len(self.passthrough_leaf_indices) + len(self.vector_leaf_indices) + + def leaves_by_vector(self, leaves: Sequence[Any]) -> list[list[Any]]: + """Returns, for each vector, the entries of `leaves` that go into it.""" + return [ + [leaves[i] for i in leaf_indices] + for leaf_indices in self.vector_leaf_indices + ] + + def concat_vectors( + self, leaves_by_vector: Sequence[Sequence[jax.Array]] + ) -> list[jax.Array]: + """Returns one raveled 1-D vector per entry of `leaves_by_vector`.""" + return [ + flatten_util.ravel_pytree(list(vector_leaves))[0] + for vector_leaves in leaves_by_vector + ] + + def split_vectors( + self, vectors: Sequence[jax.Array] + ) -> list[list[jax.Array]]: + """Returns the leaves of each vector, reshaped back. Inverts the above.""" + return [ + unravel(vector) + for unravel, vector in zip(self._unravel_vectors, vectors) + ] + + def fold_list( + self, leaves: Sequence[Any], vectors: Sequence[Any] + ) -> list[Any]: + """Returns the folded list of `leaves`, given its `vectors`.""" + return [leaves[i] for i in self.passthrough_leaf_indices] + list(vectors) + + def vectors_in_folded_list(self, folded_list: Sequence[Any]) -> list[Any]: + """Returns the vectors at the end of a folded list.""" + return list(folded_list[len(self.passthrough_leaf_indices) :]) + + def unfold_list( + self, + folded_list: Sequence[Any], + leaves_by_vector: Sequence[Sequence[Any]], + ) -> list[Any]: + """Returns the unfolded leaves, given the split vectors of `folded_list`.""" + leaves = [None] * self.num_leaves + for i, leaf in zip(self.passthrough_leaf_indices, folded_list): + leaves[i] = leaf + for vector_leaves, leaf_indices in zip( + leaves_by_vector, self.vector_leaf_indices + ): + for i, leaf in zip(leaf_indices, vector_leaves): + leaves[i] = leaf + return leaves + + +class _StateLayout: + """How the 4 lists of the training state are folded.""" + + def __init__(self, state: Sequence[Sequence[jax.Array]], max_leaf_size: int): + self.list_layouts = [_ListLayout(leaves, max_leaf_size) for leaves in state] + # `fold` and `unfold` pass only the small leaves and vectors through these: + # passing the large leaves to `jit` outside a trace would copy them. + self._concat_vectors_jit = jax.jit( + self._concat_all_vectors, + out_shardings=[layout.vector_shardings for layout in self.list_layouts], + ) + self._split_vectors_jit = jax.jit( + self._split_all_vectors, + out_shardings=[ + layout.leaves_by_vector(layout.leaf_shardings) + for layout in self.list_layouts + ], + ) + logging.info( + "Folding the training state from %s to %s arrays.", + [layout.num_leaves for layout in self.list_layouts], + [layout.folded_list_length for layout in self.list_layouts], + ) + + def _concat_all_vectors( + self, leaves_by_vector_by_list: Sequence[Sequence[Sequence[jax.Array]]] + ) -> list[list[jax.Array]]: + return [ + layout.concat_vectors(leaves_by_vector) + for layout, leaves_by_vector in zip( + self.list_layouts, leaves_by_vector_by_list + ) + ] + + def _split_all_vectors( + self, vectors_by_list: Sequence[Sequence[jax.Array]] + ) -> list[list[Any]]: + return [ + layout.split_vectors(vectors) + for layout, vectors in zip(self.list_layouts, vectors_by_list) + ] + + def fold( + self, state: Sequence[Sequence[jax.Array]] + ) -> tuple[list[jax.Array], ...]: + """Returns `state` with its small replicated leaves folded into vectors.""" + leaves_by_vector_by_list = [ + layout.leaves_by_vector(leaves) + for layout, leaves in zip(self.list_layouts, state) + ] + vectors_by_list = self._concat_vectors_jit(leaves_by_vector_by_list) + return tuple( + layout.fold_list(leaves, vectors) + for layout, leaves, vectors in zip( + self.list_layouts, state, vectors_by_list + ) + ) + + def unfold( + self, folded: Sequence[Sequence[jax.Array]] + ) -> tuple[list[jax.Array], ...]: + """Returns the unfolded `state` of `folded`.""" + vectors_by_list = [ + layout.vectors_in_folded_list(folded_list) + for layout, folded_list in zip(self.list_layouts, folded) + ] + leaves_by_vector_by_list = self._split_vectors_jit(vectors_by_list) + return tuple( + layout.unfold_list(folded_list, leaves_by_vector) + for layout, folded_list, leaves_by_vector in zip( + self.list_layouts, folded, leaves_by_vector_by_list + ) + ) + + def fold_shardings( + self, shardings: Sequence[Sequence[jax.sharding.Sharding]] + ) -> tuple[list[jax.sharding.Sharding], ...]: + """Returns the shardings of the folded state.""" + return tuple( + layout.fold_list(list_shardings, layout.vector_shardings) + for layout, list_shardings in zip(self.list_layouts, shardings) + ) + + +class _Folding: + """The folded train step of a model and the state conversions around it.""" + + def __init__(self, model: keras.Model, max_leaf_size: int): + self._model = model + self._max_leaf_size = max_leaf_size + self._layout: _StateLayout | None = None + self.original_train_step = model.train_step + self.original_make_train_function = model.make_train_function + self.original_state_sync = model.state_sync + # pylint: disable=protected-access + self.original_get_jax_state = model._get_jax_state + self.original_get_state_sharding_spec = model._get_state_sharding_spec + # pylint: enable=protected-access + + def _get_layout(self, state: Any = None) -> _StateLayout: + """Returns the layout, built on first use from the unfolded `state`.""" + if self._layout is None: + if state is None: + state = self.original_get_jax_state( + trainable_variables=True, + non_trainable_variables=True, + optimizer_variables=True, + metrics_variables=True, + ) + self._layout = _StateLayout(state, self._max_leaf_size) + return self._layout + + def folded_train_step( + self, state: Sequence[Sequence[jax.Array]], data: Any + ) -> tuple[Any, tuple[list[jax.Array], ...]]: + layout = self._get_layout() + logs, state = self.original_train_step(layout.unfold(state), data) + return logs, layout.fold(state) + + def folded_make_train_function(self, force: bool = False) -> None: + """Builds `model.train_function` with folded state shardings.""" + if self._model.train_function is not None and not force: + return + # Only the train function sees folded shardings; `evaluate` and + # `validation_data` keep the unfolded ones. + with _override_attribute( + self._model, + "_get_state_sharding_spec", + self.folded_get_state_sharding_spec, + ): + self.original_make_train_function(force=force) + + def folded_get_jax_state(self, **kwargs: bool) -> Any: + state = self.original_get_jax_state(**kwargs) + # `fit` and `train_on_batch` ask for all 4 lists. `evaluate`, `predict` and + # the checkpoint helpers ask for fewer and keep getting unfolded state. + if not all(kwargs.get(key, False) for key in _STATE_KEYS): + return state + return self._get_layout(state).fold(state) + + def folded_get_state_sharding_spec( + self, + ) -> tuple[list[jax.sharding.Sharding], ...]: + unfolded_shardings = self.original_get_state_sharding_spec() + return self._get_layout().fold_shardings(unfolded_shardings) + + def folded_state_sync(self) -> None: + """Unfolds `_jax_state` before syncing back to the model variables.""" + model = self._model + jax_state = getattr(model, "_jax_state", None) + if ( + jax_state + and not model._jax_state_synced # pylint: disable=protected-access + and all(jax_state.get(key) is not None for key in _STATE_KEYS) + ): + state = [jax_state[key] for key in _STATE_KEYS] + unfolded = self._get_layout().unfold(state) + model._jax_state = dict( # pylint: disable=protected-access + zip(_STATE_KEYS, unfolded) + ) + self.original_state_sync() + + +@contextlib.contextmanager +def model_state_folding_scope( + model: keras.Model, max_leaf_size: int +) -> Iterator[None]: + """Makes `model.fit` pass small replicated arrays of the state folded. + + Args: + model: A Keras model on the JAX backend. + max_leaf_size: Arrays with at most this many elements are folded. `0` leaves + `model` untouched. + + Yields: + None. + """ + if max_leaf_size == 0: + yield + return + folding = _Folding(model, max_leaf_size) + model.train_function = None + try: + with ( + _override_attribute(model, "train_step", folding.folded_train_step), + _override_attribute( + model, "make_train_function", folding.folded_make_train_function + ), + _override_attribute(model, "state_sync", folding.folded_state_sync), + _override_attribute( + model, "_get_jax_state", folding.folded_get_jax_state + ), + ): + yield + finally: + model.train_function = None diff --git a/recml/core/training/state_folding_test.py b/recml/core/training/state_folding_test.py new file mode 100644 index 0000000..657e466 --- /dev/null +++ b/recml/core/training/state_folding_test.py @@ -0,0 +1,296 @@ +# 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. +from absl.testing import absltest +import jax +import jax.numpy as jnp +import keras +import numpy as np +from recml.core.training import state_folding + + +def _data() -> tuple[np.ndarray, np.ndarray]: + rng = np.random.default_rng(0) + x = rng.random([32, 8]).astype(np.float32) + y = rng.random([32, 1]).astype(np.float32) + return x, y + + +class _ExtraLogModel(keras.Model): + """Adds a non-metric scalar key to `train_step` logs.""" + + def __init__(self): + super().__init__() + self._dense1 = keras.layers.Dense(4, activation="relu") + self._bn = keras.layers.BatchNormalization() + self._dense2 = keras.layers.Dense(1) + + def call(self, x, training=False): + return self._dense2(self._bn(self._dense1(x), training=training)) + + def train_step(self, state, data): + logs, state = super().train_step(state, data) + logs["extra_scalar"] = jnp.float32(7.5) + return logs, state + + +def _model(build: bool = True) -> keras.Model: + keras.utils.set_random_seed(0) + model = _ExtraLogModel() + model.compile( + optimizer=keras.optimizers.Adam(0.01), + loss="mse", + metrics=["mae", "mse", keras.metrics.RootMeanSquaredError()], + ) + if build: + model.build((None, 8)) + return model + + +class StateFoldingTest(absltest.TestCase): + + def _fit( + self, + max_leaf_size: int, + steps_per_execution: int = 1, + build: bool = True, + ) -> tuple[keras.Model, dict[str, list[float]], list[int]]: + x, y = _data() + model = _model(build=build) + model.steps_per_execution = steps_per_execution + num_args = [] + with state_folding.model_state_folding_scope(model, max_leaf_size): + model.make_train_function() + train_function = model.train_function + + def counting_train_function(state, iterator): + num_args.append(len(jax.tree.leaves(state))) + return train_function(state, iterator) + + model.train_function = counting_train_function + history = model.fit( + x, y, batch_size=8, epochs=3, shuffle=False, verbose=0 + ) + return model, history.history, num_args + + def _assert_same_training(self, model, history, expected, expected_history): + self.assertNotEmpty(model.get_weights()) + for actual, want in zip( + model.get_weights(), expected.get_weights(), strict=True + ): + np.testing.assert_allclose(actual, want, rtol=1e-6, atol=1e-7) + self.assertCountEqual(history, expected_history) + for key, values in expected_history.items(): + np.testing.assert_allclose(history[key], values, rtol=1e-6, err_msg=key) + + def test_same_weights_and_logs(self): + expected, expected_history, expected_args = self._fit(max_leaf_size=0) + model, history, num_args = self._fit(max_leaf_size=16) + + self.assertLess(max(num_args), min(expected_args)) + self.assertLen(num_args, len(expected_args)) + self._assert_same_training(model, history, expected, expected_history) + for actual, want in zip( + model.optimizer.variables, expected.optimizer.variables, strict=True + ): + np.testing.assert_allclose( + np.asarray(actual), np.asarray(want), rtol=1e-6, atol=1e-7 + ) + self.assertEqual(history["extra_scalar"], [7.5, 7.5, 7.5]) + + def test_same_weights_with_steps_per_execution(self): + expected, expected_history, _ = self._fit( + max_leaf_size=0, steps_per_execution=2 + ) + model, history, _ = self._fit(max_leaf_size=16, steps_per_execution=2) + + self._assert_same_training(model, history, expected, expected_history) + + def test_unbuilt_model_folds(self): + expected, expected_history, expected_args = self._fit( + max_leaf_size=0, build=False + ) + model, history, num_args = self._fit(max_leaf_size=16, build=False) + + self.assertLess(max(num_args), min(expected_args)) + self._assert_same_training(model, history, expected, expected_history) + + def test_validation_data_under_distribution_works(self): + x, y = _data() + distribution = keras.distribution.DataParallel() + with distribution.scope(): + expected = _model() + want_history = expected.fit( + x, + y, + validation_data=(x[:8], y[:8]), + batch_size=8, + epochs=2, + shuffle=False, + verbose=0, + ).history + + model = _model() + with state_folding.model_state_folding_scope(model, max_leaf_size=16): + got_history = model.fit( + x, + y, + validation_data=(x[:8], y[:8]), + batch_size=8, + epochs=2, + shuffle=False, + verbose=0, + ).history + + self.assertCountEqual(got_history, want_history) + for key, values in want_history.items(): + np.testing.assert_allclose( + got_history[key], values, rtol=1e-6, err_msg=key + ) + + def test_sharded_small_leaves_pass_through(self): + devices = jax.devices() + self.assertLen(devices, 4) + mesh = jax.sharding.Mesh(np.array(devices), ("model",)) + sharded = jax.sharding.NamedSharding( + mesh, jax.sharding.PartitionSpec("model") + ) + replicated = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()) + leaves = [ + jax.device_put(jnp.zeros([4], jnp.float32), sharded), + jax.device_put(jnp.zeros([4], jnp.float32), replicated), + jax.device_put(jnp.zeros([], jnp.float32), replicated), + ] + + layout = state_folding._ListLayout(leaves, max_leaf_size=16) + + # The sharded leaf is small but passes through, keeping its sharding. + self.assertEqual(layout.passthrough_leaf_indices, [0]) + self.assertEqual(layout.vector_leaf_indices, [[1, 2]]) + + def test_same_weights_under_model_parallel(self): + x, y = _data() + mesh = keras.distribution.DeviceMesh( + (2, 2), ["batch", "model"], devices=keras.distribution.list_devices() + ) + layout_map = keras.distribution.LayoutMap(mesh) + layout_map[".*dense/kernel"] = (None, "model") + # Small but sharded: passes through. + layout_map[".*dense/bias"] = ("model",) + distribution = keras.distribution.ModelParallel( + layout_map=layout_map, batch_dim_name="batch" + ) + results = [] + for max_leaf_size in (0, 16): + with distribution.scope(): + keras.utils.set_random_seed(0) + inputs = keras.Input((8,)) + hidden = keras.layers.Dense(4, name="dense")(inputs) + model = keras.Model(inputs, keras.layers.Dense(1)(hidden)) + model.compile(optimizer="adam", loss="mse", metrics=["mae"]) + with state_folding.model_state_folding_scope(model, max_leaf_size): + history = model.fit( + x, y, batch_size=8, epochs=2, shuffle=False, verbose=0 + ).history + results.append((model, history)) + + expected, expected_history = results[0] + model, history = results[1] + self._assert_same_training(model, history, expected, expected_history) + for variable, want in zip(model.weights, expected.weights, strict=True): + self.assertEqual( + variable.value.sharding, want.value.sharding, msg=variable.path + ) + + def test_restores_the_model_on_exit(self): + model, _, _ = self._fit(max_leaf_size=16) + + for name in ( + "train_step", + "make_train_function", + "state_sync", + "_get_jax_state", + "_get_state_sharding_spec", + ): + self.assertNotIn(name, vars(model), msg=name) + x, y = _data() + model.fit(x, y, batch_size=8, epochs=1, verbose=0) + model.evaluate(x, y, batch_size=8, verbose=0) + + custom_train_step = lambda state, data: _ExtraLogModel.train_step( + model, state, data + ) + model.train_step = custom_train_step + with state_folding.model_state_folding_scope(model, max_leaf_size=16): + self.assertIsNot(model.train_step, custom_train_step) + self.assertIs(model.train_step, custom_train_step) + + def test_state_layout_preserves_replicated_named_shardings(self): + devices = jax.devices() + self.assertLen(devices, 4) + mesh = jax.sharding.Mesh( + np.array(devices).reshape(2, 2), ("batch", "model") + ) + replicated_2d = jax.sharding.NamedSharding( + mesh, jax.sharding.PartitionSpec(None, None) + ) + replicated_0d = jax.sharding.NamedSharding( + mesh, jax.sharding.PartitionSpec() + ) + state = [[ + jax.device_put(jnp.ones([2, 2], jnp.float32), replicated_2d), + jax.device_put(jnp.array(3.0, jnp.float32), replicated_0d), + ]] + + layout = state_folding._StateLayout(state, max_leaf_size=4) + folded = layout.fold(state) + unfolded = layout.unfold(folded) + + self.assertEqual(folded[0][0].sharding, replicated_0d) + self.assertEqual(unfolded[0][0].sharding, replicated_2d) + self.assertEqual(unfolded[0][1].sharding, replicated_0d) + np.testing.assert_array_equal(unfolded[0][0], np.ones([2, 2], np.float32)) + np.testing.assert_array_equal(unfolded[0][1], 3.0) + + def test_zero_max_leaf_size_leaves_the_model_untouched(self): + model = _model() + + with state_folding.model_state_folding_scope(model, max_leaf_size=0): + self.assertNotIn("train_step", vars(model)) + + def test_list_layout_folds_small_leaves_by_dtype(self): + leaves = [ + jnp.zeros([3], jnp.int32), # 3 elements: too large, passes through. + jnp.zeros([], jnp.float32), + jnp.zeros([], jnp.int32), + jnp.zeros([1, 1], jnp.float32), + jnp.zeros([2], jnp.int32), + ] + + layout = state_folding._ListLayout(leaves, max_leaf_size=2) + + self.assertEqual(layout.passthrough_leaf_indices, [0]) + self.assertEqual(layout.vector_leaf_indices, [[1, 3], [2, 4]]) + self.assertEqual(layout.folded_list_length, 3) + float_leaves, int_leaves = layout.split_vectors([ + jnp.array([5.0, 6.0], jnp.float32), + jnp.array([7, 8, 9], jnp.int32), + ]) + np.testing.assert_array_equal(float_leaves[0], np.float32(5.0)) + np.testing.assert_array_equal(float_leaves[1], [[6.0]]) + np.testing.assert_array_equal(int_leaves[0], np.int32(7)) + np.testing.assert_array_equal(int_leaves[1], [8, 9]) + + +if __name__ == "__main__": + absltest.main() diff --git a/recml/core/training/zero_copy_batches.py b/recml/core/training/zero_copy_batches.py new file mode 100644 index 0000000..9c14d7f --- /dev/null +++ b/recml/core/training/zero_copy_batches.py @@ -0,0 +1,99 @@ +# 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. +"""Hands `tf.data` batches to the Keras JAX training loop without a copy. + +Keras converts each leaf of a `tf.data` batch with `np.array`, which copies the +leaf on the Python thread that feeds the device. For a batch of tens of MB this +copy takes longer than the transfer to the device. + +`zero_copy_adapter` wraps a dataset in a Keras data adapter that passes NumPy +views of the dense integer, floating-point and boolean `tf.Tensor` leaves +instead. Other leaves, including bfloat16 and complex tensors, are converted +(and copied) as Keras converts them. A view keeps its tensor alive, and +`tf.data` never writes to a tensor that it has produced, so the views are safe +to read for as long as they exist. The views are read-only. +""" + +from collections.abc import Iterator +from typing import Any + +import jax.numpy as jnp +import keras +from keras.src.backend.tensorflow.ops import core as keras_tf_core +from keras.src.trainers.data_adapters import data_adapter_utils +from keras.src.trainers.data_adapters import tf_dataset_adapter +import numpy as np +import tensorflow as tf + + +def _to_jax_or_numpy(x: Any) -> Any: + """Returns a zero-copy NumPy view of a dense tensor, or Keras' conversion.""" + if isinstance(x, tf.SparseTensor): + return data_adapter_utils.tf_sparse_to_jax_sparse(x) + if ( + isinstance(x, tf.Tensor) + and (x.dtype.is_integer or x.dtype.is_floating or x.dtype.is_bool) + and x.dtype != tf.bfloat16 + ): + return np.asarray(x) + return keras_tf_core.convert_to_numpy(x) + + +class _ZeroCopyAdapter(tf_dataset_adapter.TFDatasetAdapter): + """A `TFDatasetAdapter` that does not copy dense leaves for JAX.""" + + def get_jax_iterator(self, super_batch: int | None = None) -> Iterator[Any]: + dataset = self._dataset + if super_batch and self.batch_size is not None: + dataset = dataset.batch(super_batch).prefetch(tf.data.AUTOTUNE) + iterator = ( + keras.tree.map_structure(_to_jax_or_numpy, batch, none_is_leaf=False) + for batch in dataset + ) + if super_batch and self.batch_size is None: + iterator = data_adapter_utils.super_batch_iterator( + iterator, super_batch, stack_fn=jnp.stack + ) + yield from iterator + + +def zero_copy_adapter( + dataset: tf.data.Dataset | tf.distribute.DistributedDataset, +) -> _ZeroCopyAdapter: + """Returns a Keras data adapter for `dataset` that does not copy batches. + + Pass the result to `model.fit` in place of `dataset`. Dense integer, + floating-point and boolean tensors become read-only NumPy views. Other + leaves, including bfloat16 and complex tensors, are converted (and copied) + as Keras converts them. + + Args: + dataset: A `tf.data.Dataset` or `tf.distribute.DistributedDataset`. + + Returns: + A Keras data adapter that yields NumPy views of the dense leaves on the JAX + backend and behaves as the Keras `tf.data` adapter otherwise. + + Raises: + TypeError: If `dataset` is not a `tf.data.Dataset` or `DistributedDataset`. + """ + if not isinstance( + dataset, (tf.data.Dataset, tf.distribute.DistributedDataset) + ): + raise TypeError( + f"zero_copy_batches needs a `tf.data.Dataset`, got {type(dataset)}." + ) + return _ZeroCopyAdapter( + dataset, distribution=keras.distribution.distribution() + ) diff --git a/recml/core/training/zero_copy_batches_test.py b/recml/core/training/zero_copy_batches_test.py new file mode 100644 index 0000000..af4a1db --- /dev/null +++ b/recml/core/training/zero_copy_batches_test.py @@ -0,0 +1,153 @@ +# 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. +from typing import Any + +from absl.testing import absltest +import jax.experimental.sparse as jax_sparse +import keras +from keras.src.trainers.data_adapters import data_adapter_utils +from keras.src.trainers.data_adapters import tf_dataset_adapter +import numpy as np +from recml.core.training import zero_copy_batches +import tensorflow as tf + + +def _dataset(drop_remainder: bool = True, size: int = 64) -> tf.data.Dataset: + x = np.arange(size * 3, dtype=np.float32).reshape(size, 3) / 100.0 + y = (x.sum(axis=1, keepdims=True) > 1.0).astype(np.float32) + ids = np.arange(size, dtype=np.int64).reshape(size, 1) + return tf.data.Dataset.from_tensor_slices(({"x": x, "ids": ids}, y)).batch( + 8, drop_remainder=drop_remainder + ) + + +def _leaves(batches: list[Any]) -> list[Any]: + """Flattens batches, including Keras' unstacked remainders, into leaves.""" + leaves = [] + for batch in batches: + if isinstance(batch, data_adapter_utils.PartialBatchList): + for part in batch: + leaves.extend(keras.tree.flatten(part)) + else: + leaves.extend(keras.tree.flatten(batch)) + return leaves + + +def _model() -> keras.Model: + keras.utils.set_random_seed(0) + x = keras.Input((3,), name="x") + ids = keras.Input((1,), name="ids", dtype="int32") + h = keras.layers.Concatenate()([x, keras.ops.cast(ids, "float32") / 64.0]) + model = keras.Model({"x": x, "ids": ids}, keras.layers.Dense(1)(h)) + model.compile(optimizer=keras.optimizers.SGD(0.1), loss="mse") + return model + + +class ZeroCopyBatchesTest(absltest.TestCase): + + def test_yields_views_with_the_same_values(self): + dataset = _dataset() + adapter = zero_copy_batches.zero_copy_adapter(dataset) + + batches = list(adapter.get_jax_iterator()) + + expected = list(dataset.as_numpy_iterator()) + self.assertLen(batches, len(expected)) + for (features, labels), (want_features, want_labels) in zip( + batches, expected + ): + for got, want in ( + (features["x"], want_features["x"]), + (features["ids"], want_features["ids"]), + (labels, want_labels), + ): + self.assertIsInstance(got, np.ndarray) + self.assertFalse(got.flags.owndata) + self.assertEqual(got.dtype, want.dtype) + np.testing.assert_array_equal(got, want) + + def test_matches_the_keras_tf_dataset_adapter(self): + # 68 examples leave a remainder batch; bfloat16 takes Keras' conversion. + for drop_remainder in (True, False): + dataset = _dataset(drop_remainder, size=68).map( + lambda features, labels: ( + {**features, "x_bf16": tf.cast(features["x"], tf.bfloat16)}, + labels, + ) + ) + for super_batch in (None, 2, 3): + with self.subTest( + drop_remainder=drop_remainder, super_batch=super_batch + ): + keras_adapter = tf_dataset_adapter.TFDatasetAdapter(dataset) + want = list(keras_adapter.get_jax_iterator(super_batch)) + adapter = zero_copy_batches.zero_copy_adapter(dataset) + got = list(adapter.get_jax_iterator(super_batch)) + + self.assertEqual([type(b) for b in got], [type(b) for b in want]) + got_leaves, want_leaves = _leaves(got), _leaves(want) + self.assertLen(got_leaves, len(want_leaves)) + for got_leaf, want_leaf in zip(got_leaves, want_leaves, strict=True): + self.assertIs(type(got_leaf), type(want_leaf)) + self.assertEqual(got_leaf.dtype, want_leaf.dtype) + np.testing.assert_array_equal(got_leaf, want_leaf) + + def test_sparse_tensor_converts_to_jax_sparse(self): + sparse = tf.SparseTensor( + indices=[[0, 1], [1, 2]], values=[3.0, 4.0], dense_shape=[2, 4] + ) + dataset = tf.data.Dataset.from_tensors(sparse) + adapter = zero_copy_batches.zero_copy_adapter(dataset) + + [got] = list(adapter.get_jax_iterator()) + + self.assertIsInstance(got, jax_sparse.BCOO) + np.testing.assert_array_equal( + got.todense(), np.array([[0.0, 3.0, 0.0, 0.0], [0.0, 0.0, 4.0, 0.0]]) + ) + + def test_fit_matches_fit_on_the_dataset(self): + reference = _model() + want = reference.fit(_dataset(), epochs=2, verbose=0).history["loss"] + + model = _model() + got = model.fit( + zero_copy_batches.zero_copy_adapter(_dataset()), epochs=2, verbose=0 + ).history["loss"] + + self.assertEqual(got, want) + for variable, want_variable in zip(model.weights, reference.weights): + np.testing.assert_array_equal( + keras.ops.convert_to_numpy(variable), + keras.ops.convert_to_numpy(want_variable), + ) + + def test_preserves_none_elements_in_batch(self): + dataset = tf.data.Dataset.from_tensors( + (tf.constant([[1.0, 2.0]], dtype=tf.float32), None) + ) + adapter = zero_copy_batches.zero_copy_adapter(dataset) + + [(features, sample_weight)] = list(adapter.get_jax_iterator()) + + np.testing.assert_array_equal(features, [[1.0, 2.0]]) + self.assertIsNone(sample_weight) + + def test_rejects_non_tf_dataset(self): + with self.assertRaisesRegex(TypeError, "tf.data.Dataset"): + zero_copy_batches.zero_copy_adapter([1, 2, 3]) + + +if __name__ == "__main__": + absltest.main()