NNX: save checkpoints in the Linen on-disk layout (interchangeable with Linen)

[ecnal-cienet]

Summary

A pure_nnx run now writes the same on-disk checkpoint layout as a Linen run, so Linen and NNX checkpoints are interchangeable. After the NNX migration this lets existing Linen checkpoints keep working, and lets a checkpoint saved by either framework be loaded by the other.

The conversion lives on TrainStateNNX (no separate load-time adapter): the state serializes to the Linen layout on save and reads it back on load.

Why

NNX's natural nnx.state(...).to_pure_dict() differs from Linen's checkpoint in three ways, none of which are in the model/decoder modules — the weight names already match. They come from the train state and optimizer:

1. top-level keys: NNX {model, optimizer} vs Linen {params, step, opt_state}
2. weights: NNX model/... vs Linen params/params/... (the Linen params collection)
3. opt_state: NNX serializes the optax chain as an int-keyed dict; Linen as a list with None for each EmptyState, and wraps mu/nu in the params collection

Plus step dtype (NNX uint32 vs Linen int32) and NNX-only rngs/dropout state that Linen never had.

What changes

Two source files plus their tests:

• src/maxtext/layers/train_state_nnx.py — to_linen_checkpoint_dict (save) and from_linen_checkpoint_dict (load) reshape between the two layouts, handling the three differences above, the step dtype, and dropping NNX-only rngs/dropout.
• src/maxtext/common/checkpointing.py:
  • save: maybe_save_checkpoint writes to_linen_checkpoint_dict(state.to_pure_dict()).
  • load (NNX target): all three paths reshape the Linen layout back into the NNX state and re-initialize the dropped rngs/dropout from defaults — load_state_if_possible (auto-resume), _load_full_state_from_path (load_full_state_path), and load_params_from_path (load_parameters_path).

Linen runs are unchanged; they already read and write this layout.

Test plan

Byte-exactness. A pure_nnx gemma2-2b checkpoint is structurally identical to a
Linen one — same 76 leaves, no path or shape/dtype differences.

Compatibility matrix. For each model: save with both frameworks, then load each
of the three mechanisms in both directions (Linen→NNX and NNX→Linen). All eight
phases pass on six models on v6e-8:

Model	default (auto-resume)	weight-only (load_parameters_path)	full-state (load_full_state_path)
gpt3-52k	✓	✓	✓
gpt3-6b	✓	✓	✓
gemma2-2b	✓	✓	✓
gemma2-9b	✓	✓	✓
llama2-7b	✓	✓	✓
qwen3-8b	✓	✓	✓

(3 model families, both optimizer types: adam_pax single-state and adamw 3-element chain.)

Unit tests. tests/unit/train_state_nnx_checkpoint_test.py and tests/unit/checkpointing_nnx_load_test.py cover the converters (round-trip, opt_state list/dict, step dtype, RNG strip) and a real Orbax round-trip for the weight-only NNX load. 16 tests 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 52.84553% with 58 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/common/checkpointing.py	37.73%	30 Missing and 3 partials ⚠️
src/maxtext/layers/train_state_nnx.py	64.28%	13 Missing and 12 partials ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @ecnal-cienet, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

The Pull Request successfully implements the interchangeability of Linen and NNX checkpoints by providing bi-directional conversion logic. The implementation is robust, handling differences in top-level keys, weight wrapping, optimizer state representations, and step dtypes.

🔍 General Feedback

• Well-Structured Conversion: The separation of concerns between the checkpointing logic and the layout transformation in train_state_nnx.py is excellent.
• Abstract State Handling: The use of _cast_step and _default_for_sds to handle jax.ShapeDtypeStruct ensures that the conversion works correctly during abstract restoration, which is critical for performance and multi-host compatibility.
• Comprehensive Testing: The addition of unit tests for both full-state and weight-only restoration, as well as round-trip conversion, provides high confidence in the implementation.