NNX migration prep (5/N): enable NNX by default

[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. ✅ NNX fully supported end-to-end: TrainStateNNX, model creation, gradient accumulation, checkpointing, and training loop dispatch. (PR #3500)
4. ✅ NNX sharding diagnostics and bidirectional Linen↔NNX checkpoint conversion utilities. (PR #3525)
5. 🔄 [This PR] Enable NNX by default; fix unit test failures.
6. ❌ Remove Linen-specific code paths and NNX compatibility flags.

Description

Note: This is the fifth in a series of NNX migration PRs. This PR flips all three NNX flags to True in base.yml, making NNX the default training path, and fixes the unit test failures that surface as a result.

Config change

src/maxtext/configs/base.yml — three flags flipped to True:

enable_nnx: True
pure_nnx_decoder: True
pure_nnx: True

Unit test fixes

File	Fix
src/maxtext/layers/nnx_decoders.py	Add multimodal_input=None to NNXDecoder.__call__ and unpack into individual fields — Transformer.__call__ passes a unified MultimodalInput object but NNXDecoder previously only accepted the fields individually.
src/maxtext/utils/muon_utils.py	Return mapped_state (an nnx.State with correct structure) directly instead of converting to a flat nested dict with a "params" wrapper — the old conversion broke attribute access expected by tests.
src/maxtext/trainers/post_train/distillation/distillation_utils.py	Guard optimizer_state restore on whether it exists in the checkpoint — PeftTrainer.save() only saves model_params, so restoring optimizer_state unconditionally caused a KeyError.
tests/integration/gradient_accumulation_test.py	Switch test_sft_grad_accumulate_same_loss from the deprecated SFT loop to the NNX-native train_sft.py — the deprecated loop always passed nextrng as a 3rd positional arg, mismatching the 2-element NNX in_shardings.
tests/unit/sharding_compare_test.py	Filter abstract_state.model leaves to floating-point only before asserting dtype == float32 — the NNX model state includes RNG state variables (uint32/key) that are not weight parameters.

Tests

python3 -m pytest tests/unit/ tests/integration/ -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 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 21.13238% with 989 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...xtext/checkpoint_conversion/linen_nnx_converter.py	8.97%	283 Missing and 1 partial ⚠️
...ckpoint_conversion/compare_linen_nnx_checkpoint.py	12.09%	268 Missing and 1 partial ⚠️
src/maxtext/trainers/pre_train/train.py	15.48%	119 Missing and 12 partials ⚠️
src/maxtext/utils/maxtext_utils.py	35.00%	70 Missing and 8 partials ⚠️
src/maxtext/utils/sharding.py	5.55%	51 Missing ⚠️
src/maxtext/utils/train_utils.py	14.58%	41 Missing ⚠️
src/maxtext/utils/muon_utils.py	13.04%	20 Missing ⚠️
src/maxtext/utils/gradient_accumulation.py	5.55%	17 Missing ⚠️
src/maxtext/common/checkpointing.py	6.25%	15 Missing ⚠️
src/maxtext/layers/nnx_decoders.py	33.33%	11 Missing and 3 partials ⚠️
... and 10 more

📢 Thoughts on this report? Let us know!