Merge pull request #2769 from AI-Hypercomputer:hengtaoguo-nnx-vis

PiperOrigin-RevId: 839920093

Author: Google-ML-Automation

src/MaxText/layers/encoders.py

@@ -15,55 +15,78 @@
 """"Module for encoder layers."""
 
 import jax
-from flax import linen as nn
+from flax import nnx
 from jax.sharding import Mesh
 
 from MaxText.common_types import Config
-from MaxText.layers import quantizations
+from MaxText.layers import nnx_wrappers
+from MaxText.layers import initializers
 
-# Type alias for cleaner type hints
-Quant = quantizations.AqtQuantization
 
-
-class VisionEncoder(nn.Module):
+class VisionEncoder(nnx.Module):
   """Vision encoder to encode images into soft tokens."""
 
-  config: Config
-  mesh: Mesh
-
-  def setup(self):
-    self.vision_encoder_layer = self.get_vision_encoder_layers()
+  def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs):
+    self.config = config
+    self.mesh = mesh
+    self.rngs = rngs
+    self.encoder_name, self.projector_name = self._setup_vision_encoder_layers()
 
-  def get_vision_encoder_layers(self):
-    """Get vision encoder layers specific to the model, classes of nn.Module type."""
+  def _setup_vision_encoder_layers(self):
+    """Setup vision encoder layers specific to the model, instantiate NNX modules."""
     if self.config.model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]:
       from MaxText.layers import gemma3  # pylint: disable=import-outside-toplevel
 
-      return [gemma3.gemma3visionencoder_as_linen, gemma3.visionembedder_as_linen]
+      encoder_name = "Gemma3VisionEncoderLayer_0"
+      projector_name = "VisionEmbedder_0"
+      setattr(self, encoder_name, gemma3.Gemma3VisionEncoderLayer(config=self.config, mesh=self.mesh, rngs=self.rngs))
+      setattr(self, projector_name, gemma3.VisionEmbedder(config=self.config, mesh=self.mesh, rngs=self.rngs))
+      return encoder_name, projector_name
     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_as_linen, llama4.llama4multimodalprojector_as_linen]
+      encoder_name = "Llama4VisionModel_0"
+      projector_name = "Llama4MultiModalProjector_0"
+      setattr(self, encoder_name, llama4.Llama4VisionModel(config=self.config, mesh=self.mesh, rngs=self.rngs))
+      setattr(self, projector_name, llama4.Llama4MultiModalProjector(config=self.config, mesh=self.mesh, rngs=self.rngs))
+      return encoder_name, projector_name
     elif self.config.model_name in ["qwen3-omni-30b-a3b"]:
       from MaxText.layers import qwen3  # pylint: disable=import-outside-toplevel
 
-      return [qwen3.qwen3omni_visionencoder_as_linen, qwen3.qwen3omni_visionprojector_as_linen]
+      encoder_name = "Qwen3OmniMoeVisionEncoder_0"
+      projector_name = "Qwen3OmniMoeVisionProjector_0"
+      setattr(self, encoder_name, qwen3.Qwen3OmniMoeVisionEncoder(config=self.config, mesh=self.mesh, rngs=self.rngs))
+      setattr(self, projector_name, qwen3.Qwen3OmniMoeVisionProjector(config=self.config, rngs=self.rngs))
+      return encoder_name, projector_name
     else:
       raise ValueError(f"No VisionEncoder implemented for {self.config.model_name} yet")
 
-  @nn.compact
   def __call__(self, input_images, deterministic=False):
-    cfg = self.config
-    mesh = self.mesh
     # vision encoder output, frozen params in many cases
-    embeddings = self.vision_encoder_layer[0](config=cfg, mesh=mesh)(input_images, deterministic=deterministic)
-    if cfg.model_name in ["qwen3-omni-30b-a3b"]:
-      embeddings = embeddings[0]  # todo(eitanporat) add deepstack support
+    encoder = getattr(self, self.encoder_name)
+    embeddings = encoder(input_images, deterministic=deterministic)
 
-    if cfg.freeze_vision_encoder_params:
+    if self.config.freeze_vision_encoder_params:
       embeddings = jax.lax.stop_gradient(embeddings)
 
-    if len(self.vision_encoder_layer) > 1:
-      # vision embedder / projection layer, not frozen in most cases, trained / finetuned together with main model
-      embeddings = self.vision_encoder_layer[1](config=cfg, mesh=mesh)(embeddings)
+    # vision embedder / projection layer, not frozen in most cases, trained / finetuned together with main model
+    projector = getattr(self, self.projector_name)
+    embeddings = projector(embeddings)
+
     return embeddings
+
+
+def vision_encoder_as_linen(
+    config: Config,
+    mesh: Mesh,
+):
+  """Creates a VisionEncoder module."""
+  module = nnx_wrappers.to_linen(
+      VisionEncoder,
+      config=config,
+      mesh=mesh,
+      name="vision_encoder",
+      abstract_init=False,
+      metadata_fn=initializers.variable_to_logically_partitioned,
+  )
+  return module

src/MaxText/layers/models.py

@@ -32,7 +32,7 @@
 from MaxText.layers import nnx_wrappers
 from MaxText.layers.decoders import Decoder
 from MaxText.layers.embeddings import Embed, embed_as_linen
-from MaxText.layers.encoders import VisionEncoder
+from MaxText.layers.encoders import VisionEncoder, vision_encoder_as_linen
 from MaxText.layers.quantizations import AqtQuantization as Quant
 from MaxText.layers.multi_token_prediction import MultiTokenPredictionBlock
 from MaxText.sharding import all_gather_over_fsdp
@@ -85,7 +85,7 @@ def setup(self):
         config=cfg,
         mesh=self.mesh,
     )
-    self.vision_encoder = VisionEncoder(config=cfg, mesh=mesh) if cfg.use_multimodal else None
+    self.vision_encoder = vision_encoder_as_linen(config=cfg, mesh=mesh) if cfg.use_multimodal else None
     self.decoder = Decoder(config=cfg, mesh=mesh, quant=self.quant, model_mode=self.model_mode)
     # If MTP is enabled via config, set up the MTP block.
     if self.config.mtp_num_layers > 0:
@@ -304,7 +304,7 @@ def __init__(self, config: Config, mesh: Mesh, quant: Quant, *, model_mode: str
         config=cfg,
         rngs=rngs,
     )
-    self.vision_encoder = VisionEncoder(config=cfg, mesh=mesh) if cfg.use_multimodal else None
+    self.vision_encoder = VisionEncoder(config=cfg, mesh=mesh, rngs=rngs) if cfg.use_multimodal else None
 
     decoder_linen = Decoder(config=cfg, mesh=mesh, quant=self.quant, model_mode=self.model_mode)
     self.decoder = nnx_wrappers.ToNNX(decoder_linen, rngs=rngs)

tests/integration_tests/vision_encoder_test.py

@@ -84,7 +84,7 @@ def test_image_embedding_gemma3_4b_tpu(self):
     input_images = images[jnp.newaxis, jnp.newaxis, ...]  # pytype: disable=unsupported-operands
 
     # Initialize only the vision encoder part and extract the corresponding params
-    vision_encoder_model = models.VisionEncoder(config)
+    vision_encoder_model = models.VisionEncoder(config, engine.mesh, rngs=engine.rng)
     vision_encoder_params = params["params"]["vision_encoder"]
 
     # Apply the vision encoder to get the image embeddings