Register LazyTensor with the Orbax v1 leaf-handler registry

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")

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.")

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

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(