[ecnal-cienet]

Description

This PR aims to fix issue b/530068041

The maxtext_emc_resume_from_gcs e2e job fails when a pure_nnx run resumes from its GCS checkpoint:

ValueError: User-provided restore item and on-disk value metadata tree structures do not match
  model:     - Source: {'decoder': ...}   - Target: MISSING
  opt_state: - Source: MISSING            - Target: [{'count': ...}, ...]

A pure_nnx run saves in the on-disk checkpoint layout — Linen params/opt_state/step plus an nnx_aux subtree, via to_checkpoint_dict — so pure_nnx and Linen checkpoints stay interchangeable. Both the regular and the emergency manager save that way.

The difference is on the restore side. The regular manager restores against the item it is given at restore time. The emergency manager instead bakes in the abstract it is constructed with and restores against that, ignoring the restore-time item. create_orbax_emergency_checkpoint_manager was handing it get_abstract_state's NNX-native tree (model/optimizer), so Orbax compared a checkpoint-layout tree on disk against an NNX target and raised. This is also why converting the restore-time item alone did not help: the emergency manager never looks at that item.

Fix: convert the abstract with to_checkpoint_dict in create_orbax_emergency_checkpoint_manager when it is an nnx.State, so the blueprint the manager holds matches what maybe_save_checkpoint writes. split_for_checkpoint already requires that save and restore partition the state identically; the manager's baked-in abstract is a third copy of that partition, and it now goes through the same function as the other two. The restore path reshapes back to NNX, so the state we train on is still NNX, and the on-disk format is unchanged. Six lines, in the constructor. A Linen state is not an nnx.State, so it is passed through untouched.

Tests

tests/unit/checkpointing_emergency_nnx_test.py (5) — the emergency-manager counterpart of checkpointing_nnx_roundtrip_test.py:

• the manager is handed the checkpoint layout (params/opt_state/step/nnx_aux), not the NNX model/optimizer roots. Orbax is mocked, so this one runs in normal CI and catches a regression here without an emergency-checkpoint setup.
• a Linen state is passed through unchanged.
• a full round trip through a real EmergencyCheckpointManager is leaf-exact: weights, optimizer moments, dropout rng, batch stats.
• nnx_aux is genuinely restored rather than reset to a fresh init — checked by value (rng count, batch-norm mean, optimizer step), since a reset keeps the tree structure intact and would slip past the structural check.
• caches stay ephemeral: excluded from the baked-in abstract, so they keep their init value on restore while the rng beside them still comes back from nnx_aux.

Reverting the six lines fails four of the five (the fifth is the Linen pass-through case, which correctly still passes).

End-to-end on a v6e-8 (jax 0.10.0, orbax-checkpoint 0.12.1)

1. Save checkpoint for later tests: (passed, log)
2. EMC resume WITHOUT fix (failed expected, log)
3. EMC resume WITH fix (passed, log)

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[NuojCheng]

might be better adding is_leaf=lambda...