PR #2831: Migrate Decoder to NNX

Imported from GitHub PR #2831

# Description
Migrate the Transformer decoder layer into NNX.

Note: The following models are currently not supported:

- DeepSeek
- Gemma3
- Llama4

Support for these models will be added in a follow-up PR.

Strategy:
A `pure_nnx_decoder` flag is added to control whether NNX or Linen decoder shall be used.
Initial migration doesn't include the pipeline NNX support.

# Tests
Conducted these tests. Details in the [GDoc file](https://docs.google.com/document/d/1NbUP3g5glgbC6bMyt44pwM_vQA1NR7U2rBUzfbTDwSs/edit?pli=1&resourcekey=0-9EUahtzL-hCycdu7l0grhQ&tab=t.htq5367h8au0)
1. Test with different model and compare with Linen training
2. Golden logits comparison
3. Inference
4. Checkpoint comparison (Including TreeStructure Comparison)
5. Sharding comparison

TODOs:
- NNX version of unit tests (future PRs)

# Checklist

Before submitting this PR, please make sure (put X in square brackets):
- [x] I have performed a self-review of my code. For an optional AI review, add the `gemini-review` label.
- [x] I have necessary comments in my code, particularly in hard-to-understand areas.
- [x] I have run end-to-end tests tests and provided workload links above if applicable.
- [x] 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](https://maxtext.readthedocs.io/en/latest/development.html#adding-new-documentation-files).

Copybara import of the project:

--
073e916 by hsuan-lun-chiang <hsuan-lun.chiang@cienet.com>:

Migrate Decoder to NNX

    Adding nnx_decoders.py in parallel with decoders.py

    1. Dup and modifiy decoders.py on new file nnx_decoders.py
    2. add new config pure_nnx_decoder to control if model will use NNXDecoder, default false for now
    3. modify relative code to accomodate the change
    4. add/modify unit test

Merging this change closes #2831

COPYBARA_INTEGRATE_REVIEW=#2831 from CIeNET-International:feat/Migrate-Decoder-to-NNX 073e916
PiperOrigin-RevId: 884170982

Author: hsuan-lun-chiang

src/maxtext/configs/base.yml

@@ -1087,6 +1087,7 @@ subslice_shape: ""
 
 # NNX
 enable_nnx: false
+pure_nnx_decoder: false
 
 ################################## Qwen3-Next Specific Configs ##################################
 # Kernel size for the 1D convolution in the Gated Delta Net
@@ -1152,4 +1153,4 @@ distill_temperature: 1.0
 # distill_beta is used for cosine similarity loss between intermediate activataitions of out_proj in teacher/student models.
 # 0.0 value disables this feature.
 distill_beta: 0.0
-distill_layer_indices: None
+distill_layer_indices: None

src/maxtext/configs/types.py

@@ -783,6 +783,7 @@ class HardwareAndMesh(BaseModel):
   enable_nnx: bool = Field(False, description="Whether to use NNX for model definition.")
   optimize_mesh_for_tpu_v6e: bool = Field(False, description="Apply transformations to the mesh for TPU v6e.")
   shardy: bool = Field(True, description="Whether to use shardy XLA backend.")
+  pure_nnx_decoder: bool = Field(False, description="Whether to enable pure NNX decoder.")
 
 
 class LayoutAndSharding(BaseModel):

src/maxtext/layers/multi_token_prediction.py

@@ -22,8 +22,8 @@
 import jax.numpy as jnp
 from jax.sharding import Mesh
 from maxtext.common.common_types import Config, MODEL_MODE_TRAIN
+from maxtext.layers.nnx_decoders import NNXDecoderLayer
 from maxtext.utils.globals import EPS
-from maxtext.layers import nnx_wrappers
 from maxtext.layers.decoders import DecoderLayer
 from maxtext.layers.initializers import variable_to_logically_partitioned
 from maxtext.layers.linears import DenseGeneral
@@ -70,7 +70,7 @@ def __init__(
       config: Config,
       mesh: Mesh,
       layer_number: int,
-      transformer_layer_module: Type[DecoderLayer],
+      transformer_layer_module: Type[NNXDecoderLayer],
       *,
       rngs: nnx.Rngs,
   ):
@@ -108,22 +108,12 @@ def __init__(
         rngs=rngs,
     )
     # Use MODEL_MODE_TRAIN for initialization; runtime model_mode is passed dynamically.
-    mtp_transformer_layer = transformer_layer_module(
+    self.transformer_layer = transformer_layer_module(
         config=cfg,
         mesh=mesh,
         model_mode=MODEL_MODE_TRAIN,
         name=f"mtp_{k}_transformer_layer",
-    )
-    self.transformer_layer = nnx_wrappers.ToNNX(mtp_transformer_layer, rngs=rngs)
-
-    # ToNNX requires explicit initialization with sample inputs for proper parameter setup.
-    batch_size, seq_len = max_utils.get_batch_seq_len_for_mode(config=cfg, model_mode=MODEL_MODE_TRAIN)
-    self.transformer_layer.lazy_init(
-        inputs=jnp.zeros((batch_size, seq_len, self.config.emb_dim), dtype=self.config.dtype),
-        decoder_segment_ids=None,
-        decoder_positions=jnp.zeros((batch_size, seq_len), dtype=jnp.int32),
-        deterministic=True,
-        model_mode=MODEL_MODE_TRAIN,
+        rngs=rngs,
     )
 
   @property
@@ -212,7 +202,7 @@ def __init__(
       self,
       config: Config,
       mesh: Mesh,
-      transformer_layer_module: Type[DecoderLayer],
+      transformer_layer_module: Type[NNXDecoderLayer],
       decoder: nnx.Module,
       rngs: nnx.Rngs,
   ):