Correcting the params

Author: parambole

MaxText/layers/models.py

@@ -48,60 +48,63 @@
 # ------------------------------------------------------------------------------
 
 
-def get_decoder_layers(self):
-  """Get decoder layers, one of `DecoderBlockType` discriminants or a direct `nn.Module` inheritor"""
-  if self.config.decoder_block == DecoderBlockType.DEFAULT:
+def get_decoder_layers(config: Config):
+  """
+  Helper function to get the list of decoder layer classes based on config.
+  Get decoder layers, one of `DecoderBlockType` discriminants or a direct `nn.Module` inheritor
+  """
+  if config.decoder_block == DecoderBlockType.DEFAULT:
     return [DecoderLayer]
-  elif self.config.decoder_block == DecoderBlockType.LLAMA2:
+  elif config.decoder_block == DecoderBlockType.LLAMA2:
     from MaxText.layers import llama2  # pylint: disable=import-outside-toplevel
 
     return [llama2.LlamaDecoderLayer]
-  elif self.config.decoder_block == DecoderBlockType.MISTRAL:
+  elif config.decoder_block == DecoderBlockType.MISTRAL:
     # TODO(ranran): update to Mistral with sliding window attention
     from MaxText.layers import mistral  # pylint: disable=import-outside-toplevel
 
     return [mistral.MistralDecoderLayer]
-  elif self.config.decoder_block == DecoderBlockType.MIXTRAL:
+  elif config.decoder_block == DecoderBlockType.MIXTRAL:
     from MaxText.layers import mixtral  # pylint: disable=import-outside-toplevel
 
     return [mixtral.MixtralDecoderLayer]
-  elif self.config.decoder_block == DecoderBlockType.DEEPSEEK:
+  elif config.decoder_block == DecoderBlockType.DEEPSEEK:
     from MaxText.layers import deepseek  # pylint: disable=import-outside-toplevel
 
     return [deepseek.DeepSeekDenseLayer, deepseek.DeepSeekMoELayer]
-  elif self.config.decoder_block == DecoderBlockType.GEMMA:
+  elif config.decoder_block == DecoderBlockType.GEMMA:
     from MaxText.layers import gemma  # pylint: disable=import-outside-toplevel
 
     return [gemma.GemmaDecoderLayer]
-  elif self.config.decoder_block == DecoderBlockType.GEMMA2:
+  elif config.decoder_block == DecoderBlockType.GEMMA2:
     from MaxText.layers import gemma2  # pylint: disable=import-outside-toplevel
 
     return [gemma2.Gemma2DecoderLayer]
-  elif self.config.decoder_block == DecoderBlockType.GEMMA3:
+  elif config.decoder_block == DecoderBlockType.GEMMA3:
     from MaxText.layers import gemma3  # pylint: disable=import-outside-toplevel
 
     return [gemma3.Gemma3DecoderLayer]
-  elif self.config.decoder_block == DecoderBlockType.GPT3:
+  elif config.decoder_block == DecoderBlockType.GPT3:
     from MaxText.layers import gpt3  # pylint: disable=import-outside-toplevel
 
     return [gpt3.Gpt3DecoderLayer]
-  elif self.config.decoder_block == DecoderBlockType.SIMPLE:
+  elif config.decoder_block == DecoderBlockType.SIMPLE:
     from MaxText.layers import simple_layer  # pylint: disable=import-outside-toplevel
 
     return [simple_layer.SimpleDecoderLayer]
-  elif self.config.decoder_block == DecoderBlockType.SIMPLE_MLP:
+  elif config.decoder_block == DecoderBlockType.SIMPLE_MLP:
     from MaxText.layers import simple_layer  # pylint: disable=import-outside-toplevel
 
     return [simple_layer.SimpleMlpDecoderLayer]
-  elif self.config.decoder_block == DecoderBlockType.LLAMA4:
+  elif config.decoder_block == DecoderBlockType.LLAMA4:
     from MaxText.layers import llama4  # pylint: disable=import-outside-toplevel
 
-    if self.config.scan_layers:
+    if config.scan_layers:
       return [llama4.Llama4ScannableBlock]
     else:
       return [llama4.Llama4DecoderLayer]
   else:
-    raise ValueError(f"Incorrect decoder_block name {self.config.decoder_block.value=}")
+    raise ValueError(f"Incorrect decoder_block name {config.decoder_block.value=}")
 
 
 class SequentialBlockDecoderLayers(nn.Module):
@@ -152,7 +155,7 @@ class Decoder(nn.Module):
 
   def setup(self):
     """Initialize decoder layer."""
-    self.decoder_layer = get_decoder_layers()
+    self.decoder_layer = get_decoder_layers(self.config)
     self.norm_layer = self.get_norm_layer()
     if self.config.using_pipeline_parallelism:
       pipeline_stage_module = self.get_pipeline_stage_module(self.decoder_layer)