Update multimodal docs with new models

Author: hengtaoguo

docs/tutorials/posttraining/multimodal.md

@@ -8,9 +8,13 @@ This document provides a guide to use the multimodal functionalities in MaxText
 
 We also provide a [colab](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/examples/multimodal_gemma3_demo.ipynb) for multimodal features demonstration. The following table provides a list of models and modalities we currently support:
 
-| Models                                         | Input Modalities | Output Modalities |
-| :--------------------------------------------- | :--------------- | :---------------- |
-| - Gemma3-4B/12B/27B<br>- Llama4-Scout/Maverick | Text, images     | Text              |
+| Models                      | Input: Text | Input: Image | Input: Video | Input: Audio | Output |
+| :-------------------------- | :---------: | :----------: | :----------: | :----------: | :----: |
+| **Gemma3** (4B/12B/27B)     |      ✓      |      ✓       |              |              |  Text  |
+| **Gemma4** (26B/31B)        |      ✓      |      ✓       |              |              |  Text  |
+| **Llama4** (Scout/Maverick) |      ✓      |      ✓       |              |              |  Text  |
+| **Qwen3-Omni**              |      ✓      |      ✓       |      ✓       |      ✓       |  Text  |
+| **Qwen3.5** (35B/397B)      |      ✓      |      ✓       |      ✓       |              |  Text  |
 
 ## Introduction
 
@@ -73,6 +77,8 @@ MaxText supports multimodal decoding, allowing you to input text with multiple i
 
 Since each model uses a unique native chatting template from its pretraining, we've implemented these specific templates within `multimodal_utils.py` and apply them directly to your prompt.
 
+### Decode with text+image
+
 To run a forward pass and verify the model's output, use the following command:
 
 ```shell
@@ -124,10 +130,47 @@ python -m maxtext.inference.decode \
 
 For larger models such as Llama4-Scout/Maverick, we suggest to run the decoding on a TPU cluster such as v5p-16.
 
+### Decode with text+video+audio
+
+For models that support video input (e.g., Qwen3-Omni and Qwen3.5), pass a video file via `video_path`. For Qwen3-Omni, which also supports audio, set `use_audio_in_video=true` to additionally process the embedded audio track. Since the required token budget scales with video length and resolution, set `max_prefill_predict_length` accordingly.
+
+```shell
+# Qwen3-Omni decode with video + audio
+export MAXTEXT_CKPT_PATH=<Checkpoint GCS path>  # gs://my-bucket/path for Qwen3-Omni
+python -m maxtext.inference.decode \
+    model_name=qwen3-omni-30b-a3b \
+    tokenizer_path=Qwen/Qwen3-Omni-30B-A3B-Instruct \
+    tokenizer_type=huggingface \
+    load_parameters_path=${MAXTEXT_CKPT_PATH?}/0/items \
+    per_device_batch_size=1 \
+    run_name=ht_test \
+    steps=1 \
+    async_checkpointing=false \
+    scan_layers=false \
+    use_multimodal=true \
+    use_audio_in_video=true \
+    prompt='What can you see and hear? Answer in one short sentence.' \
+    video_path='tests/assets/test_video.mp4' \
+    max_prefill_predict_length=1250 \
+    max_target_length=1280 \
+    add_bos=false \
+    attention='dot_product' \
+```
+
+The expected output will look similar to:
+
+```
+Input `<|im_start|>user
+<|vision_start|><|video_pad|><|vision_end|>What can you see and hear? Answer in one short sentence.<|im_end|>
+<|im_start|>assistant
+` -> `A roaring Tyrannosaurus rex animatronic is displayed in a museum exhibit.
+```
+
 ## Supervised Fine-Tuning
 
 Supervised Fine-Tuning (SFT) of multimodal LLMs in MaxText focuses specifically on post-training; we don't yet support pre-training multimodal models from scratch. The SFT process typically involves training on Visual Question Answering (VQA) datasets where the model learns to generate accurate text responses based on both visual and textual inputs. During this fine-tuning phase, we recommend to freeze the pre-trained encoder layers (such as vision transformers) to preserve their learned visual representations, while the projection layers and LLM decoder components remain trainable. This selective training strategy allows the model to adapt the cross-modal alignment and text generation capabilities without disrupting the robust feature extraction abilities of the encoders, ultimately leading to improved performance on multimodal understanding and reasoning tasks while maintaining computational efficiency. This is achieved by setting `freeze_vision_encoder_params=True` in [sft-vision-chartqa.yml](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/configs/post_train/sft-vision-chartqa.yml).
