[NNX] Delete Linen 2/5: collapse dispatch in trainers (pre-train, DiLoCo, GRPO)

[ecnal-cienet]

NNX Migration Route Map

1. ✅ Add NNX scaffolding: pure_nnx flag, init_state_fn, TrainStateNNX, NNX utils. Linen workflow unchanged. (PR NNX migration prep (1/N): pure_nnx flag and init_state_fn scaffolding #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 NNX migration prep (2/N): NNX utils and sharding utilities #3470)
3. ✅ NNX fully supported end-to-end: TrainStateNNX, model creation, gradient accumulation, checkpointing, and training loop dispatch. (PR NNX migration prep (3/N): TrainState, model creation, and end-to-end training loop #3500)
4. ✅ Sharding diagnostics on NNX, plus post-training bugfixes that surfaced once the NNX path got exercised end-to-end. (PR [NNX] NNX migration prep (4/N): sharding tools and post-training fixes #3652)
5. ✅ NNX correctness fixes, feature enablements, and vocab tiling on NNX.
6. ✅ NNX-native DPO.
7. ✅ NNX-native MaxEngine inference. (PR [NNX] NNX migration prep (7/N): NNX-native MaxEngine inference #3821)
8. ✅ NNX-native LoRA + GRPO. (PR [NNX] NNX migration prep (8/N): NNX native lora grpo #3824)
9. ✅ NNX-aware QK-Clip + remaining checkpoint utilities. (PR [NNX] NNX migration prep (9/N): NNX-aware QK-Clip + checkpoint utilities #3836)
   9.5. ✅ NNX + AQT in MaxEngine + serve-mode reload + gpt3 prefill fix. (PR [NNX] NNX migration prep (9.5/N): NNX + AQT in MaxEngine + serve-mode reload + gpt3 prefill fix #3844)
10. ✅ Vocab tiling custom_vjp for NNX.
11. ✅ Set NNX defaults to True; regenerate sharding goldens; flip back integration-test pure_nnx=False annotations. (PR [NNX] NNX migration (11/N): set pure_nnx / enable_nnx / pure_nnx_decoder defaults to True #3526)
12. Delete Linen-specific code paths and NNX compatibility flags — delivered as 5 stacked PRs (12.1–12.5):
    12.1. ✅ Collapse Linen dispatch in utils / quantization / checkpointing / sharding. (PR [NNX] Delete Linen 1/5: collapse dispatch in utils, quantization, checkpointing, sharding #4357)
    12.2. 🔄 [This PR] Collapse Linen dispatch in trainers — pre_train/train.py (loss/train/eval step), train_compile.py, diloco.py, GRPO trainer + tests. After this train.py no longer references the Linen Decoder. Also regenerates golden_dpo_correctness.json for the NNX decoder (forward is bit-identical to Linen; the different RNG-init scheme just reseeds weights, so the init-specific golden is refreshed — not a numerics change). Stacks on 12.1.
    12.3. ❌ Collapse Linen dispatch in inference (maxengine, kvcache, vLLM, LoRA).
    12.4. ❌ Delete the Linen decoder/attention/embedding layers + *_as_linen model wrappers.
    12.5. ❌ Remove the pure_nnx / enable_nnx / pure_nnx_decoder config flags.

Description

Second of the stacked "Delete Linen" series (see 1/5). This one collapses the Linen dispatch in the training paths — pre_train/train.py (loss/train/eval step), train_compile.py, diloco.py, and the GRPO trainer — and updates the trainer tests. After this PR train.py no longer references the Linen Decoder.

It also regenerates golden_dpo_correctness.json for the NNX decoder. The DPO correctness test used to pin the bridged Linen decoder to match its PyTorch-synced golden; on the NNX decoder the forward pass is bit-identical (checked by transplanting weights across, max|logits diff| = 0) but the layer weights are seeded from a different RNG scheme, so the init-specific golden needed regenerating. It's an init difference, not a numerics change. The test passes against the new golden on CPU.

Like 1/5 this is dead-branch removal under the NNX defaults, and the Linen Decoder class is still present (it's deleted in 4/5), so nothing dangles. Base branch is feat/nnx-del-linen-1-utils. ~13 files, +432 / −1235.

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

Files with missing lines	Patch %	Lines
src/maxtext/trainers/pre_train/train.py	69.23%	31 Missing and 13 partials ⚠️
src/maxtext/trainers/pre_train/train_compile.py	87.50%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!