Register LazyTensor with the Orbax v1 leaf-handler registry

[lokic233]

What

Registers the converter's LazyTensor leaf type with the Orbax v1 leaf-handler
registry, so HF→MaxText checkpoint conversion works again at
--simulated_cpu_devices_count=1.

Fixes #5106.

Why

checkpoint_conversion/to_maxtext.py:297 registers LazyTensor through the Orbax
v0 API (type_handlers.register_type_handler). Since c41b2ae ("Refactor
MaxText's checkpointing system to primarily use the Orbax v1 API", 2026-09-01) the
save path runs through ocp.training.Checkpointer.save_checkpointables, which
resolves handlers against the v1 registry only.

The two registries are disjoint, and the compatibility bridge runs one way: a v1
PyTreeHandler builds a per-instance StandardLeafHandlerRegistry and derives its
v0 registry from itself. So a v0 global registration is structurally unreachable from
a v1 save. Conversion dies with NoEntryError.

The error names the checkpointable ("items" / TrainState), never the leaf type,
which is why this is easy to misdiagnose — the real failure is one leaf deeper.

c41b2ae touched zero files under checkpoint_conversion/, which is why this
regressed silently.

Scope — measured, not inferred

The failure requires simulated_cpu_devices_count=1. At 2+, shard_jax_weights
materializes every leaf at checkpoint_conversion/utils/utils.py:1194-1196 before the
tree ever reaches Orbax. Same tree, same model, same lazy_load_tensors=True, one flag
apart:

simulated_cpu_devices_count	result on main
1	NoEntryError at save_checkpointables
16 (the default)	rc=0, conversion succeeds

This is also why CI is green.
tests/integration/checkpoint_conversion_test.py::test_qwen3_30b_a3b_roundtrip_conversion
is a real round-trip test collected by the tpu-integration and gpu-integration
flavors — but it never passes --simulated_cpu_devices_count, so it always runs at 16,
the one setting that cannot fail. Not "nobody runs it"; it runs constantly, at the safe
setting.

How

Purely additive — 4 files, +68 −0. The existing v0 registration at :297 is left in
place; removing it is orthogonal.

• LazyTensorLeafHandler(NumpyLeafHandler) — materializes via np.asarray in
  serialize, then defers to the numpy handler.
• LazyTensorPyTreeHandler(PyTreeHandler) — injects a per-instance
  StandardLeafHandlerRegistry carrying that handler. (registry.add takes a type,
  not an instance, hence the subclass.)
• build_lazy_tensor_checkpointables_registry() — a local_registry with
  include_global_registry=True, bound to checkpointable name "items".
• Threaded through checkpoint_context.build_context /
  checkpointing.create_orbax_checkpoint_manager /
  checkpoint_conversion/utils/utils.save_weights_to_checkpoint as an optional
  checkpointables_registry kwarg, default None — no behavior change for any
  existing caller.

The load-bearing detail: secondary_typestrs=["np.ndarray"]

Without it, the checkpoint records the custom handler's typestr, and a standard v0
restore then fails with Unknown type: "...LazyTensorLeafHandler". The "fix" would
silently write unloadable checkpoints — strictly worse than the loud failure it
replaces. An earlier prototype did exactly that, and it was caught only because the
validator restores through the v0 PyTreeCheckpointer rather than just checking that
the save returned 0.

Reviewers: this is the line to scrutinize.

Testing

CPU logits diff against the transformers reference at reduced dimensions
(hidden 256, 8 layers, vocab 512), every parameter randomized — default init leaves
RMSNorm scales at 0, which hides a wrong norm mapping. Run on a tree carrying the
patch, not a monkeypatch, at simulated_cpu_devices_count=1:

arm	rc	max|diff|	cosine	argmax agreement
lazy_load_tensors=True, scan_layers=False	0	1.66893e-06	1.00000000	100%
lazy_load_tensors=True, scan_layers=True	0	1.57952e-06	1.00000000	100%
lazy_load_tensors=False (eager reference)	0	1.66893e-06	1.00000000	100%

Both lazy arms fail with NoEntryError without this patch; the eager arm is the
unchanged reference. Lazy now matches eager to the digit.

Both scan_layers settings are exercised because they are different code paths and
scan_layers=True is what training consumes.

The harness is sensitive to real corruption, not just to crashes: a same-shape
gate_proj↔up_proj swap in the same converter scores cosine 0.729 / argmax 18.8%.
Shape and name assertions cannot see that swap; only a value diff can.

Notes

Two imports reach into _src
(orbax...experimental.v1._src.serialization.numpy_leaf_handler and
..._src.serialization.registry). If there is a public path for these I'd rather use
it — happy to switch.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Code Review

This pull request introduces Orbax v1 support for saving checkpoints with LazyTensor leaves by implementing LazyTensorLeafHandler and LazyTensorPyTreeHandler, and updating checkpointing utilities to propagate a custom checkpointables_registry. The review feedback suggests inheriting from the public ocp_v1.PyTreeHandler instead of the private _src module, explicitly passing the expected dtype to np.asarray during tensor materialization to prevent dtype mismatches, and refactoring the handler initialization to robustly register LazyTensor on any provided custom registry.

[gemini-code-assist]

Since PyTreeHandler is publicly exposed in the Orbax v1 API, we can import and inherit from ocp_v1.PyTreeHandler directly instead of reaching into the private _src module. This reduces reliance on internal implementation details.

Suggested change

	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 orbax.checkpoint import v1 as ocp_v1
	from orbax.checkpoint.experimental.v1._src.serialization import numpy_leaf_handler
	from orbax.checkpoint.experimental.v1._src.serialization import registry as leaf_handler_registry

[gemini-code-assist]

Explicitly pass the expected dtype (param.info.dtype) to np.asarray. When converting checkpoints from a different precision (e.g., float16 to bfloat16), LazyTensor.__array__ is called with dtype=None by default, which returns the original precision array. Passing the expected dtype ensures that the materialized array matches the expected dtype of the LazyTensor and avoids potential dtype mismatches or implicit casting issues during serialization.

Suggested change

	params = [dataclasses.replace(param, value=np.asarray(param.value)) for param in params]
	params = [dataclasses.replace(param, value=np.asarray(param.value, dtype=param.info.dtype)) for param in params]

[gemini-code-assist]

Inherit from the public ocp_v1.PyTreeHandler instead of the private pytree_handler.PyTreeHandler.

Suggested change

	class LazyTensorPyTreeHandler(pytree_handler.PyTreeHandler):
	class LazyTensorPyTreeHandler(ocp_v1.PyTreeHandler):

[gemini-code-assist]

To make LazyTensorPyTreeHandler more robust, retrieve the leaf_handler_registry from kwargs if it exists, and register LazyTensor on it. This ensures that LazyTensor is always supported even if a custom registry is passed to the handler.

Suggested change

	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 __init__(self, **kwargs):
	registry = kwargs.get("leaf_handler_registry")
	if registry is None:
	registry = leaf_handler_registry.StandardLeafHandlerRegistry()
	kwargs["leaf_handler_registry"] = registry
	registry.add(
	LazyTensor,
	numpy_leaf_handler.NumpyShapeDtype,
	LazyTensorLeafHandler,
	override=True,
	secondary_typestrs=["np.ndarray"],
	)
	super().__init__(**kwargs)