Merge pull request #2413 from AI-Hypercomputer:nicogrande/sft-llama4-tiling

PiperOrigin-RevId: 816013572

Author: Google-ML-Automation

src/MaxText/decode.py

@@ -99,7 +99,7 @@ def main(argv: Sequence[str]) -> None:
 
   text = config.prompt
   prefill_length = config.max_prefill_predict_length
-  # processor_output = multimodal_utils.PreprocessorOutput()
+  processor_outputs = multimodal_utils.PreprocessorOutput()
   if config.use_multimodal:
     image_path = config.image_path.split(",")
     images = [multimodal_utils.load_image_from_path(p) for p in image_path]
@@ -151,6 +151,9 @@ def main(argv: Sequence[str]) -> None:
           params=params,
           padded_tokens=tokens,
           images=np.stack([po.pixel_values for po in processor_outputs]) if config.use_multimodal else None,
+          image_masks=np.stack([po.pixel_mask for po in processor_outputs])
+          if config.use_multimodal and "llama4" in config.model_name
+          else None,
           true_length=true_length,
           rng=rng_prefill,
           slot=i,

src/MaxText/input_pipeline/_input_pipeline_utils.py

@@ -106,8 +106,10 @@ def prepare_text_for_image_fusion(example, column_name, model_name):
       example[column_name], model_name, processor_output=example["images"]
   )
   if isinstance(example["images"], list):
+    example["image_masks"] = [image.pixel_mask for image in example["images"]]
     example["images"] = [image.pixel_values for image in example["images"]]
   else:
+    example["image_masks"] = example["images"].pixel_mask
     example["images"] = example["images"].pixel_values
   return example
 
@@ -255,6 +257,7 @@ def map(self, element):
         "inputs": np.asarray(inputs[: self.max_target_length], dtype=np.int32),
         "targets": np.asarray(targets[: self.max_target_length], dtype=np.int32),
         "images": element["images"],
+        "image_masks": element["image_masks"],
     }
 
 
@@ -434,49 +437,77 @@ def _pad_text(self, x, max_length, pad_id):
     pad_amount = [(0, pad_amount)] + [(0, 0)] * (len(x.shape) - 1)
     return np.pad(x, pad_amount, constant_values=pad_id)[: self.max_length]
 
-  def _pad_image(self, images):
-    """Pads the input images array to match the maximum required number of images per example.
+  def _pad_image_or_mask(self, tensor: np.ndarray) -> np.ndarray:
+    """Pads an input tensor (image or mask) to a maximum number of items.
 
-    The function computes the maximum number of allowed images either based on model constraints
-    (determined by model name and maximum sequence length) or the value provided in
-    `self.max_num_images_per_example` (if positive), whichever is smaller. If the provided
-    `images` array contains fewer images than the desired maximum, it pads the array with
-    zeros (dummy images) to meet the requirement.
+    This function unifies padding logic for image tensors (standard or tiled) and
+    mask tensors. It determines the tensor type based on its dimensions and applies
+    the appropriate padding along the first axis.
+
+    The maximum number of items is calculated based on model constraints or a
+    user-defined limit, ensuring that sequence length limits are respected while
+    reserving space for at least one text token. If the input tensor has fewer
+    items than this maximum, it is padded with zeros.
 
     Args:
-      images (np.ndarray): A numpy array of image data. The expected shape is
-          (num_images, height, width, channels), where num_images <= maximum allowed images.
+        tensor (np.ndarray): The input numpy array to pad.
+            - For masks, the expected shape is (num_masks, num_tiles).
+            - For standard images, the shape is (num_images, H, W, C).
+            - For tiled images, the shape is (num_images, num_tiles, H, W, C).
 
     Returns:
-      np.ndarray: An array of images padded with zero images (if necessary) such that its
-      first dimension equals the computed maximum number of images.
+        np.ndarray: The tensor, padded with zeros up to the maximum number of
+        items along the first axis.
 
     Raises:
