checkpoint_conversion: save fails with Orbax NoEntryError when --lazy_load_tensors is combined with a single simulated device

### Bug report

`save_weights_to_checkpoint` in `src/maxtext/checkpoint_conversion/utils/utils.py` fails deterministically when lazily-loaded tensors are combined with a single simulated device. No checkpoint is produced, so every conversion using that configuration is blocked.

**Steps to reproduce**

On any machine that can run MaxText checkpoint conversion (the failing step is CPU-only; we reproduced it on an 8xH100 x86_64 node and on aarch64 GB200/GB300 nodes, but no accelerator is required to hit it):

1. Install MaxText at `main` (2026-09-02 or later) with `orbax-checkpoint` 0.12.4.
2. Fetch any supported Hugging Face checkpoint, e.g. `meta-llama/Llama-3.1-8B-Instruct`.
3. Run a `checkpoint_conversion` entry point that calls `save_weights_to_checkpoint`, with lazy loading on and exactly one simulated device:

```bash
JAX_PLATFORMS=cpu python3 -m maxtext.checkpoint_conversion.to_maxtext \
src/MaxText/configs/base.yml \
model_name=llama3.1-8b \
base_output_directory=/tmp/llama31-8b-orbax \
--hf_model_path=/path/to/Llama-3.1-8B-Instruct \
--lazy_load_tensors=true \
--simulated_cpu_devices_count=1 \
--save_dtype=bfloat16
```

The default `--simulated_cpu_devices_count=16` hides the bug, because it routes through the sharding path. Setting it to `1` is the only change needed to trigger the failure.

**Expected behavior**

The conversion completes and writes the converted Orbax checkpoint to `base_output_directory`, exactly as it does today with `--simulated_cpu_devices_count > 1` and exactly as it did before `common/checkpointing.py` moved to the Orbax v1 API. `--simulated_cpu_devices_count` is documented as controlling how many CPU devices are simulated for sharding; it is not expected to change whether a checkpoint can be saved at all. The run should log `Conversion complete. Checkpoint saved to ...` and exit 0, and the resulting checkpoint should be byte-equivalent in content to one produced with more simulated devices, since the single-device path differs only in that it skips sharding.

**Actual behavior**

The save aborts with an Orbax `NoEntryError` raised from `resolve_handler_for_save` while resolving a handler for the `"items"` checkpointable, before any file is written. The process exits non-zero and `base_output_directory` is left empty. In our runs this happened about 15 seconds into the conversion step, deterministically, on every attempt and on both x86_64 and aarch64.

**Root cause**

The lazy proxies are never materialized on the single-device path:

- `to_maxtext.py` places `LazyTensor` proxies as weight-tree leaves under `--lazy_load_tensors`, and registers a handler for them through the **Orbax v0** API (`orbax.checkpoint.type_handlers.register_type_handler`).
- The only surviving code that converts those proxies to NumPy is inside `shard_jax_weights` (`if not isinstance(arr, (np.ndarray, jax.Array)): arr = np.array(arr)`).
- `save_weights_to_checkpoint` skips `shard_jax_weights` when `device_count == 1` (logging `Single device: Skip sharding`), so unmaterialized proxies end up in the `TrainState` passed to `save_checkpoint`.
- `common/checkpointing.py` now uses **Orbax v1** (`from orbax.checkpoint import v1 as ocp`) and calls `save_checkpointables` with `{"items": state}`. Orbax v1 resolves a handler by checkpointable name or `handler.is_handleable(value)`, never by the value's top-level type. `"items"` is unbound in the default global registry, so `PyTreeHandler.is_handleable` tree-maps over the **leaves** against `StandardLeafHandlerRegistry` (`jax.Array`, `np.ndarray`, `int`, `float`, `bytes`, `str`) with `issubclass`. `LazyTensor` is a plain proxy implementing `__array__` and is not an `np.ndarray` subclass, so no leaf handler matches. The underlying `UnregisteredTypeError` is swallowed by a bare `except Exception` and re-surfaces as an opaque `NoEntryError`.

The Orbax v0 and v1 leaf registries are disjoint, so the v1 migration of `common/checkpointing.py` orphaned the v0 `LazyTensor` registration in `to_maxtext.py` — the registration still runs (it still logs its `np.ndarray` collision warning) but is invisible to the v1 save path.

**Suggested fix**

Materialize the lazy leaves on the single-device branch, reusing the conversion `shard_jax_weights` already performs. We validated this experimentally as an in-image patch on 8xH100 with Llama-3.1-8B: the previously-failing conversion completes, the checkpoint is written, downstream decode passes, and inference throughput and step times are unchanged within run-to-run noise (about 0.1-0.4%) against a pre-regression baseline. The diff is below.

