[NNX] Delete Linen (12.1a): state setup and train loop

[ecnal-cienet]

Description

First of the pre-train Delete-Linen PRs. It removes the Linen branch of every pure_nnx conditional in pre-train state setup and the training loop. With the NNX defaults those branches never run, so this is pure dead-branch removal — the three flags (pure_nnx / enable_nnx / pure_nnx_decoder) still exist (they come out in a later PR), so nothing about the config surface changes here.

Source changes:

• maxtext_utils.py
  • get_functional_train_with_signature / get_functional_eval_with_signature: drop the Linen (state, batch, rng) in-shardings; the NNX step takes no rng, so it is always (state, batch).
  • load_compiled: drop the Linen example_rng extra input.
  • get_abstract_state: collapse to always call get_abstract_state_nnx, deleting the Linen abstract-state path (jax.eval_shape + nn.get_partition_spec + logical_to_mesh_sharding + the optimizer/param host-offload handling).
• train_utils.py setup_train_loop: always build the NNX abstract model and the TrainStateNNX init function, and drop the Linen from_config model, the init_initial_state partial, and the Linen else branches for param extraction, the DiLoCo step mesh, debug_sharding logical annotations, and the final train-state/model merge.

Test changes (drop the now-dead Linen cases):

• setup_train_loop_nnx_test.py: only the NNX path remains, so the test_pure_nnx_* cases lose their pure_nnx qualifier.
• maxtext_utils_test.py: remove the Linen model construction and test_linen_in_shardings_includes_rng / pure_nnx=False mock configs.
• sharding_compare_test.py: remove the Linen-golden test_sharding_dump_for_model comparison (those goldens were Linen-only).
• state_dtypes_test.py / train_compile_test.py: drop the pure_nnx=False / linen_scanned variants, keeping the NNX cases.
• correctness_tests_nnx_dispatch_test.py: drop test_sft_logits_linen_path and collapse _sft_config() to NNX-only — the Linen SFT path is deleted here, so that CPU case can no longer run.

~8 files, +115 / −628.

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

Files with missing lines	Patch %	Lines
src/maxtext/utils/train_utils.py	83.33%	2 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[bvandermoon]

Do we need to get rid of the flags to enable/disable the NNX/Linen versions with this PR? The flags break once the Linen paths start getting removed

[ecnal-cienet]

Do we need to get rid of the flags to enable/disable the NNX/Linen versions with this PR? The flags break once the Linen paths start getting removed

Hi @bvandermoon,

You're right that the flags break as the Linen paths come out. But removing them in this PR means deleting every if pure_nnx branch at the same time. The whole Linen removal in one PR is 83 files, +1011 / −6692, which is too big to review.

So I split it into 7 stacked PRs by subsystem:

• 12.1a (this one): pre-train state setup + train loop, 8 files, −637
• 12.1b: sharding + muon, 6 files, −110
• 12.1c: quantization + model creation, 5 files, −223
• 12.2: trainers (pre-train, DiLoCo, GRPO), 12 files, −1228
• 12.3: inference (maxengine, kvcache, vLLM, LoRA), 7 files, −492
• 12.4: Linen decoder/attention layers + *_as_linen wrappers, 16 files, −3538
• 12.5: the three flags, 30 files, −464

The flags are removed in 12.5. That is the only PR that touches the config surface (types.py, base.yml, inference/vllm.yml, the distillation configs, pyconfig_deprecated.py, plus doc and script mentions). All 7 branches are ready and stacked in this order, so the flags are gone by the end of the wave.

For the gap in between: I checked what pure_nnx=False does after this PR. A pre-train run fails in maxtext_utils.setup_initial_state:

ValueError: pytree structure error: different types at key path "pjit out_shardings"
  prefix pytree has flax.nnx.statelib.State, full pytree has maxtext.common.train_state_nnx.TrainStateNNX

So it fails, not silently trains something wrong.