diff --git a/recml/core/training/keras_trainer.py b/recml/core/training/keras_trainer.py index 853c1e3..dcb5736 100644 --- a/recml/core/training/keras_trainer.py +++ b/recml/core/training/keras_trainer.py @@ -16,7 +16,7 @@ from __future__ import annotations import abc -from collections.abc import Mapping +from collections.abc import Mapping, Sequence import dataclasses import functools import gc @@ -124,6 +124,8 @@ def __init__( rng_seed: int = core.DEFAULT_RNG_SEED, checkpoint_version: keras_utils.CheckpointVersion | str | None = None, legacy_checkpoint_format: bool | None = None, + train_callbacks: Sequence[keras.callbacks.Callback] = (), + eval_callbacks: Sequence[keras.callbacks.Callback] = (), ): """Initializes the instance. @@ -143,6 +145,13 @@ def __init__( legacy_checkpoint_format: Deprecated. Use checkpoint_version instead. If set, True maps to V1, and False maps to V2. TODO(b/542602169): Remove this in v2. + train_callbacks: Additional callbacks for training. They run before the + trainer's own callbacks, so metrics that they add to `logs` are written + to TensorBoard. + eval_callbacks: Additional callbacks for evaluation. They run before the + trainer's own callbacks. Use `keras_utils.MetricsCallback` to report + metrics, because Keras does not return changes made to `logs` in + `on_test_end`. resume_training_launcher: Launcher to resume training. enable_xmanager_measurements: Whether to enable XManager measurements. enable_autoxprof: Whether to enable AutoXprof. @@ -177,6 +186,8 @@ def __init__( self._checkpoint_dir = os.path.join(model_dir, core.CHECKPOINT_DIR) self._max_checkpoints_to_keep = max_checkpoints_to_keep self._checkpoint_save_interval_epochs = checkpoint_save_interval_epochs + self._custom_train_callbacks = list(train_callbacks) + self._custom_eval_callbacks = list(eval_callbacks) if checkpoint_version is not None: self._checkpoint_version = keras_utils.CheckpointVersion( checkpoint_version @@ -216,6 +227,7 @@ def train_callbacks(self) -> list[keras.callbacks.Callback]: ) callbacks = [ + *self._custom_train_callbacks, keras_utils.EpochSummaryCallback( log_dir=os.path.join(self._model_dir, core.LOG_DIR), steps_per_epoch=self._steps_per_loop, @@ -230,6 +242,7 @@ def train_callbacks(self) -> list[keras.callbacks.Callback]: return callbacks # pyrefly: ignore[bad-return] return [ + *self._custom_train_callbacks, keras.callbacks.TensorBoard( log_dir=os.path.join(self._model_dir, core.LOG_DIR), write_steps_per_second=True, @@ -253,6 +266,7 @@ def eval_callbacks(self) -> list[keras.callbacks.Callback]: """Returns the evaluation callbacks.""" if keras.backend.backend() == "jax": return [ + *self._custom_eval_callbacks, keras_utils.EpochSummaryCallback( log_dir=os.path.join(self._model_dir, core.LOG_DIR), steps_per_epoch=self._steps_per_loop, @@ -266,6 +280,7 @@ def eval_callbacks(self) -> list[keras.callbacks.Callback]: ), ] return [ + *self._custom_eval_callbacks, keras.callbacks.TensorBoard( log_dir=os.path.join(self._model_dir, core.LOG_DIR), write_steps_per_second=True, @@ -307,6 +322,27 @@ def _maybe_get_model_kws( return kws + def _val_logs( + self, + history: Mapping[str, Any], + callbacks: Sequence[keras.callbacks.Callback], + steps_per_second: float, + ) -> dict[str, Any]: + """Returns the logs for `EpochSummaryCallback` after an evaluation.""" + val_logs = {f"val_{k}": v for k, v in history.items()} + for cbk in callbacks: + if isinstance(cbk, keras_utils.MetricsCallback): + for k, v in cbk.latest_metrics.items(): + val_key = f"val_{k}" + if val_key in val_logs: + logging.warning( + f"{type(cbk).__name__} overwrites existing evaluation log key" + f" {val_key}." + ) + val_logs[val_key] = v + val_logs["val_steps_per_second"] = steps_per_second + return val_logs + def train(self, task: KerasTask) -> core.Logs: """Trains a Keras model.""" dataset = task.create_dataset( @@ -354,9 +390,9 @@ def evaluate(self, task: KerasTask) -> core.Logs: ) epoch_dt = time.time() - epoch_start_time steps_per_second = self._steps_per_eval / epoch_dt # pyrefly: ignore[unsupported-operation] - val_logs = {"val_" + k: v for k, v in history.items()} - val_logs["val_steps_per_second"] = steps_per_second - tb_cbk.on_epoch_end(0, val_logs) + tb_cbk.on_epoch_end( + 0, self._val_logs(history, self.eval_callbacks, steps_per_second) + ) return history return model.evaluate( @@ -464,6 +500,7 @@ def on_test_begin(self, logs: Mapping[str, Any] | None = None): for cbk in self.eval_callbacks if isinstance(cbk, keras_utils.EpochSummaryCallback) ] + callbacks = [restore_callback, *self.eval_callbacks] try: logging.info(f"eval | epoch: {epoch: 6d} | {steps_msg}") epoch_start_time = time.time() @@ -475,7 +512,7 @@ def on_test_begin(self, logs: Mapping[str, Any] | None = None): history = model.evaluate( eval_dataset, steps=self._steps_per_eval, - callbacks=[restore_callback] + self.eval_callbacks, + callbacks=callbacks, return_dict=True, ) @@ -491,9 +528,10 @@ def on_test_begin(self, logs: Mapping[str, Any] | None = None): epoch_dt = time.time() - epoch_start_time steps_per_second = self._steps_per_eval / epoch_dt # pyrefly: ignore[unsupported-operation] - val_logs = {"val_" + k: v for k, v in history.items()} - val_logs["val_steps_per_second"] = steps_per_second - tb_cbk.on_epoch_end(epoch, val_logs) + tb_cbk.on_epoch_end( + epoch, + self._val_logs(history, self.eval_callbacks, steps_per_second), + ) except FileNotFoundError: logging.info("Checkpoint epoch: %s did not finish writing...", epoch) diff --git a/recml/core/training/keras_trainer_test.py b/recml/core/training/keras_trainer_test.py index bfd1f7a..f396e85 100644 --- a/recml/core/training/keras_trainer_test.py +++ b/recml/core/training/keras_trainer_test.py @@ -52,6 +52,16 @@ def create_model(self) -> keras.Model: return model +class _MarkerCallback(keras_utils.MetricsCallback): + """Reports a fixed metric at the end of each training epoch and eval pass.""" + + def on_epoch_end(self, epoch, logs=None): + self._record(logs, "custom/marker", 7.0) + + def on_test_end(self, logs=None): + self._record(logs, "custom/marker", 7.0) + + class KerasTrainerTest(parameterized.TestCase): def setUp(self): @@ -63,12 +73,22 @@ def setUp(self): @parameterized.named_parameters( {"testcase_name": "train", "mode": core.Trainer.Mode.TRAIN}, + { + "testcase_name": "train_with_callbacks", + "mode": core.Trainer.Mode.TRAIN, + "with_train_callbacks": True, + }, {"testcase_name": "eval", "mode": core.Trainer.Mode.EVAL}, { "testcase_name": "train_and_eval_v3", "mode": core.Trainer.Mode.TRAIN_AND_EVAL, "checkpoint_version": "v3", }, + { + "testcase_name": "train_and_eval_with_callbacks", + "mode": core.Trainer.Mode.TRAIN_AND_EVAL, + "with_train_callbacks": True, + }, { "testcase_name": "train_and_eval_v2", "mode": core.Trainer.Mode.TRAIN_AND_EVAL, @@ -96,7 +116,10 @@ def setUp(self): }, ) def test_keras_task_and_trainer( - self, mode: str, checkpoint_version: str = "v3" + self, + mode: str, + checkpoint_version: str = "v3", + with_train_callbacks: bool = False, ): if keras.backend.backend() == "jax": distribution = keras.distribution.DataParallel() @@ -113,6 +136,7 @@ def test_keras_task_and_trainer( model_dir=self.create_tempdir().full_path, continuous_eval_timeout=5, checkpoint_version=checkpoint_version, + train_callbacks=[_MarkerCallback()] if with_train_callbacks else (), ) experiment = core.Experiment(_KerasTask(), trainer) @@ -128,6 +152,72 @@ def test_keras_task_and_trainer( and keras.backend.backend() == "jax" ): self.assertEqual(history.history["num_params/trainable"][0], 2) + if with_train_callbacks: + self.assertEqual( + history.history["custom/marker"], + [7.0] * len(history.history["loss"]), + ) + + def _run_eval_with_callbacks(self, model_dir: str, mode: core.Trainer.Mode): + if keras.backend.backend() != "jax": + self.skipTest("`EpochSummaryCallback` is only used on the Jax backend.") + trainer = keras_trainer.KerasTrainer( + distribution=keras.distribution.DataParallel(), + train_steps=2, + steps_per_eval=1, + steps_per_loop=2, + model_dir=model_dir, + continuous_eval_timeout=1, + eval_callbacks=[_MarkerCallback()], + ) + with mock.patch.object( + keras_utils.EpochSummaryCallback, "on_epoch_end", autospec=True + ) as on_epoch_end: + core.run_experiment(core.Experiment(_KerasTask(), trainer), mode) + + val_logs = on_epoch_end.call_args.args[2] + self.assertEqual(val_logs["val_custom/marker"], 7.0) + self.assertIn("val_loss", val_logs) + self.assertIn("val_steps_per_second", val_logs) + + def test_eval_metrics_callbacks_are_written_in_eval(self): + self._run_eval_with_callbacks( + self.create_tempdir().full_path, core.Trainer.Mode.EVAL + ) + + def test_eval_metrics_callbacks_are_written_in_continuous_eval(self): + if keras.backend.backend() != "jax": + self.skipTest("`EpochSummaryCallback` is only used on the Jax backend.") + model_dir = self.create_tempdir().full_path + # Produce one checkpoint so there is something to evaluate. + train_trainer = keras_trainer.KerasTrainer( + distribution=keras.distribution.DataParallel(), + train_steps=2, + steps_per_eval=1, + steps_per_loop=2, + model_dir=model_dir, + ) + core.run_experiment( + core.Experiment(_KerasTask(), train_trainer), core.Trainer.Mode.TRAIN + ) + keras.backend.clear_session() + + self._run_eval_with_callbacks(model_dir, core.Trainer.Mode.CONTINUOUS_EVAL) + + def test_val_logs_warns_when_metrics_callback_overwrites_key(self): + trainer = keras_trainer.KerasTrainer( + model_dir=self.create_tempdir().full_path + ) + callback = _MarkerCallback() + callback.on_test_end() + + with self.assertLogs(level="WARNING") as log: + val_logs = trainer._val_logs( + {"custom/marker": 1.0}, [callback], steps_per_second=2.0 + ) + + self.assertIn("overwrites existing evaluation log key", log.output[0]) + self.assertEqual(val_logs["val_custom/marker"], 7.0) def test_eval_name(self): if keras.backend.backend() != "jax": diff --git a/recml/core/utils/keras_utils.py b/recml/core/utils/keras_utils.py index ab6af89..1051ad0 100644 --- a/recml/core/utils/keras_utils.py +++ b/recml/core/utils/keras_utils.py @@ -19,6 +19,7 @@ import enum import os import re +import time from typing import Any from absl import logging @@ -1791,5 +1792,298 @@ def _collect_learning_rate(self, logs: Any) -> Any: ) return logs - def on_test_end(self, logs=None): + def on_test_end(self, logs=None) -> None: self._pop_writer() + + # This callback only writes summaries in `on_epoch_end`. The inherited + # `TensorBoard` batch hooks still call `summary.scalar` for every key in + # `logs` on every step, even though `update_freq="epoch"` means there is no + # default writer and nothing is recorded. Use the no-op `Callback` hooks + # instead so they do no per-step summary work. + on_train_batch_begin = keras.callbacks.Callback.on_train_batch_begin + on_train_batch_end = keras.callbacks.Callback.on_train_batch_end + on_test_batch_begin = keras.callbacks.Callback.on_test_batch_begin + + +class MetricsCallback(keras.callbacks.Callback): + """Base class for callbacks that add scalar metrics to the epoch summary. + + In training, a subclass can add metrics to `logs` in `on_epoch_end`. It must + run before `EpochSummaryCallback`, which writes `logs` to TensorBoard. + + In evaluation, Keras passes `on_test_end` a copy of `logs`, so changes made + there do not reach the dict that `model.evaluate` returns. A subclass must + also save its metrics in `latest_metrics`. `KerasTrainer` reads that property + after `model.evaluate` and adds the values, with a `val_` prefix, to the + validation logs that it writes. + """ + + def __init__(self): + super().__init__() + self._latest_metrics: dict[str, float] = {} + + @property + def latest_metrics(self) -> dict[str, float]: + """Returns a copy of the metrics from the most recent window.""" + return dict(self._latest_metrics) + + def _record( + self, logs: dict[str, Any] | None, name: str, value: float + ) -> None: + """Saves a metric and adds it to `logs` if `logs` is not None.""" + self._latest_metrics[name] = value + if logs is not None: + logs[name] = value + + +def _require_positive(name: str, value: int) -> int: + """Returns `value`, or raises a `ValueError` if it is not positive. + + A silent fallback to 1 would inflate the per-device metrics by up to the pod + size and nobody would notice on the dashboard. + + Args: + name: Name of the value, for the error message. + value: The value to check. + + Returns: + `value`, when it is at least 1. + + Raises: + ValueError: If `value` is less than 1. + """ + if value < 1: + raise ValueError(f"{name} must be positive, got {value}.") + return value + + +class _TokenThroughputCallback(MetricsCallback): + """Adds token throughput metrics for one timed window. + + Metrics, all under `throughput/`: + + * `k_valid_tokens_per_sec_per_device`: non-padding tokens per second per + accelerator, in thousands. It goes up when packing improves and does not + depend on the number of chips. + * `k_tokens_per_sec_per_device`: the same, but padding tokens in the padded + sequence are counted too (`examples_per_sec_per_device * padded_seq_len / + 1000`). + * `examples_per_sec_per_device`: examples per second per accelerator + (`global_batch_size * steps_per_sec / device_count`). + * `packing_efficiency_pct`: `valid_seq_len / padded_seq_len * 100`. + + `k_valid / k_tokens == packing_efficiency_pct / 100` by construction. So the + token and packing metrics have only two degrees of freedom, and the fact that + they agree does not prove that they are correct. + + The metric names are the same for training and evaluation. `KerasTrainer` + adds a `val_` prefix to evaluation metrics and `EpochSummaryCallback` removes + it again, so both series show in the same TensorBoard cards. + + The timer covers the full epoch or evaluation pass, matching + `EpochSummaryCallback`'s `steps_per_second`. The first window of a process (at + startup or upon recovery after a preemption) therefore includes JIT + compilation time and reads lower than steady-state windows. + + The callback never reads `logs` in its batch hooks and sets `async_safe`. + Thus Keras can keep dispatching batch callbacks asynchronously. Without + `async_safe`, one callback with a batch hook makes Keras convert all the + step logs to Python floats on the main thread on every step, which blocks + the accelerator. + + Subclasses select the Keras hooks that open and close the window. + """ + + def __init__( + self, + *, + global_batch_size: int, + padded_seq_len: int, + valid_seq_len_key: str, + device_count: int | None = None, + ): + """Initializes the callback. + + Args: + global_batch_size: Number of examples per step, across all devices. + padded_seq_len: Padded sequence length of one example. + valid_seq_len_key: Key of the model metric that holds the mean number of + non-padding tokens per example. A per-batch sum gives wrong values. If + the result is out of range, the token metrics are skipped and an error + is logged. + device_count: Global accelerator count. If None, it is read from + `jax.device_count()` when the first metrics are computed; pass + `device_count` explicitly on non-JAX backends. + """ + super().__init__() + self.async_safe = True + self._global_batch_size = _require_positive( + "global_batch_size", global_batch_size + ) + self._padded_seq_len = _require_positive("padded_seq_len", padded_seq_len) + self._valid_seq_len_key = valid_seq_len_key + self._device_count = ( + None + if device_count is None + else _require_positive("device_count", device_count) + ) + self._window_start = 0.0 + self._pause_start: float | None = None + self._paused_seconds = 0.0 + self._steps = 0 + self._warned_missing_key = False + + def _open_window(self) -> None: + """Starts the timer and resets the step count.""" + self._window_start = time.time() + self._pause_start = None + self._paused_seconds = 0.0 + self._steps = 0 + + def _pause_window(self) -> None: + """Pauses the timer while an in-epoch validation pass runs.""" + if self._pause_start is None: + self._pause_start = time.time() + + def _resume_window(self) -> None: + """Resumes the timer after an in-epoch validation pass finishes.""" + if self._pause_start is not None: + self._paused_seconds += max(0.0, time.time() - self._pause_start) + self._pause_start = None + + def _count_step(self, batch: int) -> None: + """Records progress from a batch-end hook. + + Args: + batch: Index of the last step that completed. Keras gives the end step, + not a count, so this stays correct when `steps_per_execution > 1`. + Because `CallbackList` dispatches async batch hooks to a multi-worker + thread pool, hooks can finish out of order; keeping the maximum + preserves the final step count. + """ + self._steps = max(self._steps, batch + 1) + + def _get_device_count(self) -> int: + if self._device_count is None: + self._device_count = _require_positive( + "jax.device_count()", jax.device_count() + ) + return self._device_count + + def _close_window(self, logs: dict[str, Any] | None) -> None: + """Computes the metrics for the window and adds them to `logs`.""" + self._resume_window() + elapsed = time.time() - self._window_start - self._paused_seconds + self._latest_metrics = {} + if not logs or self._valid_seq_len_key not in logs: + if not self._warned_missing_key: + logging.warning( + "%s: %s is not in logs; throughput metrics are skipped.", + type(self).__name__, + self._valid_seq_len_key, + ) + self._warned_missing_key = True + return + valid_seq_len = float(logs[self._valid_seq_len_key]) + packing_efficiency_pct = valid_seq_len / self._padded_seq_len * 100.0 + if not 0.0 <= packing_efficiency_pct <= 100.0: + # Do not stop a long training job because of a logging metric, and do + # not write values that are wrong by a factor of the batch size. + logging.error( + "%s: packing_efficiency_pct=%.1f is not in [0, 100]. %s is probably" + " a per-batch sum, not a per-example mean. Throughput metrics are" + " skipped.", + type(self).__name__, + packing_efficiency_pct, + self._valid_seq_len_key, + ) + return + # This metric has no time term, so record it even without a step rate. + self._record( + logs, "throughput/packing_efficiency_pct", packing_efficiency_pct + ) + + # Fall back to `params["steps"]` only when no batch hook fired in the + # window, so an epoch that stops early still uses its actual step count. + steps = self._steps or (self.params or {}).get("steps") or 0 + if not steps or elapsed <= 0: + logging.warning( + "%s: no step rate is available; tokens/sec metrics are skipped.", + type(self).__name__, + ) + return + examples_per_sec_per_device = ( + self._global_batch_size * steps / elapsed + ) / self._get_device_count() + self._record( + logs, + "throughput/examples_per_sec_per_device", + examples_per_sec_per_device, + ) + self._record( + logs, + "throughput/k_valid_tokens_per_sec_per_device", + examples_per_sec_per_device * valid_seq_len / 1000.0, + ) + self._record( + logs, + "throughput/k_tokens_per_sec_per_device", + examples_per_sec_per_device * self._padded_seq_len / 1000.0, + ) + + +class TrainThroughputCallback(_TokenThroughputCallback): + """Adds token throughput metrics once per training epoch. + + The window starts in `on_epoch_begin` and ends in `on_epoch_end`, i.e. it + covers `steps_per_loop` steps. When `model.fit` runs in-epoch validation + (as in `KerasTrainer.train_and_evaluate`), the timer pauses between + `on_test_begin` and `on_test_end` so validation time is excluded. + + See `_TokenThroughputCallback.__init__` for constructor arguments + (`global_batch_size`, `padded_seq_len`, `valid_seq_len_key`, `device_count`). + """ + + def on_epoch_begin( + self, epoch: int, logs: dict[str, Any] | None = None + ) -> None: + self._open_window() + + def on_train_batch_end( + self, batch: int, logs: dict[str, Any] | None = None + ) -> None: + self._count_step(batch) + + def on_test_begin(self, logs: dict[str, Any] | None = None) -> None: + self._pause_window() + + def on_test_end(self, logs: dict[str, Any] | None = None) -> None: + self._resume_window() + + def on_epoch_end( + self, epoch: int, logs: dict[str, Any] | None = None + ) -> None: + self._close_window(logs) + + +class EvalThroughputCallback(_TokenThroughputCallback): + """Adds token throughput metrics once per evaluation pass. + + `model.evaluate` does not call `on_epoch_end`, so the window starts in + `on_test_begin` and ends in `on_test_end`. `KerasTrainer` writes the result + from `latest_metrics` (see `MetricsCallback`). + + See `_TokenThroughputCallback.__init__` for constructor arguments + (`global_batch_size`, `padded_seq_len`, `valid_seq_len_key`, `device_count`). + """ + + def on_test_begin(self, logs: dict[str, Any] | None = None) -> None: + self._open_window() + + def on_test_batch_end( + self, batch: int, logs: dict[str, Any] | None = None + ) -> None: + self._count_step(batch) + + def on_test_end(self, logs: dict[str, Any] | None = None) -> None: + self._close_window(logs) diff --git a/recml/core/utils/keras_utils_test.py b/recml/core/utils/keras_utils_test.py index d130f55..be5012f 100644 --- a/recml/core/utils/keras_utils_test.py +++ b/recml/core/utils/keras_utils_test.py @@ -2401,5 +2401,405 @@ def call(self, x): ) +_SEQ_LEN_KEY = "labels/count_seq_len" +# With 4 devices, a global batch of 8 and 2 steps/sec, each device processes 4 +# examples/sec. This keeps the expected values exact. +_GLOBAL_BATCH_SIZE = 8 +_PADDED_SEQ_LEN = 1024 +_DEVICE_COUNT = 4 + + +def _throughput_callback( + callback_cls, + *, + global_batch_size: int = _GLOBAL_BATCH_SIZE, + padded_seq_len: int = _PADDED_SEQ_LEN, + device_count: int | None = _DEVICE_COUNT, +): + return callback_cls( + global_batch_size=global_batch_size, + padded_seq_len=padded_seq_len, + valid_seq_len_key=_SEQ_LEN_KEY, + device_count=device_count, + ) + + +def _run_train_epoch(callback, logs, *, steps=10, seconds=5.0): + """Runs the hooks of one training epoch that takes `seconds`.""" + with mock.patch.object(keras_utils.time, "time", side_effect=[0.0]): + callback.on_epoch_begin(0) + for batch in range(steps): + callback.on_train_batch_end(batch) + with mock.patch.object(keras_utils.time, "time", side_effect=[seconds]): + callback.on_epoch_end(0, logs) + + +def _run_eval_pass(callback, logs, *, steps=10, seconds=5.0): + """Runs the hooks of one evaluation pass that takes `seconds`.""" + with mock.patch.object(keras_utils.time, "time", side_effect=[0.0]): + callback.on_test_begin() + for batch in range(steps): + callback.on_test_batch_end(batch) + with mock.patch.object(keras_utils.time, "time", side_effect=[seconds]): + callback.on_test_end(logs) + + +class EpochSummaryCallbackBatchHooksTest(absltest.TestCase): + + def test_batch_hooks_are_keras_defaults(self): + # Default hooks skip per-step summary work for this callback. + callback = keras_utils.EpochSummaryCallback( + log_dir=self.create_tempdir().full_path, steps_per_epoch=10 + ) + + for name in ( + "on_train_batch_begin", + "on_train_batch_end", + "on_test_batch_begin", + ): + self.assertIs( + getattr(type(callback), name), + getattr(keras.callbacks.Callback, name), + name, + ) + self.assertTrue(callback.async_safe) + + +class MetricsCallbackTest(absltest.TestCase): + + def test_record_adds_metric_to_logs_and_latest_metrics(self): + callback = keras_utils.MetricsCallback() + logs = {"loss": 1.0} + + callback._record(logs, "custom/metric", 3.0) + + self.assertEqual(logs, {"loss": 1.0, "custom/metric": 3.0}) + self.assertEqual(callback.latest_metrics, {"custom/metric": 3.0}) + + def test_record_without_logs_only_saves_latest_metrics(self): + # Subclass hooks receive `logs=None` by default, e.g. when a test or a + # custom loop calls them directly. + callback = keras_utils.MetricsCallback() + + callback._record(None, "custom/metric", 3.0) + + self.assertEqual(callback.latest_metrics, {"custom/metric": 3.0}) + + def test_latest_metrics_returns_a_copy(self): + callback = keras_utils.MetricsCallback() + callback._record(None, "custom/metric", 3.0) + + callback.latest_metrics["custom/metric"] = 0.0 + + self.assertEqual(callback.latest_metrics, {"custom/metric": 3.0}) + + +class TrainThroughputCallbackTest(absltest.TestCase): + + def test_reports_expected_values(self): + # 10 steps in 5s = 2 steps/sec. 8 examples/step / 4 devices = 4 examples + # per second per device. 4 * 512 valid tokens = 2048, i.e. 2.048k. + logs = {_SEQ_LEN_KEY: 512.0} + callback = _throughput_callback(keras_utils.TrainThroughputCallback) + + _run_train_epoch(callback, logs) + + self.assertAlmostEqual(logs["throughput/examples_per_sec_per_device"], 4.0) + self.assertAlmostEqual( + logs["throughput/k_valid_tokens_per_sec_per_device"], 2.048 + ) + self.assertAlmostEqual( + logs["throughput/k_tokens_per_sec_per_device"], 4.096 + ) + self.assertAlmostEqual(logs["throughput/packing_efficiency_pct"], 50.0) + self.assertEqual( + callback.latest_metrics, + {k: v for k, v in logs.items() if k.startswith("throughput/")}, + ) + + def test_is_async_safe(self): + # Otherwise Keras runs all batch callbacks synchronously on the main thread + # and converts the step logs to Python floats on every step. + callback = _throughput_callback(keras_utils.TrainThroughputCallback) + + self.assertTrue(callback.async_safe) + + def test_step_count_comes_from_keras_end_step(self): + # One batch end with end step 99 means that 100 steps completed. This is + # what happens when `steps_per_execution` is more than one. + logs = {_SEQ_LEN_KEY: 512.0} + callback = _throughput_callback(keras_utils.TrainThroughputCallback) + + with mock.patch.object(keras_utils.time, "time", side_effect=[0.0]): + callback.on_epoch_begin(0) + callback.on_train_batch_end(99) + with mock.patch.object(keras_utils.time, "time", side_effect=[10.0]): + callback.on_epoch_end(0, logs) + + # 10 steps/sec => 20 examples/sec/device * 512 tokens. + self.assertAlmostEqual( + logs["throughput/k_valid_tokens_per_sec_per_device"], 10.24 + ) + + def test_uses_params_steps_when_no_batch_hook_fires(self): + logs = {_SEQ_LEN_KEY: 512.0} + callback = _throughput_callback(keras_utils.TrainThroughputCallback) + callback.set_params({"steps": 10}) + + _run_train_epoch(callback, logs, steps=0) + + self.assertAlmostEqual( + logs["throughput/k_valid_tokens_per_sec_per_device"], 2.048 + ) + + def test_out_of_order_batch_hooks_and_early_stop_override_params_steps(self): + # Async batch hooks can finish out of order, and an epoch can stop before + # `params["steps"]`. The max end step observed must win over both earlier + # batch indices and `params["steps"]`. + logs = {_SEQ_LEN_KEY: 512.0} + callback = _throughput_callback(keras_utils.TrainThroughputCallback) + callback.set_params({"steps": 20}) + + with mock.patch.object(keras_utils.time, "time", side_effect=[0.0]): + callback.on_epoch_begin(0) + callback.on_train_batch_end(9) + callback.on_train_batch_end(8) + with mock.patch.object(keras_utils.time, "time", side_effect=[5.0]): + callback.on_epoch_end(0, logs) + + # 10 steps (not 9 or 20) in 5s => 2.048k valid tokens/sec/device. + self.assertAlmostEqual( + logs["throughput/k_valid_tokens_per_sec_per_device"], 2.048 + ) + + def test_excludes_in_epoch_validation_time(self): + # In `train_and_evaluate`, Keras runs `on_test_begin`..`on_test_end` inside + # the training epoch before `on_epoch_end`. + logs = {_SEQ_LEN_KEY: 512.0} + callback = _throughput_callback(keras_utils.TrainThroughputCallback) + + with mock.patch.object(keras_utils.time, "time", side_effect=[0.0]): + callback.on_epoch_begin(0) + callback.on_train_batch_end(9) + with mock.patch.object( + keras_utils.time, "time", side_effect=[5.0, 8.0, 8.0] + ): + callback.on_test_begin() + callback.on_test_end() + callback.on_epoch_end(0, logs) + + # 3s of validation inside an 8s epoch leaves 5s of training time. + self.assertAlmostEqual( + logs["throughput/k_valid_tokens_per_sec_per_device"], 2.048 + ) + + def test_step_count_does_not_leak_into_next_epoch(self): + # The second epoch has no batch hooks and no `params["steps"]`, so it has + # no step rate. A step count left over from the first epoch would make it + # report tokens/sec anyway. + callback = _throughput_callback(keras_utils.TrainThroughputCallback) + _run_train_epoch(callback, {_SEQ_LEN_KEY: 512.0}, steps=10) + second = {_SEQ_LEN_KEY: 256.0} + + with self.assertLogs(level="WARNING") as log: + _run_train_epoch(callback, second, steps=0) + + self.assertIn("no step rate", log.output[0]) + self.assertEqual( + second, {_SEQ_LEN_KEY: 256.0, "throughput/packing_efficiency_pct": 25.0} + ) + self.assertEqual( + callback.latest_metrics, {"throughput/packing_efficiency_pct": 25.0} + ) + + def test_latest_metrics_do_not_leak_into_next_epoch(self): + # The second epoch has no sequence length, so it reports nothing. Metrics + # of the first epoch must not stay in `latest_metrics`, or `KerasTrainer` + # would write them again for the second window. + callback = _throughput_callback(keras_utils.TrainThroughputCallback) + _run_train_epoch(callback, {_SEQ_LEN_KEY: 512.0}) + self.assertLen(callback.latest_metrics, 4) + second = {"loss": 1.0} + + with self.assertLogs(level="WARNING") as log: + _run_train_epoch(callback, second) + + self.assertIn(f"{_SEQ_LEN_KEY} is not in logs", log.output[0]) + self.assertEqual(second, {"loss": 1.0}) + self.assertEmpty(callback.latest_metrics) + + def test_ignores_val_prefixed_keys(self): + logs = {f"val_{_SEQ_LEN_KEY}": 512.0} + + _run_train_epoch( + _throughput_callback(keras_utils.TrainThroughputCallback), logs + ) + + self.assertEqual(logs, {f"val_{_SEQ_LEN_KEY}": 512.0}) + + +class EvalThroughputCallbackTest(absltest.TestCase): + + def test_reports_unprefixed_keys_in_latest_metrics(self): + # `KerasTrainer` adds the `val_` prefix, so the names are the same as in + # training. + callback = _throughput_callback(keras_utils.EvalThroughputCallback) + + _run_eval_pass(callback, {_SEQ_LEN_KEY: 256.0}) + + self.assertIsInstance(callback, keras_utils.MetricsCallback) + self.assertEqual( + set(callback.latest_metrics), + { + "throughput/examples_per_sec_per_device", + "throughput/k_valid_tokens_per_sec_per_device", + "throughput/k_tokens_per_sec_per_device", + "throughput/packing_efficiency_pct", + }, + ) + self.assertAlmostEqual( + callback.latest_metrics["throughput/examples_per_sec_per_device"], + 4.0, + ) + self.assertAlmostEqual( + callback.latest_metrics["throughput/k_valid_tokens_per_sec_per_device"], + 1.024, + ) + self.assertAlmostEqual( + callback.latest_metrics["throughput/packing_efficiency_pct"], 25.0 + ) + + def test_ignores_training_hooks(self): + logs = {_SEQ_LEN_KEY: 256.0} + callback = _throughput_callback(keras_utils.EvalThroughputCallback) + + callback.on_epoch_begin(0) + callback.on_epoch_end(0, logs) + + self.assertEqual(logs, {_SEQ_LEN_KEY: 256.0}) + self.assertEmpty(callback.latest_metrics) + + def test_uses_its_own_batch_size(self): + small = _throughput_callback(keras_utils.EvalThroughputCallback) + large = _throughput_callback( + keras_utils.EvalThroughputCallback, + global_batch_size=2 * _GLOBAL_BATCH_SIZE, + ) + + _run_eval_pass(small, {_SEQ_LEN_KEY: 256.0}) + _run_eval_pass(large, {_SEQ_LEN_KEY: 256.0}) + + key = "throughput/k_valid_tokens_per_sec_per_device" + self.assertAlmostEqual( + large.latest_metrics[key], 2 * small.latest_metrics[key] + ) + + +class ThroughputCallbackEdgeCaseTest(absltest.TestCase): + + def test_missing_seq_len_leaves_logs_untouched(self): + logs = {"loss": 1.0} + callback = _throughput_callback(keras_utils.TrainThroughputCallback) + + with self.assertLogs(level="WARNING") as log: + _run_train_epoch(callback, logs) + _run_train_epoch(callback, logs) + + # Warn only once per callback instance so tasks without the metric do not + # spam the logs on every epoch. + self.assertLen(log.output, 1) + self.assertIn(f"{_SEQ_LEN_KEY} is not in logs", log.output[0]) + self.assertEqual(logs, {"loss": 1.0}) + self.assertEmpty(callback.latest_metrics) + + def test_none_logs_warns_and_reports_nothing(self): + callback = _throughput_callback(keras_utils.TrainThroughputCallback) + + with self.assertLogs(level="WARNING") as log: + _run_train_epoch(callback, None) + + self.assertIn( + f"TrainThroughputCallback: {_SEQ_LEN_KEY} is not in logs", + log.output[0], + ) + self.assertEmpty(callback.latest_metrics) + + def test_packing_efficiency_without_a_step_rate(self): + logs = {_SEQ_LEN_KEY: 512.0} + + _run_train_epoch( + _throughput_callback(keras_utils.TrainThroughputCallback), + logs, + steps=0, + ) + + self.assertAlmostEqual(logs["throughput/packing_efficiency_pct"], 50.0) + self.assertNotIn("throughput/examples_per_sec_per_device", logs) + self.assertNotIn("throughput/k_valid_tokens_per_sec_per_device", logs) + self.assertNotIn("throughput/k_tokens_per_sec_per_device", logs) + + def test_zero_elapsed_time_skips_tokens_per_sec(self): + logs = {_SEQ_LEN_KEY: 512.0} + + _run_train_epoch( + _throughput_callback(keras_utils.TrainThroughputCallback), + logs, + seconds=0.0, + ) + + self.assertIn("throughput/packing_efficiency_pct", logs) + self.assertNotIn("throughput/k_valid_tokens_per_sec_per_device", logs) + + def test_out_of_range_seq_len_is_skipped_not_raised(self): + # A per-batch sum instead of a per-example mean. The job must not stop. + logs = {_SEQ_LEN_KEY: _PADDED_SEQ_LEN * 8.0} + + with self.assertLogs(level="ERROR") as log: + _run_train_epoch( + _throughput_callback(keras_utils.TrainThroughputCallback), logs + ) + + self.assertEqual(logs, {_SEQ_LEN_KEY: _PADDED_SEQ_LEN * 8.0}) + self.assertIn("per-batch sum", log.output[0]) + + def test_rejects_non_positive_arguments(self): + with self.assertRaisesRegex(ValueError, "device_count"): + _throughput_callback(keras_utils.TrainThroughputCallback, device_count=0) + with self.assertRaisesRegex(ValueError, "global_batch_size"): + _throughput_callback( + keras_utils.TrainThroughputCallback, global_batch_size=0 + ) + with self.assertRaisesRegex(ValueError, "padded_seq_len"): + _throughput_callback( + keras_utils.TrainThroughputCallback, padded_seq_len=0 + ) + + def test_reads_device_count_from_jax(self): + logs = {_SEQ_LEN_KEY: 512.0} + callback = _throughput_callback( + keras_utils.TrainThroughputCallback, device_count=None + ) + + with mock.patch.object(keras_utils.jax, "device_count", return_value=2): + _run_train_epoch(callback, logs) + + # Half the devices of the default, so twice the per-device rate. + self.assertAlmostEqual( + logs["throughput/k_valid_tokens_per_sec_per_device"], 4.096 + ) + + def test_casts_scalar_values_to_float(self): + logs = {_SEQ_LEN_KEY: np.float32(512.0)} + + _run_train_epoch( + _throughput_callback(keras_utils.TrainThroughputCallback), logs + ) + + self.assertIsInstance(logs["throughput/packing_efficiency_pct"], float) + self.assertAlmostEqual( + logs["throughput/k_valid_tokens_per_sec_per_device"], 2.048 + ) + + if __name__ == "__main__": absltest.main()