diff --git a/src/maxtext/training_engine/inflight_throttler.py b/src/maxtext/training_engine/inflight_throttler.py index 488040fd2..a809fe5e5 100644 --- a/src/maxtext/training_engine/inflight_throttler.py +++ b/src/maxtext/training_engine/inflight_throttler.py @@ -26,14 +26,15 @@ class InflightThrottler: """Rate limits the number of inflight computations on TPU.""" - def __init__(self, config: pyconfig.HyperParameters): + def __init__(self, config: pyconfig.HyperParameters, metrics_logger: metrics_module.MetricsLogger) -> None: """Initializes the inflight throttler. Args: config: The training configuration. + metrics_logger: The metrics logger to use for writing metrics. """ self._inflight_queue = queue.Queue[Any](maxsize=config.max_inflight_computations) - self._metrics_logger = metrics_module.MetricsLogger(config=config) + self._metrics_logger = metrics_logger self._pending_metrics: abstract_engine.MetricsBuffer | None = None def add_computation(self, computation: Any, metrics: abstract_engine.MetricsBuffer | None) -> None: @@ -77,6 +78,5 @@ def wait_for_all(self) -> None: self._flush_pending_metrics() def cleanup(self) -> None: - """Closes the underlying metrics logger and releases resources.""" + """Wait for all inflight computations to finish and log their metrics.""" self.wait_for_all() - self._metrics_logger.cleanup() diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py index 356b07eeb..abce4032e 100644 --- a/src/maxtext/training_engine/maxtext_engine.py +++ b/src/maxtext/training_engine/maxtext_engine.py @@ -157,6 +157,53 @@ def _batch_signature(dynamic_batch: Any, static_batch: dict[str, Any]) -> Any: ) +def _normalize_loss_output(out: Any, has_aux: bool) -> abstract_engine.LossOutput: + """Normalizes whatever a loss function returned a `LossOutput`. + + Shared by the training and evaluation kernels. + + Args: + out: The loss function's return value. + has_aux: Whether the caller considers a 2-tuple's second element to be auxiliary + output worth recording. + + Returns: + The equivalent `LossOutput`. + + Raises: + TypeError: If `out` matches none of the accepted shapes. + """ + if isinstance(out, abstract_engine.LossOutput): + return out + if isinstance(out, abstract_engine.WeightedMetric): + return abstract_engine.LossOutput(primary_loss=out, aux_metrics={}) + if isinstance(out, (tuple, list)) and len(out) == 2: + loss_val, aux = out + if isinstance(loss_val, abstract_engine.WeightedMetric): + primary_loss = loss_val + elif isinstance(aux, dict) and "xent_sum" in aux and "total_weights" in aux: + primary_loss = abstract_engine.WeightedMetric( + unreduced_sum=aux["xent_sum"], + denominator=aux["total_weights"], + ) + else: + raise TypeError( + f"Cannot construct WeightedMetric from 2-tuple loss return with elements " + f"of type ({type(loss_val).__name__}, {type(aux).__name__}). Expected first element to be a " + "WeightedMetric, or second element to be a dict containing 'xent_sum' and 'total_weights'." + ) + + return abstract_engine.LossOutput( + primary_loss=primary_loss, + aux_metrics=aux if (has_aux and isinstance(aux, dict)) else {}, + ) + raise TypeError( + f"Unsupported return type from loss function: {type(out)}. " + "Expected abstract_engine.LossOutput, abstract_engine.WeightedMetric, " + "or a 2-element tuple/list: (loss, aux_metrics)." + ) + + @struct.dataclass(frozen=True, kw_only=True) class RouterReplayTrainerPayload(abstract_engine.TrainerPayload): """A TrainerPayload extension carrying forced router-replay expert decisions. @@ -341,15 +388,14 @@ def __init__( # keeps recording its aux metrics. `with_loss_fn` overrides this per its own default. self._has_aux: bool = True self._gen_model_input_fn: Callable[[Any], dict[str, Any]] | None = None - # Tracked per instance rather than via logging.log_first_n, which is process-wide and - # would make the warning depend on whether some earlier engine already triggered it. - self._eval_step_warned: bool = False self._compiled = False # Set by `compile()`, including when it defers for want of a dummy payload. `_compiled` # alone cannot express "wanted, not yet built", and conflating them would either make # every eager caller compile or make a deferred compile never happen. self._compile_requested = False self._compiled_signature: Any = None + self._compiled_eval: Any = None + self._compiled_eval_signature: Any = None self._signature_compare_warned: bool = False if not training_config.model_name: raise ValueError("training_config.model_name must be specified") @@ -397,7 +443,9 @@ def __init__( config=self._config, ) self._metrics_recorder = metrics_module.MetricsRecorder() - self._throttler = inflight_throttler.InflightThrottler(config=self._config) + self._eval_metrics_recorder = metrics_module.MetricsRecorder(mode=metrics_module.Mode.EVAL) + self._metrics_logger = metrics_module.MetricsLogger(config=self._config) + self._throttler = inflight_throttler.InflightThrottler(config=self._config, metrics_logger=self._metrics_logger) self._raiden_sync: Any = None @property @@ -413,6 +461,8 @@ def model(self, new_model: Any) -> None: self._compiled_fwd_bwd = None self._compiled_fwd_bwd_accum = None self._compiled_update = None + self._compiled_eval = None + self._compiled_eval_signature = None self._model_graphdef = None self._invalidate_pure_state() @@ -431,6 +481,8 @@ def optimizer(self, new_optimizer: Any) -> None: self._compiled_update = None self._state_graphdef = None self._invalidate_pure_state() + self._compiled_eval = None + self._compiled_eval_signature = None @property def train_step(self) -> int: @@ -459,6 +511,8 @@ def state(self, new_state: Any) -> None: self._compiled_update = None self._state_graphdef = None self._invalidate_pure_state() + self._compiled_eval = None + self._compiled_eval_signature = None @property def micro_step_count(self) -> int: @@ -487,6 +541,8 @@ def with_loss_fn(self, customized_fn: Callable[..., Any], has_aux: bool = False) self._loss_fn = customized_fn self._has_aux = has_aux self._compiled = False + self._compiled_eval = None + self._compiled_eval_signature = None return self def with_gen_model_input_fn(self, gen_model_input_fn: Callable[[Any], dict[str, Any]]) -> "MaxTextTrainingEngine": @@ -508,6 +564,8 @@ def with_gen_model_input_fn(self, gen_model_input_fn: Callable[[Any], dict[str, # The adapter decides which batch entries are traced and which are baked into the # executable, so a compiled kernel built against the previous one is stale. self._compiled = False + self._compiled_eval = None + self._compiled_eval_signature = None return self @contextlib.contextmanager @@ -680,44 +738,8 @@ def diff_wrapper(p, r, b): out = loss_callable(mdl, self._config, b, None, None, is_train=True) _, _, new_r = nnx.split(mdl, nnx.Param, ...) - if isinstance(out, abstract_engine.LossOutput): - return out.primary_loss.unreduced_sum, (out, new_r) - elif isinstance(out, abstract_engine.WeightedMetric): - loss_out = abstract_engine.LossOutput( - primary_loss=out, - aux_metrics={}, - ) - return out.unreduced_sum, (loss_out, new_r) - elif isinstance(out, (tuple, list)) and len(out) == 2: - loss_val, aux = out - if isinstance(loss_val, abstract_engine.WeightedMetric): - primary_loss = loss_val - elif isinstance(aux, dict) and "xent_sum" in aux and "total_weights" in aux: - primary_loss = abstract_engine.WeightedMetric( - unreduced_sum=aux["xent_sum"], - denominator=aux["total_weights"], - ) - else: - raise TypeError( - f"Cannot construct WeightedMetric from 2-tuple loss return with elements " - f"of type ({type(loss_val).__name__}, {type(aux).__name__}). Expected first element to be a " - "WeightedMetric, or second element to be a dict containing 'xent_sum' and 'total_weights'." - ) - - # `has_aux=False` means the caller does not consider the second element to be - # auxiliary output, so it is not recorded -- even though it may have been read - # above to build `primary_loss`. - loss_out = abstract_engine.LossOutput( - primary_loss=primary_loss, - aux_metrics=aux if (self._has_aux and isinstance(aux, dict)) else {}, - ) - return primary_loss.unreduced_sum, (loss_out, new_r) - else: - raise TypeError( - f"Unsupported return type from loss function: {type(out)}. " - "Expected abstract_engine.LossOutput, abstract_engine.WeightedMetric, " - "or a 2-element tuple/list: (loss, aux_metrics)." - ) + loss_out = _normalize_loss_output(out, self._has_aux) + return loss_out.primary_loss.unreduced_sum, (loss_out, new_r) grad_func = jax.value_and_grad(diff_wrapper, argnums=0, has_aux=True) # Every non-raising branch of `diff_wrapper` builds a LossOutput, so `loss_out` is always @@ -784,6 +806,26 @@ def _update_kernel(self, state_pure, accumulated_grads, accumulated_denominator, return new_state_pure, grad_norm, is_skipped_val return state_pure, grad_norm, is_skipped_val + def _eval_kernel(self, params, rest, batch): + """Executes a single forward pass, returning the loss and its aux metrics. + + Returns: + `(primary_loss, aux_metrics)` -- a `WeightedMetric` and a dict. + """ + loss_callable = self._loss_fn if self._loss_fn is not None else maxtext_train.loss_fn + mdl = nnx.merge(self._model_graphdef, params, rest, copy=True) + if self._gen_model_input_fn is not None: + if not isinstance(batch, dict): + raise TypeError( + "gen_model_input_fn must return a dict of loss-fn keyword arguments, got " f"{type(batch).__name__}." + ) + out = loss_callable(mdl, **batch) + else: + out = loss_callable(mdl, self._config, batch, None, None, is_train=False) + + loss_out = _normalize_loss_output(out, self._has_aux) + return loss_out.primary_loss, loss_out.aux_metrics + def _warn_uncomparable(self, what: str, hint: str, exc: Exception) -> None: """Warns once per instance that a signature half could not be compared. @@ -796,7 +838,7 @@ def _warn_uncomparable(self, what: str, hint: str, exc: Exception) -> None: self._signature_compare_warned = True logging.warning(_UNCOMPARABLE_SIGNATURE_WARNING, what, exc, hint) - def _needs_recompile(self, signature: Any) -> bool: + def _needs_recompile(self, signature: Any, previous: Any) -> bool: """Returns whether the compiled kernel is stale for `signature`. An unanswerable comparison counts as stale. That is the right direction for @@ -809,7 +851,6 @@ def _needs_recompile(self, signature: Any) -> bool: Comparing the signature as a whole would route a badly-behaved treedef or shape entry into a message blaming the caller's static loss arguments. """ - previous = self._compiled_signature if previous is None: return True @@ -964,6 +1005,30 @@ def accum_kernel(params, rest, dynamic, acc_grads, acc_denom): self._compiled_signature = _batch_signature(dynamic_batch, static_batch) self._compiled = True + def _compile_eval_for_batch(self, dynamic_batch: Any, static_batch: dict[str, Any]) -> None: + """JIT-compiles the forward-only eval kernel for one batch structure.""" + self._model_graphdef, params_pure, rest_pure = nnx.split(self._model, nnx.Param, ...) + + def kernel(params, rest, dynamic): + batch = {**dynamic, **static_batch} if isinstance(dynamic, dict) else dynamic + return self._eval_kernel(params, rest, batch) + + if self._mesh is not None: + params_shardings = jax.tree.map(self._mesh_sharding, params_pure) + rest_shardings = jax.tree.map(self._mesh_sharding, rest_pure) + eval_in_shardings = (params_shardings, rest_shardings, self._batch_data_shardings(dynamic_batch)) + eval_out_shardings = (None, None) + else: + eval_in_shardings = None + eval_out_shardings = None + + self._compiled_eval = jax.jit( + kernel, + in_shardings=eval_in_shardings, + out_shardings=eval_out_shardings, + ) + self._compiled_eval_signature = _batch_signature(dynamic_batch, static_batch) + def compile(self, dummy_data: abstract_engine.TrainerPayload) -> None: """Triggers SPMD JIT compilation of fwd_bwd and update steps. @@ -1022,7 +1087,7 @@ def fwd_bwd(self, payload: abstract_engine.TrainerPayload, **kwargs: Any) -> Non # changing needs a fresh kernel -- reusing it would raise an in_shardings mismatch # for the first, and silently use stale values for the second. signature = _batch_signature(dynamic_batch, static_batch) - if not self._compiled or self._needs_recompile(signature): + if not self._compiled or self._needs_recompile(signature, self._compiled_signature): self._compile_for_batch(dynamic_batch, static_batch) # After any recompile, not before: reading first would hand the new kernel a pure # state split against the old graph. @@ -1142,29 +1207,70 @@ def update(self, **kwargs: Any) -> int: return self.train_step - def eval_step(self, payload: abstract_engine.TrainerPayload, **kwargs: Any) -> None: - """Warns once that evaluation is not implemented, then does nothing. + @contextlib.contextmanager + def eval_context(self): + """Brackets a sequence of `eval_step` calls and writes their metrics on exit. - A silent no-op lets `TrainerWorker.run_eval` report success having evaluated nothing, - so any eval metrics for the run are meaningless rather than absent. Warning makes that - audible; warning only once keeps a loop that evaluates every step from flooding the - log. Implementing this properly means a forward-only pass plus deciding how eval - metrics bucket via `MetricsBuffer.mode`, which is tracked separately. + Usage: - Mutates no trainer state -- in particular not `_accumulated_grads` or - `_micro_step_count` -- as `AbstractTrainer.eval_step` requires. + with trainer.eval_context(): + for micro_batch in eval_ds: + trainer.eval_step(micro_batch) + """ + logging.info("Running evaluation on train step %d.", self.train_step) + # Drain the training queue so that the eval metrics are logged after all training metrics for this step. + self._throttler.wait_for_all() + try: + yield + finally: + for buffer in self._eval_metrics_recorder.get_metrics_history(clear_cache=True): + logging.info("Writing buffered eval metrics for train step %d.", buffer.id) + self._metrics_logger.write_metrics(buffer, mode=metrics_module.Mode.EVAL) + self._eval_metrics_recorder.cleanup() + + def eval_step(self, payload: abstract_engine.TrainerPayload, **kwargs: Any) -> None: + """Prepares inputs, runs the evaluation, and logs eval metrics. + + Callers should bracket a sequence of eval_step calls with eval_context() so that the metrics + mode is set to EVAL and buffered metrics are written on exit. Args: - payload: Packed micro-batch evaluation input. Currently unused. - **kwargs: Additional keyword arguments for evaluation. Currently unused. + payload: Packed micro-batch evaluation input. + **kwargs: Additional keyword arguments for evaluation. """ - if not self._eval_step_warned: - self._eval_step_warned = True - logging.warning( - "MaxTextTrainingEngine.eval_step is not implemented: it evaluates nothing and " - "records no metrics, so any eval result reported for this run is meaningless. " - "Logged once per engine instance." - ) + batch = self._prepare_batch(payload) + + model = getattr(self._state, "model", self._model) if self._state is not None else self._model + if not isinstance(model, nnx.Module): + raise TypeError("MaxTextTrainingEngine requires an NNX model (flax.nnx.Module), got" f" {type(model).__name__}") + + self._model_graphdef, params, rest = nnx.split(model, nnx.Param, ...) + + # Wait for previous computations to finish before dispatching the next one to TPU. + self._throttler.wait_for_next() + + if self._compile_requested: + dynamic_batch, static_batch = _split_static_and_dynamic(batch) + signature = _batch_signature(dynamic_batch, static_batch) + if self._compiled_eval is None or self._needs_recompile(signature, self._compiled_eval_signature): + self._compile_eval_for_batch(dynamic_batch, static_batch) + loss, aux = self._compiled_eval(params, rest, dynamic_batch) + else: + loss, aux = self._eval_kernel(params, rest, batch) + + # No metrics attached: eval metrics are buffered by `_eval_metrics_recorder` and written + # in EVAL mode when `eval_context` exits. + self._throttler.add_computation(computation=loss, metrics=None) + + if isinstance(loss, abstract_engine.WeightedMetric): + self.record_metrics("loss", loss, mode=metrics_module.Mode.EVAL) + else: + logging.warning("Eval loss is not a WeightedMetric, so it will not be logged. Got %s.", type(loss).__name__) + + if isinstance(aux, dict): + for key, value in aux.items(): + if value is not None: + self.record_metrics(key, value, mode=metrics_module.Mode.EVAL) def save_checkpoint(self, metadata: Any, **kwargs: Any) -> None: """Forces asynchronous Orbax checkpoint serialization. @@ -1340,6 +1446,7 @@ def record_metrics( name: str, metric: abstract_engine.WeightedMetric | jax.Array | float | int | dict[str, Any], aggregation_fn: Callable[[jax.Array], Any] | None = None, + mode: metrics_module.Mode = metrics_module.Mode.TRAIN, ) -> None: """Records a metric into the buffer, appending to JAX arrays. @@ -1357,14 +1464,23 @@ def record_metrics( f"{name}/{sub_k}" if name else sub_k, sub_v, aggregation_fn=aggregation_fn, + mode=mode, ) else: - self._metrics_recorder.buffer_metrics( - train_step=self.train_step, - name=name, - metric=metric, - aggregation_fn=aggregation_fn, - ) + if mode == metrics_module.Mode.TRAIN: + self._metrics_recorder.buffer_metrics( + train_step=self.train_step, + name=name, + metric=metric, + aggregation_fn=aggregation_fn, + ) + else: + self._eval_metrics_recorder.buffer_metrics( + train_step=self.train_step, + name=name, + metric=metric, + aggregation_fn=aggregation_fn, + ) def get_metrics(self, clear_cache: bool = True) -> abstract_engine.MetricsBuffer: """Returns the most recent step's metrics as an on-device MetricsBuffer. @@ -1536,3 +1652,4 @@ def close(self) -> None: # Cleanup metrics recorder resources after saving the checkpoint, ensuring all buffered metrics are saved properly self._metrics_recorder.cleanup() + self._metrics_logger.cleanup() diff --git a/src/maxtext/training_engine/metrics.py b/src/maxtext/training_engine/metrics.py index 25adc54dc..1992cee3b 100644 --- a/src/maxtext/training_engine/metrics.py +++ b/src/maxtext/training_engine/metrics.py @@ -18,6 +18,7 @@ from collections.abc import Callable import dataclasses +import enum import os from typing import Any @@ -46,6 +47,14 @@ _DEFAULT_MAX_BUFFERED_STEPS = 128 +class Mode(str, enum.Enum): + TRAIN = "train" + EVAL = "eval" + + def __str__(self): + return self.value + + class MetricsRecorder: """Synchronous frontend for buffering and aggregating step metrics on-device. @@ -58,7 +67,7 @@ class MetricsRecorder: for processing. """ - def __init__(self, max_buffered_steps: int = _DEFAULT_MAX_BUFFERED_STEPS): + def __init__(self, max_buffered_steps: int = _DEFAULT_MAX_BUFFERED_STEPS, mode: Mode = Mode.TRAIN) -> None: """Initializes the recorder. Args: @@ -66,6 +75,7 @@ def __init__(self, max_buffered_steps: int = _DEFAULT_MAX_BUFFERED_STEPS): as new steps start. Zero or less retains everything. """ self._metrics_buffer: list[abstract_engine.MetricsBuffer] = [] + self._mode = mode self._max_buffered_steps = max_buffered_steps self._dropped_buffer_count = 0 @@ -85,7 +95,7 @@ def buffer_metrics( aggregation_fn: Optional aggregation function to apply to the metric. """ if not self._metrics_buffer or self._metrics_buffer[-1].id != train_step: - new_buffer = abstract_engine.MetricsBuffer(id=train_step, mode="train") + new_buffer = abstract_engine.MetricsBuffer(id=train_step, mode=self._mode) self._metrics_buffer.append(new_buffer) self._evict_old_buffers() @@ -195,7 +205,7 @@ class MetricsLogger: results to TensorBoard and console stdout. """ - def __init__(self, config: pyconfig.HyperParameters): + def __init__(self, config: pyconfig.HyperParameters) -> None: """Initializes the metrics logger. Args: @@ -219,48 +229,58 @@ def write_setup_info_to_tensorboard(self, params: Any) -> None: max_utils.add_text_to_summary_writer("libtpu_init_args", os.getenv("LIBTPU_INIT_ARGS", ""), self._tb_writer) maxtext_utils.add_config_to_summary_writer(self._config, self._tb_writer) - def _log_metrics(self, step: int, metrics: dict[str, Any]) -> None: + def _log_metrics(self, step: int, metrics: dict[str, Any], mode: Mode) -> None: """Logs the metrics to the console. Args: step: The train step for which to log the metrics. metrics: Dictionary mapping metric names to reduced Python floats/numpy arrays. + mode: The mode of the training engine. """ - log_message = [f"Completed step: {step}"] - for k in _METRICS_TO_LOG: - if k in metrics: - if k == "learning_rate": - log_message.append(f"{k}: {metrics[k]:.3e}") - else: - log_message.append(f"{k}: {metrics[k]:.3f}") - - logging.info(", ".join(log_message)) + if mode == Mode.TRAIN: + log_message = [f"Train step: {step}"] + for k in ["loss", "perplexity"]: + if k in metrics: + val = metrics[k] + log_message.append(f"{k}: {val:.3f}") + if len(log_message) > 1: + logging.info(", ".join(log_message)) + elif mode == Mode.EVAL: + loss = metrics.get("loss") + if loss is not None: + logging.info( + "Eval step %d evaluation loss: %f", + step, + loss, + ) - def write_metrics(self, metrics: abstract_engine.MetricsBuffer) -> None: + def write_metrics(self, metrics: abstract_engine.MetricsBuffer, mode: Mode = Mode.TRAIN) -> None: """Write metrics to the console and TensorBoard. Args: metrics: MetricsBuffer containing the metrics to write. + mode: The mode of the training engine. """ processed_metrics = self.process_metrics(metrics) - self._log_metrics(metrics.id, processed_metrics) - self._write_metrics_to_tensorboard(metrics.id, processed_metrics) + self._log_metrics(step=metrics.id, metrics=processed_metrics, mode=mode) + self._write_metrics_to_tensorboard(step=metrics.id, metrics=processed_metrics, mode=mode) - def _write_metrics_to_tensorboard(self, step: int, metrics: dict[str, Any]) -> None: + def _write_metrics_to_tensorboard(self, step: int, metrics: dict[str, Any], mode: Mode = Mode.TRAIN) -> None: """Write metrics to TensorBoard. Args: step: The train step for which to write the metrics. metrics: Dictionary mapping metric names to reduced Python floats/numpy arrays. + mode: The mode of the training engine. """ if self._tb_writer is None: return if jax.process_index() == 0: for metric_name, value in metrics.items(): - self._tb_writer.add_scalar(metric_name, value, step) + self._tb_writer.add_scalar(f"{mode}/{metric_name}", value, step) if step % self._config.log_period == 0: logging.info( diff --git a/tests/post_training/integration/maxtext_engine_grpo_loss_test.py b/tests/post_training/integration/maxtext_engine_grpo_loss_test.py index 6284277cd..6356fefb5 100644 --- a/tests/post_training/integration/maxtext_engine_grpo_loss_test.py +++ b/tests/post_training/integration/maxtext_engine_grpo_loss_test.py @@ -26,6 +26,7 @@ """ import dataclasses +from unittest import mock from absl.testing import absltest from flax import nnx @@ -36,6 +37,7 @@ from maxtext.configs import pyconfig from maxtext.training_engine import maxtext_engine +from maxtext.training_engine import metrics as metrics_module from maxtext.utils import maxtext_utils from tests.utils.test_helpers import get_test_config_path @@ -169,6 +171,31 @@ def _param_leaves(model) -> list[jax.Array]: return jax.tree.leaves(nnx.to_pure_dict(nnx.state(model, nnx.Param))) +def _grpo_model_input(algo_config: _GrpoConfig): + """Returns grpo_loss_fn's keyword arguments for a payload. + + Tunix's loss is used directly, with no adapter closure. Setting this as the engine's + `gen_model_input_fn` tells it to invoke the loss Tunix's way, `loss_fn(model, **inputs)`, + and the dict below is already exactly grpo_loss_fn's keyword arguments -- which is what + the orchestrator's own `_grpo_model_input` produces in the real pipeline. + """ + return lambda payload: { + "train_example": payload, + "algo_config": algo_config, + "pad_id": _PAD_ID, + "eos_id": _EOS_ID, + } + + +def _mean(metric) -> float: + """Reduces a WeightedMetric to one number the way `process_metrics` does. + + `WeightedMetric.compute()` is elementwise -- `unreduced_sum * compute_scale()` -- so a + buffer holding several micro-batches returns one value per micro-batch, not their mean. + """ + return float(np.mean(np.asarray(metric.compute()))) + + class MaxTextEngineGrpoLossTest(absltest.TestCase): """The engine driven by Tunix's own GRPO loss, not a stand-in.""" @@ -194,17 +221,8 @@ def test_grpo_loss_drives_a_training_step(self): ) self.assertEqual(type(engine.model).__name__, "TunixMaxTextAdapter") - # Tunix's loss is used directly, with no adapter closure. Setting a gen_model_input_fn - # tells the engine to invoke the loss Tunix's way, `loss_fn(model, **inputs)`, and the - # dict below is already exactly grpo_loss_fn's keyword arguments -- which is what the - # orchestrator's own `_grpo_model_input` produces in the real pipeline. returned = engine.with_loss_fn(algo_core.grpo_loss_fn, has_aux=True).with_gen_model_input_fn( - lambda payload: { - "train_example": payload, - "algo_config": algo_config, - "pad_id": _PAD_ID, - "eos_id": _EOS_ID, - } + _grpo_model_input(algo_config) ) self.assertIs(returned, engine) @@ -244,6 +262,67 @@ def test_grpo_loss_drives_a_training_step(self): self.assertIn("learning_rate", buf.scalar_metrics) self.assertNotIn("learning_rate", buf.weighted_metrics) + def test_eval_step_matches_fwd_bwd_and_leaves_training_untouched(self): + """`eval_step` against the real GRPO loss.""" + cfg = _config() + mesh = jax.sharding.Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + engine = maxtext_engine.MaxTextTrainingEngine( + cfg, + mesh=mesh, + # grpo_loss_fn calls the model with Tunix's signature, so the adapter wrap is needed + # for the loss itself, not only for weight sync. + wrap_with_tunix_adapter=True, + tokenizer_pad_id=_PAD_ID, + ) + + algo_config = _GrpoConfig() + + engine.with_loss_fn(algo_core.grpo_loss_fn, has_aux=True).with_gen_model_input_fn(_grpo_model_input(algo_config)) + + payload = _train_example(engine.model, algo_config) + + # Open a training step. This is also the reference forward pass eval is compared against. + engine.fwd_bwd(payload) + train_loss_metric = engine.get_metrics(clear_cache=False).weighted_metrics["loss"] + train_loss = _mean(train_loss_metric) + per_call_entries = train_loss_metric.unreduced_sum.size + + before = _param_leaves(engine.model) + micro_steps_before = engine.micro_step_count + self.assertEqual(engine.train_step, 0) + + with mock.patch.object(engine._metrics_logger, "write_metrics") as write_metrics: # pylint: disable=protected-access + with engine.eval_context(): + engine.eval_step(payload) + engine.eval_step(payload) + + for a, b in zip(_param_leaves(engine.model), before): + self.assertTrue(jnp.array_equal(a, b), "eval_step mutated a model parameter") + self.assertEqual(engine.train_step, 0) + self.assertEqual(engine.micro_step_count, micro_steps_before) + + # Eval must not reach the train recorder + train_buf = engine.get_metrics(clear_cache=False) + self.assertEqual(train_buf.weighted_metrics["loss"].unreduced_sum.size, per_call_entries) + self.assertEqual(train_buf.mode, metrics_module.Mode.TRAIN) + + # Leaving the context writes the metrics once + self.assertEqual(write_metrics.call_count, 1) + eval_buf = write_metrics.call_args.args[0] + self.assertEqual(write_metrics.call_args.kwargs["mode"], metrics_module.Mode.EVAL) + self.assertEqual(eval_buf.mode, metrics_module.Mode.EVAL) + self.assertEqual(eval_buf.id, 0) + + # Both micro-batches accumulated into that single buffer. + eval_loss_metric = eval_buf.weighted_metrics["loss"] + self.assertEqual(eval_loss_metric.unreduced_sum.size, 2 * per_call_entries) + + eval_loss = _mean(eval_loss_metric) + self.assertAlmostEqual(eval_loss, train_loss, places=4) + + # Eval never records a learning rate + self.assertNotIn("learning_rate", eval_buf.scalar_metrics) + if __name__ == "__main__": absltest.main() diff --git a/tests/post_training/unit/maxtext_engine_test.py b/tests/post_training/unit/maxtext_engine_test.py index 5ee877bbb..60d937218 100644 --- a/tests/post_training/unit/maxtext_engine_test.py +++ b/tests/post_training/unit/maxtext_engine_test.py @@ -1066,13 +1066,8 @@ def _loss_fn(model, *args, **kwargs): self.assertIn("loss", buf.weighted_metrics) self.assertNotIn("loss", buf.scalar_metrics) - def test_eval_step_warns_once_and_mutates_no_state(self): - """eval_step is an unimplemented no-op, but an audible one, and it disturbs nothing. - - `AbstractTrainer.eval_step` forbids mutating trainer state, so this asserts against a - populated engine -- gradients accumulated and a micro step counted -- rather than a - fresh one, where "unchanged" would be trivially true. - """ + def test_eval_step_records_eval_metrics_and_mutates_no_training_state(self): + """eval_step scores a batch without disturbing training, and its metrics stay separate.""" t = maxtext_engine.MaxTextTrainingEngine(self.mock_config) t.with_loss_fn( lambda *args, **kwargs: ( @@ -1086,19 +1081,42 @@ def test_eval_step_warns_once_and_mutates_no_state(self): micro_steps_before = t._micro_step_count train_step_before = t.train_step self.assertEqual(micro_steps_before, 1) + train_entries_before = t._metrics_recorder.get_step_metrics(train_step_before).weighted_metrics["loss"] + train_entries_before = train_entries_before.unreduced_sum.size - with self.assertLogs(level="WARNING") as logs: - t.eval_step(DummyPayload()) - t.eval_step(DummyPayload()) - t.eval_step(DummyPayload()) - - eval_warnings = [line for line in logs.output if "eval_step is not implemented" in line] - self.assertLen(eval_warnings, 1) + with mock.patch.object(t._metrics_logger, "write_metrics") as write_metrics: + with t.eval_context(): + t.eval_step(DummyPayload()) + t.eval_step(DummyPayload()) + t.eval_step(DummyPayload()) + # Nothing about the in-flight training step moved. self.assertEqual(t._micro_step_count, micro_steps_before) self.assertEqual(t.train_step, train_step_before) jax.tree.map(np.testing.assert_array_equal, grads_before, t._accumulated_grads) + # The train buffer still holds exactly the one fwd_bwd loss: no eval leaked into it. + train_buf = t._metrics_recorder.get_step_metrics(train_step_before) + self.assertEqual(train_buf.weighted_metrics["loss"].unreduced_sum.size, train_entries_before) + self.assertEqual(train_buf.mode, metrics_module.Mode.TRAIN) + + # Leaving the context writes the pass once, tagged eval, against the step it ran at -- + # not once per micro-batch, which would put three points on the curve at one x. + self.assertEqual(write_metrics.call_count, 1) + eval_buf = write_metrics.call_args.args[0] + self.assertEqual(write_metrics.call_args.kwargs["mode"], metrics_module.Mode.EVAL) + self.assertEqual(eval_buf.mode, metrics_module.Mode.EVAL) + self.assertEqual(eval_buf.id, train_step_before) + + # All three micro-batches accumulated into that one buffer. `compute()` is elementwise, + # so it stays per-micro-batch here; `process_metrics` is what averages it down. + eval_loss = eval_buf.weighted_metrics["loss"] + self.assertEqual(eval_loss.unreduced_sum.size, 3) + self.assertAlmostEqual(float(np.mean(np.asarray(eval_loss.compute()))), 0.5, places=4) + + # The recorder is drained, so a later pass cannot re-write this one's numbers. + self.assertEmpty(t._eval_metrics_recorder.get_metrics_history(clear_cache=False)) + def test_get_metrics_returns_one_buffer_and_a_sentinel_when_empty(self): """`get_metrics` returns a single buffer, matching both ABCs.