End2End Video SFT with Varied Size/Duration

src/maxtext/configs/post_train/sft-vision-llava-video-178k.yml

@@ -0,0 +1,23 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+base_config: "sft-vision-chartqa.yml"
+
+# Dataset
+hf_path: "parquet"
+hf_train_files: "/mounted/LLaVA-Video-178K/0_30_s_academic_v0_1/*.parquet"
+train_data_columns: ["query", "response"]
+train_image_column: "video"
+eval_data_columns: ["query", "response"]
+eval_image_column: "video"

src/maxtext/input_pipeline/hf_data_processing.py

@@ -42,6 +42,21 @@ def _get_pad_id(tokenizer):
   return pad_id
 
 
+def format_sharegpt(example):
+  """Convert a ShareGPT conversation into query and response fields."""
+  query = ""
+  response = ""
+  for turn in example["conversations"]:
+    if turn["from"] == "human":
+      query = turn["value"]
+    elif turn["from"] == "gpt":
+      response = turn["value"]
+      break
+  example["query"] = query
+  example["response"] = response
+  return example
+
+
 def vision_sft_preprocessing_pipeline(
     dataset,
     config,
@@ -54,6 +69,9 @@ def vision_sft_preprocessing_pipeline(
 ):
   """pipeline for multimodal SFT with HF dataset"""
 
+  if "conversations" in dataset.features and len(text_columns) == 2 and text_columns[0] not in dataset.features:
+    dataset = dataset.map(format_sharegpt)
+
   assert len(text_columns) == 2, f"Need two text_columns for query and response, received {text_columns=}"
   # Tunix GA requires per-micro-batch slicing at the data level,
   # whereas Native GA processes the full batch and splits it internally.
@@ -99,6 +117,7 @@ def vision_sft_preprocessing_pipeline(
           "column": text_columns[0],
           "image_placeholder": config.image_placeholder,
           "model_name": config.model_name,
+          "video_placeholder": getattr(config, "video_placeholder", "<|video|>"),
       },
   )
   dataset = dataset.map(
@@ -181,7 +200,7 @@ def vision_sft_preprocessing_pipeline(
       data_source=dataset,
       operations=operations,
       sampler=dummy_index_sampler,
-      worker_count=1,  # only supports <=1 for now, more workers results in duplicated data
+      worker_count=1,
       worker_buffer_size=1,
       read_options=grain.ReadOptions(num_threads=1, prefetch_buffer_size=batch_size * 4),
   )

src/maxtext/input_pipeline/input_pipeline_utils.py

@@ -90,19 +90,56 @@ def _process_string(string_tensor):
 ########## Functions used by HF pipeline
 
 
-def reformat_prompt(example, column, image_placeholder, model_name):
+def is_video_file(path):
+  if not isinstance(path, str):
+    return False
+  video_extensions = (".mp4", ".avi", ".mkv", ".webm", ".mov", ".gif")
+  return path.lower().endswith(video_extensions)
+
+
+def reformat_prompt(
+    example,
+    column,
+    image_placeholder,
+    model_name,
+    video_placeholder="<|video|>",
+    num_videos=0,
+    num_image_tokens=None,
+    num_video_tokens=None,
+):
   """reformat prompt for multimodal SFT"""
   if isinstance(example["images"], list):
-    num_images = len(example["images"])
+    num_images = sum(1 for img in example["images"] if not is_video_file(img))
+    num_videos_in_example = sum(1 for img in example["images"] if is_video_file(img))
+    if num_videos_in_example > 0:
+      num_videos = num_videos_in_example
   else:
-    num_images = 1
-  example[column] = mm_processor.reformat_prompt(example[column], image_placeholder, model_name, num_images)
+    if is_video_file(example["images"]):
+      num_images = 0
+      num_videos = 1
+    else:
+      num_images = 1
+      num_videos = 0
+  example[column] = mm_processor.reformat_prompt(
+      example[column],
+      image_placeholder=image_placeholder,
+      model_name=model_name,
+      num_images=num_images,
+      video_placeholder=video_placeholder,
+      num_videos=num_videos,
+      num_image_tokens=num_image_tokens,
+      num_video_tokens=num_video_tokens,
+  )
   return example
 
 
 def reformat_response(example, column, model_name):
   """reformat response for multimodal SFT"""
-  example[column] = mm_processor.reformat_response(example[column][0], model_name)
+  if isinstance(example[column], (list, np.ndarray)):
+    response = example[column][0]
+  else:
+    response = example[column]
+  example[column] = mm_processor.reformat_response(response, model_name)
   return example
 
 
@@ -123,12 +160,16 @@ def pre_process_image_sft(example, image_column, config):
   """pre-process image for multimodal SFT"""
 
   def _process_image_fn(image):
-    if isinstance(image, list):
+    if isinstance(image, str) and is_video_file(image):
+      image = mm_processor.preprocess_image_for_training(image, config)
+    elif isinstance(image, list) and len(image) > 0 and isinstance(image[0], str) and is_video_file(image[0]):
+      image = mm_processor.preprocess_image_for_training(image, config)
+    elif isinstance(image, list):
       image = [np.array(mm_utils.convert_to_RGB(img)) for img in image]
+      image = mm_processor.preprocess_image_for_training(image, config)
     else:
       image = np.array(mm_utils.convert_to_RGB(image))
-
-    image = mm_processor.preprocess_image_for_training(image, config)
+      image = mm_processor.preprocess_image_for_training(image, config)
     return image
 
   example[image_column] = _process_image_fn(example[image_column])
@@ -453,7 +494,10 @@ def __getitem__(self, index):
     The next item in the iterator is returned."""
     if not self.data_iters:
       self.data_iters = [iter(x) for x in self.datasets]
-    idx = int(current_thread().name.split("_")[1])
+    try:
+      idx = int(current_thread().name.split("_")[1])
+    except (IndexError, ValueError):
+      idx = 0
 
     while True:
       try:
@@ -777,6 +821,9 @@ def _pad_image_and_mask(self, preprocessed_image: mm_utils.PreprocessorOutput) -
     if not isinstance(preprocessed_image, mm_utils.PreprocessorOutput):
       raise TypeError(f"Input must be multimodal_utils.PreprocessorOutput, but got {type(preprocessed_image)}")
 
+    if getattr(preprocessed_image, "video_values", None) is not None:
+      return preprocessed_image
+
     if preprocessed_image.pixel_values is None:
       raise ValueError("Input preprocessed_image must have pixel_values to pad images.")
 
@@ -870,16 +917,7 @@ def map(
 
 @dataclasses.dataclass
 class ExtractImagesAndMasks(grain.MapTransform):
-  """Extracts images and masks from a PreprocessorOutput object.
-
-  This transform is used in multi-modal data pipelines to extract the image
-  tensors and their corresponding masks from a PreprocessorOutput object.
-  The extracted images and masks are then added to the data element under
-  the keys 'images' and 'image_masks', respectively.
-
-  If the 'images' key is not present in the input element, the transform
-  returns the element unchanged.
-  """
+  """Extracts images and masks from a PreprocessorOutput object."""
 
   def map(self, element: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
     """Applies the extraction transformation to the 'images' field if present."""
@@ -891,28 +929,25 @@ def map(self, element: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
       raise TypeError(f"'images' must be of type PreprocessorOutput, but got {type(preprocessed_image)}")
 
     output = element.copy()
-    output["images"] = preprocessed_image.pixel_values  # pyrefly: ignore[unsupported-operation]
-    if preprocessed_image.pixel_mask is not None:
-      output["image_masks"] = preprocessed_image.pixel_mask
+    if getattr(preprocessed_image, "video_values", None) is not None:
+      output["images"] = preprocessed_image.video_values
+      if getattr(preprocessed_image, "video_mask", None) is not None:
+        output["image_masks"] = preprocessed_image.video_mask
+      if getattr(preprocessed_image, "video_grid_thw", None) is not None:
+        output["video_grid_thw"] = preprocessed_image.video_grid_thw
+      if getattr(preprocessed_image, "video_second_per_grid", None) is not None:
+        output["second_per_grids"] = preprocessed_image.video_second_per_grid
+    else:
+      output["images"] = preprocessed_image.pixel_values
+      if preprocessed_image.pixel_mask is not None:
+        output["image_masks"] = preprocessed_image.pixel_mask
 
     return output
 
 
 @dataclasses.dataclass
 class FoldImagesIntoBatch(grain.MapTransform):
-  """Folds the 'image' dimension into the batch dimension.
-
-  This transform is used in multi-modal data pipelines where each data example
-  might have multiple associated images. For model processing, it's often
-  efficient to treat each image as a separate item in a larger batch.
-
-  This operation reshapes the 'images' tensor from a shape like
-  (B, N, T, H, W, C) to (B * N, T, H, W, C), where B is the batch size, N is
-  the number of images per example, and T is the number of image tiles.
-
-  The transformation is triggered only if the input 'images' tensor has more
-  dimensions than the expected batched image tensor.
-  """
+  """Folds the 'image' dimension into the batch dimension."""
 
   model_name: str | None = None
 
@@ -926,16 +961,21 @@ def map(self, element: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
     if images is None:
       return element
 
-    # If ndim is greater than the expected ndim for a batched image tensor,
-    # it implies an extra dimension (e.g., number of images per example)
-    # that needs to be folded into the batch dimension.
+    image_masks = element.get("image_masks")
+    video_grid_thw = element.get("video_grid_thw")
+    second_per_grids = element.get("second_per_grids")
+
     if images.ndim > len(self.target_shape):
-      # Compute the new shape by merging the batch and image count dimensions.
-      trailing_dims = self.target_shape[1:]
+      element["images"] = images.reshape(-1, *images.shape[2:])
+
+      if image_masks is not None:
+        element["image_masks"] = image_masks.reshape(-1, *image_masks.shape[2:])
+
+      if video_grid_thw is not None:
+        element["video_grid_thw"] = video_grid_thw.reshape(-1, video_grid_thw.shape[-1])
 
-      # Reshape merges the leading dimensions (B, N) into one (-1) and
-      # appends the correct trailing dimensions.
-      element["images"] = images.reshape(-1, *trailing_dims)
+      if second_per_grids is not None:
+        element["second_per_grids"] = second_per_grids.reshape(-1)
 
     return element
 

src/maxtext/multimodal/processor.py

@@ -98,7 +98,16 @@ def get_image_offsets(config, processor_output: mm_utils.PreprocessorOutput | No
     return 0
 
 
-def reformat_prompt(prompt, image_placeholder, model_name, num_images, video_placeholder="<|video|>", num_videos=0):
+def reformat_prompt(
+    prompt,
+    image_placeholder,
+    model_name,
+    num_images,
+    video_placeholder="<|video|>",
+    num_videos=0,
+    num_image_tokens=None,
+    num_video_tokens=None,
+):
   """Reformat prompt for different models."""
   if model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]:
     from maxtext.multimodal.processor_gemma3 import reformat_prompt_gemma3  # pylint: disable=import-outside-toplevel
@@ -121,6 +130,8 @@ def reformat_prompt(prompt, image_placeholder, model_name, num_images, video_pla
         num_images=num_images,
         video_placeholder=video_placeholder,
         num_videos=num_videos,
+        num_image_tokens=num_image_tokens,
+        num_video_tokens=num_video_tokens,
     )
   else:
     return prompt