From 2fcc1b0ffb330d2f38a8c314ee7f268e85999c10 Mon Sep 17 00:00:00 2001 From: Buddh Prakash Date: Tue, 28 Jul 2026 06:42:33 -0700 Subject: [PATCH] Cache compiled JAX init HLOs in MaxEngine. This CL caches the compiled init_cache HLO for init state and init cache on the MaxEngine instance. By storing the compiled functions and reusing them, we avoid redundant compilation overhead on subsequent calls. FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/AI-Hypercomputer/maxtext/pull/4585 from AI-Hypercomputer:zxhe/eval_start_step 25702cd498557fc23717d541b33c5362e05c6f6b PiperOrigin-RevId: 955227312 --- .github/workflows/gemini_invoke.yml | 2 +- .github/workflows/gemini_review.yml | 4 +- .../inspect_checkpoint.py | 46 ++- .../convert_gpt3_ckpt_from_paxml.py | 81 +++--- src/maxtext/common/checkpointing.py | 112 +++++--- src/maxtext/common/data_loader.py | 46 ++- src/maxtext/configs/base.yml | 3 +- src/maxtext/configs/types.py | 11 +- src/maxtext/inference/maxengine/maxengine.py | 37 ++- src/maxtext/models/deepseek.py | 9 +- src/maxtext/models/deepseek4.py | 1 + src/maxtext/trainers/pre_train/train.py | 7 +- src/maxtext/utils/elastic_utils.py | 21 +- tests/unit/configs_value_test.py | 11 + tests/unit/decoder_layer_model_mode_test.py | 180 ++++++++++++ .../unit/deepseek_decode_consistency_test.py | 268 ++++++++++++++++++ tests/unit/elastic_utils_test.py | 84 +++++- tests/unit/pyconfig_test.py | 9 + tests/unit/train_state_nnx_checkpoint_test.py | 38 ++- 19 files changed, 824 insertions(+), 146 deletions(-) create mode 100644 tests/unit/decoder_layer_model_mode_test.py create mode 100644 tests/unit/deepseek_decode_consistency_test.py diff --git a/.github/workflows/gemini_invoke.yml b/.github/workflows/gemini_invoke.yml index c3c6b9bebd..16e50eebd5 100644 --- a/.github/workflows/gemini_invoke.yml +++ b/.github/workflows/gemini_invoke.yml @@ -86,7 +86,7 @@ jobs: settings: |- { "model": { - "maxSessionTurns": 25 + "maxSessionTurns": 100 }, "telemetry": { "enabled": false, diff --git a/.github/workflows/gemini_review.yml b/.github/workflows/gemini_review.yml index 1babf822ba..0e518de050 100644 --- a/.github/workflows/gemini_review.yml +++ b/.github/workflows/gemini_review.yml @@ -36,7 +36,7 @@ defaults: jobs: review: runs-on: 'ubuntu-latest' - timeout-minutes: 10 + timeout-minutes: 30 permissions: contents: 'read' id-token: 'write' @@ -103,7 +103,7 @@ jobs: settings: |- { "model": { - "maxSessionTurns": 25 + "maxSessionTurns": 100 }, "telemetry": { "enabled": false, diff --git a/src/maxtext/checkpoint_conversion/inspect_checkpoint.py b/src/maxtext/checkpoint_conversion/inspect_checkpoint.py index 7e9784e516..93294454dc 100644 --- a/src/maxtext/checkpoint_conversion/inspect_checkpoint.py +++ b/src/maxtext/checkpoint_conversion/inspect_checkpoint.py @@ -48,7 +48,7 @@ [Mode 2: MaxText Architecture] python -m maxtext.checkpoint_conversion.inspect_checkpoint maxtext \ - model_name= scan_layers= + model_name= scan_layers= enable_nnx= (Optional: other maxtext config) [Mode 3: Orbax] @@ -68,6 +68,8 @@ import absl from maxtext.inference.inference_utils import str2bool from maxtext.checkpoint_conversion.utils.utils import print_peak_memory +from maxtext.utils.model_creation_utils import create_nnx_abstract_model +from flax import nnx def natural_sort_key(s: str): @@ -238,15 +240,19 @@ def inspect_maxtext(args, remaining_args): print(argv) config = pyconfig.initialize(argv) - print(f"\n--- Inspecting MaxText Architecture: {config.model_name} (Scan: {config.scan_layers}) ---") + print( + f"\n--- Inspecting MaxText Architecture: {config.model_name} " + f"(scan_layers: {config.scan_layers}, enable_nnx: {config.enable_nnx}) ---" + ) devices_array = maxtext_utils.create_device_mesh(config) mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) - quant = quantizations.configure_quantization(config) - model = Transformer(config, mesh=mesh, quant=quant) - - # Extract abstract parameters. This returns a PyTree of `ShapeDtypeStruct` - # objects, without materializing weights. - abstract_param = maxtext_utils.get_abstract_param(model, config) + if config.enable_nnx: + _, abstract_model = create_nnx_abstract_model(config, mesh=mesh) + _, abstract_param, _ = nnx.split(abstract_model, nnx.Param, ...) + else: + quant = quantizations.configure_quantization(config) + model = Transformer(config, mesh=mesh, quant=quant) + abstract_param = maxtext_utils.get_abstract_param(model, config) # Calculate and display the total parameter count based purely on abstract shapes. num_params = max_utils.calculate_num_params_from_pytree(abstract_param) @@ -261,8 +267,14 @@ def inspect_maxtext(args, remaining_args): for path_tuple, abstract_leaf_value in abstract_params_flat: key_parts = param_key_parts_from_path(path_tuple) - # Construct a MaxText-style parameter key (e.g., "params.params.layer.weight"). - param_key = "params." + ".".join(key_parts) + # Construct a MaxText-style parameter key. Examples: + # "params.params.decoder.decoder_norm.scale" (for standard model weights) + # "params.Tid2EidVar.decoder.layers_0.mlp.MoeBlock_0.tid2eid" (for legacy custom collections) + key_str = ".".join(key_parts) + if config.enable_nnx and not key_str.startswith(("params", "Tid2EidVar")): + param_key = "params.params." + key_str + else: + param_key = "params." + key_str shape = abstract_leaf_value.shape param_dict[param_key] = f"shape: {shape}" @@ -294,6 +306,7 @@ def inspect_orbax(args): # Defer imports to avoid overhead when running in other modes. import orbax.checkpoint as ocp from etils import epath + from maxtext.checkpoint_conversion.utils.utils import param_key_parts_from_path path = epath.Path(args.path) @@ -309,9 +322,13 @@ def inspect_orbax(args): # Filter strictly for parameter keys and format them. param_dict = {} for k, v in dictionary.items(): - # `k` is a tuple representing the path hierarchy; join it into a single string. - # `v` is a metadata object containing `.shape` and `.dtype`. - param_key = ".".join(k) + # `k` is a tuple representing the path hierarchy. + # `v` is a shape-dtype struct. + # Pass it through param_key_parts_from_path to normalize NNX structures + # (folding list indices like "0" into "layers_0" and dropping "value"). + key_parts = param_key_parts_from_path(k) + param_key = ".".join(key_parts) + if not param_key.startswith("params"): continue @@ -372,7 +389,8 @@ def main(): if args.mode == "hf": inspect_hf(args) elif args.mode == "maxtext": - # remaining_args accepts maxtext config, like `model_name= scan_layers=` + # remaining_args accepts maxtext config, like `model_name= + # scan_layers= enable_nnx=` inspect_maxtext(args, remaining_args) elif args.mode == "orbax": inspect_orbax(args) diff --git a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py index 2f3e26db7b..138d5be4a5 100644 --- a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py +++ b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py @@ -41,6 +41,7 @@ import sys from flax import nnx +from flax import serialization import jax from jax import random from jax.sharding import Mesh @@ -118,6 +119,12 @@ def init_state_fn(): ) state, _, _, _, _ = maxtext_utils.setup_training_state(None, cfg, mesh, checkpoint_manager, init_state_fn) + if cfg.pure_nnx: + state = train_state_nnx.to_checkpoint_dict(state) + state.pop("nnx_aux", None) + else: + state = serialization.to_state_dict(state) + max_logging.log("start") max_utils.print_mem_stats("After params initialized") @@ -201,24 +208,10 @@ def init_state_fn(): "['decoder']['decoder_norm']['bias']": (".params.lm.final_ln.bias", None), } - if cfg.pure_nnx: - # NNX state-tree paths after `nnx.split(TrainStateNNX)`. The state is a - # nested `nnx.State` (dict-like Mapping) with `nnx.Variable` leaves, so - # `jax.tree_util.keystr` produces dict-style entries (`['key']`) plus - # `.value` for the Variable leaf, plus `[idx]` for the optax tuple: - # model params -> ['model'].value - # adam mu / nu -> ['optimizer']['opt_state'][0]['mu' | 'nu'].value - # step -> ['optimizer']['step'].value - # opt count -> ['optimizer']['opt_state'][0]['count'].value - state_map = { - "['optimizer']['step'].value": ("step", None), - "['optimizer']['opt_state'][0]['count'].value": ("opt_states_0.no_prefix_0.count", None), - } - else: - state_map = { - ".step": ("step", None), - ".opt_state.count": ("opt_states_0.no_prefix_0.count", None), - } + state_map = { + "['step']": ("step", None), + "['opt_state']['count']": ("opt_states_0.no_prefix_0.count", None), + } def get_layer_prefix(keystr_pax): # different path format between decoder_layer variable @@ -231,26 +224,18 @@ def get_layer_prefix(keystr_pax): for keystr_maxtext, (keystr_pax, transform_fn) in keystr_map.items(): prefix_pax_opt_state = get_layer_prefix(keystr_pax) - if cfg.pure_nnx: - state_map[f"['model']{keystr_maxtext}.value"] = (f"mdl_vars{keystr_pax}", transform_fn) - state_map[f"['optimizer']['opt_state'][0]['mu']{keystr_maxtext}.value"] = ( - f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}", - transform_fn, - ) - state_map[f"['optimizer']['opt_state'][0]['nu']{keystr_maxtext}.value"] = ( - f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}", - transform_fn, - ) - else: - state_map[f".params['params']{keystr_maxtext}"] = (f"mdl_vars{keystr_pax}", transform_fn) - state_map[f".opt_state.mu['params']{keystr_maxtext}"] = ( - f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}", - transform_fn, - ) - state_map[f".opt_state.nu['params']{keystr_maxtext}"] = ( - f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}", - transform_fn, - ) + state_map[f"['params']['params']{keystr_maxtext}"] = ( + f"mdl_vars{keystr_pax}", + transform_fn, + ) + state_map[f"['opt_state']['mu']['params']{keystr_maxtext}"] = ( + f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}", + transform_fn, + ) + state_map[f"['opt_state']['nu']['params']{keystr_maxtext}"] = ( + f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}", + transform_fn, + ) def verify_fn(key_path, _): keystr = jax.tree_util.keystr(key_path) @@ -302,9 +287,9 @@ def map_fn(key_path, value): max_logging.log("converted state finished") max_utils.print_mem_stats("converted state finished") - step_value = int(converted_state.optimizer.step.value) if cfg.pure_nnx else converted_state.step - if checkpointing.save_checkpoint(checkpoint_manager, step_value, converted_state): - max_logging.log(f"saved a checkpoint at step {step_value}") + step_value = int(converted_state["step"]) if isinstance(converted_state, dict) else int(converted_state.step) + if checkpointing.save_checkpoint(checkpoint_manager, step_value, converted_state, config=cfg): + max_logging.log(f"Saved a checkpoint at step {step_value}.") # Upon preemption, exit when and only when all ongoing saves are complete. if checkpointing.reached_preemption(checkpoint_manager, step_value): checkpointing.wait_until_finished(checkpoint_manager) @@ -322,7 +307,12 @@ def map_fn(key_path, value): default="gs://mlperf-llm-public2/gpt3_spmd1x64x24_tpuv4-3072_v84_20221101/checkpoints/checkpoint_00004000", required=True, ) - parser.add_argument("--maxtext-model-name", choices=["gpt3-175b", "gpt3-52k"], type=str, required=True) + parser.add_argument( + "--maxtext-model-name", + choices=["gpt3-175b", "gpt3-52k"], + type=str, + required=True, + ) parser.add_argument("--base-output-directory", type=str, required=True) parser.add_argument("--run-name", type=str, required=True) @@ -330,4 +320,9 @@ def map_fn(key_path, value): if not args.paxml_ckpt_path.startswith("gs://"): raise ValueError("--paxml-ckpt-path should be a gcs path starting with gs://") - convert(args.paxml_ckpt_path, args.maxtext_model_name, args.base_output_directory, args.run_name) + convert( + args.paxml_ckpt_path, + args.maxtext_model_name, + args.base_output_directory, + args.run_name, + ) diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index 4827cbc1a4..d0d3bd7028 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -15,6 +15,7 @@ """Create an Orbax CheckpointManager with specified (Async or not) Checkpointer.""" +import contextlib import datetime import importlib import time @@ -716,7 +717,35 @@ def _handle_post_checkpoint_preemption(checkpoint_manager, step, force_ckpt_save if force_ckpt_save or is_preempted: wait_until_finished(checkpoint_manager) if is_preempted: - raise exceptions.StopTraining("Job is preempted.") + raise exceptions.StopTraining("Job received termination signal (SIGTERM).") + + +@contextlib.contextmanager +def checkpoint_exception_guard(config, checkpoint_manager, handler_fn=None): + """Context manager that wraps checkpointing save Exception handling. + + On block success (checkpoint written without errors): runs the scale-up check + if elastic training is active. + On block failure: bubbles up JAX/ScaleUp errors if elastic training is active; + otherwise delegates to `handler_fn`. + + Args: + config: maxtext configuration object. + checkpoint_manager: The CheckpointManager instance. + handler_fn: Optional callback function(Exception) that handles/wraps + non-elastic exceptions. If this handler raises a new exception, that new + exception is propagated. If it returns normally (returns None) or is not + provided, the original exception is re-raised (preserving its traceback). + """ + try: + yield + elastic_utils.maybe_elastic_scale_up(config, checkpoint_manager) + except Exception as e: # pylint: disable=broad-except + elastic_utils.maybe_bubble_elastic_exception(config, e) + if handler_fn: + handler_fn(e) + else: + raise def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step=None): @@ -753,38 +782,21 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step max_logging.log(f"Checkpoint for step {actual_step} already exists, skipping save.") return - if config.pure_nnx: - # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable. - if config.enable_diloco: - # DiLoCoTrainState: persist the synchronized global model (outer params). - # The per-replica inner optimizer / outer-momentum state is not checkpointed. - step_value = state.step.get_value() if hasattr(state.step, "get_value") else state.step - state = train_state_nnx.to_linen_checkpoint_dict({"model": state.params, "optimizer": {"step": step_value}}) - else: - # rngs/dropout/batch-stats are packed under items/nnx_aux so the RNG/dropout - # stream continues across resumes instead of resetting to a base key. - state = train_state_nnx.to_checkpoint_dict(state) - - try: - checkpoint_saved = save_checkpoint(checkpoint_manager, actual_step, state, config, data_iterator, force_ckpt_save) - if checkpoint_saved: - print_save_message(actual_step, config.async_checkpointing) - if config.elastic_enabled: - elastic_utils.maybe_elastic_scale_up(config, checkpoint_manager) - except elastic_utils.manager.ScaleUpSignalError as e: - if config.elastic_enabled: - max_logging.log(f"Elastic event detected, letting exception bubble up: {e}") - raise - else: - raise exceptions.StopTraining("Job is preempted.") from e - except jax.errors.JaxRuntimeError as e: - if config.elastic_enabled: - max_logging.log(f"Elastic event detected, letting exception bubble up: {e}") - raise - else: - raise exceptions.StopTraining("Job is preempted.") from e - except Exception as e: - raise exceptions.StopTraining(f"Checkpointing failed. {str(e)}") from e + def _checkpoint_error_handler(err): + """Handles checkpointing errors, when not in an elastic context.""" + raise exceptions.StopTraining(f"Checkpointing failed. {str(err)}") from err + + with checkpoint_exception_guard(config, checkpoint_manager, _checkpoint_error_handler): + checkpoint_saved = save_checkpoint( + checkpoint_manager, + actual_step, + state, + config, + data_iterator, + force_ckpt_save, + ) + if checkpoint_saved: + print_save_message(actual_step, config.async_checkpointing) # Wait for any pending checkpoint save to finish during preemption or final # step save, then raise upon preemption. @@ -793,12 +805,26 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=None, force=False): """Wrapper for saving checkpoint.""" - if config and config.enable_checkpointing: + if not isinstance(state, (dict, nnx.State, train_state.TrainState)): + if isinstance(state, train_state_nnx.TrainStateNNX): + state = nnx.state(state) + elif not isinstance(state, (dict, nnx.State)): + state = {} + + if config and getattr(config, "pure_nnx", False) and isinstance(state, nnx.State): + # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable. + if getattr(config, "enable_diloco", False): + step_value = state.step.get_value() if hasattr(state.step, "get_value") else state.step + state = train_state_nnx.to_linen_checkpoint_dict({"model": state.params, "optimizer": {"step": step_value}}) + else: + state = train_state_nnx.to_checkpoint_dict(state) + + if config and getattr(config, "enable_checkpointing", False): if ( force - or (step % config.checkpoint_period == 0 and not config.enable_continuous_checkpointing) + or (step % config.checkpoint_period == 0 and not getattr(config, "enable_continuous_checkpointing", False)) or (_uses_local_checkpoint_period(config) and step % config.local_checkpoint_period == 0) - or (config.enable_autocheckpoint and reached_preemption(checkpoint_manager, step)) + or (getattr(config, "enable_autocheckpoint", False) and reached_preemption(checkpoint_manager, step)) ): blocking_until_ready_start = time.time() max_logging.log(f"Waiting for step {step} to finish before checkpoint...") @@ -812,7 +838,13 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator= # specify chunk_byte_size to force orbax to control maximum file size in checkpoint chunk_byte_size = ( - config.checkpoint_storage_target_data_file_size_bytes if config else DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE + getattr( + config, + "checkpoint_storage_target_data_file_size_bytes", + DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE, + ) + if config + else DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE ) checkpoint_args = ocp.args.PyTreeSave( @@ -822,7 +854,11 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator= ) save_args_composite = {"items": checkpoint_args} - if config and config.dataset_type == "grain" and not isinstance(data_iterator, PlaceHolderDataIterator): + if ( + config + and getattr(config, "dataset_type", None) == "grain" + and not isinstance(data_iterator, PlaceHolderDataIterator) + ): if isinstance(data_iterator, RemoteIteratorWrapper): # Pass the wrapper directly; GrainCheckpointHandler will call save_state with the step save_args_composite["iter"] = grain_utility.GrainCheckpointSave( diff --git a/src/maxtext/common/data_loader.py b/src/maxtext/common/data_loader.py index 21bd870bc8..2bd04c527b 100644 --- a/src/maxtext/common/data_loader.py +++ b/src/maxtext/common/data_loader.py @@ -15,19 +15,40 @@ # pytype: disable=unsupported-operands """Module to load data for training.""" +import contextlib import jax -import jax.numpy as jnp from jax.experimental import checkify - +import jax.numpy as jnp from maxtext.common.goodput import ( GoodputEvent, maybe_record_goodput, ) from maxtext.trainers.diloco import diloco +from maxtext.utils import elastic_utils from maxtext.utils import exceptions from maxtext.utils.sharding import get_input_data_sharding +@contextlib.contextmanager +def loader_exception_guard(config): + """Context manager that wraps data loading Exception handling. + + On block failure: bubbles up JAX/ScaleUp errors if elastic training is active; + otherwise raises a StopTraining exception. + + Args: + config: maxtext configuration object. + """ + try: + yield + except Exception as e: # pylint: disable=broad-except + elastic_utils.maybe_bubble_elastic_exception(config, e) + if isinstance(e, StopIteration): + raise exceptions.StopTraining(f"You may have run out of training data. Received {type(e)}" f" exception: ({e})") + else: + raise exceptions.StopTraining(f"`next(self.data_iterator)` failed with {type(e)} exception: ({e}).") + + class DataLoader: """ Loads preprocessed data for training. @@ -54,23 +75,18 @@ def update_data_iterator(self): def load_next_batch_pre_sharding(self): """Loads the next batch w/o sharding. Can keep reusing the same batch for performance reasons.""" with maybe_record_goodput(self.goodput_recorder, GoodputEvent.DATA_LOADING): - try: - if self.config.reuse_example_batch and self.last_batch: - example_batch = self.last_batch - else: + if self.config.reuse_example_batch and self.last_batch: + example_batch = self.last_batch + else: + with loader_exception_guard(self.config): example_batch = next(self.data_iterator) - self.update_data_iterator() - self.last_batch = example_batch - self.check_example_batch() - except Exception as e: # pylint: disable=broad-except - if isinstance(e, StopIteration): - raise exceptions.StopTraining(f"You may have run out of training data. Received {type(e)} exception: ({e})") - else: - raise exceptions.StopTraining(f"`load_next_batch()` failed with {type(e)} exception: ({e}).") + self.update_data_iterator() + self.last_batch = example_batch + self.check_example_batch() return self.last_batch def load_next_batch(self, *args, **kwargs): - """Loads the next batch with sharding hint""" + """Loads the next batch with sharding hint.""" example_batch = self.load_next_batch_pre_sharding() if self.config.enable_diloco: example_batch = diloco.reshape_first_axis_with_diloco(self.config.num_diloco_replicas, example_batch) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index b0843a9939..73f6ddc4a6 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -988,7 +988,7 @@ decode_sampling_nucleus_p: -1 # set if you're doing nucleus / top-p decode_sampling_top_k: 0 # set if you're doing top-k decode_sampling_temperature: 1. -eval_start_step: 0 # start eval after train step is >= eval_start_step +eval_start_step: 0 # start eval when train step is >= eval_start_step eval_interval: -1 # the specific number of train step between eval_step eval_steps: -1 # run this number of steps for eval, recommend setting this to prevent error due to running out of evel data target_eval_loss: 0. # early stop once reaching target eval_loss @@ -1339,6 +1339,7 @@ distill_beta_schedule: "constant" ##### Elastic training parameters # Elastic training is Pathways-specific and does not work on McJAX. elastic_enabled: false +elastic_backup_kind: "snapshot" elastic_timeout_seconds: 300 elastic_max_retries: 10 elastic_min_slice_count: -1 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index edc7d3112e..2d6ae37b12 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1600,7 +1600,8 @@ class TrainingLoop(BaseModel): log_period: int = Field(100, description="Frequency (in steps) to log metrics and flush Tensorboard.") eval_start_step: int = Field( 0, - description="Start evaluation after training step is >= eval_start_step.", + ge=0, + description="Start evaluation when training step is >= eval_start_step.", ) eval_interval: int = Field( -1, @@ -2035,6 +2036,10 @@ class ElasticTraining(BaseModel): """ elastic_enabled: bool = Field(False, description="Whether to enable elastic training.") + elastic_backup_kind: str = Field( + "snapshot", + description=("The kind of backup to use for elastic training: 'snapshot' or 'checkpoint'."), + ) elastic_timeout_seconds: int = Field( 300, description=( @@ -3227,6 +3232,10 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de ) if self.elastic_enabled and not self.enable_single_controller: raise ValueError("Elastic training is only supported with Pathways (`enable_single_controller=True`).") + if self.elastic_backup_kind not in ("snapshot", "checkpoint"): + raise ValueError( + "elastic_backup_kind must be one of 'snapshot' or 'checkpoint', got" f" '{self.elastic_backup_kind}'." + ) if self.colocated_python_data_input and not self.enable_single_controller: raise ValueError( "Colocated python data input is only supported with Pathways (single" diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index b13d72d89e..c4c7d9f225 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -145,11 +145,16 @@ def __init__(self, config: Any, devices: Any | None = None): # nnx.merge. `rest` (RNG state etc.) is materialized in load_params. graphdef, _, _, _ = nnx.split(abstract_model, nnx.Param, nnx.Cache, ...) self.graphdef = graphdef + # Layers may bake their construction-time model_mode into static attributes, + # so a call must be merged with the graphdef built for that same mode. + graphdef_ar, _, _, _ = nnx.split(abstract_model_ar, nnx.Param, nnx.Cache, ...) + self.graphdef_ar = graphdef_ar self._create_model_fn = _create_model self._nnx_rest_state = None else: self.model = models.transformer_as_linen(config, mesh=self._mesh, quant=quant, model_mode=MODEL_MODE_PREFILL) self.graphdef = None + self.graphdef_ar = None self._create_model_fn = None self.replicated_sharding = jax.sharding.NamedSharding(self._mesh, P(None)) @@ -164,6 +169,8 @@ def __init__(self, config: Any, devices: Any | None = None): self.decode_state_layouts = None self.param_layouts = None self.rng = None + self._compiled_initialize_fn = None + self._compiled_init_cache_fn = None def print_stats(self, label: str): max_utils.print_mem_stats(label) @@ -218,10 +225,14 @@ def _nnx_run_model( """NNX equivalent of `model.apply(..., mutable=["cache"])`. Returns (logits, new_cache_dict).""" cache_state = self._nnx_cache_state_template(mode=model_mode) nnx.replace_by_pure_dict(cache_state, cache_dict) + # Merge with the graphdef built for this mode. Layers that captured their + # model_mode at construction (e.g. the DeepSeek layers) would otherwise run + # the prefill attention path during autoregressive decode. + graphdef = self.graphdef_ar if model_mode == MODEL_MODE_AUTOREGRESSIVE else self.graphdef # copy=True avoids reusing Variable objects across traces (TraceContextError), # mirroring the workaround in train.py's diff_wrapper. model = nnx.merge( - self.graphdef, params, cache_state, self._nnx_rest_state, copy=True + graphdef, params, cache_state, self._nnx_rest_state, copy=True ) # pyrefly: ignore[no-matching-overload] logits = model( decoder_input_tokens, @@ -1974,11 +1985,15 @@ def init(abstract_params): mesh_annotations, ) - @functools.partial(jax.jit, out_shardings=shardings) - def initialize(): - return jax.tree_util.tree_map(lambda x: jnp.zeros(x.shape, x.dtype), abstract_outputs) + if self._compiled_initialize_fn is None: - init_state = initialize() + @functools.partial(jax.jit, out_shardings=shardings) + def initialize(): + return jax.tree_util.tree_map(lambda x: jnp.zeros(x.shape, x.dtype), abstract_outputs) + + self._compiled_initialize_fn = initialize + + init_state = self._compiled_initialize_fn() cache = init_state["cache"] def is_lp(k): @@ -2002,11 +2017,15 @@ def _init_decode_state_nnx(self) -> DecodeState: # AR-mode cache so the batch dim matches generate's input shape. cache_dict_abs = self._nnx_init_cache_dict(mode=MODEL_MODE_AUTOREGRESSIVE) - @functools.partial(jax.jit, out_shardings=(self.kv_cache_shardings,)) - def _init_cache(): - return (jax.tree.map(lambda x: jnp.zeros(x.shape, x.dtype), cache_dict_abs),) + if self._compiled_init_cache_fn is None: + + @functools.partial(jax.jit, out_shardings=(self.kv_cache_shardings,)) + def _init_cache(): + return (jax.tree.map(lambda x: jnp.zeros(x.shape, x.dtype), cache_dict_abs),) + + self._compiled_init_cache_fn = _init_cache - (cache,) = _init_cache() + (cache,) = self._compiled_init_cache_fn() # Per-leaf logical axes for bulk_insert's "cache_batch" lookup. Use model_ar # so segment_id leaves carry CACHE_BATCH (under PREFILL they'd carry diff --git a/src/maxtext/models/deepseek.py b/src/maxtext/models/deepseek.py index 53166f00ee..e0e09f4175 100644 --- a/src/maxtext/models/deepseek.py +++ b/src/maxtext/models/deepseek.py @@ -211,6 +211,7 @@ def attention_op( decoder_segment_ids, decoder_positions, deterministic, + model_mode, previous_chunk=None, slot: None | int = None, ): @@ -221,7 +222,7 @@ def attention_op( decoder_positions, decoder_segment_ids=decoder_segment_ids, deterministic=deterministic, - model_mode=self.model_mode, + model_mode=model_mode, out_sharding=self.out_sharding, previous_chunk=previous_chunk, slot=slot, @@ -270,6 +271,7 @@ def self_attention_with_norm_op( decoder_segment_ids, decoder_positions, deterministic, + model_mode, previous_chunk=None, slot: None | int = None, ): @@ -283,7 +285,7 @@ def self_attention_with_norm_op( decoder_segment_ids=decoder_segment_ids, inputs_positions=decoder_positions, deterministic=deterministic, - model_mode=self.model_mode, + model_mode=model_mode, out_sharding=self.out_sharding, previous_chunk=previous_chunk, slot=slot, @@ -295,6 +297,7 @@ def self_attention_with_norm_op( decoder_segment_ids, decoder_positions, deterministic, + model_mode, previous_chunk, slot, ) @@ -368,6 +371,7 @@ def __call__( decoder_segment_ids, decoder_positions, deterministic, + model_mode, previous_chunk, slot, ) @@ -581,6 +585,7 @@ def extract_fn(x): decoder_segment_ids, decoder_positions, deterministic, + model_mode, previous_chunk, slot, ) diff --git a/src/maxtext/models/deepseek4.py b/src/maxtext/models/deepseek4.py index c259db8c66..6181d684e0 100644 --- a/src/maxtext/models/deepseek4.py +++ b/src/maxtext/models/deepseek4.py @@ -151,6 +151,7 @@ def __call__( decoder_segment_ids, decoder_positions, deterministic, + model_mode, previous_chunk, slot, ) diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index e28316ef15..367695211f 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -718,7 +718,12 @@ def training_loop_iteration( all_host_upload=dump_hlo_upload_all, ) - if eval_interval > 0 and step >= start_step and step >= eval_start_step and (step + 1) % eval_interval == 0: + if ( + eval_interval > 0 + and step >= start_step + and step >= eval_start_step + and (step - eval_start_step) % eval_interval == 0 + ): assert eval_data_iterator # Explicitly reset the eval iterator and counters before starting the eval loop eval_data_iterator.reset() diff --git a/src/maxtext/utils/elastic_utils.py b/src/maxtext/utils/elastic_utils.py index 1cba750774..4ebfb97571 100644 --- a/src/maxtext/utils/elastic_utils.py +++ b/src/maxtext/utils/elastic_utils.py @@ -14,8 +14,8 @@ """Utility functions for Elastic Training.""" -import functools from collections import Counter +import functools from types import SimpleNamespace import jax @@ -64,6 +64,22 @@ def elastic_enabled(config) -> bool: return pathwaysutils.is_pathways_backend_used() and config.elastic_enabled +def elastic_snapshot(config) -> bool: + """Returns whether elastic snapshot mode is enabled.""" + return elastic_enabled(config) and config.elastic_backup_kind == "snapshot" + + +def maybe_bubble_elastic_exception(config, e: Exception) -> None: + """Checks JAX/ScaleUp elastic errors and re-raises them if elasticity is enabled. + + Args: + config: Maxtext configuration object. + e: The exception currently being evaluated. + """ + if elastic_enabled(config) and isinstance(e, (jax.errors.JaxRuntimeError, manager.ScaleUpSignalError)): + raise e + + def should_use_elastic(config) -> bool: """Returns whether elastic training should be used.""" return config is not None and elastic_enabled(config) @@ -218,6 +234,9 @@ def is_scale_up_event(config) -> bool: def maybe_elastic_scale_up(config, checkpoint_manager): """Waits for a checkpoint to finish before interrupting for scale up.""" + if not should_use_elastic(config): + max_logging.log("maybe_elastic_scale_up: Elastic training is not enabled.") + return if is_scale_up_event(config): max_logging.log( "Started a checkpoint and a new slice is available. Waiting for current" diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py index 465adbc307..78dfab1da9 100644 --- a/tests/unit/configs_value_test.py +++ b/tests/unit/configs_value_test.py @@ -312,6 +312,17 @@ def test_safetensors_dynamic_disallows_single_controller(self): with self.assertRaises(pydantic.ValidationError): pyconfig.initialize(argv) + def test_elastic_backup_kind_validation(self): + """Tests that elastic_backup_kind must be either 'snapshot' or 'checkpoint'.""" + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "elastic_backup_kind=invalid_backup_kind", + ] + with self.assertRaises(pydantic.ValidationError): + pyconfig.initialize(argv) + if __name__ == "__main__": absltest.main() diff --git a/tests/unit/decoder_layer_model_mode_test.py b/tests/unit/decoder_layer_model_mode_test.py new file mode 100644 index 0000000000..1d60a3b69d --- /dev/null +++ b/tests/unit/decoder_layer_model_mode_test.py @@ -0,0 +1,180 @@ +# Copyright 2023-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 +# +# https://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. + +"""A decoder layer must attend in the mode it is called with, not the one it was built in. + +Under NNX a model object is built once and reused across modes — MaxEngine builds it +in PREFILL and then passes MODEL_MODE_AUTOREGRESSIVE per call while decoding. A layer +that forwards its construction-time `self.model_mode` to attention instead of the +`model_mode` argument therefore runs the prefill KV-cache path during decode and +produces wrong tokens from the first generated one, with no error anywhere. + +The check stops at the attention boundary and asserts on the mode that arrives there, +so it stays cheap and does not depend on any particular attention implementation. +""" + +import sys +import unittest +from unittest import mock + +import jax +import jax.numpy as jnp +from absl.testing import parameterized +from flax import nnx +from flax.linen import partitioning as nn_partitioning + +from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL +from maxtext.configs import pyconfig +from maxtext.layers import attention_mla, attentions +from maxtext.models import deepseek, gemma, gemma2, llama2, mistral, qwen2, qwen3 +from maxtext.utils import maxtext_utils +from tests.utils.test_helpers import get_test_config_path + +# Config knobs shared by every family; small enough that construction is the only cost. +_COMMON = { + "base_emb_dim": 64, + "base_mlp_dim": 64, + "base_num_query_heads": 4, + "base_num_kv_heads": 4, + "base_num_decoder_layers": 3, + "vocab_size": 64, + "max_prefill_predict_length": 8, + "max_target_length": 16, + "per_device_batch_size": 1, + "scan_layers": False, + "attention": "dot_product", + "sparse_matmul": False, + "dtype": "float32", + "weight_dtype": "float32", + "enable_checkpointing": False, + "skip_jax_distributed_system": True, + "pure_nnx": True, +} + +_DEEPSEEK = { + "model_name": "deepseek2-16b", + "override_model_config": True, + "mla_naive_kvcache": False, + "first_num_dense_layers": 1, + "num_experts": 4, + "num_experts_per_tok": 2, + "shared_experts": 1, +} + +# Each entry is (case name, layer class, extra config, extra constructor kwargs). +_LAYERS = [ + ("deepseek_dense", deepseek.DeepSeekDenseLayer, _DEEPSEEK, {}), + ("deepseek_moe", deepseek.DeepSeekMoELayer, _DEEPSEEK, {}), + ("llama2", llama2.LlamaDecoderLayer, {}, {}), + ("mistral", mistral.MistralDecoderLayer, {"decoder_block": "mistral"}, {}), + ("gemma", gemma.GemmaDecoderLayer, {"decoder_block": "gemma"}, {}), + ("gemma2", gemma2.Gemma2DecoderLayer, {"decoder_block": "gemma2"}, {}), + ("qwen2", qwen2.Qwen2DecoderLayer, {"decoder_block": "qwen2"}, {"quant": None}), + ("qwen3", qwen3.Qwen3DecoderLayer, {"decoder_block": "qwen3"}, {"quant": None}), +] + + +def _mesh(cfg): + """Builds the mesh over every available device. + + ici_fsdp_parallelism keeps its base.yml default of -1, so the fsdp axis absorbs + however many devices the host has. Pinning the ICI axes instead would fix the mesh + at one device and fail everywhere except a single-device runner. + + Args: + cfg: Model config. + + Returns: + The device mesh. + """ + return jax.sharding.Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + + +class _StopAtAttention(Exception): + """Raised by the spy so the forward pass ends once the mode has been observed.""" + + +class DecoderLayerModelModeTest(parameterized.TestCase): + """Every decoder layer must forward the call-time model_mode to its attention.""" + + def _mode_seen_by_attention(self, layer_cls, extra_config, ctor_kwargs, build_mode, call_mode): + """Builds a layer in one mode and calls it in another. + + Args: + layer_cls: Decoder layer class to instantiate. + extra_config: Config overrides this family needs, on top of _COMMON. + ctor_kwargs: Extra constructor arguments this family requires. + build_mode: Model mode the layer is constructed with. + call_mode: Model mode the layer is called with. + + Returns: + The model mode that reached the attention module. + """ + cfg = pyconfig.initialize([sys.argv[0], get_test_config_path()], **(_COMMON | extra_config)) + mesh = _mesh(cfg) + seen = [] + + def spy(_self, *args, **kwargs): + # model_mode is the 5th positional parameter of Attention.__call__, but every + # in-tree caller passes it by keyword. Accept both so the spy is not brittle. + seen.append(kwargs.get("model_mode", args[4] if len(args) > 4 else None)) + raise _StopAtAttention() + + with nn_partitioning.axis_rules(cfg.logical_axis_rules), mesh: + layer = layer_cls( + config=cfg, + model_mode=build_mode, + mesh=mesh, + rngs=nnx.Rngs(params=0, dropout=0), + **ctor_kwargs, + ) + batch = cfg.micro_batch_size_to_train_on + inputs = jnp.zeros((batch, 1, cfg.emb_dim), dtype=jnp.float32) + positions = jnp.zeros((batch, 1), dtype=jnp.int32) + segment_ids = jnp.ones((batch, 1), dtype=jnp.int32) + with ( + mock.patch.object(attentions.Attention, "__call__", spy), + mock.patch.object(attention_mla.MLA, "__call__", spy), + ): + try: + layer(inputs, segment_ids, positions, True, call_mode) + except _StopAtAttention: + pass + + self.assertEqual(len(seen), 1, f"expected exactly one attention call, saw {len(seen)}") + return seen[0] + + @parameterized.named_parameters(*_LAYERS) + def test_decode_call_overrides_prefill_construction(self, layer_cls, extra_config, ctor_kwargs): + """Checks a layer built for prefill and called for decode, as MaxEngine does.""" + seen = self._mode_seen_by_attention( + layer_cls, extra_config, ctor_kwargs, MODEL_MODE_PREFILL, MODEL_MODE_AUTOREGRESSIVE + ) + self.assertEqual( + seen, + MODEL_MODE_AUTOREGRESSIVE, + "attention ran in the layer's construction-time mode; decode would use the prefill KV-cache path", + ) + + @parameterized.named_parameters(*_LAYERS) + def test_prefill_call_overrides_decode_construction(self, layer_cls, extra_config, ctor_kwargs): + """Checks the reverse direction, so a layer cannot pass by hardcoding either mode.""" + seen = self._mode_seen_by_attention( + layer_cls, extra_config, ctor_kwargs, MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL + ) + self.assertEqual(seen, MODEL_MODE_PREFILL, "attention ran in the layer's construction-time mode") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/deepseek_decode_consistency_test.py b/tests/unit/deepseek_decode_consistency_test.py new file mode 100644 index 0000000000..181636d9a8 --- /dev/null +++ b/tests/unit/deepseek_decode_consistency_test.py @@ -0,0 +1,268 @@ +# Copyright 2023-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 +# +# https://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. + +"""Greedy decode through MaxEngine must match a teacher-forced forward pass. + +Golden-logit tests only cover the forward pass, so nothing catches a broken +prefill -> autoregressive handoff. Here a tiny model is decoded greedily through +prefill/insert/generate and compared against the argmax rollout of a plain forward +pass over the same growing prefix, which is the same sequence by definition. That +holds for any model, so new families can be added to the parameter list. + +DeepSeek is the case that motivated this: MaxEngine's NNX path builds the model once +in PREFILL mode and passes the mode per call, so a layer that reads its +construction-time model_mode instead of the argument runs the prefill attention path +during decode and silently produces wrong tokens. +""" + +import sys +import unittest +from unittest import mock + +import jax +import jax.numpy as jnp +import pytest +from absl.testing import parameterized +from flax import nnx +from flax.linen import partitioning as nn_partitioning + +from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN +from maxtext.configs import pyconfig +from maxtext.models import deepseek +from maxtext.utils import maxtext_utils, model_creation_utils + +pytest.importorskip("jetstream", reason="jetstream not installed") +from maxtext.inference.maxengine import maxengine +from tests.utils.test_helpers import get_test_config_path + +PROMPT = [3, 17, 42, 5, 9] +STEPS = 4 + +_COMMON = { + "base_emb_dim": 64, + "base_mlp_dim": 64, + "base_num_query_heads": 4, + "base_num_kv_heads": 4, + "base_num_decoder_layers": 3, + "vocab_size": 64, + "max_prefill_predict_length": 8, + "max_target_length": 16, + "per_device_batch_size": 1, + "scan_layers": False, + "attention": "dot_product", + "sparse_matmul": False, + "dtype": "float32", + "weight_dtype": "float32", + "matmul_precision": "highest", + "decode_sampling_strategy": "greedy", + "enable_checkpointing": False, + "skip_jax_distributed_system": True, + "pure_nnx": True, +} + +_DEEPSEEK = { + "model_name": "deepseek2-16b", + "override_model_config": True, + "num_experts": 4, + "num_experts_per_tok": 2, + "shared_experts": 1, + # The naive MLA kv cache sizes its key cache with the value head dim, so prefill + # writes a wider key than the cache holds. Unrelated to decode consistency and it + # fails the same way on the Linen path. + "mla_naive_kvcache": False, +} + +# MoE and dense DeepSeek layers take different paths to attention, so cover both. +_DEEPSEEK_MOE = _DEEPSEEK | {"first_num_dense_layers": 1} +_DEEPSEEK_DENSE = _DEEPSEEK | {"first_num_dense_layers": 3} + + +def _make_config(**overrides): + return pyconfig.initialize([sys.argv[0], get_test_config_path()], **(_COMMON | overrides)) + + +def _mesh(cfg): + """Builds the mesh over every available device. + + ici_fsdp_parallelism keeps its base.yml default of -1, so the fsdp axis absorbs + however many devices the host has. Pinning the ICI axes instead would fix the mesh + at one device and fail everywhere except a single-device runner. + + Args: + cfg: Model config. + + Returns: + The device mesh. + """ + return jax.sharding.Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + + +class DecodeConsistencyTest(parameterized.TestCase): + """Decode must agree with the forward pass it is supposed to be replaying.""" + + def _forward_rollout(self, cfg, mesh, params_state): + """Rolls out greedily, recomputing a full forward pass over the prefix each step. + + Args: + cfg: Model config. + mesh: Device mesh the model is built on. + params_state: nnx.Param state to load into the model. + + Returns: + The generated token ids. + """ + with nn_partitioning.axis_rules(cfg.logical_axis_rules), mesh: + model = model_creation_utils.create_model( + cfg, mesh, model_mode=MODEL_MODE_TRAIN, rngs=nnx.Rngs(params=0, dropout=0) + ) + nnx.update(model, params_state) + + tokens = list(PROMPT) + generated = [] + pad = cfg.max_target_length + batch = cfg.micro_batch_size_to_train_on + for _ in range(STEPS + 1): + ids = jnp.tile(jnp.asarray([tokens + [0] * (pad - len(tokens))], dtype=jnp.int32), (batch, 1)) + positions = jnp.tile(jnp.asarray([list(range(pad))], dtype=jnp.int32), (batch, 1)) + segment_ids = jnp.tile(jnp.asarray([[1] * len(tokens) + [0] * (pad - len(tokens))], dtype=jnp.int32), (batch, 1)) + with nn_partitioning.axis_rules(cfg.logical_axis_rules), mesh: + logits = model( + ids, + positions, + decoder_segment_ids=segment_ids, + enable_dropout=False, + model_mode=MODEL_MODE_TRAIN, + ) + next_token = int(jnp.argmax(logits[0, len(tokens) - 1])) + generated.append(next_token) + tokens.append(next_token) + return generated + + def _engine_rollout(self, engine, params): + """Rolls out greedily through prefill, insert and generate. + + Args: + engine: MaxEngine with params already loaded. + params: Params returned by load_params. + + Returns: + The generated token ids. + """ + padded = jnp.asarray( + PROMPT + [0] * (engine.config.max_prefill_predict_length - len(PROMPT)), + dtype=jnp.int32, + ) + prefix, first_token = engine.prefill(params=params, padded_tokens=padded, true_length=len(PROMPT)) + generated = [int(first_token.data[0, 0])] + + decode_state = engine.init_decode_state() + decode_state = engine.insert(prefix, decode_state, slot=0) + for _ in range(STEPS): + decode_state, result = engine.generate(params, decode_state) + generated.append(int(result.data[0, 0])) + return generated + + def _build(self, cfg): + """Builds a freshly initialized model. + + Args: + cfg: Model config. + + Returns: + A tuple of (mesh, params_state). + """ + mesh = _mesh(cfg) + with nn_partitioning.axis_rules(cfg.logical_axis_rules), mesh: + model = model_creation_utils.create_model( + cfg, mesh, model_mode=MODEL_MODE_PREFILL, rngs=nnx.Rngs(params=0, dropout=0) + ) + _, params_state, _ = nnx.split(model, nnx.Param, ...) + return mesh, params_state + + @parameterized.named_parameters( + ("deepseek_moe", _DEEPSEEK_MOE), + ("deepseek_dense", _DEEPSEEK_DENSE), + ("generic", {}), + ) + def test_greedy_decode_matches_forward_pass(self, overrides): + cfg = _make_config(**overrides) + mesh, params_state = self._build(cfg) + + expected = self._forward_rollout(cfg, mesh, params_state) + + engine = maxengine.MaxEngine(cfg, jax.devices()) + params = engine.load_params(params=params_state) + actual = self._engine_rollout(engine, params) + + self.assertEqual( + expected, + actual, + f"decode diverged from the forward pass: expected {expected}, got {actual}", + ) + + +class EngineGraphdefModeTest(unittest.TestCase): + """MaxEngine must merge the graphdef built for the mode it is running. + + Layers are supposed to honor the call-time model_mode, but the engine should not + depend on that: merging a prefill graphdef for an autoregressive step leaves every + layer's `self.model_mode` reading "prefill". This guards the engine side on its own, + so the two defenses cannot regress together silently. + """ + + def test_generate_merges_autoregressive_graphdef(self): + cfg = _make_config(**_DEEPSEEK_MOE) + mesh = _mesh(cfg) + with nn_partitioning.axis_rules(cfg.logical_axis_rules), mesh: + model = model_creation_utils.create_model( + cfg, mesh, model_mode=MODEL_MODE_PREFILL, rngs=nnx.Rngs(params=0, dropout=0) + ) + _, params_state, _ = nnx.split(model, nnx.Param, ...) + + engine = maxengine.MaxEngine(cfg, jax.devices()) + params = engine.load_params(params=params_state) + + # The mode each merged layer was constructed with, recorded from inside the + # jitted prefill / generate bodies. + built_modes = [] + original = deepseek.DeepSeekGenericLayer.attention_op + + def recording_attention_op(self, *args, **kwargs): + built_modes.append(self.model_mode) + return original(self, *args, **kwargs) + + padded = jnp.asarray( + PROMPT + [0] * (cfg.max_prefill_predict_length - len(PROMPT)), + dtype=jnp.int32, + ) + with mock.patch.object(deepseek.DeepSeekGenericLayer, "attention_op", recording_attention_op): + prefix, _ = engine.prefill(params=params, padded_tokens=padded, true_length=len(PROMPT)) + self.assertTrue(built_modes, "prefill did not reach the DeepSeek attention") + self.assertEqual(set(built_modes), {MODEL_MODE_PREFILL}) + + decode_state = engine.init_decode_state() + decode_state = engine.insert(prefix, decode_state, slot=0) + built_modes.clear() + engine.generate(params, decode_state) + + self.assertTrue(built_modes, "generate did not reach the DeepSeek attention") + self.assertEqual( + set(built_modes), + {MODEL_MODE_AUTOREGRESSIVE}, + "generate merged a graphdef built for another mode", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/elastic_utils_test.py b/tests/unit/elastic_utils_test.py index 634ec2f992..9307c7dd30 100644 --- a/tests/unit/elastic_utils_test.py +++ b/tests/unit/elastic_utils_test.py @@ -15,16 +15,19 @@ """Unit tests for Elastic Training utility functions.""" import unittest -from unittest.mock import create_autospec, Mock +from unittest.mock import Mock, create_autospec from absl.testing import parameterized - - +from maxtext.common import checkpointing from maxtext.utils import elastic_utils from maxtext.utils import gcs_utils import pathwaysutils from pathwaysutils.elastic.manager import ScaleUpSignalError +class MockJaxRuntimeError(Exception): + """Fake JAX Runtime Error class for unit tests.""" + + class FakeDevice: """Fake Device object.""" @@ -77,6 +80,7 @@ def setUp(self): # Inject fakes into elastic_utils namespace elastic_utils.pathwaysutils = self.fake_pathwaysutils elastic_utils.jax = self.fake_jax + self.fake_jax.errors.JaxRuntimeError = MockJaxRuntimeError elastic_utils.gcs_utils = self.fake_gcs_utils elastic_utils.max_logging = self.fake_logging @@ -467,6 +471,80 @@ def test_is_scale_up_event_with_set(self, available_inactive_slices, expected): self.fake_manager.available_inactive_slices = available_inactive_slices self.assertEqual(elastic_utils.is_scale_up_event(config), expected) + def test_maybe_bubble_elastic_exception_bubbles_on_elastic_errors(self): + """Tests that elastic exceptions are bubbled up, while other exceptions are returned normally.""" + config = FakeConfig() + config.elastic_enabled = True + self.fake_pathwaysutils.is_pathways_backend_used.return_value = True + + # Scenario 1: Elastic JAX error propagates + with self.assertRaises(MockJaxRuntimeError): + elastic_utils.maybe_bubble_elastic_exception(config, MockJaxRuntimeError("TPU offline")) + + # Scenario 2: ScaleUpSignalError propagates + with self.assertRaises(ScaleUpSignalError): + elastic_utils.maybe_bubble_elastic_exception(config, ScaleUpSignalError()) + + # Scenario 3: Non-elastic error is returned/ignored + elastic_utils.maybe_bubble_elastic_exception(config, ValueError("Disk full")) + + def test_maybe_bubble_elastic_exception_disabled_does_not_bubble(self): + """If elasticity is disabled, JaxRuntimeError should not bubble.""" + config = FakeConfig() + config.elastic_enabled = False # Disabled + self.fake_pathwaysutils.is_pathways_backend_used.return_value = True + + # Executes normally (no exception raised) + elastic_utils.maybe_bubble_elastic_exception(config, MockJaxRuntimeError("JAX error but elasticity disabled")) + + def test_checkpoint_exception_guard_checks_scale_up_on_success(self): + """Signals ScaleUpSignalError if scale-up is active when save completes.""" + config = FakeConfig() + config.elastic_enabled = True + elastic_utils.elastic_manager = self.fake_manager + self.fake_pathwaysutils.is_pathways_backend_used.return_value = True + self.fake_manager.available_inactive_slices = {1} # Trigger scale-up + + mock_checkpoint_manager = Mock(spec=["wait_until_finished"]) + + # Successful checkpoint save block raises ScaleUpSignalError to trigger restart + with self.assertRaises(ScaleUpSignalError): + with checkpointing.checkpoint_exception_guard(config, mock_checkpoint_manager): + pass + + mock_checkpoint_manager.wait_until_finished.assert_called_once() + + def test_checkpoint_exception_guard_none_manager(self): + """Checks that checkpoint_manager=None doesn't raise AttributeError on scale-up.""" + config = FakeConfig() + config.elastic_enabled = True + elastic_utils.elastic_manager = self.fake_manager + self.fake_pathwaysutils.is_pathways_backend_used.return_value = True + self.fake_manager.available_inactive_slices = {1} # Trigger scale-up + + with self.assertRaises(ScaleUpSignalError): + with checkpointing.checkpoint_exception_guard(config, checkpoint_manager=None): + pass + + def test_checkpoint_exception_guard_skips_scale_up_on_failure(self): + """If checkpoint save fails, scale-up check should be skipped, and exception handled.""" + config = FakeConfig() + config.elastic_enabled = True + elastic_utils.elastic_manager = self.fake_manager + self.fake_pathwaysutils.is_pathways_backend_used.return_value = True + self.fake_manager.available_inactive_slices = {1} + + handler_called = False + + def handler(_err): + nonlocal handler_called + handler_called = True + + with checkpointing.checkpoint_exception_guard(config, self.fake_manager, handler): + raise ValueError("Save failed") + + self.assertTrue(handler_called) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/pyconfig_test.py b/tests/unit/pyconfig_test.py index d3afcb612e..9b7862bc5a 100644 --- a/tests/unit/pyconfig_test.py +++ b/tests/unit/pyconfig_test.py @@ -332,6 +332,15 @@ def test_eval_start_step_config(self): ) self.assertEqual(config_override.eval_start_step, 50) + def test_eval_start_step_negative_raises_error(self): + """Verifies that eval_start_step < 0 raises a validation error.""" + with self.assertRaises((ValueError, Exception)): + pyconfig.initialize( + [os.path.join(MAXTEXT_PKG_DIR, "train.py"), get_test_config_path()], + skip_jax_distributed_system=True, + eval_start_step=-1, + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/train_state_nnx_checkpoint_test.py b/tests/unit/train_state_nnx_checkpoint_test.py index ef19ac427e..46ee0b73e9 100644 --- a/tests/unit/train_state_nnx_checkpoint_test.py +++ b/tests/unit/train_state_nnx_checkpoint_test.py @@ -434,13 +434,31 @@ def test_nnx_and_linen_agree_on_actual_step(self): ) def test_nnx_state_is_saved_in_linen_layout(self): - """For pure_nnx=True, maybe_save_checkpoint reshapes the NNX state to the Linen on-disk layout.""" + """For pure_nnx=True, save_checkpoint reshapes the NNX state to the Linen on-disk layout.""" state = self._build_nnx_state(self.N_STEPS) self.assertIsInstance(state, nnx.State) # precondition: NNX train_step returns an nnx.State - captured = self._invoke_maybe_save(state, pure_nnx=True) + config = self._config(pure_nnx=True, enable_checkpointing=True, checkpoint_period=1) + mgr = mock.MagicMock() + mgr.reached_preemption.return_value = False + + captured = {} - # save_checkpoint should receive a plain dict in Linen layout, not the nnx.State. + def fake_save(_step, *args, **kwargs): + composite = kwargs.get("args") + if composite: + for key in ["items", "state"]: + if hasattr(composite, "_items") and key in composite._items: # pylint: disable=protected-access + val = composite[key] + if val is not None and hasattr(val, "item"): + captured["state"] = val.item + break + return True + + mgr.save.side_effect = fake_save + checkpointing.save_checkpoint(mgr, self.N_STEPS - 1, state, config=config, force=True) + + # save_checkpoint should pass a plain dict in Linen layout to Orbax, not the nnx.State. self.assertIsInstance(captured["state"], dict) self.assertNotIsInstance(captured["state"], nnx.State) # Linen layout: {params: {params: ...}, step, opt_state}; not the NNX {model, optimizer}. @@ -558,16 +576,12 @@ def test_maybe_save_checkpoint_allows_local_checkpoint_period(self): mgr.reached_preemption.return_value = False save_checkpoint_mock = mock.MagicMock(return_value=False) - with ( - mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock), - mock.patch.object(train_state_nnx, "to_checkpoint_dict") as to_checkpoint_dict_mock, - ): + with mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock): checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5) mgr.latest_step.assert_called_once_with() mgr.reached_preemption.assert_called_once_with(5) mgr.wait_until_finished.assert_not_called() - to_checkpoint_dict_mock.assert_called_once_with(state) save_checkpoint_mock.assert_called_once() def test_maybe_save_checkpoint_allows_mtc_period_with_continuous_policy( @@ -587,16 +601,12 @@ def test_maybe_save_checkpoint_allows_mtc_period_with_continuous_policy( mgr.reached_preemption.return_value = False save_checkpoint_mock = mock.MagicMock(return_value=False) - with ( - mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock), - mock.patch.object(train_state_nnx, "to_checkpoint_dict") as to_checkpoint_dict_mock, - ): + with mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock): checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5) mgr.should_save.assert_called_once_with(5) mgr.latest_step.assert_called_once_with() mgr.reached_preemption.assert_called_once_with(5) - to_checkpoint_dict_mock.assert_called_once_with(state) save_checkpoint_mock.assert_called_once() def test_maybe_save_checkpoint_checks_scale_up_after_unsaved_dispatch(self): @@ -610,12 +620,10 @@ def test_maybe_save_checkpoint_checks_scale_up_after_unsaved_dispatch(self): with ( mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock), - mock.patch.object(train_state_nnx, "to_checkpoint_dict", return_value={}) as to_checkpoint_dict_mock, mock.patch.object(checkpointing.elastic_utils, "maybe_elastic_scale_up") as mock_maybe_scale_up, ): checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5) - to_checkpoint_dict_mock.assert_called_once_with(state) save_checkpoint_mock.assert_called_once() mock_maybe_scale_up.assert_called_once_with(config, mgr)