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

[janbernloehr]

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:

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.

--- 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 "<class '__main__.LazyTensor'>" 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.