Migrate Pipeline to use NNX module.

Author: ecnal-cienet

src/MaxText/layers/decoders.py

@@ -260,6 +260,138 @@ def __call__(
       return inputs
 
 
+class SequentialNNXWrapper(nnx.Module):
+  """Wrapper that creates sequential decoder layers for pipeline stages.
+
+  This wrapper matches the decoder layer signature expected by Pipeline.
+  """
+
+  def __init__(
+      self,
+      decoder_layer_class: type,
+      num_decoder_layers: int,
+      config: Config,
+      mesh: Mesh,
+      model_mode: str,
+      rngs: nnx.Rngs,
+      quant: None | Quant = None,
+  ):
+    """Initialize wrapper with sequential decoder layers.
+
+    Args:
+      decoder_layer_class: NNX decoder layer class to instantiate
+      num_decoder_layers: Number of layers to create
+      config: Model configuration
+      mesh: Device mesh
+      model_mode: 'train', 'eval', etc.
+      rngs: RNG state
+      quant: Quantization config
+    """
+    self.sequential = SequentialBlockNNXDecoderLayers(
+        decoder_layer_class=decoder_layer_class,
+        num_decoder_layers=num_decoder_layers,
+        config=config,
+        mesh=mesh,
+        model_mode=model_mode,
+        rngs=rngs,
+        quant=quant
+    )
+
+  def __call__(self, *args, **kwargs):
+    """Forward pass through sequential layers."""
+    return self.sequential(*args, **kwargs)
+
+
+class SequentialBlockNNXDecoderLayers(nnx.Module):
+  """Sequential unscanned series of NNX decoder layers."""
+
+  def __init__(
+      self,
+      decoder_layer_class: type,
+      num_decoder_layers: int,
+      config: Config,
+      mesh: Mesh,
+      model_mode: str,
+      rngs: nnx.Rngs,
+      quant: None | Quant = None,
+  ):
+    """Initialize multiple NNX decoder layer instances.
+
+    Args:
+      decoder_layer_class: The NNX decoder layer class to instantiate
+      num_decoder_layers: Number of decoder layers to create
+      config: Model configuration
+      mesh: Device mesh for sharding
+      model_mode: 'train', 'eval', etc.
+      rngs: RNG state for initialization
+      quant: Quantization configuration
+    """
+    self.config = config
+    self.num_decoder_layers = num_decoder_layers
+
+    # Create multiple independent decoder layer instances
+    # IMPORTANT: Store as individual attributes so NNX tracks them as pytree nodes
+    # Regular Python lists are not tracked by NNX!
+    for lyr in range(num_decoder_layers):
+      layer = decoder_layer_class(
+          config=config,
+          mesh=mesh,
+          model_mode=model_mode,
+          rngs=rngs,
+          quant=quant,
+      )
+      # Store as attribute with unique name so NNX can track it
+      setattr(self, f'layer_{lyr}', layer)
+
+  def __call__(
+      self,
+      inputs: jnp.ndarray,
+      decoder_segment_ids,
+      decoder_positions,
+      deterministic: bool,
+      model_mode,
+      slot: None | int = None,
+      page_state: None | page_manager.PageState = None,
+  ) -> jnp.ndarray:
+    """Sequentially apply all decoder layers.
+
+    Args:
+      inputs: Input tensor
+      decoder_segment_ids: Segment IDs for attention masking
+      decoder_positions: Position indices
+      deterministic: Whether to use deterministic mode (no dropout)
+      model_mode: 'train', 'eval', etc.
+      slot: Optional slot index for paged attention
+      page_state: Optional page state for paged attention
+
+    Returns:
+      Output tensor after all layers, or (output, None) if scan_layers is True
+    """
+    # Iterate over layer attributes (layer_0, layer_1, ...)
+    for lyr in range(self.num_decoder_layers):
+      layer = getattr(self, f'layer_{lyr}')
+      outputs = layer(
+          inputs,
+          decoder_segment_ids,
+          decoder_positions,
+          deterministic,
+          model_mode,
+          slot=slot,
+          page_state=page_state,
+      )
+      # Handle tuple outputs (e.g., from scan_layers)
+      if self.config.scan_layers:
+        inputs = outputs[0]  # When scan_layers is True the decoder layers return (outputs, None).
+      else:
+        inputs = outputs
+
+    # Return format matching scan_layers configuration
+    if self.config.scan_layers:
+      return inputs, None
+    else:
+      return inputs
+
+
 class Decoder(nn.Module):
   """A stack of decoder layers as a part of an encoder-decoder architecture."""
 
@@ -273,10 +405,17 @@ 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)
+      # Try to get pure NNX decoder classes for pipeline parallelism
+      nnx_decoder_classes = self.get_nnx_decoder_layers()
+      if nnx_decoder_classes is not None:
+        # Use pure NNX classes for pipeline - pass the class, not instance
+        pipeline_stage_module = self.get_pipeline_stage_module(nnx_decoder_classes, use_nnx=True)
+      else:
+        # Fallback to Linen-wrapped classes
+        pipeline_stage_module = self.get_pipeline_stage_module(self.decoder_layer, use_nnx=False)
       remat_policy = self.get_remat_policy()
-      self.pipeline_module = pipeline.Pipeline(
-          config=self.config, mesh=self.mesh, layers=pipeline_stage_module, remat_policy=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)
       )
 
   def minimal_policy(self, with_context=False):