-      AssertionError: If the number of images in the input exceeds the allowed maximum.
+        ValueError: If the input tensor's dimension is not 2, 4, or 5.
+        ValueError: If the number of items in the input tensor exceeds the
+        allowed maximum.
 
     Notes:
       - The computation of maximum images ensures that space is reserved in the sequence
         for at least one text token.
       - The dummy images used for padding are based on the image shape for initialization
         of this model (ignoring batch size).
     """
+    if not isinstance(tensor, np.ndarray):
+      raise TypeError(f"Input must be a numpy array, but got {type(tensor)}")
+
+    # Determine the maximum number of images/masks allowed.
     image_offsets = multimodal_utils.get_image_offsets(self.model_name, None)
-    max_num_images = (self.max_length // image_offsets) - 1  # -1 to reserve space for at least one text token
+    # Reserve space for at least one text token.
+    max_num_items = (self.max_length // image_offsets) - 1
     if self.max_num_images_per_example > 0:
-      max_num_images = min(self.max_num_images_per_example, max_num_images)
-    image_shape = multimodal_utils.get_dummy_image_shape_for_init(self.model_name)[2:]
-    assert (
-        images.shape[0] <= max_num_images
-    ), f"Number of images {images.shape[0]} exceeds the maximum allowed {max_num_images}"
-    if images.shape[0] < max_num_images:
-      pad_size = max_num_images - images.shape[0]
-      pad_shape = (pad_size,) + image_shape
-      pad_images = np.zeros(pad_shape, dtype=images.dtype)
-      if images is not None and images.size > 0:
-        images = np.concatenate([images, pad_images], axis=0)
+      max_num_items = min(self.max_num_images_per_example, max_num_items)
+
+    # Validate tensor dimensions.
+    if tensor.ndim in (4, 5):  # Standard or Tiled Image
+      tensor_type = "images"
+    elif tensor.ndim == 2:  # Mask
+      tensor_type = "masks"
+    else:
+      raise ValueError(
+          "Input tensor must be 2D (mask), 4D (image), or 5D (tiled image), " f"but got {tensor.ndim} dimensions."
+      )
+
+    # Assert that the input tensor does not exceed the maximum size.
+    if tensor.shape[0] > max_num_items:
+      raise ValueError(f"Number of {tensor_type} ({tensor.shape[0]}) exceeds the maximum allowed ({max_num_items}).")
+
+    # Apply padding if the tensor is smaller than the maximum size.
+    if tensor.shape[0] < max_num_items:
+      pad_size = max_num_items - tensor.shape[0]
+      pad_shape_suffix = tensor.shape[1:]
+      pad_shape = (pad_size,) + pad_shape_suffix
+      pad_tensor = np.zeros(pad_shape, dtype=tensor.dtype)
+
+      if tensor.size > 0:
+        tensor = np.concatenate([tensor, pad_tensor], axis=0)
       else:
-        images = pad_images
-    return images
+        # If the input tensor is empty, the result is just the padding.
+        tensor = pad_tensor
+
+    return tensor
 
   def map(self, element: dict[str, np.ndarray]):
     """map to each element"""
@@ -487,13 +518,27 @@ def map(self, element: dict[str, np.ndarray]):
         element[f"{data_column}_position"] = np.arange(element[data_column].shape[0], dtype=np.int32)
         if self.add_true_length:
           element[f"{data_column}_true_length"] = np.array([element[data_column].shape[0]], dtype=np.int32)
+
     for key, _ in element.items():
       if key == "images":
-        if isinstance(element["images"], list):
-          assert self.model_name is not None, "model_name must be provided when padding images"
-          element["images"] = self._pad_image(np.asarray(element["images"]))
-        else:
+        if isinstance(element["images"], list) and self.model_name is None:
+          raise ValueError("model_name must be provided when padding images")
+        elif isinstance(element["images"], list):
+          element["images"] = self._pad_image_or_mask(np.asarray(element["images"]))
+        elif element["images"].ndim == 3:
           element["images"] = np.asarray(element["images"])[None, ...]
+        else:
+          # Do not add extra image dimension for image tiling case
+          element["images"] = np.asarray(element["images"])
+
+      elif key == "image_masks" and element["image_masks"] is not None:
+        if isinstance(element["image_masks"], list) and self.model_name is None:
+          raise ValueError("model_name must be provided when padding image masks")
+        elif isinstance(element["image_masks"], list):
+          element["image_masks"] = self._pad_image_or_mask(np.asarray(element["image_masks"]))
+        else:
+          element["image_masks"] = np.asarray(element["image_masks"])
+
       elif "true_length" not in key:
         element[key] = self._pad_text(element[key], self.max_length, self.pad_id)
     return element

src/MaxText/layers/decoders.py

@@ -286,6 +286,7 @@ def get_remat_policy(self):
         # save all
         if cfg.remat_policy == "minimal_flash":
           max_logging.log("WARNING: 'minimal_flash' will be deprecated soon, please use 'minimal_with_context' instead.")
+          max_logging.log("WARNING: 'minimal_flash' will be deprecated soon, please use 'minimal_with_context' instead.")
         policy = self.minimal_policy(with_context=True)
       elif cfg.remat_policy == "minimal":
         # save all except context
@@ -528,6 +529,7 @@ def _apply_embedding(
       model_mode,
       image_embeddings=None,
       bidirectional_mask=None,
+      image_masks=None,
   ):
     """Applies token and positional embeddings to the input tokens."""
     cfg = self.config
@@ -541,6 +543,7 @@ def _apply_embedding(
             text_embeddings=y,
             vision_embeddings=image_embeddings,
             mask=bidirectional_mask,
+            image_masks=image_masks,
         )
       # TODO(hengtaoguo): Add support for other multimodal models such as Llama4, refactor if needed
       else:
@@ -635,6 +638,7 @@ def __call__(
       page_state: None | page_manager.PageState = None,
       bidirectional_mask: None | Any = None,
       image_embeddings: None | jnp.ndarray = None,
+      image_masks: None | jnp.ndarray = None,
   ):
     cfg = self.config
     mesh = self.mesh
@@ -649,6 +653,7 @@ def __call__(
         model_mode,
         image_embeddings,
         bidirectional_mask,
+        image_masks,
     )
 
     policy = self.get_remat_policy()

src/MaxText/layers/llama4.py

@@ -295,8 +295,12 @@ def __call__(self, image_features: Array) -> Array:
       Tensor of shape [batch_size, num_patches, (pixel_shuffle_ratio**2), vision_hidden_size]
     """
     b, t, c, d = image_features.shape
