diff --git a/src/maxtext/checkpoint_conversion/to_maxtext.py b/src/maxtext/checkpoint_conversion/to_maxtext.py index 17ced7f095..5911af0153 100644 --- a/src/maxtext/checkpoint_conversion/to_maxtext.py +++ b/src/maxtext/checkpoint_conversion/to_maxtext.py @@ -50,6 +50,7 @@ """ import argparse +import dataclasses from functools import partial import json import os @@ -76,6 +77,10 @@ from maxtext.utils.globals import HF_IDS import numpy as np from orbax.checkpoint import type_handlers +from orbax.checkpoint import v1 as ocp_v1 +from orbax.checkpoint.experimental.v1._src.handlers import pytree_handler +from orbax.checkpoint.experimental.v1._src.serialization import numpy_leaf_handler +from orbax.checkpoint.experimental.v1._src.serialization import registry as leaf_handler_registry from safetensors import safe_open try: @@ -297,6 +302,51 @@ async def serialize(self, value, *args, **kwargs): type_handlers.register_type_handler(LazyTensor, LazyTensorHandler(), override=True) +class LazyTensorLeafHandler(numpy_leaf_handler.NumpyLeafHandler): + """Orbax v1 leaf handler for LazyTensor. + + The v0 registration above cannot serve the v1 save path. A v1 ``PyTreeHandler`` + resolves leaves through a per-instance ``LeafHandlerRegistry`` and *derives* its + v0 registry from itself, so the compatibility bridge only runs v1 -> v0; a v0 + global registration is unreachable from v1. + + Like its v0 counterpart this masquerades as the standard numpy handler -- + ``secondary_typestrs`` below writes ``np.ndarray`` as the leaf's typestr -- so + the checkpoint is indistinguishable from one a plain ``NumpyLeafHandler`` + produced and restores in a standard MaxText instance. + """ + + async def serialize(self, params, serialization_context): + # MATERIALIZE: trigger the lazy load (__array__) explicitly before saving. + # This ensures the parent NumpyLeafHandler receives real np.ndarrays. + params = [dataclasses.replace(param, value=np.asarray(param.value)) for param in params] + return await super().serialize(params, serialization_context) + + +class LazyTensorPyTreeHandler(pytree_handler.PyTreeHandler): + """``PyTreeHandler`` whose leaf registry additionally accepts ``LazyTensor``.""" + + def __init__(self, **kwargs): + if "leaf_handler_registry" not in kwargs: + registry = leaf_handler_registry.StandardLeafHandlerRegistry() + registry.add( + LazyTensor, + numpy_leaf_handler.NumpyShapeDtype, + LazyTensorLeafHandler, + override=True, + secondary_typestrs=["np.ndarray"], + ) + kwargs["leaf_handler_registry"] = registry + super().__init__(**kwargs) + + +def build_lazy_tensor_checkpointables_registry(): + """Registry that saves MaxText's "items" checkpointable with LazyTensor support.""" + registry = ocp_v1.handlers.local_registry(include_global_registry=True) + registry.add(LazyTensorPyTreeHandler, checkpointable_name="items") + return registry + + def get_maxtext_model_info(config): """Initializes the abstract MaxText model and returns parameter mapping information. @@ -1093,6 +1143,7 @@ def _eager_getter(key): config.checkpoint_storage_use_ocdbt, config.checkpoint_storage_use_zarr3, config=config, + checkpointables_registry=build_lazy_tensor_checkpointables_registry(), ) print_ram_usage("Program Ends") diff --git a/src/maxtext/checkpoint_conversion/utils/utils.py b/src/maxtext/checkpoint_conversion/utils/utils.py index 17209a037b..7f887cb060 100644 --- a/src/maxtext/checkpoint_conversion/utils/utils.py +++ b/src/maxtext/checkpoint_conversion/utils/utils.py @@ -53,6 +53,7 @@ from maxtext.checkpoint_conversion.utils.tensor_handling import nesting_depth, stacked_axes from maxtext.utils import max_logging import orbax.checkpoint as ocp +from orbax.checkpoint import v1 as ocp_v1 _storage = gcs_storage() Client = _storage.Client @@ -1238,6 +1239,7 @@ def save_weights_to_checkpoint( use_ocdbt: bool, use_zarr3: bool, config=None, + checkpointables_registry: ocp_v1.handlers.CheckpointableHandlerRegistry | None = None, ): """Saves model weights to a MaxText-compatible checkpoint with optional sharding. @@ -1254,6 +1256,10 @@ def save_weights_to_checkpoint( (OCDBT) format for improved metadata handling. use_zarr3: If True, uses the Zarr3 storage format for the underlying array data. config: Optional config to save along with checkpoint metadata. + checkpointables_registry: Optional Orbax v1 handler registry, for converters whose + weight pytrees hold leaf types Orbax does not know natively. Only matters when + device_count is 1; above that shard_jax_weights has already materialized every + leaf into a sharded jax.Array. """ mem_info = psutil.Process() logging.debug("Memory usage: %f GB", mem_info.memory_info().rss / (1024**3)) @@ -1281,6 +1287,7 @@ def save_weights_to_checkpoint( save_interval_steps, use_ocdbt=use_ocdbt, use_zarr3=use_zarr3, + checkpointables_registry=checkpointables_registry, ) if checkpoint_manager is None: raise RuntimeError("Failed to create Orbax checkpoint manager.") diff --git a/src/maxtext/common/checkpoint_context.py b/src/maxtext/common/checkpoint_context.py index 433868fb87..6ca12e2a42 100644 --- a/src/maxtext/common/checkpoint_context.py +++ b/src/maxtext/common/checkpoint_context.py @@ -94,6 +94,7 @@ def build_context( colocated_python_checkpointing: bool = False, partial_load: bool = False, checkpoint_layout: ocp.options.CheckpointLayout | None = None, + checkpointables_registry: ocp.handlers.CheckpointableHandlerRegistry | None = None, ) -> ocp.Context: """Builds an Orbax v1 ``Context`` from MaxText checkpoint flags. @@ -119,6 +120,10 @@ def build_context( equivalent of v0 ``partial_restore=True``). checkpoint_layout: On-disk layout (``ORBAX`` or ``SAFETENSORS``) for loading. + checkpointables_registry: Handler registry deciding which + ``CheckpointableHandler`` serves each named checkpointable. Callers whose + pytrees hold leaf types Orbax does not know natively pass a registry + carrying a handler for them; ``None`` leaves the global registry in place. Returns: A configured, unfrozen ``ocp_v1.Context``. @@ -172,4 +177,7 @@ def build_context( if checkpoint_layout is not None: ctx.checkpoint_layout = checkpoint_layout + if checkpointables_registry is not None: + ctx.checkpointables.registry = checkpointables_registry + return ctx diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index dec2634bed..26bf1a4f60 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -375,6 +375,7 @@ def create_orbax_checkpoint_manager( todelete_subdir: str | None = None, todelete_full_path: str | None = None, ocdbt_target_data_file_size_bytes: int | None = None, + checkpointables_registry: ocp.handlers.CheckpointableHandlerRegistry | None = None, ): """Returns an Orbax v1 training ``Checkpointer``, or None if checkpointing is disabled.""" if not enable_checkpointing: @@ -411,6 +412,7 @@ def create_orbax_checkpoint_manager( todelete_full_path=todelete_full_path, todelete_subdir=todelete_subdir, partial_load=True, + checkpointables_registry=checkpointables_registry, ) manager = ocp.training.Checkpointer(