@@ -431,6 +570,63 @@ 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_layers(self):
+    """Retrieves pure NNX decoder layer classes (without Linen wrappers) for pipeline parallelism.
+
+    Returns:
+        A list containing one or more NNX Module classes for the decoder.
+    """
+    match self.config.decoder_block:
+      case DecoderBlockType.DEFAULT:
+        # DecoderLayer is Linen-only, no NNX version available
+        return None
+      case DecoderBlockType.LLAMA2:
+        return [llama2.LlamaDecoderLayer]  # Pure NNX version
+      case DecoderBlockType.MISTRAL:
+        return [mistral.MistralDecoderLayer] if hasattr(mistral, 'MistralDecoderLayer') else None
+      case DecoderBlockType.MIXTRAL:
+        return [mixtral.MixtralDecoderLayer] if hasattr(mixtral, 'MixtralDecoderLayer') else None
+      case DecoderBlockType.DEEPSEEK:
+        # DeepSeek uses specific dense/MoE layers
+        if self.config.use_batch_split_schedule:
+          return [deepseek_batchsplit.DeepSeekDenseLayer, deepseek_batchsplit.DeepSeekMoELayer]
+        else:
+          return [deepseek.DeepSeekDenseLayer, deepseek.DeepSeekMoELayer]
+      case DecoderBlockType.GEMMA:
+        return [gemma.GemmaDecoderLayer] if hasattr(gemma, 'GemmaDecoderLayer') else None
+      case DecoderBlockType.GEMMA2:
+        return [gemma2.Gemma2DecoderLayer] if hasattr(gemma2, 'Gemma2DecoderLayer') else None
+      case DecoderBlockType.GEMMA3:
+        return [gemma3.Gemma3DecoderLayer] if hasattr(gemma3, 'Gemma3DecoderLayer') else None
+      case DecoderBlockType.GPT3:
+        return [gpt3.Gpt3DecoderLayer]
+      case DecoderBlockType.GPT_OSS:
+        # Check if pure NNX version exists
+        if self.config.scan_layers:
+          return [gpt_oss.GptOssScannableBlock] if hasattr(gpt_oss, 'GptOssScannableBlock') else None
+        else:
+          return [gpt_oss.GptOssDecoderLayer] if hasattr(gpt_oss, 'GptOssDecoderLayer') else None
+      case DecoderBlockType.QWEN3:
+        return [qwen3.Qwen3DecoderLayer] if hasattr(qwen3, 'Qwen3DecoderLayer') else None
+      case DecoderBlockType.QWEN3_MOE:
+        return [qwen3.Qwen3MoeDecoderLayer] if hasattr(qwen3, 'Qwen3MoeDecoderLayer') else None
+      case DecoderBlockType.QWEN3_NEXT:
+        if self.config.scan_layers:
+          return [qwen3.Qwen3NextScannableBlock] if hasattr(qwen3, 'Qwen3NextScannableBlock') else None
+        else:
+          return [qwen3.Qwen3NextDecoderLayer] if hasattr(qwen3, 'Qwen3NextDecoderLayer') else None
+      case DecoderBlockType.SIMPLE:
+        return [simple_layer.SimpleDecoderLayer]  # Pure NNX version
+      case DecoderBlockType.SIMPLE_MLP:
+        return [simple_layer.SimpleMlpDecoderLayer]  # Pure NNX version
+      case DecoderBlockType.LLAMA4:
+        if self.config.scan_layers:
+          return [llama4.Llama4ScannableBlock] if hasattr(llama4, 'Llama4ScannableBlock') else None
+        else:
+          return [llama4.Llama4DecoderLayer] if hasattr(llama4, 'Llama4DecoderLayer') else None
+      case _:
+        return None
+
   def set_remat_policy(self, block_layers, policy):
     """Set remat policy"""
     RemattedBlockLayers = []
@@ -510,8 +706,14 @@ 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_pipeline_stage_module(self, decoder_blocks, use_nnx=False):
+    """get pipeline stage module
+
+    Args:
+      decoder_blocks: List of decoder layer classes (either Linen or NNX)
+      use_nnx: If True, decoder_blocks are NNX classes and should be passed to Pipeline
+               without instantiation. Pipeline will handle NNX instantiation with proper rngs.
+    """
 
     def get_layer_to_pipeline(blocks, cfg):
       if cfg.decoder_block == DecoderBlockType.DEEPSEEK:
@@ -521,6 +723,30 @@ def get_layer_to_pipeline(blocks, cfg):
 
     cfg = self.config
     base_stage = get_layer_to_pipeline(decoder_blocks, cfg)
+
+    # For NNX classes, return a class that Pipeline can instantiate
+    if use_nnx:
+      if cfg.num_layers_per_pipeline_stage == 1:
+        # Return the NNX class itself, Pipeline will instantiate it
+        return base_stage
+      else:
+        # For multiple layers per stage, return a partial wrapper
+        max_logging.log(
+            f"Pipeline: Creating sequential NNX wrapper with {cfg.num_layers_per_pipeline_stage} layers per stage"
+        )
+        # Return a lambda that creates the wrapper with the right parameters
+        # Pipeline will call this with (config, mesh, model_mode, rngs, quant)
+        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
+        )
+
+    # For Linen classes, instantiate as before
     if cfg.set_remat_policy_on_layers_per_stage:
       policy = self.get_remat_policy()
       base_stage = self.set_remat_policy([base_stage], policy)[0]