diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 1569f8436b..769a548422 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -3982,9 +3982,32 @@ class RLConfig( Decoding, IciParallelism, DcnParallelism, + PipelineParallelism, + DilocoParams, HardwareAndMesh, ModelArchitecture, + MTP, MoBa, + # Advanced Architectures, Tuning, and Optimizers + Muon, + FineTuning, + Distillation, + # Datasets and Loading Compatibility + DatasetGeneral, + TfdsDataset, + HfDataset, + GrainDataset, + OlmoGrainDataset, + # Inference, Checkpointing, and Monitoring + EmergencyCheckpointing, + ElasticTraining, + InferenceServer, + InferenceBenchmark, + PrefixCaching, + HloDump, + Goodput, + GcpMonitoring, + ManagedMLDiagnostics, # Positional Embeddings PositionalEmbedding, Rope, @@ -4011,9 +4034,12 @@ class RLConfig( AttentionIndexer, SplashAttention, Qwen3Next, - # Debugging and Profiling + # Debugging, Profiling, and Telemetry + AOT, DevelopmentAndDebugging, Profiling, + Metrics, + Tensorboard, # For compatibility with trainer in post_train/rl RL, RLCluster, @@ -4022,6 +4048,8 @@ class RLConfig( RLReward, RLSpecialTokens, VLLM, + TrainingLoop, + DerivedValues, ): """ Configuration for Reinforcement Learning in MaxText. @@ -4207,19 +4235,20 @@ def set_derived_values_and_validate(self) -> "RLConfig": # Dynamically inject model dimensions. emb_scale, num_head_scale, mlp_dim_scale, layer_scale = get_individual_scales(self.global_parameter_scale) - object.__setattr__(self, "emb_dim", int((2**emb_scale) * self.base_emb_dim)) - object.__setattr__(self, "num_query_heads", int((2**num_head_scale) * self.base_num_query_heads)) - object.__setattr__(self, "num_kv_heads", int((2**num_head_scale) * self.base_num_kv_heads)) - object.__setattr__(self, "mlp_dim", int((2**mlp_dim_scale) * self.base_mlp_dim)) - object.__setattr__(self, "moe_mlp_dim", int((2**mlp_dim_scale) * getattr(self, "base_moe_mlp_dim", 0))) - object.__setattr__(self, "num_decoder_layers", int((2**layer_scale) * self.base_num_decoder_layers)) + self.emb_dim = int((2**emb_scale) * self.base_emb_dim) + self.num_query_heads = int((2**num_head_scale) * self.base_num_query_heads) + self.num_kv_heads = int((2**num_head_scale) * self.base_num_kv_heads) + self.mlp_dim = int((2**mlp_dim_scale) * self.base_mlp_dim) + self.moe_mlp_dim = int((2**mlp_dim_scale) * getattr(self, "base_moe_mlp_dim", 0)) + self.num_decoder_layers = int((2**layer_scale) * self.base_num_decoder_layers) # Mirror into internal MaxText fields for backward compatibility. train_micro_batch_size = getattr(self.dataset, "train_micro_batch_size", -1) batch_size = getattr(self.dataset, "batch_size", 1) if train_micro_batch_size <= 0: train_micro_batch_size = batch_size - object.__setattr__(self, "micro_batch_size_to_train_on", train_micro_batch_size) + self.micro_batch_size_to_train_on = train_micro_batch_size + self.steps = getattr(self, "train_steps", getattr(self, "num_batches", 10)) if self.remat_policy == "custom": tensors = [ @@ -4244,7 +4273,7 @@ def set_derived_values_and_validate(self) -> "RLConfig": "attention_out", "out_proj", ] - object.__setattr__(self, "tensors_on_device", [t for t in tensors if getattr(self, t) == "device"]) - object.__setattr__(self, "tensors_to_offload", [t for t in tensors if getattr(self, t) == "offload"]) + self.tensors_on_device = [t for t in tensors if getattr(self, t) == "device"] + self.tensors_to_offload = [t for t in tensors if getattr(self, t) == "offload"] return self diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py index 52ba94a6a9..2a03a6df0d 100644 --- a/src/maxtext/training_engine/maxtext_engine.py +++ b/src/maxtext/training_engine/maxtext_engine.py @@ -21,6 +21,7 @@ from __future__ import annotations from collections.abc import Callable +import dataclasses from typing import Any from absl import logging @@ -28,16 +29,17 @@ import jax import jax.numpy as jnp from maxtext.common import common_types +from maxtext.common import train_state_nnx from maxtext.configs import pyconfig from maxtext.trainers.pre_train import train as maxtext_train from maxtext.training_engine import abstract_engine from maxtext.training_engine import checkpointing from maxtext.training_engine import inflight_throttler from maxtext.training_engine import metrics as metrics_module -from maxtext.utils import gradient_accumulation from maxtext.utils import max_utils from maxtext.utils import maxtext_utils from maxtext.utils import model_creation_utils +from maxtext.utils import sharding from maxtext.utils import train_utils @@ -85,7 +87,7 @@ def __init__( self._train_step: int = 0 self._checkpoint_manager = checkpointing.CheckpointManager( - checkpoint_dir=getattr(self._config, "checkpoint_directory", ""), + checkpoint_dir=getattr(self._config, "checkpoint_dir", getattr(self._config, "checkpoint_directory", "")), config=self._config, ) self._metrics_recorder = metrics_module.MetricsRecorder() @@ -100,6 +102,10 @@ def model(self) -> Any: def model(self, new_model: Any) -> None: """Sets the NNX model instance.""" self._model = new_model + self._compiled = False + self._compiled_fwd_bwd = None + self._compiled_update = None + self._model_graphdef = None @property def optimizer(self) -> Any: @@ -110,6 +116,10 @@ def optimizer(self) -> Any: def optimizer(self, new_optimizer: Any) -> None: """Sets the NNX optimizer instance.""" self._optimizer = new_optimizer + self._compiled = False + self._compiled_fwd_bwd = None + self._compiled_update = None + self._state_graphdef = None @property def train_step(self) -> int: @@ -121,6 +131,32 @@ def train_step(self, step: int) -> None: """Sets the current step integer.""" self._train_step = step + @property + def state(self) -> Any: + """Returns the current train state, initializing it if necessary.""" + if self._state is None and self._model is not None and self._optimizer is not None: + self._state = train_state_nnx.TrainStateNNX(self._model, self._optimizer) + return self._state + + @state.setter + def state(self, new_state: Any) -> None: + """Sets the current train state.""" + self._state = new_state + self._compiled = False + self._compiled_fwd_bwd = None + self._compiled_update = None + self._state_graphdef = None + + @property + def micro_step_count(self) -> int: + """Returns the current micro-batch count in gradient accumulation.""" + return self._micro_step_count + + @property + def has_accumulated_grads(self) -> bool: + """Returns True if accumulated gradients are present.""" + return self._accumulated_grads is not None + def with_loss_fn(self, customized_fn: Callable[..., Any]) -> None: """Overrides the default autoregressive loss function with a custom RL loss. @@ -135,12 +171,124 @@ def with_gen_model_input_fn(self, gen_model_input_fn: Callable[[Any], dict[str, self._gen_model_input_fn = gen_model_input_fn return self + def _fwd_bwd_kernel(self, params, rest, batch): + """Executes a single forward and backward pass to compute gradients.""" + loss_callable = self._loss_fn if self._loss_fn is not None else maxtext_train.loss_fn + + def diff_wrapper(p, r, b): + mdl = nnx.merge(self._model_graphdef, p, r, copy=True) + loss, aux = loss_callable(mdl, self._config, b, None, None, is_train=True) + _, _, new_r = nnx.split(mdl, nnx.Param, ...) + return loss, (aux, new_r) + + grad_func = jax.value_and_grad(diff_wrapper, argnums=0, has_aux=True) + (loss, (aux, new_rest)), micro_grads = grad_func(params, rest, batch) + micro_grads = jax.tree.map( + lambda x: ( + x.astype(getattr(self._config, "grad_dtype", jnp.float32)) + if hasattr(x, "dtype") and x.dtype == jnp.float32 + else x + ), + micro_grads, + ) + return loss, aux, new_rest, micro_grads + + def _update_kernel(self, state_pure, accumulated_grads, micro_step_count, mean_loss): + """Applies accumulated gradients to update the NNX model state.""" + grad_norm = None + is_skipped_val = None + if state_pure is not None: + if micro_step_count <= 1: + grads = accumulated_grads + else: + grads = jax.tree.map( + lambda g: g / micro_step_count, + accumulated_grads, + ) + if getattr(self._config, "gradient_clipping_threshold", 0.0) > 0: + grads = maxtext_utils.apply_gradient_clipping(grads, None, self._config.gradient_clipping_threshold) + local_state = nnx.merge(self._state_graphdef, state_pure, copy=True) + if hasattr(local_state, "apply_gradients"): + if getattr(self._config, "skip_step_on_spikes", False): + grad_norm = max_utils.l2norm_pytree(grads) + local_state.apply_gradients(grads, loss=mean_loss, grad_norm=grad_norm) + opt_obj = getattr(local_state, "optimizer", self._optimizer) + if opt_obj is not None: + opt_state = nnx.to_pure_dict(nnx.state(opt_obj)).get("opt_state", {}) + is_skipped = opt_state.get("is_skipped") if isinstance(opt_state, dict) else None + if is_skipped is not None: + is_skipped_val = is_skipped.astype(jnp.float32) + else: + local_state.apply_gradients(grads) + _, new_state_pure = nnx.split(local_state) + return new_state_pure, grad_norm, is_skipped_val + return state_pure, grad_norm, is_skipped_val + def compile(self, dummy_data: abstract_engine.TrainerPayload) -> None: - """Triggers SPMD JIT compilation of fwd_bwd, update, and eval steps. + """Triggers SPMD JIT compilation of fwd_bwd and update steps. Args: dummy_data: Sample TrainerPayload providing representative tensor shapes. """ + if self._compiled: + return + + if self._state is None: + self._state = train_state_nnx.TrainStateNNX(self._model, self._optimizer) + + self._state_graphdef, state_pure = nnx.split(self._state) + self._model_graphdef, params_pure, rest_pure = nnx.split(self._model, nnx.Param, ...) + + if self._mesh is not None: + data_sharding = sharding.get_input_data_sharding(self._config, self._mesh) + state_mesh_shardings = jax.tree.map( + lambda x: getattr( + x, + "sharding", + jax.sharding.NamedSharding(self._mesh, jax.sharding.PartitionSpec()), + ), + state_pure, + ) + params_shardings = jax.tree.map( + lambda x: getattr( + x, + "sharding", + jax.sharding.NamedSharding(self._mesh, jax.sharding.PartitionSpec()), + ), + params_pure, + ) + rest_shardings = jax.tree.map( + lambda x: getattr( + x, + "sharding", + jax.sharding.NamedSharding(self._mesh, jax.sharding.PartitionSpec()), + ), + rest_pure, + ) + fwd_bwd_in_shardings = (params_shardings, rest_shardings, data_sharding) + fwd_bwd_out_shardings = (None, None, rest_shardings, params_shardings) + update_in_shardings = (state_mesh_shardings, params_shardings, None) + update_out_shardings = (state_mesh_shardings, None, None) + else: + fwd_bwd_in_shardings = None + fwd_bwd_out_shardings = None + update_in_shardings = None + update_out_shardings = None + + # 1. JIT Compile Micro FWD/BWD Pass + self._compiled_fwd_bwd = jax.jit( + self._fwd_bwd_kernel, + in_shardings=fwd_bwd_in_shardings, + out_shardings=fwd_bwd_out_shardings, + ) + + # 2. JIT Compile Optimizer Update Pass + self._compiled_update = jax.jit( + self._update_kernel, + in_shardings=update_in_shardings, + out_shardings=update_out_shardings, + static_argnums=(2,), + ) self._compiled = True def fwd_bwd(self, payload: abstract_engine.TrainerPayload) -> None: @@ -151,9 +299,10 @@ def fwd_bwd(self, payload: abstract_engine.TrainerPayload) -> None: """ if self._gen_model_input_fn is not None: batch = self._gen_model_input_fn(payload) + elif dataclasses.is_dataclass(payload): + batch = {k: getattr(payload, k) for k in payload.__dataclass_fields__ if getattr(payload, k) is not None} else: batch = payload - loss_callable = self._loss_fn if self._loss_fn is not None else maxtext_train.loss_fn model = getattr(self._state, "model", None) if self._state is not None else self._model if not isinstance(model, nnx.Module): @@ -162,28 +311,31 @@ def fwd_bwd(self, payload: abstract_engine.TrainerPayload) -> None: # Wait for previous computations to finish before dispatching the next one to TPU. self._throttler.wait_for_next() - # TODO(mazumdera): This function call should be pre-compiled. - loss, aux, micro_grads = gradient_accumulation.gradient_accumulation_loss_and_grad( - loss_callable, - self._config, - model, - None, - None, - batch, - None, - ) + if self._state is None: + self._state = train_state_nnx.TrainStateNNX(self._model, self._optimizer) + model = getattr(self._state, "model", self._model) + self._model_graphdef, params, rest = nnx.split(model, nnx.Param, ...) + + if self._compiled and hasattr(self, "_compiled_fwd_bwd"): + loss, aux, new_rest, micro_grads = self._compiled_fwd_bwd(params, rest, batch) + else: + loss, aux, new_rest, micro_grads = self._fwd_bwd_kernel(params, rest, batch) + nnx.update(model, new_rest) # Don't add metrics to the throttler queue because metrics are logged after # the update step. self._throttler.add_computation(computation=loss, metrics=None) - if isinstance(loss, abstract_engine.WeightedMetric): + if loss is not None: + # TODO(mazumdera): This needs to be modified to become + # if isinstance(loss, abstract_engine.WeightedMetric): self.record_metrics("loss", loss) # Record auxiliary metrics. if isinstance(aux, dict): for key, value in aux.items(): - self.record_metrics(key, value) + if value is not None: + self.record_metrics(key, value) self._cached_losses.append(loss) if self._accumulated_grads is None: @@ -201,40 +353,32 @@ def update(self) -> None: return if self._learning_rate_schedule is not None: - try: - lr = self._learning_rate_schedule(self.train_step) - self.record_metrics("learning_rate", lr) - except Exception: # pylint: disable=broad-except - pass + lr = self._learning_rate_schedule(self.train_step) + self.record_metrics("learning_rate", lr) # Wait for previous computations to finish before dispatching the update step to TPU. self._throttler.wait_for_next() # TODO(mazumdera): The logic below should be pre-compiled. - if self._state is not None: - # TODO(mazumdera): Figure out how exactly we should normalize the losses - # (if at all). Given that inputs are varying in size, it not correct to - # simply divide by the number of micro-steps. - grads = jax.tree.map( - lambda g: g / max(self._micro_step_count, 1), - self._accumulated_grads, + if self._state is None: + self._state = train_state_nnx.TrainStateNNX(self._model, self._optimizer) + self._state_graphdef, state_pure = nnx.split(self._state) + + mean_loss = jnp.mean(jnp.array(self._cached_losses)) if self._cached_losses else jnp.array(0.0) + if self._compiled and hasattr(self, "_compiled_update"): + new_state_pure, grad_norm, is_skipped = self._compiled_update( + state_pure, self._accumulated_grads, self._micro_step_count, mean_loss ) - if getattr(self._config, "gradient_clipping_threshold", 0.0) > 0: - grads = maxtext_utils.apply_gradient_clipping(grads, None, self._config.gradient_clipping_threshold) - if hasattr(self._state, "apply_gradients"): - if getattr(self._config, "skip_step_on_spikes", False): - grad_norm = max_utils.l2norm_pytree(grads) - self.record_metrics("gradient_norm", grad_norm) - mean_loss = jnp.mean(jnp.array(self._cached_losses)) if self._cached_losses else jnp.array(0.0) - self._state.apply_gradients(grads, loss=mean_loss, grad_norm=grad_norm) - opt_obj = getattr(self._state, "optimizer", self._optimizer) - if opt_obj is not None: - opt_state = nnx.to_pure_dict(nnx.state(opt_obj)).get("opt_state", {}) - is_skipped = opt_state.get("is_skipped") if isinstance(opt_state, dict) else None - if is_skipped is not None: - self.record_metrics("step_skipped", is_skipped.astype(jnp.float32)) - else: - self._state.apply_gradients(grads) + else: + new_state_pure, grad_norm, is_skipped = self._update_kernel( + state_pure, self._accumulated_grads, self._micro_step_count, mean_loss + ) + nnx.update(self._state, new_state_pure) + + if grad_norm is not None: + self.record_metrics("gradient_norm", grad_norm) + if is_skipped is not None: + self.record_metrics("step_skipped", is_skipped) # Add the state to the throttler queue so jax.block_until_ready() waits # for the optimizer update to complete before logging the metrics. @@ -319,15 +463,35 @@ def restore_checkpoint(self, **kwargs: Any) -> Any: checkpoint_state=checkpoint_state, step=step, ) - if not restored_step: + if restored_step is None: return None logging.info("Checkpoint restored from step %d.", restored_step) self.train_step = restored_step if restored_checkpoint_state.accumulated_metrics: + buffers = [] + for b in restored_checkpoint_state.accumulated_metrics: + if isinstance(b, dict): + wms = {} + for k, wm in b.get("weighted_metrics", {}).items(): + if isinstance(wm, dict): + wms[k] = abstract_engine.WeightedMetric(**wm) + else: + wms[k] = wm + buffers.append( + abstract_engine.MetricsBuffer( + id=b.get("id", 0), + mode=b.get("mode", "train"), + weighted_metrics=wms, + scalar_metrics=b.get("scalar_metrics", {}), + aggregation_fns=b.get("aggregation_fns", {}), + ) + ) + else: + buffers.append(b) # pylint: disable-next=protected-access - self._metrics_recorder._metrics_buffer = restored_checkpoint_state.accumulated_metrics + self._metrics_recorder._metrics_buffer = buffers restored_additional_metadata = None if restored_metadata: @@ -360,7 +524,7 @@ def restore_checkpoint(self, **kwargs: Any) -> Any: def record_metrics( self, name: str, - metric: abstract_engine.WeightedMetric | jax.Array | float | int, + metric: abstract_engine.WeightedMetric | jax.Array | float | int | dict[str, Any], aggregation_fn: Callable[[jax.Array], Any] | None = None, ) -> None: """Records a metric into the buffer, appending to JAX arrays. @@ -370,12 +534,23 @@ def record_metrics( metric: The metric to record. aggregation_fn: The aggregation function to apply to the metric. """ - self._metrics_recorder.buffer_metrics( - train_step=self.train_step, - name=name, - metric=metric, - aggregation_fn=aggregation_fn, - ) + if metric is None: + return + if isinstance(metric, dict): + for sub_k, sub_v in metric.items(): + if sub_v is not None: + self.record_metrics( + f"{name}/{sub_k}" if name else sub_k, + sub_v, + aggregation_fn=aggregation_fn, + ) + else: + self._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 accumulated step metrics as an on-device MetricsBuffer. diff --git a/src/maxtext/utils/gradient_accumulation.py b/src/maxtext/utils/gradient_accumulation.py index 35fdf65503..106baee07d 100644 --- a/src/maxtext/utils/gradient_accumulation.py +++ b/src/maxtext/utils/gradient_accumulation.py @@ -176,7 +176,12 @@ def reshape_to_microbatch_accumulations(batch_arr): unreduced_shardings = jax.tree.map(update_sharding_for_unreduced, params_shardings) raw_grads = jax.tree.map(_maybe_shard_with_name, raw_grads, unreduced_shardings) raw_grads = jax.tree.map(_maybe_shard_with_name, raw_grads, params_shardings) - raw_grads = jax.tree_util.tree_map(lambda arr: arr / grad_and_loss["total_weights"], raw_grads) + divisor = ( + config.gradient_accumulation_steps + if getattr(config, "use_tunix_gradient_accumulation", False) + else grad_and_loss["total_weights"] + ) + raw_grads = jax.tree_util.tree_map(lambda arr: arr / divisor, raw_grads) aux = jax.tree.map(lambda x: jnp.sum(x, axis=0), aux) # pytype: disable=module-attr if is_nnx: diff --git a/tests/maxtext_engine_test.py b/tests/maxtext_engine_test.py new file mode 100644 index 0000000000..f19efd3ab9 --- /dev/null +++ b/tests/maxtext_engine_test.py @@ -0,0 +1,393 @@ +# Copyright 2026 Google LLC +# +# 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. + +"""Unit tests for MaxText training engine.""" +# pylint: disable=protected-access + +import dataclasses +from typing import Any +from unittest import mock + +from absl.testing import absltest +from flax import nnx +import jax +import jax.numpy as jnp +from maxtext.configs import pyconfig +from maxtext.training_engine import abstract_engine +from maxtext.training_engine import maxtext_engine +from tests.utils.test_helpers import get_test_config_path +import numpy as np +import optax +import orbax.checkpoint as ocp + + +class DummyNNXModel(nnx.Module): + + def __init__(self): + self.weights = nnx.Param(jnp.array([1.0, 2.0])) + + +@dataclasses.dataclass(kw_only=True) +class DummyPayload(abstract_engine.TrainerPayload): + token_ids: Any = dataclasses.field(default_factory=lambda: jnp.ones((2, 2))) + token_mask: Any = dataclasses.field(default_factory=lambda: jnp.ones((2, 2))) + + +class MaxTextTrainingEngineTest(absltest.TestCase): + + def setUp(self): + """Sets up test dependencies and mocks.""" + super().setUp() + dummy_model = DummyNNXModel() + dummy_opt = nnx.Optimizer(dummy_model, optax.sgd(0.01), wrt=nnx.Param) + patcher = mock.patch.object( + maxtext_engine.train_utils, + "create_training_optimizer", + return_value=(lambda step: jnp.array(0.001), dummy_opt), + ) + self.addCleanup(patcher.stop) + patcher.start() + + from_pretrained_patcher = mock.patch.object( + maxtext_engine.model_creation_utils, + "from_pretrained", + return_value=dummy_model, + ) + self.addCleanup(from_pretrained_patcher.stop) + self.mock_from_pretrained = from_pretrained_patcher.start() + self.mock_config = self.setup_config() + + def setup_config(self, enable_checkpointing: bool = False, **kwargs): + """Sets up a MaxText config via pyconfig.initialize.""" + overrides = { + "model_name": "llama3.1-8b", + "run_name": "test_run", + "base_output_directory": self.create_tempdir().full_path, + "init_weights_seed": 42, + "micro_batch_size_to_train_on": 2, + "gradient_accumulation_steps": 1, + "enable_dropout": False, + "record_internal_nn_metrics": False, + "enable_tensorboard": False, + "tensorboard_dir": self.create_tempdir().full_path, + "skip_jax_distributed_system": True, + "enable_checkpointing": enable_checkpointing, + } + if enable_checkpointing: + overrides.update( + { + "checkpoint_dir": self.create_tempdir().full_path, + "checkpoint_period": 1, + "max_num_checkpoints_to_keep": 10, + "async_checkpointing": False, + } + ) + overrides.update(kwargs) + return pyconfig.initialize([None, get_test_config_path()], **overrides) + + def test_raises_type_error_for_non_pyconfig(self): + invalid_config = abstract_engine.TrainingConfig() + with self.assertRaises(TypeError): + maxtext_engine.MaxTextTrainingEngine(invalid_config) # pytype: disable=wrong-arg-types + + def test_raises_value_error_for_missing_model_name(self): + with self.assertRaises(ValueError): + self.setup_config(model_name="") + + def test_max_text_trainer_instantiation_with_pyconfig(self): + t = maxtext_engine.MaxTextTrainingEngine(self.mock_config) + t.with_loss_fn(lambda *args, **kwargs: (jnp.array(0.5), {})) + self.assertIsInstance(t, abstract_engine.AbstractTrainingEngine) + self.mock_from_pretrained.assert_called_once() + + for step in range(2): + self.assertEqual(t.train_step, step) + payload = DummyPayload( + token_ids=jnp.ones((2, 2)), + token_mask=jnp.ones((2, 2)), + ) + t.compile(payload) + self.assertTrue(t._compiled) + t.with_loss_fn(lambda *args, **kwargs: (jnp.array(0.5), {})) + self.assertFalse(t._compiled) + t.fwd_bwd(payload) + self.assertEqual(t._micro_step_count, 1) + t.update() + self.assertEqual(t._micro_step_count, 0) + self.assertIsNone(t._accumulated_grads) + self.assertEqual(t.train_step, 2) + + @mock.patch("orbax.checkpoint.CheckpointManager") + def test_max_text_trainer_checkpoint_manager_init(self, mock_create_mgr): + mock_config = self.setup_config(enable_checkpointing=True) + + _ = maxtext_engine.MaxTextTrainingEngine(mock_config) + mock_create_mgr.assert_called_once_with( + directory=mock_config.checkpoint_dir, + options=ocp.CheckpointManagerOptions( + save_interval_steps=mock_config.checkpoint_period, + max_to_keep=mock_config.max_num_checkpoints_to_keep, + enable_async_checkpointing=mock_config.async_checkpointing, + ), + ) + + def test_save_checkpoint_called_after_update(self): + mock_config = self.setup_config(enable_checkpointing=True) + + t = maxtext_engine.MaxTextTrainingEngine(mock_config) + mock_orbax_mgr = mock.MagicMock() + mock_orbax_mgr.latest_step.return_value = None + mock_orbax_mgr.save.return_value = True + t._checkpoint_manager._checkpoint_manager = mock_orbax_mgr + + dummy_metadata = mock.MagicMock() + t.save_checkpoint(metadata=dummy_metadata) + + # Verify orbax save was called + mock_orbax_mgr.save.assert_called_once() + call_kwargs = mock_orbax_mgr.save.call_args.kwargs + self.assertNotIn("micro_step_count", call_kwargs["custom_metadata"]) + self.assertEqual(call_kwargs["custom_metadata"]["additional_metadata"], dummy_metadata) + args_dict = ( + dict(call_kwargs["args"].items()) + if hasattr(call_kwargs["args"], "items") and callable(call_kwargs["args"].items) + else call_kwargs["args"].__dict__ + ) + self.assertIn("model_params", args_dict) + self.assertIn("accumulated_metrics", args_dict) + self.assertNotIn("accumulated_grads", args_dict) + + def test_save_checkpoint_skips_if_already_saved(self): + mock_config = self.setup_config(enable_checkpointing=True) + + t = maxtext_engine.MaxTextTrainingEngine(mock_config) + mock_orbax_mgr = mock.MagicMock() + mock_orbax_mgr.latest_step.return_value = 10 + t._checkpoint_manager._checkpoint_manager = mock_orbax_mgr + t.train_step = 10 + + t.save_checkpoint(metadata={"key": "val"}) + mock_orbax_mgr.save.assert_not_called() + + def test_save_checkpoint_drains_inflight_throttler(self): + mock_config = self.setup_config(enable_checkpointing=True) + t = maxtext_engine.MaxTextTrainingEngine(mock_config) + mock_orbax_mgr = mock.MagicMock() + mock_orbax_mgr.latest_step.return_value = None + mock_orbax_mgr.save.return_value = True + t._checkpoint_manager._checkpoint_manager = mock_orbax_mgr + + # Add a dummy item to the throttler queue. + dummy_computation = jnp.array(1.0) + t._throttler.add_computation(computation=dummy_computation, metrics=None) + self.assertEqual(t._throttler._inflight_queue.qsize(), 1) + + t.save_checkpoint(metadata={"test": "val"}) + + # Checkpoint should be saved and throttler queue should be drained. + mock_orbax_mgr.save.assert_called_once() + self.assertTrue(t._throttler._inflight_queue.empty()) + + def test_save_checkpoint_called_after_fwd_bwd_before_update(self): + mock_config = self.setup_config(enable_checkpointing=True) + t = maxtext_engine.MaxTextTrainingEngine(mock_config) + mock_orbax_mgr = mock.MagicMock() + mock_orbax_mgr.latest_step.return_value = None + mock_orbax_mgr.save.return_value = True + t._checkpoint_manager._checkpoint_manager = mock_orbax_mgr + + t._micro_step_count = 1 + t._accumulated_grads = {"params": {"w": jnp.array([0.5, 0.5])}} + + dummy_metadata = mock.MagicMock() + t.save_checkpoint(metadata=dummy_metadata) + + # Verify orbax save was called + mock_orbax_mgr.save.assert_called_once() + call_kwargs = mock_orbax_mgr.save.call_args.kwargs + self.assertEqual(call_kwargs["custom_metadata"]["micro_step_count"], 1) + self.assertEqual(call_kwargs["custom_metadata"]["additional_metadata"], dummy_metadata) + args_dict = ( + dict(call_kwargs["args"].items()) + if hasattr(call_kwargs["args"], "items") and callable(call_kwargs["args"].items) + else call_kwargs["args"].__dict__ + ) + self.assertIn("model_params", args_dict) + self.assertIn("accumulated_metrics", args_dict) + self.assertIn("accumulated_grads", args_dict) + + def test_restore_checkpoint_no_checkpoint_returns_defaults(self): + mock_config = self.setup_config(enable_checkpointing=True) + + t = maxtext_engine.MaxTextTrainingEngine(mock_config) + mock_orbax_mgr = mock.MagicMock() + mock_orbax_mgr.latest_step.return_value = None + t._checkpoint_manager._checkpoint_manager = mock_orbax_mgr + + restored_metadata = t.restore_checkpoint() + self.assertIsNone(restored_metadata) + + def test_restore_checkpoint_restores_ckpt_metadata(self): + mock_config = self.setup_config(enable_checkpointing=True) + t = maxtext_engine.MaxTextTrainingEngine(mock_config) + mock_orbax_mgr = mock.MagicMock() + mock_orbax_mgr.latest_step.return_value = 10 + + # Mock metadata with item_metadata and custom_metadata attributes + dummy_metadata = mock.MagicMock() + mock_metadata = mock.MagicMock() + mock_metadata.item_metadata = {"model_params": {}, "optimizer_state": {}} + mock_metadata.custom_metadata = {"additional_metadata": dummy_metadata} + mock_orbax_mgr.metadata.return_value = mock_metadata + + # Return dummy model and optimizer state from orbax restore + dummy_model = DummyNNXModel() + dummy_opt = nnx.Optimizer(dummy_model, optax.sgd(0.01), wrt=nnx.Param) + dummy_opt_state = nnx.state(dummy_opt, nnx.optimizer.OptState) + mock_orbax_mgr.restore.return_value = { + "model_params": nnx.state(dummy_model), + "optimizer_state": dummy_opt_state, + } + t._checkpoint_manager._checkpoint_manager = mock_orbax_mgr + + restored_metadata = t.restore_checkpoint(step=10) + self.assertEqual(t.train_step, 10) + self.assertEqual(restored_metadata, dummy_metadata) + mock_orbax_mgr.restore.assert_called_once() + + def test_restore_intra_step_checkpoint(self): + mock_config = self.setup_config(enable_checkpointing=True) + t = maxtext_engine.MaxTextTrainingEngine(mock_config) + mock_orbax_mgr = mock.MagicMock() + mock_orbax_mgr.latest_step.return_value = 5 + + # Mock metadata with item_metadata and custom_metadata attributes + dummy_metadata = mock.MagicMock() + mock_metadata = mock.MagicMock() + mock_metadata.item_metadata = {"model_params": {}, "optimizer_state": {}} + mock_metadata.custom_metadata = {"micro_step_count": 2, "additional_metadata": dummy_metadata} + mock_orbax_mgr.metadata.return_value = mock_metadata + + metrics_buf = abstract_engine.MetricsBuffer(id=5, mode="train") + metrics_buf.weighted_metrics["loss"] = abstract_engine.WeightedMetric( + unreduced_sum=jnp.array([4.0, 6.0]), + denominator=jnp.array([2.0, 2.0]), + ) + dummy_grads = {"params": {"w": jnp.array([0.5, 0.5])}} + dummy_model = DummyNNXModel() + dummy_opt = nnx.Optimizer(dummy_model, optax.sgd(0.01), wrt=nnx.Param) + dummy_opt_state = nnx.state(dummy_opt, nnx.optimizer.OptState) + mock_orbax_mgr.restore.return_value = { + "model_params": nnx.state(dummy_model), + "optimizer_state": dummy_opt_state, + "accumulated_metrics": [metrics_buf], + "accumulated_grads": dummy_grads, + } + t._checkpoint_manager._checkpoint_manager = mock_orbax_mgr + + _ = t.restore_checkpoint(step=5) + self.assertEqual(t._micro_step_count, 2) + self.assertEqual(t._accumulated_grads, dummy_grads) + self.assertEqual(len(t._cached_losses), 2) + self.assertTrue(isinstance(t._cached_losses[0], abstract_engine.WeightedMetric)) + self.assertAlmostEqual(float(t._cached_losses[0].unreduced_sum), 4.0) + self.assertAlmostEqual(float(t._cached_losses[1].unreduced_sum), 6.0) + + def test_record_and_get_metrics(self): + t = maxtext_engine.MaxTextTrainingEngine(self.mock_config) + # Record WeightedMetric + t.record_metrics( + name="loss", + metric=abstract_engine.WeightedMetric(unreduced_sum=jnp.array(20.0), denominator=jnp.array(4.0)), + ) + t.record_metrics( + name="loss", + metric=abstract_engine.WeightedMetric(unreduced_sum=jnp.array(30.0), denominator=jnp.array(6.0)), + ) + + # Record scalar + t.record_metrics( + name="lr", + metric=0.002, + aggregation_fn=lambda x: np.round(np.asarray(x), 4), + ) + + metrics_buffer: Any = t.get_metrics(clear_cache=True) + self.assertLen(metrics_buffer, 1) + step0_metrics = metrics_buffer[0] + self.assertIn("loss", step0_metrics.weighted_metrics) + np.testing.assert_array_equal( + step0_metrics.weighted_metrics["loss"].unreduced_sum, + jnp.array([20.0, 30.0]), + ) + np.testing.assert_array_equal( + step0_metrics.weighted_metrics["loss"].denominator, + jnp.array([4.0, 6.0]), + ) + self.assertIn("lr", step0_metrics.scalar_metrics) + np.testing.assert_array_equal(step0_metrics.scalar_metrics["lr"], jnp.array([0.002])) + self.assertIn("lr", step0_metrics.aggregation_fns) + self.assertEqual(step0_metrics.aggregation_fns["lr"](jnp.array([0.002])), 0.002) + + def test_update_with_inflight_throttling(self): + t = maxtext_engine.MaxTextTrainingEngine(self.mock_config) + t.with_loss_fn(lambda *args, **kwargs: (jnp.array(0.5), {})) + + payload = DummyPayload() + t.compile(payload) + + # train_step=0: fwd_bwd + fwd_bwd + update + t.fwd_bwd(payload) + # Loss for micro_step_count=0 is queued. qsize=1. + self.assertEqual(t._throttler._inflight_queue.qsize(), 1) + t.fwd_bwd(payload) + # Loss for micro_step_count=1 is also queued. qsize=2 (full). + self.assertEqual(t._throttler._inflight_queue.qsize(), 2) + t.update() + self.assertEqual(t.train_step, 1) + # wait_for_next() in update() sees qsize=2 (full), so it pops + # index 0 (loss for micro_step_count=0), leaving qsize=1. + # Then add_computation() queues the updated model state and step 0 metrics. + # Since we removed the trailing wait_for_next() from update(), qsize + # remains 2. + self.assertEqual(t._throttler._inflight_queue.qsize(), 2) + expected_state_leaves = jax.tree.leaves(t._state if t._state else t._model) + for idx, (computation, metrics) in enumerate(t._throttler._inflight_queue.queue): + if idx == 0: + # Loss for micro_step_count=0. + self.assertIsNone(metrics) + if idx == 1: + # Metrics for train_step=0. + self.assertIsNotNone(metrics) + self.assertEqual(computation, expected_state_leaves) + + # train_step=1: fwd_bwd + update + # Calling fwd_bwd() while queue is full (qsize=2) triggers wait_for_next(), + # popping index 0 (loss from micro_step_count=1) before adding the new loss. + t.fwd_bwd(payload) + self.assertEqual(t._throttler._inflight_queue.qsize(), 2) + # When update() runs for train_step=1, wait_for_next() pops the metrics + # for train_step=0. This blocks on expected_state_leaves and logs + # train_step=0 metrics. + t.update() + self.assertEqual(t.train_step, 2) + self.assertEqual(t._throttler._inflight_queue.qsize(), 2) + + # Closing trainer drains remaining inflight items. + t.close() + self.assertTrue(t._throttler._inflight_queue.empty()) + + +if __name__ == "__main__": + absltest.main()