NNX migration prep (3/N): TrainState, model creation, and end-to-end training loop

[ecnal-cienet]

NNX Migration Route Map

1. ✅ Add NNX scaffolding: pure_nnx flag, init_state_fn, TrainStateNNX, NNX utils. Linen workflow unchanged. (PR #3427)
2. ✅ NNX sharding utilities: get_abstract_state_nnx, get_named_sharding_nnx, set_named_sharding_nnx, get_partition_spec_nnx, get_mesh_from_config. (PR #3470)
3. 🔄 [This PR] NNX fully supported end-to-end. pure_nnx=True enables full NNX training; default remains False.
4. ❌ NNX unit tests and performance verification complete. Set pure_nnx=True as default.
5. ❌ Remove Linen-specific code paths and NNX compatibility flags.

Description

Note: This is the third in a series of NNX migration PRs. With this PR, pure_nnx=True runs a complete NNX training loop — initialization, sharding, gradient accumulation, eval, and checkpointing — without hitting any NotImplementedError. The pure_nnx flag still defaults to False, preserving the existing Linen workflow unchanged.

TrainStateNNX and unit tests

src/maxtext/layers/train_state_nnx.py implements the TrainStateNNX container, which holds an NNX model and its Optax optimizer as a single composable unit. Unit tests cover state creation, optimizer step, and Orbax checkpoint round-trip:

• tests/unit/train_state_nnx_test.py
• tests/unit/train_state_nnx_checkpoint_test.py

Muon optimizer and model creation utilities

• muon_utils.py — updated to support NNX models alongside Linen.
• model_creation_utils.py — refactored to expose create_nnx_abstract_model and from_config, which create and initialize an NNX model from a config without running a full forward pass.

End-to-end training loop (train.py + supporting modules)

The core training loop in train.py now dispatches on pure_nnx at every major decision point:

• Sharding (sharding.py) — maybe_update_params_sharding_with_opt dispatches to a new maybe_update_params_sharding_with_opt_nnx, which extracts nnx.Param-only shardings from the flat nnx.State without accessing .params.
• Gradient accumulation (gradient_accumulation.py) — NNX path uses nnx.value_and_grad with nnx.split / nnx.merge per microbatch inside jax.lax.scan, carrying non-Param rest state (RNGs) through the loop.
• Train/eval step JIT (maxtext_utils.py) — get_functional_train_with_signature and get_functional_eval_with_signature use a 2-element in_shardings tuple (state, batch) for NNX (no rng argument), vs. 3-element for Linen.
• Checkpointing (checkpointing.py) — maybe_save_checkpoint converts nnx.State to a plain dict via state.to_pure_dict() before Orbax save; load_state_if_possible restores via nnx.replace_by_pure_dict(abstract_state, dict).

Tests

Unit tests:

python3 -m pytest tests/unit/train_state_nnx_test.py -v
python3 -m pytest tests/unit/train_state_nnx_checkpoint_test.py -v
python3 -m pytest tests/unit/sharding_nnx_test.py -v
python3 -m pytest tests/unit/maxtext_utils_nnx_test.py -v
python3 -m pytest tests/unit/train_compile_test.py -v

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 45.23364% with 293 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/trainers/pre_train/train.py	30.32%	88 Missing and 20 partials ⚠️
src/maxtext/utils/sharding.py	5.55%	50 Missing and 1 partial ⚠️
src/maxtext/utils/maxtext_utils.py	69.79%	20 Missing and 9 partials ⚠️
src/maxtext/utils/train_utils.py	47.91%	21 Missing and 4 partials ⚠️
src/maxtext/trainers/pre_train/train_compile.py	33.33%	14 Missing and 4 partials ⚠️
src/maxtext/utils/model_creation_utils.py	52.77%	15 Missing and 2 partials ⚠️
src/maxtext/utils/gradient_accumulation.py	27.77%	8 Missing and 5 partials ⚠️
src/maxtext/common/checkpointing.py	43.75%	6 Missing and 3 partials ⚠️
...rc/maxtext/utils/generate_param_only_checkpoint.py	30.00%	5 Missing and 2 partials ⚠️
src/maxtext/utils/muon_utils.py	66.66%	4 Missing and 3 partials ⚠️
... and 2 more

📢 Thoughts on this report? Let us know!