```diff
--- a/src/maxtext/checkpoint_conversion/utils/utils.py
+++ b/src/maxtext/checkpoint_conversion/utils/utils.py
@@ def shard_jax_weights(...)
return jax_weights


+def materialize_lazy_weights(jax_weights, mem_info):
+ """Materializes lazily-loaded weight proxies into NumPy arrays."""
+ start = time.time()
+ jax_weights_flat, jax_weights_struct = tree.flatten(jax_weights)
+ del jax_weights
+ gc.collect()
+
+ jax_weights_new = []
+ jax_weights_flat.reverse()
+ num_weights = len(jax_weights_flat)
+ for _ in tqdm(range(num_weights)):
+ jax_weight = jax_weights_flat.pop()
+ if not isinstance(jax_weight, (np.ndarray, jax.Array)):
+ # materialize lazy tensor
+ jax_weight = np.array(jax_weight)
+ jax_weights_new.append(jax_weight)
+ del jax_weight
+ gc.collect()
+ logging.debug("Memory usage: %f GB", mem_info.memory_info().rss / (1024**3))
+
+ jax_weights = tree.unflatten(jax_weights_struct, jax_weights_new)
+ max_logging.log(f"Elapse for lazy weight materialization: {(time.time() - start) / 60:.2f} min")
+
+ return jax_weights
+
+
def save_weights_to_checkpoint(
@@ in save_weights_to_checkpoint
else:
# If number of simulated devices is 1, SKIP sharding and SKIP jax conversion.
max_logging.log("Single device: Skip sharding")
+ # Sharding is what would otherwise materialize lazily-loaded weights, so do
+ # it here: the save path below has no handler for lazy proxy leaves.
+ max_logging.log("Single device: Materializing lazy weights")
+ jax_weights = materialize_lazy_weights(jax_weights, mem_info)
```

It mirrors `shard_jax_weights`' reverse-pop plus `gc.collect()` loop, so peak memory stays roughly one tensor above the converted tree, and checkpoint contents are unchanged: this is the same conversion the multi-device path already performs and the same conversion the pre-migration `LazyTensorHandler.serialize` performed.

Alternatively, port `LazyTensorHandler` to an Orbax v1 `LeafHandler` registered through `PyTreeOptions`/`Context`, which would preserve streaming into serialization but pins MaxText to the v1 leaf-handler API and only helps `to_maxtext.py`.

The blast radius is wider than one entry point: `save_weights_to_checkpoint` serves several conversion entry points, and any of them would hit this if they produced lazy leaves. Today only `to_maxtext.py` constructs `LazyTensor`, so only it is affected.

### Logs/Output

Relevant log lines, in order, from a failing conversion run:

```
Type handler registry overriding type "" collision on np.ndarray
Lazy loading ENABLED
Single device: Skip sharding
```

followed by:

```
File ".../checkpoint_conversion/utils/utils.py", line 1302, in save_weights_to_checkpoint
checkpointing.save_checkpoint(checkpoint_manager, step, state_new, config=config)
File ".../common/checkpointing.py", line 992, in save_checkpoint
checkpoint_manager.save_checkpointables(...)
File ".../orbax/checkpoint/experimental/v1/_src/saving/saving.py", line 321, in ...
resolve_handler_for_save(...)
File ".../orbax/checkpoint/experimental/v1/_src/handlers/registration.py", line 623, in resolve_handler_for_save
raise NoEntryError(
```

No checkpoint files are produced. After applying the suggested fix, the same run instead logs `Single device: Materializing lazy weights`, then `Elapse for lazy weight materialization: 1.04 min`, then `Conversion complete. Checkpoint saved to ...`, and downstream decode succeeds with no traceback anywhere in the log.

### Environment Information

- **MaxText version / commit:** `main` as of 2026-09-02. Relevant blobs: `src/maxtext/checkpoint_conversion/utils/utils.py` = `26677a0`, `src/maxtext/common/checkpointing.py` = `a11681b`, `src/maxtext/checkpoint_conversion/to_maxtext.py` = `ce2eecb`.
- **orbax-checkpoint:** 0.12.4
- **jax:** 0.11.2.dev20260901
- **Model:** `meta-llama/Llama-3.1-8B-Instruct` (bf16 safetensors)
- **Conversion step:** CPU-only (`JAX_PLATFORMS=cpu`), `--simulated_cpu_devices_count=1`, `--lazy_load_tensors=true`, `--save_dtype=bfloat16`, `async_checkpointing=false`
- **Operating system / hardware:** Linux; reproduced on `x86_64` (8xH100 node) and `aarch64` (GB200/GB300 nodes). The failure is platform-independent — it happens before any device work.

### Additional Context

Reported by NVIDIA. This is a regression relative to the pre-Orbax-v1 behaviour: the same configuration worked before `common/checkpointing.py` moved to the v1 API, because the v0 `LazyTensorHandler` was still consulted at save time. It began failing in our nightly runs on the first day that picked up the v1 migration commit, and the last passing run used MaxText from the day before.

---

_This issue was drafted with assistance from the `opus` AI model._