Merge pull request #2885 from CIeNET-International:test/pipeline-scan-nnx

PiperOrigin-RevId: 933577624

Author: Google-ML-Automation

src/maxtext/layers/decoders.py

@@ -33,6 +33,7 @@
 from maxtext.layers import mhc
 from maxtext.layers import normalizations
 from maxtext.layers import pipeline
+from maxtext.layers.nnx_decoders import NNXDecoderLayer, NNXSequentialPipelineStage, NNXScannedPipelineStage
 from maxtext.layers import quantizations
 from maxtext.layers.attentions import attention_as_linen
 from maxtext.layers.embeddings import attend_on_embedding, embed_as_linen, positional_embedding_as_linen
@@ -260,7 +261,7 @@ def __call__(
           slot=slot,
       )
       if self.config.scan_layers:
-        inputs = inputs[0]  #  When scan_layers is True the decoder layers return (outputs, None).
+        inputs = inputs[0]  # When scan_layers is True the decoder layers return (outputs, None).
     if self.config.scan_layers:
       return inputs, None  # pytype: disable=bad-return-type
     else:
@@ -305,10 +306,19 @@ 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:
-      pipeline_stage_module = self.get_pipeline_stage_module(self.decoder_layer)
       remat_policy = self.get_remat_policy()
+      nnx_blocks = self._get_nnx_decoder_block_classes()
+
+      # Per-stage builder handed to create_pipeline: the pipeline transform invokes it
+      # once per pipeline stage, passing that stage's rngs, to construct the stage's
+      # decoder block(s). nnx_blocks (the selected decoder block classes) is captured from
+      # this setup scope; remat_policy is applied separately by create_pipeline.
+      def build_pipeline_stage_layers(rngs):
+        """Builds one pipeline stage module from the selected NNX decoder block classes."""
+        return self._build_nnx_pipeline_stage(nnx_blocks, rngs)
+
       self.pipeline_module = pipeline.create_pipeline(
-          config=self.config, mesh=self.mesh, layers=pipeline_stage_module, remat_policy=remat_policy
+          config=self.config, layers=build_pipeline_stage_layers, mesh=self.mesh, remat_policy=remat_policy
       )
 
   def minimal_policy(self, with_context=False, with_quantization=False):
@@ -498,6 +508,42 @@ def get_decoder_layers(self):
         # Default case to handle any unknown decoder block types.
         raise ValueError(f"Incorrect decoder_block name {self.config.decoder_block.value=}")
 
+  def _get_nnx_decoder_block_classes(self):
+    """Returns NNX decoder block classes for pipeline stage creation."""
+    cfg = self.config
+
+    def get_scannable(normal_cls, scannable_cls):
+      return [scannable_cls] if cfg.scan_layers else [normal_cls]
+
+    layer_map = {
+        DecoderBlockType.DEFAULT: [NNXDecoderLayer],
+        DecoderBlockType.LLAMA2: [llama2.LlamaDecoderLayer],
+        DecoderBlockType.MISTRAL: [mistral.MistralDecoderLayer],
+        DecoderBlockType.MIXTRAL: [mixtral.MixtralDecoderLayer],
+        DecoderBlockType.GEMMA: [gemma.GemmaDecoderLayer],
+        DecoderBlockType.GEMMA2: [gemma2.Gemma2DecoderLayer],
+        DecoderBlockType.GEMMA3: [gemma3.Gemma3DecoderLayer],
+        DecoderBlockType.GEMMA4: get_scannable(gemma4.Gemma4DecoderLayer, gemma4.Gemma4ScannableBlock),
+        DecoderBlockType.GEMMA4_SMALL: [gemma4_small.Gemma4SmallDecoderLayer],
+        DecoderBlockType.GPT3: [gpt3.Gpt3DecoderLayer],
+        DecoderBlockType.GPT_OSS: get_scannable(gpt_oss.GptOssDecoderLayer, gpt_oss.GptOssScannableBlock),
+        DecoderBlockType.QWEN2: [qwen2.Qwen2DecoderLayer],
+        DecoderBlockType.QWEN3: [qwen3.Qwen3DecoderLayer],
+        DecoderBlockType.QWEN3_MOE: [qwen3.Qwen3MoeDecoderLayer],
+        DecoderBlockType.QWEN3_CUSTOM_MOE: [qwen3_custom.Qwen3CustomMoeDecoderLayer],
+        DecoderBlockType.QWEN3_5: get_scannable(qwen3_5.Qwen3_5DecoderLayer, qwen3_5.Qwen3_5ScannableBlock),
+        DecoderBlockType.QWEN3_NEXT: get_scannable(qwen3.Qwen3NextDecoderLayer, qwen3.Qwen3NextScannableBlock),
+        DecoderBlockType.SIMPLE: [simple_layer.SimpleDecoderLayer],
+        DecoderBlockType.SIMPLE_MLP: [simple_layer.SimpleMlpDecoderLayer],
+        DecoderBlockType.DEEPSEEK: [deepseek.DeepSeekDenseLayer, deepseek.DeepSeekMoELayer],
+        DecoderBlockType.LLAMA4: get_scannable(llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock),
+        DecoderBlockType.OLMO3: get_scannable(olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock),
+    }
+
+    if cfg.decoder_block not in layer_map:
+      raise ValueError(f"Incorrect decoder_block name {cfg.decoder_block.value=}")
+    return layer_map[cfg.decoder_block]
+
   def set_remat_policy(self, block_layers, policy):
     """Set remat policy"""
     RemattedBlockLayers = []
@@ -526,6 +572,58 @@ def map_fn(path, value):
       RemattedBlockLayers.append(layer)
     return RemattedBlockLayers
 
