Migrate Decoder (Gemma3/Deepseek/Llama4) and utils to NNX

Author: hsuan-lun-chiang

src/maxtext/configs/base.yml

@@ -1074,7 +1074,8 @@ position_id_per_seconds: 25
 subslice_shape: ""
 
 # NNX
-enable_nnx: false
+enable_nnx: True
+pure_nnx_decoder: True
 
 ################################## Qwen3-Next Specific Configs ##################################
 # Kernel size for the 1D convolution in the Gated Delta Net

src/maxtext/configs/types.py

@@ -777,6 +777,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

@@ -113,7 +113,9 @@ def __init__(
         mesh=mesh,
         model_mode=MODEL_MODE_TRAIN,
         name=f"mtp_{k}_transformer_layer",
+        rngs=rngs,
     )
+
     self.transformer_layer = nnx_wrappers.ToNNX(mtp_transformer_layer, rngs=rngs)
 
     # ToNNX requires explicit initialization with sample inputs for proper parameter setup.