-Here, we use [ChartQA](https://huggingface.co/datasets/HuggingFaceM4/ChartQA) as an example to demonstrate SFT functionality:
+
+**Text+image SFT is supported for all models listed above.** The following example uses Gemma3-4B with the [ChartQA](https://huggingface.co/datasets/HuggingFaceM4/ChartQA) dataset:
 
 ```shell
 export MAXTEXT_CKPT_PATH=<your-checkpoints-path>  # either set to an already available MaxText ckpt or to the one we just converted in the previous step

tests/end_to_end/tpu/gemma3/Run_Gemma3.md

@@ -20,6 +20,8 @@
 
 We provide examples for checkpoint conversion and decoding/training/finetuning Gemma3 in test scripts at [tests/end_to_end/tpu/gemma3](https://github.com/AI-Hypercomputer/maxtext/tree/main/tests/end_to_end/tpu/gemma3). 
 
+For multimodal functionality, see the [Multimodal Support guide](../../../../../docs/tutorials/posttraining/multimodal.md).
+
 
 ## Pre-training
 You can train from scratch to generate a new checkpoint. One example command to run pretraining Gemma3-4B model is as follows:

tests/end_to_end/tpu/llama4/Run_Llama4.md

@@ -22,6 +22,7 @@
 * LLama4 Maverick (17B-128E)
 * Llama4 Maverick (17B-128E-Instruct)
 
+For multimodal functionality, see the [Multimodal Support guide](../../../../docs/tutorials/posttraining/multimodal.md).
 
 ## Checkpoint conversion
 Currently, we support converting both [PyTorch](https://www.llama.com/) and [HuggingFace](https://huggingface.co/collections/meta-llama/llama-4-67f0c30d9fe03840bc9d0164) checkpoints.  Note that we recommend using the `huggingface-cli download` command with environment variable

tests/end_to_end/tpu/qwen/moe/qwen3-omni-30b-a3b/1_test_qwen3_omni_30b_a3b.sh

@@ -0,0 +1,114 @@
+#!/bin/bash
+
+# This file is documentation for how to get started with Qwen3-Omni-30B-A3B.
+# This e2e script was tested on a v5p-8 TPU VM.
+
+# This file runs Step 1 on CPU/TPU.
+# 1. Convert the HuggingFace checkpoint (bf16) to MaxText-compatible checkpoint (bf16):
+#    Unscanned format is used here as it is better suited for decoding.
+# 2. Run multimodal decoding: text+image, and text+video+audio.
+# ---
+# Example Usage:
+#
+# export HF_TOKEN=<your_hf_token>
+# export BASE_OUTPUT_PATH=gs://your-gcs-bucket/qwen3-omni-30b-a3b_maxtext_ckpt
+# bash tests/end_to_end/tpu/qwen/moe/qwen3-omni-30b-a3b/1_test_qwen3_omni_30b_a3b.sh
+# ---
+
+set -ex
+
+export MODEL_NAME="${MODEL_NAME:-qwen3-omni-30b-a3b}"
+export TOKENIZER_PATH="${TOKENIZER_PATH:-Qwen/Qwen3-Omni-30B-A3B-Instruct}"
+
+# (Optional) Path to your local Hugging Face checkpoint
+export HF_MODEL_PATH="${HF_MODEL_PATH:-}"
+
+# Base output path for MaxText checkpoint and SFT output.
+export BASE_OUTPUT_PATH="${BASE_OUTPUT_PATH:-gs://your-gcs-bucket/qwen3-omni-30b-a3b_maxtext_ckpt}"
+
+if [ -z "${HF_TOKEN}" ]; then
+  echo "Error: HF_TOKEN environment variable is not set. Please export your Hugging Face token."
+  echo "Example: export HF_TOKEN=hf_..."
+  exit 1
+fi
+
+# Strip trailing slash from base path to avoid malformed URIs
+BASE_OUTPUT_PATH=${BASE_OUTPUT_PATH%/}
+echo "Using BASE_OUTPUT_PATH = ${BASE_OUTPUT_PATH}"
+
+# Install torch for checkpoint conversion
+python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
+
+# Setup local HF path argument if one was provided
+HF_LOCAL_ARG=""
+if [ -n "${HF_MODEL_PATH}" ]; then
+  HF_LOCAL_ARG="hf_model_path=${HF_MODEL_PATH}"
+fi
+
+# ---
+# Step 1: Checkpoint Conversion
+# Convert HuggingFace checkpoint to MaxText unscanned format (better for decoding).
+# use_multimodal=true is required to include vision/audio encoder weights.
+# ---
+JAX_PLATFORMS=cpu python3 -m maxtext.checkpoint_conversion.to_maxtext src/maxtext/configs/base.yml \
+    model_name=${MODEL_NAME} \
+    base_output_directory=${BASE_OUTPUT_PATH}/unscanned \
+    hf_access_token=${HF_TOKEN} \
+    scan_layers=false \
+    use_multimodal=true \
+    --lazy_load_tensors=False \
+    ${HF_LOCAL_ARG}
+
+UNSCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/unscanned/0/items
+
+# ---
+# Step 2a: Multimodal Decode — text + image
+# Uses a test image from the repo assets.
+# max_prefill_predict_length accounts for image tokens (~256) + text prompt tokens.
+# ---
+python3 -m maxtext.inference.decode src/maxtext/configs/base.yml \
+    model_name=${MODEL_NAME} \
+    tokenizer_path=${TOKENIZER_PATH} \
+    tokenizer_type=huggingface \
+    load_parameters_path=${UNSCANNED_CKPT_PATH} \
+    hf_access_token=${HF_TOKEN} \
+    per_device_batch_size=1 \
+    run_name=qwen3_omni_decode_image \
+    steps=1 \
+    async_checkpointing=false \
+    scan_layers=false \
+    use_multimodal=true \
+    prompt='Describe this image in one sentence.' \
+    image_path='tests/assets/test_image.jpg' \
+    max_prefill_predict_length=512 \
+    max_target_length=542 \
+    add_bos=false \
+    attention=dot_product \
+    ici_tensor_parallelism=4
+
+# ---
+# Step 2b: Multimodal Decode — text + video + audio
+# Passes a video file and enables audio processing from the video track.
+# max_prefill_predict_length is set higher to accommodate video frame tokens (~1126)
+# plus audio tokens (~77) plus text prompt tokens.
+# ---
+python3 -m maxtext.inference.decode src/maxtext/configs/base.yml \
+    model_name=${MODEL_NAME} \
+    tokenizer_path=${TOKENIZER_PATH} \
+    tokenizer_type=huggingface \
+    load_parameters_path=${UNSCANNED_CKPT_PATH} \
+    hf_access_token=${HF_TOKEN} \
+    per_device_batch_size=1 \
+    run_name=qwen3_omni_decode_video \
+    steps=1 \
+    async_checkpointing=false \
+    scan_layers=false \
+    use_multimodal=true \
+    use_audio_in_video=true \
+    prompt='What can you see and hear? Answer in one short sentence.' \
+    video_path='tests/assets/test_video.mp4' \
+    max_prefill_predict_length=1240 \
+    max_target_length=1280 \
+    add_bos=false \
+    attention=dot_product \
+    ici_tensor_parallelism=4

tests/end_to_end/tpu/qwen/moe/run_qwen_moe.md

@@ -9,6 +9,8 @@ Qwen3 is a family of open-source large language models from the Qwen team at Ali
 
 -   **Qwen3-480B-A35B**
 
+-   **Qwen3-Omni-30B-A3B**
+
 -   **Qwen3.5-397B-A17B**
 
 -   **Qwen3.5-35B-A3B**
@@ -17,6 +19,8 @@ For more details on Qwen3 architecture, see the [Qwen3 Technical Report](https:/
 
 For more details on Qwen3.5 architecture, see the [Qwen3.5 Blog](https://qwen.ai/blog?id=qwen3.5)
 
+For multimodal functionality (image, video, and audio input), see the [Multimodal Support guide](../../../../../docs/tutorials/posttraining/multimodal.md).
+
 * * * * *
 
 Checkpoint Conversion
@@ -153,6 +157,22 @@ export MAXTEXT_CHECKPOINT_PATH=gs://your-gcs-bucket/qwen3-480b-a35b_maxtext_ckpt
 bash tests/end_to_end/tpu/qwen/moe/qwen3-480b-a35b/1_test_qwen3_480b_a35b.sh
 ```
 
+### Qwen3-Omni-30B-A3B
+
+```bash
+# 1. Export your Hugging Face token
+export HF_TOKEN="your_hf_token_here"
+
+# 2. Set the base path for conversion and SFT outputs
+export BASE_OUTPUT_PATH=gs://<YOUR-GCS-BUCKET>/qwen3-omni-30b-a3b_maxtext_ckpt
+
+# (Optional) Set the path if you are using a local Hugging Face checkpoint instead of downloading
+# export HF_MODEL_PATH=/path/to/local/qwen3-omni-30b-a3b_hf_checkpoint
+
+# 3. Execute the conversion and multimodal decode verification
+bash tests/end_to_end/tpu/qwen/moe/qwen3-omni-30b-a3b/1_test_qwen3_omni_30b_a3b.sh
+```
+
 ### Qwen3.5-35B-A3B
 
 ```bash