[NNX] Emergency checkpoint: build the manager with a Linen-layout abstract

[ecnal-cienet]

Description

The EMC resume-from-GCS test (maxtext_emc_resume_from_gcs) fails when a pure_nnx run restores from step 99:

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

What's happening: pure_nnx saves checkpoints in the on-disk checkpoint layout (params/opt_state/step plus an nnx_aux subtree, via to_checkpoint_dict) so old Linen checkpoints stay compatible. Both the regular and the emergency manager save this way. The difference is that when we build the emergency manager we hand it the NNX abstract state (model/optimizer), and the emergency manager keeps that abstract and restores against it — it does not use the item we pass at restore time. So it ends up matching a checkpoint-layout tree on disk against an NNX tree, and Orbax complains.

This is also why the earlier fix (0da53c3) didn't help: it converts the restore-time item, but the emergency manager never looks at that item.

Fix: convert the abstract with to_checkpoint_dict inside create_orbax_emergency_checkpoint_manager when it's an nnx.State, so what the manager holds matches what's on disk. Restore still reshapes back to NNX afterwards, so the state we train on is NNX. On-disk format is unchanged, so regular and EMC checkpoints are still interchangeable. 5 lines, in the constructor.

Tests

Reproduced and verified on our v6e-8 (jax 0.10.0, same as the failing job): before the fix the save→restore hits the exact same ValueError, after the fix it restores fine.

• TestEmergencyManagerAbstractLayout — mocks the Orbax manager and checks it gets a Linen abstract (params/opt_state/step), not the NNX model/optimizer. Runs anywhere, so CI catches it if someone breaks this again.
• TestEmergencySaveRestoreRoundTrip — full save→restore through the real emergency manager, checks it comes back as NNX with the right weights. Skips if the EMC manager can't be built.
• Both fail if I revert the fix. 3 pass on the TPU, existing checkpoint tests still pass.

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

❌ Patch coverage is 89.28571% with 9 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/utils/maxtext_utils.py	58.82%	4 Missing and 3 partials ⚠️
src/maxtext/common/checkpointing.py	96.42%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!