Video/audio SFT for Qwen3-Omni

src/maxtext/common/common_types.py

@@ -84,6 +84,7 @@ class MultimodalInput:
   video_masks: Array | None = None
   audio_embeddings: Array | None = None
   audio_masks: Array | None = None
+  audio_token_masks: Array | None = None
   bidirectional_mask: Array | None = None
   bidirectional_mask_video: Array | None = None
 

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

@@ -0,0 +1,38 @@
+# 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: "base.yml"
+
+use_sft: true
+use_tunix_gradient_accumulation: true
+use_multimodal: true
+sft_train_on_completion_only: true
+packing: false  # packing is not supported yet
+freeze_vision_encoder_params: true
+learning_rate: 2.e-5
+
+# -------------- Model --------------
+model_name: "qwen3-omni-30b-a3b"
+tokenizer_path: "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+
+# -------------- HF pipeline --------------
+dataset_type: "hf"
+hf_path: "parquet"
+hf_train_files: "gs://YOUR_BUCKET/path/to/parquet/*.parquet"
+train_split: "train"
+train_data_columns: ["query", "label"]
+train_image_column: "video"
+
+# Local SSD path for videos on the TPU VM
+video_directory: "/path/to/video_directory"

src/maxtext/configs/types.py

@@ -1872,13 +1872,18 @@ class MultimodalGeneral(BaseModel):
       description="Maximum number of images per example for training with image lists. -1 means no limit.",
   )
   video_path: PathStr = Field("", description="Path to a video for decoding.")
+  video_directory: PathStr = Field("", description="Local directory path containing video files for SFT.")
   audio_path: PathStr = Field("", description="Path to an audio file for decoding.")
   video_placeholder: str = Field("<|video|>", description="Placeholder string for video in text prompts.")
   audio_placeholder: str = Field("<|audio|>", description="Placeholder string for audio in text prompts.")
   use_audio_in_video: bool = Field(False, description="Extract and use audio from video files.")
   use_mrope: bool = Field(False, description="Enable Multi-dimensional RoPE for Qwen3-Omni models.")
   mrope_section: list[int] = Field([24, 20, 20], description="Dimensions for temporal, height, width in MRoPE.")
   position_id_per_seconds: int = Field(25, description="Temporal granularity for MRoPE (tokens per second).")
+  filter_sft_sequences_by_length: bool = Field(
+      False,
+      description="Filter out multimodal SFT sequences that exceed max_prefill_predict_length or max_target_length.",
+  )
 
 
 class VisionTower(BaseModel):

src/maxtext/input_pipeline/hf_data_processing.py

