Test/mesa migrate pipeline to nnx

src/MaxText/layers/decoders.py

@@ -399,14 +399,10 @@ def setup(self):
     self.decoder_layer = self.get_decoder_layers()
     self.norm_layer = self.get_norm_layer(num_features=self.config.emb_dim)
     if self.config.using_pipeline_parallelism:
-      nnx_decoder_classes = self.get_nnx_decoder_layers()
-      if nnx_decoder_classes is not None:
-        pipeline_stage_module = self.get_pipeline_stage_module(nnx_decoder_classes, use_nnx=True)
-      else:
-        pipeline_stage_module = self.get_pipeline_stage_module(self.decoder_layer, use_nnx=False)
+      pipeline_stage_module = self.get_pipeline_stage_module(self.decoder_layer)
       remat_policy = self.get_remat_policy()
       self.pipeline_module = pipeline.create_pipeline(
-          config=self.config, mesh=self.mesh, layers=pipeline_stage_module, remat_policy=remat_policy, use_nnx=(nnx_decoder_classes is not None)
+          config=self.config, mesh=self.mesh, layers=pipeline_stage_module, remat_policy=remat_policy,
       )
 
   def minimal_policy(self, with_context=False):
@@ -685,7 +681,7 @@ def scan_decoder_layers(self, cfg, decoder_layer, length, metadata_axis_name, me
         config=cfg, mesh=mesh, name=metadata_axis_name, quant=self.quant, **kwargs  # pytype: disable=wrong-keyword-args
     )
 
-  def get_pipeline_stage_module(self, decoder_blocks, use_nnx=False):
+  def get_pipeline_stage_module(self, decoder_blocks):
     """get pipeline stage module
 
     Args:
@@ -703,20 +699,6 @@ def get_layer_to_pipeline(blocks, cfg):
     cfg = self.config
     base_stage = get_layer_to_pipeline(decoder_blocks, cfg)
 
-    if use_nnx:
-      if cfg.num_layers_per_pipeline_stage == 1:
-        return base_stage
-      else:
-        return lambda config, mesh, model_mode, rngs, quant=None: SequentialNNXWrapper(
-            decoder_layer_class=base_stage,
-            num_decoder_layers=cfg.num_layers_per_pipeline_stage,
-            config=config,
-            mesh=mesh,
-            model_mode=model_mode,
-            rngs=rngs,
-            quant=quant
-        )
-
     if cfg.set_remat_policy_on_layers_per_stage:
       policy = self.get_remat_policy()
       base_stage = self.set_remat_policy([base_stage], policy)[0]
@@ -730,6 +712,7 @@ def get_layer_to_pipeline(blocks, cfg):
           "layers_per_stage",
           self.mesh,
           in_axes_tuple=(nn.broadcast,) * 4,
+          model_mode=self.model_mode,
       )
     else:
       stage_module = SequentialBlockDecoderLayers(

src/MaxText/layers/embeddings.py

@@ -342,6 +342,7 @@ def __call__(
           "The embedding dims of the rotary position embedding" "must match the hidden dimension of the inputs."
       )
 
+
     position = position[:, :, jnp.newaxis, jnp.newaxis]
     sinusoid_inp = position / self.timescale
     sin = jnp.sin(sinusoid_inp).astype(inputs.dtype)