diff --git a/src/dependencies/extra_deps/cuda12_overrides.txt b/src/dependencies/extra_deps/cuda12_overrides.txt index 647a1ca427..5f837af0a5 100644 --- a/src/dependencies/extra_deps/cuda12_overrides.txt +++ b/src/dependencies/extra_deps/cuda12_overrides.txt @@ -1,3 +1,4 @@ datasets>=4.8.5 fsspec==2026.2.0 gcsfs==2026.2.0 +orbax-checkpoint>=0.12.1 diff --git a/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt b/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt index 266e7b3ff9..4b2ab63041 100644 --- a/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt +++ b/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt @@ -30,7 +30,7 @@ cloudpickle>=3.1.2 clu>=0.0.12 colorama>=0.4.6 contourpy>=1.3.3 -coverage>=7.14.3 +coverage>=7.15.0 cryptography>=49.0.0 cycler>=0.12.1 datasets>=5.0.0 @@ -84,11 +84,11 @@ grpcio-status>=1.80.0 gviz-api>=1.10.0 h11>=0.16.0 h5py>=3.14.0 -hf-xet>=1.5.1 ; platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64' +hf-xet>=1.5.1 ; platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'x86_64' httpcore>=1.0.9 httplib2>=0.32.0 httpx>=0.28.1 -huggingface-hub>=1.21.0 +huggingface-hub>=1.22.0 humanize>=4.16.0 hypothesis>=6.142.1 identify>=2.6.19 @@ -158,7 +158,7 @@ opt-einsum>=3.4.0 optax>=0.2.8 optree>=0.19.0 optype>=0.17.0 -orbax-checkpoint>=0.11.39 +orbax-checkpoint>=0.12.1 packaging>=26.1 pandas>=3.0.3 parameterized>=0.9.0 @@ -169,6 +169,7 @@ platformdirs>=4.10.0 pluggy>=1.6.0 portpicker>=1.6.0 pre-commit>=4.6.0 +prometheus-client>=0.20.0 promise>=2.3 propcache>=0.5.2 proto-plus>=1.28.0 @@ -195,7 +196,7 @@ pyproject-hooks>=1.2.0 pytest>=8.4.2 pytest-xdist>=3.8.0 python-dateutil>=2.9.0.post0 -python-discovery>=1.4.2 +python-discovery>=1.4.3 pytokens>=0.4.1 pytype>=2024.10.11 pyyaml>=6.0.3 @@ -229,7 +230,7 @@ tensorflow>=2.20.0 tensorflow-datasets>=4.9.10 tensorflow-metadata>=1.21.0 tensorflow-text>=2.20.1 -tensorstore>=0.1.82 +tensorstore>=0.1.84 termcolor>=3.3.0 tiktoken>=0.13.0 tokamax>=0.0.12 @@ -241,24 +242,24 @@ tqdm>=4.68.3 transformer-engine>=2.16.1 transformer-engine-cu12>=2.16.1 transformer-engine-jax>=2.16.1 -transformers>=5.12.1 +transformers>=5.13.0 treescope>=0.1.10 typeguard>=2.13.3 -typer>=0.25.1 +typer>=0.26.8 typing-extensions>=4.15.0 typing-inspection>=0.4.2 tzdata>=2026.2 ; sys_platform == 'emscripten' or sys_platform == 'win32' uritemplate>=4.2.0 urllib3>=2.6.3 -uvicorn>=0.49.0 +uvicorn>=0.50.2 uvloop>=0.22.1 -virtualenv>=21.5.1 +virtualenv>=21.5.2 wadler-lindig>=0.1.7 websockets>=16.0 werkzeug>=3.1.8 wheel>=0.46.3 wrapt>=2.2.2 -xxhash>=3.8.0 +xxhash>=3.8.1 yarl>=1.24.2 zipp>=3.23.1 zstandard>=0.25.0 diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index c527427144..4a7b3cf35e 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -15,28 +15,27 @@ """Create an Orbax CheckpointManager with specified (Async or not) Checkpointer.""" +import datetime import time -from typing import Any, Optional +from typing import Any from absl import flags -import datetime from etils import epath from flax import nnx from flax.training import train_state +from grain.experimental import ElasticIterator import jax -import jax.numpy as jnp -from maxtext.utils.globals import DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE +from maxtext.checkpoint_conversion.utils.load_dynamic import load_safetensors_dynamic_state +from maxtext.common import grain_utility +from maxtext.common import train_state_nnx from maxtext.input_pipeline.multihost_dataloading import MultiHostDataLoadIterator from maxtext.input_pipeline.multihost_dataloading import RemoteIteratorWrapper from maxtext.input_pipeline.synthetic_data_processing import PlaceHolderDataIterator -from maxtext.common import train_state_nnx +from maxtext.utils import elastic_utils from maxtext.utils import exceptions -from maxtext.utils import max_logging from maxtext.utils import gcs_utils -from maxtext.utils import elastic_utils -from maxtext.checkpoint_conversion.utils.load_dynamic import load_safetensors_dynamic_state - -import numpy as np +from maxtext.utils import max_logging +from maxtext.utils.globals import DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE import orbax.checkpoint as ocp from orbax.checkpoint import v1 as ocp_v1 from orbax.checkpoint._src.arrays import sharding as sharding_utils @@ -44,13 +43,7 @@ from orbax.checkpoint._src.checkpoint_managers import save_decision_policy as save_decision_policy_lib import orbax.checkpoint.experimental.emergency.checkpoint_manager as emergency_checkpoint_manager import orbax.checkpoint.experimental.emergency.replicator_checkpoint_manager as emergency_replicator_checkpoint_manager -# pylint: disable=too-many-positional-arguments -import dataclasses -import json -import grain -from grain.python import PyGrainCheckpointHandler -from grain.experimental import ElasticIterator CheckpointManager = ocp.CheckpointManager CheckpointManagerOptions = ocp.CheckpointManagerOptions @@ -62,158 +55,6 @@ EmergencyReplicatorCheckpointManager = emergency_replicator_checkpoint_manager.ReplicatorCheckpointManager -class GrainCheckpointHandler(PyGrainCheckpointHandler, ocp.CheckpointHandler): - """A CheckpointHandler that allows specifying process_index and process_count.""" - - def save( - self, - directory: epath.Path, - # `item` is for backwards compatibility with older Orbax API, see - # https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html. - item: Optional[Any] = None, - args: Any = None, - ): - """Saves the given iterator to the checkpoint in `directory`.""" - item = item or args.item # pytype:disable=attribute-error - - # RemoteIteratorWrapper handles checkpointing via colocated python - if isinstance(item, RemoteIteratorWrapper): - step = int(directory.parent.name) - item.save_state(step) - return - - # ElasticIterator state is a single global scalar shared by all shards, - # so we write one fixed `process_0.json` from process 0 only. This file - # layout survives changes in `jax.process_count()`. - if isinstance(item, ElasticIterator): - if jax.process_index() == 0: - directory.mkdir(parents=True, exist_ok=True) - filename = directory / "process_0.json" - filename.write_text(json.dumps(item.get_state(), indent=4)) - return - - def save_single_process(item, process_index, process_count): - filename = directory / f"process_{process_index}-of-{process_count}.json" - if isinstance(item, grain.DatasetIterator): - state = json.dumps(item.get_state(), indent=4) - else: - state = item.get_state().decode() - filename.write_text(state) - - if isinstance(item, list): - for local_iterator, process_index, process_count in item: - save_single_process(local_iterator, process_index, process_count) - else: - process_index, process_count = jax.process_index(), jax.process_count() - save_single_process(item, process_index, process_count) - - def restore( - self, - directory: epath.Path, - item: Optional[Any] = None, - args: Any = None, - ) -> Any: - """Restores the given iterator from the checkpoint in `directory`.""" - item = item or args.item - process_index = getattr(args, "process_index", None) - process_count = getattr(args, "process_count", None) - - # In Pathways + colocated_python environment, RemoteIteratorWrapper handles checkpointing - if isinstance(item, RemoteIteratorWrapper): - step = int(directory.parent.name) - item.restore_state(step) - return item - - # McJax and Pathways through controller cases - # ElasticIterator: every process reads the same shared `process_0.json`. - if isinstance(item, ElasticIterator): - filename = directory / "process_0.json" - if not filename.exists(): - raise ValueError(f"File {filename} does not exist.") - item.set_state(json.loads(filename.read_text())) - return item - - def restore_single_process(item, process_index, process_count): - filename = directory / f"process_{process_index}-of-{process_count}.json" - if not filename.exists(): - raise ValueError(f"File {filename} does not exist.") - state = filename.read_text() - if isinstance(item, grain.DatasetIterator): - state = json.loads(state) - else: - state = state.encode() - item.set_state(state) - return item - - if isinstance(item, list): - restored_items = [] - for data_iter, process_idx in zip(item, process_index): - restored_items.append(restore_single_process(data_iter, process_idx, process_count)) - return restored_items - else: - if process_index is None or process_count is None: - process_index, process_count = jax.process_index(), jax.process_count() - return restore_single_process(item, process_index, process_count) - - -@ocp.args.register_with_handler(GrainCheckpointHandler, for_save=True) -@dataclasses.dataclass -class GrainCheckpointSave(ocp.args.CheckpointArgs): - item: Any - - -@ocp.args.register_with_handler(GrainCheckpointHandler, for_restore=True) -@dataclasses.dataclass -class GrainCheckpointRestore(ocp.args.CheckpointArgs): - item: Any - process_index: Optional[int | list[int]] = None - process_count: Optional[int] = None - - -def _default_for_sds(sds): - """Returns a deterministic value matching `sds` shape/dtype/sharding. - - Used to fill NNX-only state (rngs/dropout) that the Linen on-disk layout never - carried. Materializes under jit with the target out_shardings so it works on - multi-host meshes (device_put can't place a global sharding whose devices - aren't - all addressable from this process). - """ - if not (hasattr(sds, "dtype") and hasattr(sds, "shape")): - return sds - - def _make(): - if "key" in str(sds.dtype): - base = jax.random.key(0) - return base if sds.shape == () else jax.random.split(base, int(np.prod(sds.shape))).reshape(sds.shape) - return jnp.zeros(sds.shape, dtype=sds.dtype) - - sharding = getattr(sds, "sharding", None) - if sharding is None: - return _make() - return jax.jit(_make, out_shardings=sharding)() - - -def _populate_pure_dict_from_partial(abstract_pure, partial_concrete): - """Fills `abstract_pure` with values from `partial_concrete` (by path), defaulting the rest. - - Paths present in `partial_concrete` take the restored value; paths absent from - it - (NNX-only state the Linen checkpoint never had) get `_default_for_sds`. - """ - if isinstance(abstract_pure, dict): - return { - k: _populate_pure_dict_from_partial( - v, - partial_concrete.get(k) if isinstance(partial_concrete, dict) else None, - ) - for k, v in abstract_pure.items() - } - if partial_concrete is not None and not isinstance(partial_concrete, dict): - return partial_concrete - return _default_for_sds(abstract_pure) - - def _load_linen_checkpoint_into_nnx( path, abstract_nnx_state, @@ -241,7 +82,7 @@ def _load_linen_checkpoint_into_nnx( restored = ocp.args.PyTreeRestore(item=linen_abstract, restore_args=restore_args, partial_restore=True) restored = ckptr.restore(epath.Path(path), args=restored) partial_nnx = train_state_nnx.from_linen_checkpoint_dict(restored) - return _populate_pure_dict_from_partial(nnx_abstract_pure, partial_nnx) + return train_state_nnx.populate_pure_dict_from_partial(nnx_abstract_pure, partial_nnx) def _restore_emergency_linen_checkpoint_into_nnx( @@ -262,19 +103,7 @@ def _restore_emergency_linen_checkpoint_into_nnx( ) restored = checkpoint_manager.restore(step, args=Composite(state=checkpoint_args)).state partial_nnx = train_state_nnx.from_linen_checkpoint_dict(restored) - return _populate_pure_dict_from_partial(nnx_abstract_pure, partial_nnx) - - -def _rebuild_nnx_with_values(abstract_nnx_state, concrete_weights): - """Fills each Variable in `abstract_nnx_state` with the matching restored array.""" - leaves, treedef = jax.tree_util.tree_flatten(abstract_nnx_state, is_leaf=lambda x: isinstance(x, nnx.Variable)) - concrete = jax.tree_util.tree_leaves(concrete_weights) - if len(leaves) != len(concrete): - raise ValueError( - f"Params load leaf-count mismatch: {len(leaves)} abstract Variables vs" f" {len(concrete)} restored." - ) - new_leaves = [v.replace(value=a) if isinstance(v, nnx.Variable) else a for v, a in zip(leaves, concrete)] - return jax.tree_util.tree_unflatten(treedef, new_leaves) + return train_state_nnx.populate_pure_dict_from_partial(nnx_abstract_pure, partial_nnx) def _load_linen_params_into_nnx( @@ -306,7 +135,7 @@ def _load_linen_params_into_nnx( epath.Path(path), args=ocp.args.PyTreeRestore(item=linen_abstract, restore_args=restore_args, partial_restore=True), ) - return _rebuild_nnx_with_values(nnx_params_abstract, restored["params"]["params"]) + return train_state_nnx.rebuild_nnx_with_values(nnx_params_abstract, restored["params"]["params"]) def _load_full_state_from_path( @@ -431,7 +260,7 @@ def create_orbax_checkpoint_manager( if dataset_type is not None and dataset_type == "grain": item_names += ("iter",) - item_handlers["iter"] = GrainCheckpointHandler() + item_handlers["iter"] = grain_utility.GrainCheckpointHandler() # local storage checkpoint needs parent directory created p = gcs_utils.mkdir_and_check_permissions(checkpoint_dir) @@ -586,137 +415,6 @@ def print_save_message(step, async_checkpointing): max_logging.log(f"Saved a checkpoint at step {step}.") -def _find_idx(array: np.ndarray, replica_axis_idx: int): - """Returns the index along given dimension that the current host belongs to.""" - idx = None - for idx, val in np.ndenumerate(array): - if val.process_index == jax.process_index(): - break - return idx[replica_axis_idx] - - -def _replica_devices(device_array: np.ndarray, replica_axis_idx: int): - """Returns the devices from the replica that current host belongs to. - - Replicas are assumed to be restricted to the first axis. - - Args: - device_array: devices of the mesh that can be obtained by mesh.devices() - replica_axis_idx: axis dimension along which replica is taken - - Returns: - devices inside the replica that current host is in - """ - idx = _find_idx(device_array, replica_axis_idx) - replica_result = np.take(device_array, idx, axis=replica_axis_idx) - return np.expand_dims(replica_result, axis=replica_axis_idx) - - -def _prepare_scaled_down_grain_restore_args( - data_iterator: list, process_count_jax: int, process_count_stored: int, directory: epath.Path -) -> GrainCheckpointRestore: - """ - Prepares the restore arguments for a scaled-up (list) data iterator. - - This is used when restoring a checkpoint saved with more processes than - the current run (e.g., 64 files onto 32 JAX processes). - """ - # 1. Validation Assertions - assert isinstance(data_iterator, list), ( - f"{process_count_stored} processes found in Grain checkpoint directory {directory}, but only " - f"{process_count_jax} jax processes in this run, please set expansion_factor_real_data accordingly." - ) - - scaling_factor = len(data_iterator) - expected_process_count = process_count_stored / process_count_jax - assert scaling_factor == expected_process_count, ( - f"Found {process_count_stored} processes in checkpoint and {process_count_jax} " - f"JAX processes, implying a scaling factor of {expected_process_count}. " - f"However, the data_iterator list has {scaling_factor} items." - ) - - # 2. Prepare Arguments - local_iterator_list = [x.local_iterator for x in data_iterator] - # Each JAX process calculates the global indices it's responsible for. - # e.g., process 0 with scaling_factor=2 handles checkpoints from processes [0, 32] - # e.g., process 1 with scaling_factor=2 handles checkpoints from processes [1, 33] - process_index_list = [jax.process_index() + i * process_count_jax for i in range(scaling_factor)] - - return GrainCheckpointRestore(local_iterator_list, process_index=process_index_list, process_count=process_count_stored) - - -def _restore_grain_iterator( - checkpoint_manager, - step: int, - data_iterator, - checkpoint_args, - expansion_factor_real_data: int, # This must be defined in the outer scope -) -> tuple[Any, None]: - """ - Handles the complex logic for restoring a Grain data iterator checkpoint. - This function dispatches to the correct restore strategy based on - the number of stored checkpoint files vs. current JAX processes. - """ - if isinstance(data_iterator, RemoteIteratorWrapper): - grain_restore_args = GrainCheckpointRestore(item=data_iterator) - restored_state = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args, iter=grain_restore_args)) - return (restored_state, None) - - # ElasticIterator: one shared `process_0.json` regardless of shard count. - if not isinstance(data_iterator, list) and isinstance(data_iterator.local_iterator, ElasticIterator): - grain_restore_args = GrainCheckpointRestore(item=data_iterator.local_iterator) - restored_state = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args, iter=grain_restore_args)) - return (restored_state, None) - - directory = checkpoint_manager.directory / str(step) / "iter" - process_count_jax = jax.process_count() - - # Count the number of checkpoint files - process_count_stored = len(list(directory.glob("process_*-of-*.json"))) - - grain_restore_args = None - - if process_count_stored > process_count_jax: - # Scaling down from a larger number of hosts. (e.g., 128 files -> 64 processes) - # In this case, each host restores a list of data iterators. - grain_restore_args = _prepare_scaled_down_grain_restore_args( - data_iterator, process_count_jax, process_count_stored, directory - ) - - elif process_count_stored == process_count_jax: - # Normal case: number of hosts is the same. (e.g., 64 files -> 64 processes) - assert not isinstance(data_iterator, list), ( - f"{process_count_stored} processes found in Grain checkpoint directory {directory}, matching the number of " - "jax process, please do not set expansion_factor_real_data." - ) - grain_restore_args = GrainCheckpointRestore(data_iterator.local_iterator) - - elif expansion_factor_real_data > 1 and process_count_stored == process_count_jax // expansion_factor_real_data: - # Scaling up to a larger number of hosts.(e.g., 32 files -> 64 processes) - # In this case, a subset of hosts restore the data iterator. - assert not isinstance( - data_iterator, list - ), "when expansion_factor_real_data > 1, the data iterator should not be a list." - grain_restore_args = GrainCheckpointRestore( - data_iterator.local_iterator, process_index=jax.process_index(), process_count=process_count_stored - ) - - else: - # Case 4: Mismatch - raise ValueError( - f"Error restoring Grain checkpoint in {directory}: " - f"The number of stored checkpoint files ({process_count_stored}) " - f"is incompatible with the number of JAX processes ({process_count_jax}). " - "If you are resuming training with a different number of chips, see instructions in " - "https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/data_input_pipeline/" - "data_input_grain.md#using-grain" - ) - - # Call restore once with the composed arguments - restored_state = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args, iter=grain_restore_args)) - return (restored_state, None) - - def load_state_if_possible( checkpoint_manager: CheckpointManager | None, data_iterator: MultiHostDataLoadIterator | list[MultiHostDataLoadIterator] | None, @@ -774,7 +472,7 @@ def map_to_pspec(data): pspec = data.sharding.spec mesh = data.sharding.mesh replica_axis_index = 0 - replica_devices = _replica_devices(mesh.devices, replica_axis_index) + replica_devices = grain_utility.replica_devices(mesh.devices, replica_axis_index) replica_mesh = jax.sharding.Mesh(replica_devices, mesh.axis_names) single_replica_sharding = jax.sharding.NamedSharding(replica_mesh, pspec) @@ -865,7 +563,7 @@ def map_to_pspec(data): and not isinstance(data_iterator, PlaceHolderDataIterator) and (checkpoint_manager.directory / str(step) / "iter").exists() ): - return _restore_grain_iterator( + return grain_utility.restore_grain_iterator( checkpoint_manager, step, data_iterator, @@ -1113,10 +811,10 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator= if config and config.dataset_type == "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"] = GrainCheckpointSave(item=data_iterator) + save_args_composite["iter"] = grain_utility.GrainCheckpointSave(item=data_iterator) elif not isinstance(data_iterator, list) and isinstance(data_iterator.local_iterator, ElasticIterator): # ElasticIterator checkpoints a single global scalar shared by all shards. - save_args_composite["iter"] = GrainCheckpointSave(item=data_iterator.local_iterator) + save_args_composite["iter"] = grain_utility.GrainCheckpointSave(item=data_iterator.local_iterator) else: if not isinstance(data_iterator, list): data_iterator = [data_iterator] @@ -1127,7 +825,7 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator= for i, data_iter in enumerate(data_iterator): process_index = jax.process_index() + i * jax.process_count() grain_iters_to_save.append((data_iter.local_iterator, process_index, process_count_total)) - save_args_composite["iter"] = GrainCheckpointSave(item=grain_iters_to_save) + save_args_composite["iter"] = grain_utility.GrainCheckpointSave(item=grain_iters_to_save) custom_metadata = {} if config: diff --git a/src/maxtext/common/grain_utility.py b/src/maxtext/common/grain_utility.py new file mode 100644 index 0000000000..1d8d7caf3f --- /dev/null +++ b/src/maxtext/common/grain_utility.py @@ -0,0 +1,273 @@ +# 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 +# +# 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. + +"""Grain utility functions for checkpointing.""" + +import dataclasses +import json +from typing import Any, Optional + +from etils import epath +import grain +from grain import experimental as grain_experimental +from grain import python +import jax +from maxtext.input_pipeline import multihost_dataloading +import numpy as np +import orbax.checkpoint as ocp + +PyGrainCheckpointHandler = python.PyGrainCheckpointHandler +Composite = ocp.args.Composite +RemoteIteratorWrapper = multihost_dataloading.RemoteIteratorWrapper + +ElasticIterator = grain_experimental.ElasticIterator + + +class GrainCheckpointHandler(PyGrainCheckpointHandler, ocp.CheckpointHandler): + """A CheckpointHandler that allows specifying process_index and process_count.""" + + def save( + self, + directory: epath.Path, + # `item` is for backwards compatibility with older Orbax API, see + # https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html. + item: Optional[Any] = None, + args: Any = None, + ): + """Saves the given iterator to the checkpoint in `directory`.""" + item = item or args.item # pytype:disable=attribute-error + + # RemoteIteratorWrapper handles checkpointing via colocated python + if isinstance(item, RemoteIteratorWrapper): + step = int(directory.parent.name) + item.save_state(step) + return + + # ElasticIterator state is a single global scalar shared by all shards, + # so we write one fixed `process_0.json` from process 0 only. This file + # layout survives changes in `jax.process_count()`. + if isinstance(item, ElasticIterator): + if jax.process_index() == 0: + directory.mkdir(parents=True, exist_ok=True) + filename = directory / "process_0.json" + filename.write_text(json.dumps(item.get_state(), indent=4)) + return + + def save_single_process(item, process_index, process_count): + filename = directory / f"process_{process_index}-of-{process_count}.json" + if isinstance(item, grain.DatasetIterator): + state = json.dumps(item.get_state(), indent=4) + else: + state = item.get_state().decode() + filename.write_text(state) + + if isinstance(item, list): + for local_iterator, process_index, process_count in item: + save_single_process(local_iterator, process_index, process_count) + else: + process_index, process_count = jax.process_index(), jax.process_count() + save_single_process(item, process_index, process_count) + + def restore( + self, + directory: epath.Path, + item: Optional[Any] = None, + args: Any = None, + ) -> Any: + """Restores the given iterator from the checkpoint in `directory`.""" + item = item or args.item + process_index = getattr(args, "process_index", None) + process_count = getattr(args, "process_count", None) + + # In Pathways + colocated_python environment, RemoteIteratorWrapper handles checkpointing + if isinstance(item, RemoteIteratorWrapper): + step = int(directory.parent.name) + item.restore_state(step) + return item + + # McJax and Pathways through controller cases + # ElasticIterator: every process reads the same shared `process_0.json`. + if isinstance(item, ElasticIterator): + filename = directory / "process_0.json" + if not filename.exists(): + raise ValueError(f"File {filename} does not exist.") + item.set_state(json.loads(filename.read_text())) + return item + + def restore_single_process(item, process_index, process_count): + filename = directory / f"process_{process_index}-of-{process_count}.json" + if not filename.exists(): + raise ValueError(f"File {filename} does not exist.") + state = filename.read_text() + if isinstance(item, grain.DatasetIterator): + state = json.loads(state) + else: + state = state.encode() + item.set_state(state) + return item + + if isinstance(item, list): + restored_items = [] + for data_iter, process_idx in zip(item, process_index): + restored_items.append(restore_single_process(data_iter, process_idx, process_count)) + return restored_items + else: + if process_index is None or process_count is None: + process_index, process_count = jax.process_index(), jax.process_count() + return restore_single_process(item, process_index, process_count) + + +@ocp.args.register_with_handler(GrainCheckpointHandler, for_save=True) +@dataclasses.dataclass +class GrainCheckpointSave(ocp.args.CheckpointArgs): + item: Any + + +@ocp.args.register_with_handler(GrainCheckpointHandler, for_restore=True) +@dataclasses.dataclass +class GrainCheckpointRestore(ocp.args.CheckpointArgs): + item: Any + process_index: Optional[int | list[int]] = None + process_count: Optional[int] = None + + +def find_idx(array: np.ndarray, replica_axis_idx: int): + """Returns the index along given dimension that the current host belongs to.""" + idx = None + for idx, val in np.ndenumerate(array): + if val.process_index == jax.process_index(): + break + return idx[replica_axis_idx] + + +def replica_devices(device_array: np.ndarray, replica_axis_idx: int): + """Returns the devices from the replica that current host belongs to. + + Replicas are assumed to be restricted to the first axis. + + Args: + device_array: devices of the mesh that can be obtained by mesh.devices() + replica_axis_idx: axis dimension along which replica is taken + + Returns: + devices inside the replica that current host is in + """ + idx = find_idx(device_array, replica_axis_idx) + replica_result = np.take(device_array, idx, axis=replica_axis_idx) + return np.expand_dims(replica_result, axis=replica_axis_idx) + + +def prepare_scaled_down_grain_restore_args( + data_iterator: list, process_count_jax: int, process_count_stored: int, directory: epath.Path +) -> GrainCheckpointRestore: + """ + Prepares the restore arguments for a scaled-up (list) data iterator. + + This is used when restoring a checkpoint saved with more processes than + the current run (e.g., 64 files onto 32 JAX processes). + """ + # 1. Validation Assertions + assert isinstance(data_iterator, list), ( + f"{process_count_stored} processes found in Grain checkpoint directory {directory}, but only " + f"{process_count_jax} jax processes in this run, please set expansion_factor_real_data accordingly." + ) + + scaling_factor = len(data_iterator) + expected_process_count = process_count_stored / process_count_jax + assert scaling_factor == expected_process_count, ( + f"Found {process_count_stored} processes in checkpoint and {process_count_jax} " + f"JAX processes, implying a scaling factor of {expected_process_count}. " + f"However, the data_iterator list has {scaling_factor} items." + ) + + # 2. Prepare Arguments + local_iterator_list = [x.local_iterator for x in data_iterator] + # Each JAX process calculates the global indices it's responsible for. + # e.g., process 0 with scaling_factor=2 handles checkpoints from processes [0, 32] + # e.g., process 1 with scaling_factor=2 handles checkpoints from processes [1, 33] + process_index_list = [jax.process_index() + i * process_count_jax for i in range(scaling_factor)] + + return GrainCheckpointRestore(local_iterator_list, process_index=process_index_list, process_count=process_count_stored) + + +def restore_grain_iterator( + checkpoint_manager, + step: int, + data_iterator, + checkpoint_args, + expansion_factor_real_data: int, # This must be defined in the outer scope +) -> tuple[Any, None]: + """ + Handles the complex logic for restoring a Grain data iterator checkpoint. + This function dispatches to the correct restore strategy based on + the number of stored checkpoint files vs. current JAX processes. + """ + if isinstance(data_iterator, RemoteIteratorWrapper): + grain_restore_args = GrainCheckpointRestore(item=data_iterator) + restored_state = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args, iter=grain_restore_args)) + return (restored_state, None) + + # ElasticIterator: one shared `process_0.json` regardless of shard count. + if not isinstance(data_iterator, list) and isinstance(data_iterator.local_iterator, ElasticIterator): + grain_restore_args = GrainCheckpointRestore(item=data_iterator.local_iterator) + restored_state = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args, iter=grain_restore_args)) + return (restored_state, None) + + directory = checkpoint_manager.directory / str(step) / "iter" + process_count_jax = jax.process_count() + + # Count the number of checkpoint files + process_count_stored = len(list(directory.glob("process_*-of-*.json"))) + + grain_restore_args = None + + if process_count_stored > process_count_jax: + # Scaling down from a larger number of hosts. (e.g., 128 files -> 64 processes) + # In this case, each host restores a list of data iterators. + grain_restore_args = prepare_scaled_down_grain_restore_args( + data_iterator, process_count_jax, process_count_stored, directory + ) + + elif process_count_stored == process_count_jax: + # Normal case: number of hosts is the same. (e.g., 64 files -> 64 processes) + assert not isinstance(data_iterator, list), ( + f"{process_count_stored} processes found in Grain checkpoint directory {directory}, matching the number of " + "jax process, please do not set expansion_factor_real_data." + ) + grain_restore_args = GrainCheckpointRestore(data_iterator.local_iterator) + + elif expansion_factor_real_data > 1 and process_count_stored == process_count_jax // expansion_factor_real_data: + # Scaling up to a larger number of hosts.(e.g., 32 files -> 64 processes) + # In this case, a subset of hosts restore the data iterator. + assert not isinstance( + data_iterator, list + ), "when expansion_factor_real_data > 1, the data iterator should not be a list." + grain_restore_args = GrainCheckpointRestore( + data_iterator.local_iterator, process_index=jax.process_index(), process_count=process_count_stored + ) + + else: + # Case 4: Mismatch + raise ValueError( + f"Error restoring Grain checkpoint in {directory}: " + f"The number of stored checkpoint files ({process_count_stored}) " + f"is incompatible with the number of JAX processes ({process_count_jax}). " + "If you are resuming training with a different number of chips, see instructions in " + "https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/data_input_pipeline/" + "data_input_grain.md#using-grain" + ) + + # Call restore once with the composed arguments + restored_state = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args, iter=grain_restore_args)) + return (restored_state, None) diff --git a/src/maxtext/common/train_state_nnx.py b/src/maxtext/common/train_state_nnx.py index a123956a73..c19331658e 100644 --- a/src/maxtext/common/train_state_nnx.py +++ b/src/maxtext/common/train_state_nnx.py @@ -12,13 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""The NNX Unified TrainState.""" +"""The NNX Unified TrainState. + +The NNX <-> Linen on-disk checkpoint format conversions are implemented below. +""" from typing import Any from flax import nnx import jax import jax.numpy as jnp +import numpy as np class TrainStateNNX(nnx.Module): @@ -193,3 +197,59 @@ def from_linen_checkpoint_dict(linen_pure_dict): if optimizer: result["optimizer"] = optimizer return result + + +def default_for_sds(sds): + """Returns a deterministic value matching `sds` shape/dtype/sharding. + + Used to fill NNX-only state (rngs/dropout) that the Linen on-disk layout never + carried. Materializes under jit with the target out_shardings so it works on + multi-host meshes (device_put can't place a global sharding whose devices + aren't + all addressable from this process). + """ + if not (hasattr(sds, "dtype") and hasattr(sds, "shape")): + return sds + + def _make(): + if "key" in str(sds.dtype): + base = jax.random.key(0) + return base if sds.shape == () else jax.random.split(base, int(np.prod(sds.shape))).reshape(sds.shape) + return jnp.zeros(sds.shape, dtype=sds.dtype) + + sharding = getattr(sds, "sharding", None) + if sharding is None: + return _make() + return jax.jit(_make, out_shardings=sharding)() + + +def populate_pure_dict_from_partial(abstract_pure, partial_concrete): + """Fills `abstract_pure` with values from `partial_concrete` (by path), defaulting the rest. + + Paths present in `partial_concrete` take the restored value; paths absent from + it + (NNX-only state the Linen checkpoint never had) get `default_for_sds`. + """ + if isinstance(abstract_pure, dict): + return { + k: populate_pure_dict_from_partial( + v, + partial_concrete.get(k) if isinstance(partial_concrete, dict) else None, + ) + for k, v in abstract_pure.items() + } + if partial_concrete is not None and not isinstance(partial_concrete, dict): + return partial_concrete + return default_for_sds(abstract_pure) + + +def rebuild_nnx_with_values(abstract_nnx_state, concrete_weights): + """Fills each Variable in `abstract_nnx_state` with the matching restored array.""" + leaves, treedef = jax.tree_util.tree_flatten(abstract_nnx_state, is_leaf=lambda x: isinstance(x, nnx.Variable)) + concrete = jax.tree_util.tree_leaves(concrete_weights) + if len(leaves) != len(concrete): + raise ValueError( + f"Params load leaf-count mismatch: {len(leaves)} abstract Variables vs" f" {len(concrete)} restored." + ) + new_leaves = [v.replace(value=a) if isinstance(v, nnx.Variable) else a for v, a in zip(leaves, concrete)] + return jax.tree_util.tree_unflatten(treedef, new_leaves) diff --git a/src/maxtext/trainers/post_train/distillation/distillation_utils.py b/src/maxtext/trainers/post_train/distillation/distillation_utils.py index f063cdb23a..c17e1573ea 100644 --- a/src/maxtext/trainers/post_train/distillation/distillation_utils.py +++ b/src/maxtext/trainers/post_train/distillation/distillation_utils.py @@ -32,7 +32,7 @@ from maxtext.utils import max_logging from maxtext.utils import maxtext_utils # Reuse MaxText's native checkpointing logic. -from maxtext.common.checkpointing import GrainCheckpointHandler, GrainCheckpointSave, GrainCheckpointRestore +from maxtext.common.grain_utility import GrainCheckpointHandler, GrainCheckpointSave, GrainCheckpointRestore from tunix.sft import checkpoint_manager as tunix_checkpoint_manager from tunix.sft import peft_trainer diff --git a/tests/unit/train_state_nnx_test.py b/tests/unit/train_state_nnx_test.py index 59bc4a2a8f..b5f7eb1838 100644 --- a/tests/unit/train_state_nnx_test.py +++ b/tests/unit/train_state_nnx_test.py @@ -16,6 +16,7 @@ import unittest from flax import nnx +import jax import jax.numpy as jnp from maxtext.common import train_state_nnx import optax @@ -85,5 +86,51 @@ def test_apply_gradients_raises_runtime_error(self): self.assertIn("inference only", str(cm.exception)) +class TestStateReconstructionHelpers(unittest.TestCase): + """default_for_sds / populate_pure_dict_from_partial / rebuild_nnx_with_values.""" + + def test_default_for_sds_zeros_for_array(self): + sds = jax.ShapeDtypeStruct((2, 3), jnp.float32) + out = train_state_nnx.default_for_sds(sds) + self.assertEqual(out.shape, (2, 3)) + self.assertEqual(out.dtype, jnp.float32) + self.assertTrue(jnp.array_equal(out, jnp.zeros((2, 3), jnp.float32))) + + def test_default_for_sds_key_for_key_dtype(self): + key_dtype = jax.random.key(0).dtype # str(dtype) contains "key" + out = train_state_nnx.default_for_sds(jax.ShapeDtypeStruct((), key_dtype)) + self.assertIn("key", str(out.dtype)) + + def test_default_for_sds_passthrough_non_sds(self): + # A value without shape/dtype is returned unchanged. + self.assertEqual(train_state_nnx.default_for_sds("not-an-array"), "not-an-array") + + def test_populate_takes_concrete_and_defaults_missing(self): + abstract = { + "restored": jax.ShapeDtypeStruct((2,), jnp.float32), + "missing": jax.ShapeDtypeStruct((3,), jnp.float32), + } + partial = {"restored": jnp.ones((2,))} + out = train_state_nnx.populate_pure_dict_from_partial(abstract, partial) + self.assertTrue(jnp.array_equal(out["restored"], jnp.ones((2,)))) # taken from concrete + self.assertTrue(jnp.array_equal(out["missing"], jnp.zeros((3,)))) # defaulted + + def test_rebuild_nnx_with_values_binds_arrays(self): + abstract = {"w": nnx.Param(jnp.zeros((2,))), "b": nnx.Param(jnp.zeros(()))} + concrete = {"w": jnp.ones((2,)), "b": jnp.asarray(5.0)} + rebuilt = train_state_nnx.rebuild_nnx_with_values(abstract, concrete) + # Read the bound array off each Variable leaf (version-robust vs. `.value`). + w_val = jax.tree_util.tree_leaves(rebuilt["w"])[0] + b_val = jax.tree_util.tree_leaves(rebuilt["b"])[0] + self.assertTrue(jnp.array_equal(w_val, jnp.ones((2,)))) + self.assertEqual(float(b_val), 5.0) + + def test_rebuild_nnx_with_values_raises_on_count_mismatch(self): + abstract = {"w": nnx.Param(jnp.zeros((2,))), "b": nnx.Param(jnp.zeros(()))} + concrete = {"w": jnp.ones((2,))} # one leaf vs two abstract Variables + with self.assertRaises(ValueError): + train_state_nnx.rebuild_nnx_with_values(abstract, concrete) + + if __name__ == "__main__": unittest.main()