@@ -55,6 +55,25 @@ def vision_sft_preprocessing_pipeline(
   """pipeline for multimodal SFT with HF dataset"""
 
   assert len(text_columns) == 2, f"Need two text_columns for query and response, received {text_columns=}"
+
+  # Format conversations if columns are missing
+  features_keys = list(dataset.features.keys()) if dataset.features else []
+  if "conversations" in features_keys and not all(col in features_keys for col in text_columns):
+    def format_llava_video_dataset(example):
+      conversations = example["conversations"]
+      query = ""
+      label = ""
+      for turn in conversations:
+        if turn["from"] == "human" and not query:
+          query = turn["value"]
+        elif turn["from"] == "gpt" and not label:
+          label = turn["value"]
+      example[text_columns[0]] = query
+      example[text_columns[1]] = label
+      return example
+
+    dataset = dataset.map(format_llava_video_dataset)
+
   # Tunix GA requires per-micro-batch slicing at the data level,
   # whereas Native GA processes the full batch and splits it internally.
   if config.elastic_enabled:
@@ -137,6 +156,18 @@ def vision_sft_preprocessing_pipeline(
       fn_kwargs={"column_name": text_columns[0], "config": config},
   )
 
+  # Filter out sequences exceeding max_prefill_predict_length or max_target_length
+  if getattr(config, "filter_sft_sequences_by_length", False):
+    max_prefill = getattr(config, "max_prefill_predict_length", 8192)
+    max_target = getattr(config, "max_target_length", 8192 + 512)
+
+    def filter_by_length(example):
+      prefill_len = len(example[text_columns[0]])
+      response_len = len(example[text_columns[1]])
+      return (prefill_len <= max_prefill) and (prefill_len + response_len <= max_target)
+
+    dataset = dataset.filter(filter_by_length)
+
   dataset = input_pipeline_utils.HFDataSource(
       dataset=dataset,
       dataloading_host_index=dataloading_host_index,

src/maxtext/input_pipeline/input_pipeline_utils.py

@@ -92,17 +92,25 @@ def _process_string(string_tensor):
 
 def reformat_prompt(example, column, image_placeholder, model_name):
   """reformat prompt for multimodal SFT"""
-  if isinstance(example["images"], list):
-    num_images = len(example["images"])
+  if isinstance(example["images"], str):
+    example[column] = mm_processor.reformat_prompt(
+        example[column], image_placeholder, model_name, num_images=0, video_placeholder=image_placeholder, num_videos=1
+    )
   else:
-    num_images = 1
-  example[column] = mm_processor.reformat_prompt(example[column], image_placeholder, model_name, num_images)
+    if isinstance(example["images"], list):
+      num_images = len(example["images"])
+    else:
+      num_images = 1
+    example[column] = mm_processor.reformat_prompt(example[column], image_placeholder, model_name, num_images)
   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)
+  val = example[column]
+  if isinstance(val, (list, tuple)) and len(val) > 0:
+    val = val[0]
+  example[column] = mm_processor.reformat_response(val, model_name)
   return example
 
 
@@ -120,9 +128,17 @@ def merge_image_columns(example, image_columns, max_num_images_per_example):
 
 
 def pre_process_image_sft(example, image_column, config):
-  """pre-process image for multimodal SFT"""
+  """pre-process image or video for multimodal SFT"""
 
   def _process_image_fn(image):
+    if isinstance(image, str):
+      import os
+
+      video_directory = getattr(config, "video_directory", "")
+      if video_directory:
+        image = os.path.join(video_directory, image)
+      return mm_processor.preprocess_image_for_training(image, config)
+
     if isinstance(image, list):
       image = [np.array(mm_utils.convert_to_RGB(img)) for img in image]
     else:
@@ -131,7 +147,7 @@ def _process_image_fn(image):
     image = mm_processor.preprocess_image_for_training(image, config)
     return image
 
-  example[image_column] = _process_image_fn(example[image_column])
+  example[image_column] = _process_image_fn(example[image_column]) if example.get(image_column) is not None else None
   return example
 
 
@@ -702,12 +718,80 @@ 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 preprocessed_image.pixel_values is None:
-      raise ValueError("Input preprocessed_image must have pixel_values to pad images.")
-
     if self.config.model_name and self.config.model_name.startswith("qwen3-omni"):
+      # Pad video_values and audio_values to fixed shapes so grain.Batch can stack them.
+      video_values = getattr(preprocessed_image, "video_values", None)
+      video_grid_thw = getattr(preprocessed_image, "video_grid_thw", None)
+      if video_values is not None:
+        target_shape = mm_processor.get_dummy_video_shape_for_init(
+            self.config.model_name, batch_size=1, config=self.config
+        )
+        # target_shape = (1, C, max_T_px, max_H_px, max_W_px)
+        padded = np.zeros(target_shape[1:], dtype=video_values.dtype)  # (C, max_T, max_H, max_W)
+        _, c, t, h, w = video_values.shape
+        max_t_px, max_h_px, max_w_px = target_shape[2], target_shape[3], target_shape[4]
+        t_clip = min(t, max_t_px)
+        h_clip = min(h, max_h_px)
+        w_clip = min(w, max_w_px)
+        padded[:, :t_clip, :h_clip, :w_clip] = video_values[0, :, :t_clip, :h_clip, :w_clip]
+        preprocessed_image.video_values = padded
+
+        if video_grid_thw is not None:
+          from maxtext.multimodal.processor_qwen3_omni import VIDEO_MAX_GRID_T, VIDEO_MAX_GRID_H, VIDEO_MAX_GRID_W
+          merge_size = getattr(self.config, "spatial_merge_size_for_vit", 2)
+          max_t = VIDEO_MAX_GRID_T
+          max_h_merged = VIDEO_MAX_GRID_H // merge_size
+          max_w_merged = VIDEO_MAX_GRID_W // merge_size
+
+          actual_t, actual_h, actual_w = video_grid_thw[0]
+          actual_t = min(actual_t, VIDEO_MAX_GRID_T)
+          actual_h = min(actual_h, VIDEO_MAX_GRID_H)
+          actual_w = min(actual_w, VIDEO_MAX_GRID_W)
+
+          actual_h_merged = actual_h // merge_size
+          actual_w_merged = actual_w // merge_size
+
+          mask_3d = np.zeros((max_t, max_h_merged, max_w_merged), dtype=np.int32)
+          mask_3d[:actual_t, :actual_h_merged, :actual_w_merged] = 1
+          preprocessed_image.video_mask = mask_3d.flatten()
+
+          print(
+              f"[SFT_DEBUG] Padding Video: Original Pixel Shape: {video_values.shape}, Padded Pixel Shape: {padded.shape}. "
+              f"Original Grid (THW): {video_grid_thw[0]}, Clipped Grid: [{actual_t}, {actual_h}, {actual_w}]. "
+              f"Video Mask Shape (flattened): {preprocessed_image.video_mask.shape}, Valid tokens count: {np.sum(preprocessed_image.video_mask)}"
+          )
+
+      audio_values = getattr(preprocessed_image, "audio_values", None)
+      audio_lengths = getattr(preprocessed_image, "audio_lengths", None)
+      if audio_values is not None:
+        target_audio = mm_processor.get_dummy_audio_shape_for_sft(
+            self.config.model_name, batch_size=1, config=self.config
+        )
+        # target_audio = (1, num_mel_bins, AUDIO_MAX_TIME)
+        _, mel, t_audio = audio_values.shape
+        padded_audio = np.zeros(target_audio[1:], dtype=audio_values.dtype)  # (mel, max_time)
+        padded_audio[:, :t_audio] = audio_values[0]
+        preprocessed_image.audio_values = padded_audio
+
+        if audio_lengths is not None:
+          from maxtext.multimodal.processor_qwen3_omni import AUDIO_MAX_TIME, _get_feat_extract_output_lengths
+          max_audio_tokens = _get_feat_extract_output_lengths(AUDIO_MAX_TIME)
+          actual_audio_tokens = audio_lengths[0]
+
+          audio_token_mask = np.zeros(max_audio_tokens, dtype=np.int32)
+          audio_token_mask[:actual_audio_tokens] = 1
+          preprocessed_image.audio_token_mask = audio_token_mask
+
+          print(
+              f"[SFT_DEBUG] Padding Audio: Original Mel Shape: {audio_values.shape}, Padded Mel Shape: {padded_audio.shape}. "
+              f"Audio Mask Shape: {preprocessed_image.audio_token_mask.shape}, Valid audio tokens count: {np.sum(preprocessed_image.audio_token_mask)}"
+          )
+
       return preprocessed_image
 
+    if preprocessed_image.pixel_values is None:
+      raise ValueError("Input preprocessed_image must have pixel_values to pad images.")
+
     # Determine the maximum number of images/masks allowed.
     image_offsets = mm_processor.get_image_offsets(self.config, preprocessed_image)
     single_image_offset = image_offsets // preprocessed_image.pixel_values.shape[0]
@@ -812,6 +896,27 @@ def map(self, element: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
     if preprocessed_image.pixel_mask is not None:
       output["image_masks"] = preprocessed_image.pixel_mask
 
+    # Extract video and audio tensors from Qwen3OmniPreprocessorOutput.
+    video_values = getattr(preprocessed_image, "video_values", None)
+    if video_values is not None:
+      output["videos"] = video_values
+    video_grid_thw = getattr(preprocessed_image, "video_grid_thw", None)
+    if video_grid_thw is not None:
+      output["video_grid_thw"] = video_grid_thw
+    video_mask = getattr(preprocessed_image, "video_mask", None)
+    if video_mask is not None:
+      output["video_masks"] = video_mask
+
+    audio_values = getattr(preprocessed_image, "audio_values", None)
+    if audio_values is not None:
+      output["audios"] = audio_values
+    audio_lengths = getattr(preprocessed_image, "audio_lengths", None)
+    if audio_lengths is not None:
+      output["audio_lengths"] = audio_lengths
+    audio_token_mask = getattr(preprocessed_image, "audio_token_mask", None)
+    if audio_token_mask is not None:
+      output["audio_token_masks"] = audio_token_mask
+
     return output
 
 

src/maxtext/layers/decoders.py

@@ -752,7 +752,7 @@ def _apply_embedding(
               text_embeddings=y,
               multimodal_embeddings=audio_embeddings,
               mask=audio_masks,
-              token_masks=None,
+              token_masks=getattr(multimodal_input, "audio_token_masks", None),
           )
         else:
           raise ValueError(f"Unsupported model_name for audio: {cfg.model_name}")

src/maxtext/layers/encoders.py

@@ -94,6 +94,13 @@ def __call__(self, input_images, deterministic=False):
     else:
       embeddings = encoder_output
 
+    if self.config.model_name in ["qwen3-omni-30b-a3b"]:
+      jax.debug.print(
+          "[SFT_DEBUG] VisionEncoder: Input Shape: {x}, Encoder Output Shape: {y}",
+          x=input_images.shape,
+          y=embeddings.shape
+      )
+
     if self.config.freeze_vision_encoder_params:
       embeddings = jax.lax.stop_gradient(embeddings)
       if deep_feats is not None:
@@ -103,6 +110,12 @@ def __call__(self, input_images, deterministic=False):
     projector = getattr(self, self.projector_name)
     embeddings = projector(embeddings)
 
+    if self.config.model_name in ["qwen3-omni-30b-a3b"]:
+      jax.debug.print(
+          "[SFT_DEBUG] VisionEncoder: Projector Output Shape: {x}",
+          x=embeddings.shape
+      )
+
     return embeddings, deep_feats
 
 

src/maxtext/models/models.py

@@ -130,6 +130,7 @@ def __call__(
       encoder_videos: None | jnp.ndarray = None,
       encoder_video_masks: None | jnp.ndarray = None,
       encoder_audios: None | jnp.ndarray = None,
+      encoder_audio_token_masks: None | jnp.ndarray = None,
       enable_dropout=True,
       model_mode=MODEL_MODE_TRAIN,
       previous_chunk=None,
@@ -195,6 +196,7 @@ def __call__(
           video_masks=encoder_video_masks,
           audio_embeddings=audio_embeddings,
           audio_masks=audio_masks,
+          audio_token_masks=encoder_audio_token_masks,
           bidirectional_mask=bidirectional_mask_image,
           bidirectional_mask_video=bidirectional_mask_video,
       )
@@ -443,6 +445,7 @@ def __call__(
       encoder_videos: jax.Array | None = None,
       encoder_video_masks: jax.Array | None = None,
       encoder_audios: jax.Array | None = None,
+      encoder_audio_token_masks: jax.Array | None = None,
       enable_dropout=True,
       model_mode=MODEL_MODE_TRAIN,
       previous_chunk=None,
@@ -522,6 +525,7 @@ def __call__(
           video_masks=encoder_video_masks,
           audio_embeddings=audio_embeddings,
           audio_masks=audio_masks,
+          audio_token_masks=encoder_audio_token_masks,
           bidirectional_mask=bidirectional_mask_image,
           bidirectional_mask_video=bidirectional_mask_video,
       )

src/maxtext/models/qwen3.py

@@ -1906,7 +1906,7 @@ def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None):
         num_kv_heads=self.config.num_attention_heads_for_vit,
         head_dim=head_dim,
         max_target_length=self.config.num_position_embeddings_for_vit,
-        attention_kernel="dot_product",
+        attention_kernel="autoselected",
         inputs_q_shape=(1, 1, self.config.hidden_size_for_vit),
         inputs_kv_shape=(1, 1, self.config.hidden_size_for_vit),
         float32_qk_product=self.config.float32_qk_product,

src/maxtext/multimodal/processor.py

@@ -69,9 +69,13 @@ def preprocess_image_for_training(image, config):
 
     return preprocess_mm_data_llama4(image)
   elif config.model_name in ["qwen3-omni-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b"]:
-    from maxtext.multimodal.processor_qwen3_omni import preprocess_mm_data_qwen3_omni_for_training  # pylint: disable=import-outside-toplevel
+    from maxtext.multimodal.processor_qwen3_omni import preprocess_mm_data_qwen3_omni_for_training, preprocess_mm_data_qwen3_omni_for_training_video  # pylint: disable=import-outside-toplevel
 
-    return preprocess_mm_data_qwen3_omni_for_training(image, config)
+    if isinstance(image, str):
+      use_audio_in_video = getattr(config, "use_audio_in_video", False)
+      return preprocess_mm_data_qwen3_omni_for_training_video(image, config)
+    else:
+      return preprocess_mm_data_qwen3_omni_for_training(image, config)
   else:
     raise ValueError(f"Model {config.model_name} not supported for image preprocessing.")
 
@@ -188,6 +192,24 @@ def get_dummy_image_shape_for_init(model_name, batch_size=1, num_image_per_seque
   return image_shape
 
 
+def get_dummy_video_shape_for_init(model_name, batch_size=1, config=None):
+  """Return the fixed padded shape for video batch tensors used in SFT training."""
+  if model_name in ["qwen3-omni-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b"]:
+    from maxtext.multimodal.processor_qwen3_omni import get_dummy_video_shape_for_init_qwen3_omni  # pylint: disable=import-outside-toplevel
+
+    return get_dummy_video_shape_for_init_qwen3_omni(batch_size, config)
+  return ()
+
+
+def get_dummy_audio_shape_for_sft(model_name, batch_size=1, config=None):
+  """Return the fixed padded shape for audio batch tensors used in SFT training."""
+  if model_name in ["qwen3-omni-30b-a3b"]:
+    from maxtext.multimodal.processor_qwen3_omni import get_dummy_audio_shape_for_init_qwen3_omni_sft  # pylint: disable=import-outside-toplevel
+
+    return get_dummy_audio_shape_for_init_qwen3_omni_sft(batch_size, config)
+  return ()
+
+
 def get_dummy_audio_shape_for_init(config):
   """Return the shape of the dummy audio for specific model's initialization.