Merge pull request #1974 from mesakhcienet:feat/simple-decoder-layer-nnx

PiperOrigin-RevId: 820006566

Author: Google-ML-Automation

src/MaxText/layers/decoders.py

@@ -402,9 +402,9 @@ def get_decoder_layers(self):
       case DecoderBlockType.QWEN3_MOE:
         return [qwen3.Qwen3MoeDecoderLayerToLinen]
       case DecoderBlockType.SIMPLE:
-        return [simple_layer.SimpleDecoderLayer]
+        return [simple_layer.SimpleDecoderLayerToLinen]
       case DecoderBlockType.SIMPLE_MLP:
-        return [simple_layer.SimpleMlpDecoderLayer]
+        return [simple_layer.SimpleMlpDecoderLayerToLinen]
       case DecoderBlockType.LLAMA4:
         return [llama4.Llama4ScannableBlockToLinen] if self.config.scan_layers else [llama4.Llama4DecoderLayerToLinen]
       case _:

src/MaxText/layers/simple_layer.py

@@ -17,56 +17,82 @@
 from jax import numpy as jnp
 from jax.sharding import Mesh
 
-from flax import linen as nn
-
+from flax import nnx
 from MaxText.common_types import Config
-from MaxText.layers import quantizations
+from MaxText.layers import quantizations, nnx_wrappers
+from MaxText.layers.initializers import variable_to_logically_partitioned
+
 
+from typing import Optional
 # pytype: disable=attribute-error
 
 
-class SimpleDecoderLayer(nn.Module):
+class SimpleDecoderLayer(nnx.Module):
   """Decoder layer consisting of a single [embed, embed] weight matrix."""
 
-  config: Config
-  mesh: Mesh
-  model_mode: str
-  quant: None | quantizations.AqtQuantization = None
+  def __init__(
+      self,
+      config: Config,
+      mesh: Mesh,
+      model_mode: str,
+      rngs: nnx.Rngs,
+      quant: Optional[quantizations.AqtQuantization] = None,
+  ) -> None:
+
+    self.config = config
+    self.mesh = mesh
+    self.model_mode = model_mode
+    self.rngs = rngs
+    self.quant = quant
 
-  def setup(self):
-    self.weight_mat = self.param(
-        "weights",
-        nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")),
-        (self.config.emb_dim, self.config.emb_dim),
+    init_fn = nnx.with_partitioning(nnx.initializers.lecun_normal(), sharding=("embed", "mlp"), mesh=self.mesh)
+
+    self.weights = nnx.Param(
+        init_fn(self.rngs.params(), (self.config.emb_dim, self.config.emb_dim)),
     )
 
   def __call__(
       self, inputs: jnp.ndarray, positions, segmentation, deterministic, model_mode, previous_chunk=None, page_state=None
   ):
     if self.config.scan_layers:
-      return inputs @ self.weight_mat.astype(inputs.dtype), None
-    else:
-      return inputs @ self.weight_mat.astype(inputs.dtype)
+      return inputs @ self.weights.astype(inputs.dtype), None
+    return inputs @ self.weights.astype(inputs.dtype)
+
+
+SimpleDecoderLayerToLinen = nnx_wrappers.to_linen_class(
+    SimpleDecoderLayer,
+    base_metadata_fn=variable_to_logically_partitioned,
+)
 
 
-class SimpleMlpDecoderLayer(nn.Module):
+class SimpleMlpDecoderLayer(nnx.Module):
   """Decoder layer consisting of [embed,mlp] followed by an [mlp,embed] matmul."""
 
-  config: Config
-  mesh: Mesh
-  model_mode: str
-  quant: None | quantizations.AqtQuantization = None
+  def __init__(
+      self,
+      config: Config,
+      mesh: Mesh,
+      model_mode: str,
+      rngs: nnx.Rngs,
+      quant: Optional[quantizations.AqtQuantization] = None,
+  ) -> None:
+
+    self.config = config
+    self.mesh = mesh
+    self.model_mode = model_mode
+    self.rngs = rngs
+    self.quant = quant
 
-  def setup(self):
-    self.ff_1 = self.param(
-        "ff_1",
-        nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")),
-        (self.config.emb_dim, self.config.mlp_dim),
+    init_ff1_fn = nnx.with_partitioning(nnx.initializers.lecun_normal(), sharding=("embed", "mlp"), mesh=self.mesh)
+
+    self.ff_1 = nnx.Param(
+        init_ff1_fn(self.rngs.params(), (self.config.emb_dim, self.config.mlp_dim)),
     )
-    self.ff_2 = self.param(
-        "ff_2",
-        nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("mlp", "embed")),
-        (self.config.mlp_dim, self.config.emb_dim),
+
+    init_ff2_fn = nnx.with_partitioning(nnx.initializers.lecun_normal(), sharding=("mlp", "embed"), mesh=self.mesh)
+
+    self.ff_2 = nnx.Param(
+        init_ff2_fn(self.rngs.params(), (self.config.mlp_dim, self.config.emb_dim)),
     )
 
   def __call__(
@@ -84,5 +110,10 @@ def __call__(
     output = intermediate @ self.ff_2.astype(inputs.dtype)
     if self.config.scan_layers:
       return output, None
-    else:
-      return output
+    return output
+
+
+SimpleMlpDecoderLayerToLinen = nnx_wrappers.to_linen_class(
+    SimpleMlpDecoderLayer,
+    base_metadata_fn=variable_to_logically_partitioned,
+)

tests/pipeline_parallelism_test.py

@@ -27,6 +27,7 @@
 
 from flax.core import meta
 from flax import linen as nn
+from flax import nnx
 
 from MaxText import maxtext_utils
 from MaxText import pyconfig
@@ -63,7 +64,10 @@ def assert_pipeline_same_output_and_grad(self, config, single_pipeline_stage_cla
     mesh = Mesh(devices_array, config.mesh_axes)
     model_mode = MODEL_MODE_TRAIN
     if single_pipeline_stage_class is None:
-      single_pipeline_stage = simple_layer.SimpleDecoderLayer(config=config, mesh=mesh, model_mode=model_mode)
+      rngs = nnx.Rngs(params=0)
+      single_pipeline_stage = simple_layer.SimpleDecoderLayerToLinen(
+          config=config, mesh=mesh, model_mode=model_mode, rngs=rngs
+      )
     else:
       single_pipeline_stage = single_pipeline_stage_class(config=config, mesh=mesh, model_mode=model_mode)
 
@@ -90,7 +94,10 @@ def get_inputs(batch_size, sequence, features):
     )
     deterministic = True
     # We use a simpler single matmul decoder layer for fast compilation in these tests.
-    single_pipeline_stage = simple_layer.SimpleDecoderLayer(config=config, mesh=mesh, model_mode=model_mode)
+    rngs = nnx.Rngs(params=0)
+    single_pipeline_stage = simple_layer.SimpleDecoderLayerToLinen(
+        config=config, mesh=mesh, model_mode=model_mode, rngs=rngs
+    )
     my_pipeline = pipeline.Pipeline(config=config, layers=single_pipeline_stage, mesh=mesh)
     init_pipeline_params = my_pipeline.init(
         jax.random.PRNGKey(0), inputs, inputs_position, inputs_segmentation, deterministic, model_mode