Llama4 vision layers NNX migration

minor fix

rngs fix

Author: hengtaoguo

src/MaxText/configs/models/llama4-17b-128e.yml

@@ -41,4 +41,13 @@ temperature_tuning: True
 # Chunk attention is used on all RoPE layers
 # otherwise, on NoPE layers, use global attention
 chunk_attn_window_size: 8192
-image_size_for_vit: 336 
+
+# Multimodal flags (need to set use_multimodal=true)
+image_size_for_vit: 336
+num_channels_for_vit: 3
+patch_size_for_vit: 14
+hidden_size_for_vit: 1408
+intermediate_size_for_vit: 5632
+num_hidden_layers_for_vit: 34
+num_attention_heads_for_vit: 16
+image_placeholder: "<|image|>"

src/MaxText/configs/models/llama4-17b-16e.yml

@@ -41,4 +41,13 @@ temperature_tuning: True
 # Chunk attention is used on all RoPE layers
 # otherwise, on NoPE layers, use global attention
 chunk_attn_window_size: 8192
-image_size_for_vit: 336 
+
+# Multimodal flags (need to set use_multimodal=true)
+image_size_for_vit: 336
+num_channels_for_vit: 3
+patch_size_for_vit: 14
+hidden_size_for_vit: 1408
+intermediate_size_for_vit: 5632
+num_hidden_layers_for_vit: 34
+num_attention_heads_for_vit: 16
+image_placeholder: "<|image|>"

src/MaxText/layers/encoders.py

@@ -43,7 +43,7 @@ def get_vision_encoder_layers(self):
     elif self.config.model_name in ["llama4-17b-16e", "llama4-17b-128e"]:
       from MaxText.layers import llama4  # pylint: disable=import-outside-toplevel
 
-      return [llama4.Llama4VisionModel, llama4.Llama4MultiModalProjector]
+      return [llama4.llama4visionmodel_as_linen, llama4.llama4multimodalprojector_as_linen]
     else:
       raise ValueError(f"No VisionEncoder implemented for {self.config.model_name} yet")