+  def _build_nnx_pipeline_stage(self, decoder_blocks, rngs):
+    """Creates a single NNX pipeline stage module."""
+    cfg = self.config
+    base_stage_cls = decoder_blocks[1] if cfg.decoder_block == DecoderBlockType.DEEPSEEK else decoder_blocks[0]
+
+    if cfg.num_layers_per_pipeline_stage == 1:
+      return base_stage_cls(config=cfg, mesh=self.mesh, quant=self.quant, model_mode=self.model_mode, rngs=rngs)
+    elif cfg.scan_layers_per_stage:
+      return NNXScannedPipelineStage(
+          base_stage_cls, cfg.num_layers_per_pipeline_stage, cfg, self.mesh, self.quant, self.model_mode, rngs=rngs
+      )
+    return NNXSequentialPipelineStage(
+        base_stage_cls, cfg.num_layers_per_pipeline_stage, cfg, self.mesh, self.quant, self.model_mode, rngs=rngs
+    )
+
+  def get_pipeline_stage_module(self, decoder_blocks):
+    """get pipeline stage module"""
+
+    def get_layer_to_pipeline(blocks, cfg):
+      if cfg.decoder_block == DecoderBlockType.DEEPSEEK:
+        return blocks[1]  # return the sparse block
+      else:
+        return blocks[0]
+
+    cfg = self.config
+    base_stage = get_layer_to_pipeline(decoder_blocks, cfg)
+    if cfg.set_remat_policy_on_layers_per_stage:
+      policy = self.get_remat_policy()
+      base_stage = self.set_remat_policy([base_stage], policy)[0]
+    if cfg.num_layers_per_pipeline_stage == 1:
+      stage_module = base_stage(config=cfg, mesh=self.mesh, quant=self.quant, model_mode=self.model_mode)
+    elif cfg.scan_layers_per_stage:
+      stage_module = self.scan_decoder_layers(
+          cfg,
+          base_stage,
+          cfg.num_layers_per_pipeline_stage,
+          "layers_per_stage",
+          self.mesh,
+          in_axes_tuple=(nn.broadcast,) * 4,
+          model_mode=self.model_mode,
+      )
+    else:
+      stage_module = SequentialBlockDecoderLayers(
+          decoder_layer=base_stage,
+          num_decoder_layers=cfg.num_layers_per_pipeline_stage,
+          config=cfg,
+          mesh=self.mesh,
+          quant=self.quant,
+          model_mode=self.model_mode,
+      )
+    return stage_module
+
   def get_norm_layer(self, num_features: int):
     """get normalization layer (return type inherits from nn.Module)"""
     if self.config.decoder_block in (
@@ -586,42 +684,6 @@ 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):
-    """get pipeline stage module"""
-
-    def get_layer_to_pipeline(blocks, cfg):
-      if cfg.decoder_block == DecoderBlockType.DEEPSEEK:
-        return blocks[1]  # return the sparse block
-      else:
-        return blocks[0]
-
-    cfg = self.config
-    base_stage = get_layer_to_pipeline(decoder_blocks, cfg)
-    if cfg.set_remat_policy_on_layers_per_stage:
-      policy = self.get_remat_policy()
-      base_stage = self.set_remat_policy([base_stage], policy)[0]
-    if cfg.num_layers_per_pipeline_stage == 1:
-      stage_module = base_stage(config=cfg, mesh=self.mesh, quant=self.quant, model_mode=self.model_mode)
-    elif cfg.scan_layers_per_stage:
-      stage_module = self.scan_decoder_layers(
-          cfg,
-          base_stage,
-          cfg.num_layers_per_pipeline_stage,
-          "layers_per_stage",
-          self.mesh,
-          in_axes_tuple=(nn.broadcast,) * 4,
-      )
-    else:
-      stage_module = SequentialBlockDecoderLayers(
-          decoder_layer=base_stage,
-          num_decoder_layers=cfg.num_layers_per_pipeline_stage,
-          config=cfg,
-          mesh=self.mesh,
-          quant=self.quant,
-          model_mode=self.model_mode,
-      )
-    return stage_module
-
   @nn.compact
   def _apply_embedding(
       self,
@@ -1301,11 +1363,17 @@ def _apply_gemma4_scanned_blocks(
         decoder_positions,
         deterministic,
         model_mode,
-        slot,
-        previous_chunk,
-        bidirectional_mask,
     )
 
+    # Pass slot/previous_chunk/bidirectional_mask by keyword only (via layer_call_kwargs),
+    # never positionally in broadcast_args: Gemma4DecoderLayer and Gemma4ScannableBlock
+    # declare slot and previous_chunk in swapped order, so positional passing misroutes them.
+    layer_call_kwargs = {
+        "slot": slot,
+        "previous_chunk": previous_chunk,
+        "bidirectional_mask": bidirectional_mask,
+    }
+
     if num_full_blocks > 0:
       ScannableBlockToLinen = gemma4.Gemma4ScannableBlockToLinen
       policy = self.get_remat_policy()
@@ -1335,7 +1403,7 @@ def _apply_gemma4_scanned_blocks(
           num_of_layers=block_pattern_len,
           name="scanned_blocks",
       )(
-          y, *broadcast_args
+          y, *broadcast_args, **layer_call_kwargs
       )
 
     # Process any remaining layers that don't fit into a full scanned block
@@ -1349,7 +1417,7 @@ def _apply_gemma4_scanned_blocks(
           attention_type=attention_type,
           layer_idx=layer_id,
       )
-      y = layer(y, *broadcast_args)
+      y = layer(y, *broadcast_args, **layer_call_kwargs)
       if cfg.scan_layers:
         y = y[0]