Unify the NNX and Linen restore paths in load_state_if_possible

[ecnal-cienet]

Description

The NNX path had two early-return blocks that each duplicated the Linen restore flow, one for the standard manager and one for the emergency managers. Any change to restore logic had to be made in three places, and they had already drifted: the NNX grain branch dropped the iterator element that the Linen branch returns.

pure_nnx saves in the Linen on-disk layout, so both paths restore the same tree. This converts the NNX abstract to that layout going in, runs the existing match statement unchanged, and reshapes what comes back on the way out. Linen behavior is unchanged.

Bugs fixed behind that duplication

• Dynamic safetensors got the whole abstract state instead of the weights. For NNX that is fatal before any weight is read: opt_state is keyed by integer chain indices, and the loader flattens its target with a "." separator, so it died in flatten_dict with TypeError: sequence item 2: expected str instance, int found.
• A weight no mapping covered reached the model as an untrained init value. Orbax rejects a shape conflict itself, but leaves a weight the checkpoint does not carry as an unmaterialized placeholder. The weights-only Orbax load already checks for this; the dynamic load now does too.
• checkpoint_conversion_fn was unusable. It arrives from config as a dotted string but was called directly, raising TypeError for any value that was set. It is now imported, and resolved before the checkpoint is read.
• Full-state safetensors now handles NNX, previously wired only for the orbax layout.

Testing

92 unit tests pass across the six suites that touch this code:

python -m pytest tests/unit/checkpointing_nnx_load_test.py \
                 tests/unit/checkpointing_test.py \
                 tests/unit/checkpointing_nnx_roundtrip_test.py \
                 tests/unit/checkpointing_emergency_nnx_test.py \
                 tests/unit/setup_initial_state_nnx_test.py \
                 tests/unit/grain_utility_test.py

End-to-end on a v6e-8: 74 assertions across gpt3-52k / gemma2-2b / llama2-7b, 0 failures.

Case	Path	Model	Evidence	Log
1a/1b	NNX save → resume (standard manager)	llama2-7b	restoring from this run's directory step 9, resumed 10 → 19	log log
2a/2b	Linen save → resume (regression)	gemma2-2b	Resumed 10 → 19, behavior unchanged	log log
3	Linen checkpoint → NNX resume	llama2-7b	Restoring Linen-layout checkpoint into NNX state, continued from the Linen step 19 → 22	log
4	Weights-only load (load_parameters_path)	gemma2-2b	Params restored into the NNX state, training continues	log
5	Full-state load (load_full_state_path)	llama2-7b	Reshaped into NNX, continued from the restored step 9 → 12	log
9a–9d	Emergency manager save → resume	gemma2-2b	Local dir wiped between save and resume, so the restore came from the persistent GCS copy; NNX and Linen both 6 → 11	log log log log
6a	Shape mismatch is rejected	gpt3-52k	Requested shape (2048, 16) is not compatible with the stored shape (1024, 16), zero training steps	log
6b	Missing weight is named	gpt3-52k	'decoder/layers/2/self_attention/qkv_proj/kernel': missing (model expects (16, 3, 2, 8) float32), zero training steps	log
7	safetensors without a conversion fn	gpt3-52k	Names the missing config flag; no raw 'NoneType' object is not callable	log
8/8b	HuggingFace safetensors → NNX	qwen3-0.6b	Successfully mapped 13 parameters, trained 3 steps. Pre-fix, same command: TypeError, zero parameters mapped, zero steps	log log

The resume also matches Linen numerically. Two independently trained gemma2-2b runs, across the save and the restore:

step	NNX	Linen
9 (checkpoint written)	5.488	5.498
17 (after restore)	2.280	2.288
19	2.036	2.043

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.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov]

Codecov Report

❌ Patch coverage is 98.03922% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/common/checkpointing.py	97.95%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[xibinliu]

move to the very beginning of the func so it can be re-used by all branches.

[xibinliu]

set(tree) == {"params"}

Nit: make it efficient and understandable:

len(tree) == 1 and "params" in tree