+
+    # Reshape image_features to [b * t, c, d] and project to text hidden dimension
     image_features = image_features.reshape(b * t, c, d)
     hidden_states = self.linear(image_features)
+
+    # Reshape hidden states back to [b, t, c, d]
     _, c, d = hidden_states.shape
     hidden_states = hidden_states.reshape(b, t, c, d)
     return hidden_states
@@ -754,15 +758,18 @@ def __call__(
     """Forward pass of the Llama4 vision model.
 
     Args:
-      inputs: Input tensor of shape [batch_size, num_tiles, num_channels_for_vit, tile_size_for_vit, tile_size_for_vit]
+      inputs: Input tensor of shape:
+              [batch_size * num_images, num_tiles, num_channels_for_vit, tile_size_for_vit, tile_size_for_vit]
       deterministic: Whether to use deterministic mode (disables dropout)
 
     Returns:
-      Final hidden states from the vision encoder of shape [batch_size, num_tiles, num_patches, vision_output_dim_for_vit]
+      Final hidden states from the vision encoder of shape:
+      [batch_size * num_images, num_tiles, num_patches, vision_output_dim_for_vit]
     """
     cfg = self.config
     mesh = self.mesh
 
+    # Reshape pixel values to combine batch and num_tiles dimensions
     b, t, c, h, w = pixel_values.shape
     pixel_values = jnp.reshape(pixel_values, [b * t, c, h, w])
 

src/MaxText/layers/models.py

@@ -116,6 +116,7 @@ def __call__(
       decoder_positions: jnp.ndarray,
       decoder_segment_ids=None,
       encoder_images: None | jnp.ndarray = None,
+      encoder_image_masks: None | jnp.ndarray = None,
       enable_dropout=True,
       model_mode=MODEL_MODE_TRAIN,
       previous_chunk=None,
@@ -162,6 +163,7 @@ def __call__(
         page_state=page_state,
         bidirectional_mask=bidirectional_mask,
         image_embeddings=image_embeddings,
+        image_masks=encoder_image_masks,
     )
 
     # If we are initializing the model AND MTP is enabled, we must create
@@ -353,6 +355,7 @@ def __call__(
       decoder_segment_ids=None,
       cache=None,
       encoder_images: jax.Array | None = None,
+      encoder_image_masks: jax.Array | None = None,
       enable_dropout=True,
       model_mode=MODEL_MODE_TRAIN,
       previous_chunk=None,
@@ -413,6 +416,7 @@ def __call__(
         page_state=page_state,
         bidirectional_mask=bidirectional_mask,
         image_embeddings=image_embeddings,
+        image_masks=encoder_image_masks,
     )
 
     # Materialize hidden state when vocab tiling is enabled
@@ -493,6 +497,7 @@ def __call__(
       decoder_positions: jnp.ndarray,
       decoder_segment_ids=None,
       encoder_images: None | jnp.ndarray = None,
+      encoder_image_masks: None | jnp.ndarray = None,
       enable_dropout=True,
       model_mode=MODEL_MODE_TRAIN,
       previous_chunk=None,
@@ -535,6 +540,7 @@ def __call__(
           decoder_positions=decoder_positions,
           decoder_segment_ids=decoder_segment_ids,
           encoder_images=encoder_images,
+          encoder_image_masks=encoder_image_masks,
           enable_dropout=enable_dropout,
           model_mode=model_mode,
           previous_chunk=previous_chunk,
@@ -552,6 +558,7 @@ def __call__(
         decoder_positions=decoder_positions,
         decoder_segment_ids=decoder_segment_ids,
         encoder_images=encoder_images,
+        encoder_image_masks=encoder_image_masks,
         enable_dropout=enable_dropout,
         model_mode=model_mode,
         previous_chunk=previous_chunk,

src/MaxText/maxengine.py

@@ -327,17 +327,17 @@ def quantize_params(self, state, rng: PRNGKeyType | None = None):
 
     @jax.jit
     def model_apply(_p, _rng):
+      image_shape = multimodal_utils.get_dummy_image_shape_for_init(
+          self.config.model_name, batch_size=self.config.micro_batch_size_to_train_on
+      )
       return self.model.apply(
           _p | {"aqt": {}},
           jnp.ones((1, self.config.max_prefill_predict_length), dtype=jnp.int32),
           jnp.ones((1, self.config.max_prefill_predict_length), dtype=jnp.int32),
-          encoder_images=jnp.ones(
-              multimodal_utils.get_dummy_image_shape_for_init(
-                  self.config.model_name, batch_size=self.config.micro_batch_size_to_train_on
-              ),
-              dtype=jnp.float32,
-          )
-          if self.config.use_multimodal
+          encoder_images=jnp.ones(image_shape, dtype=jnp.float32) if self.config.use_multimodal else None,
+          # encoder_image_masks indicates valid tiles if image tiling + padding is used in vision encoder input.
+          encoder_image_masks=jnp.ones(image_shape[:2], dtype=jnp.int32)
+          if self.config.use_multimodal and "llama4" in self.config.model_name
           else None,
           decoder_segment_ids=jnp.zeros((1, self.config.max_prefill_predict_length), dtype=jnp.int32),
           enable_dropout=False,
@@ -412,6 +412,7 @@ def _prefill_jit(
       existing_prefix: ExistingPrefix | None = None,
       padded_tokens: jax.Array,
       images: jax.Array | None = None,
+      image_masks: jax.Array | None = None,
       true_length: int,
       sampler: Callable[[Any], Any] | None = None,  # pylint: disable=unused-argument
       rng: PRNGKeyType | None = None,
@@ -479,9 +480,11 @@ def _prefill_jit(
       if images.ndim == 3:
         # For Gemma3 single image, add batch and image count dimensions
         images = images[jnp.newaxis, jnp.newaxis, ...]
+        image_masks = image_masks[jnp.newaxis, jnp.newaxis, ...] if image_masks is not None else None
       elif images.ndim == 4:
         # add batch dimension
         images = images[jnp.newaxis, ...]
+        image_masks = image_masks[jnp.newaxis, ...] if image_masks is not None else None
 
     # sequence_indicator will be concatenated to existing_prefix decoder_segment_ids
     start_to_n = jnp.arange(start_position, start_position + input_tokens.shape[1])
@@ -496,6 +499,7 @@ def _prefill_jit(
           input_tokens,
           positions,
           encoder_images=images,
+          encoder_image_masks=image_masks,
           decoder_segment_ids=sequence_indicator,
           enable_dropout=False,
           model_mode=MODEL_MODE_PREFILL,
@@ -570,6 +574,7 @@ def prefill(
       existing_prefix: ExistingPrefix | None = None,
       padded_tokens: jax.Array,
       images: jax.Array | None = None,
+      image_masks: jax.Array | None = None,
       true_length: int,
       sampler: Callable[[Any], Any] | None = None,  # pylint: disable=unused-argument
       rng: PRNGKeyType | None = None,
@@ -602,6 +607,7 @@ def prefill(
         existing_prefix=existing_prefix,
         padded_tokens=padded_tokens,
         images=images,
+        image_masks=image_masks,
         sampler=sampler,
         true_length=true_length,
         page_state=self.page_state,  # Pass current page state

src/MaxText/maxtext_utils.py

@@ -91,6 +91,7 @@ def get_shaped_batch(config):
         config.model_name, batch_size=config.micro_batch_size_to_train_on
     )
     shaped_batch["images"] = jax.ShapeDtypeStruct(image_shape, jnp.int32)
+    shaped_batch["image_masks"] = jax.ShapeDtypeStruct(image_shape[:2], jnp.int32)
   return shaped_batch