Skip to content

Commit c5b9a18

Browse files
hengtaoguohengtaoguo
committed
Update
1 parent eb71c23 commit c5b9a18

3 files changed

Lines changed: 89 additions & 56 deletions

File tree

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

Lines changed: 2 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,10 @@
11
#!/bin/bash
22

33
# This file is documentation for how to get started with Qwen3-Omni-30B-A3B.
4-
# This e2e script was tested on a v5p-8 TPU VM.
54

6-
# This file runs Step 1 on CPU/TPU.
5+
# This file runs Step 1 on CPU.
76
# 1. Convert the HuggingFace checkpoint (bf16) to MaxText-compatible checkpoint (bf16):
87
# Unscanned format is used here as it is better suited for decoding.
9-
# 2. Run multimodal decoding: text+image, and text+video+audio.
108
# ---
119
# Example Usage:
1210
#
@@ -23,7 +21,7 @@ export TOKENIZER_PATH="${TOKENIZER_PATH:-Qwen/Qwen3-Omni-30B-A3B-Instruct}"
2321
# (Optional) Path to your local Hugging Face checkpoint
2422
export HF_MODEL_PATH="${HF_MODEL_PATH:-}"
2523

26-
# Base output path for MaxText checkpoint and SFT output.
24+
# Base output path for MaxText checkpoint.
2725
export BASE_OUTPUT_PATH="${BASE_OUTPUT_PATH:-gs://your-gcs-bucket/qwen3-omni-30b-a3b_maxtext_ckpt}"
2826

2927
if [ -z "${HF_TOKEN}" ]; then
@@ -59,56 +57,4 @@ JAX_PLATFORMS=cpu python3 -m maxtext.checkpoint_conversion.to_maxtext src/maxtex
5957
--lazy_load_tensors=False \
6058
${HF_LOCAL_ARG}
6159

62-
UNSCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/unscanned/0/items
6360

64-
# ---
65-
# Step 2a: Multimodal Decode — text + image
66-
# Uses a test image from the repo assets.
67-
# max_prefill_predict_length accounts for image tokens (~256) + text prompt tokens.
68-
# ---
69-
python3 -m maxtext.inference.decode src/maxtext/configs/base.yml \
70-
model_name=${MODEL_NAME} \
71-
tokenizer_path=${TOKENIZER_PATH} \
72-
tokenizer_type=huggingface \
73-
load_parameters_path=${UNSCANNED_CKPT_PATH} \
74-
hf_access_token=${HF_TOKEN} \
75-
per_device_batch_size=1 \
76-
run_name=qwen3_omni_decode_image \
77-
steps=1 \
78-
async_checkpointing=false \
79-
scan_layers=false \
80-
use_multimodal=true \
81-
prompt='Describe this image in one sentence.' \
82-
image_path='tests/assets/test_image.jpg' \
83-
max_prefill_predict_length=512 \
84-
max_target_length=542 \
85-
add_bos=false \
86-
attention=dot_product \
87-
ici_tensor_parallelism=4
88-
89-
# ---
90-
# Step 2b: Multimodal Decode — text + video + audio
91-
# Passes a video file and enables audio processing from the video track.
92-
# max_prefill_predict_length is set higher to accommodate video frame tokens (~1126)
93-
# plus audio tokens (~77) plus text prompt tokens.
94-
# ---
95-
python3 -m maxtext.inference.decode src/maxtext/configs/base.yml \
96-
model_name=${MODEL_NAME} \
97-
tokenizer_path=${TOKENIZER_PATH} \
98-
tokenizer_type=huggingface \
99-
load_parameters_path=${UNSCANNED_CKPT_PATH} \
100-
hf_access_token=${HF_TOKEN} \
101-
per_device_batch_size=1 \
102-
run_name=qwen3_omni_decode_video \
103-
steps=1 \
104-
async_checkpointing=false \
105-
scan_layers=false \
106-
use_multimodal=true \
107-
use_audio_in_video=true \
108-
prompt='What can you see and hear? Answer in one short sentence.' \
109-
video_path='tests/assets/test_video.mp4' \
110-
max_prefill_predict_length=1240 \
111-
max_target_length=1280 \
112-
add_bos=false \
113-
attention=dot_product \
114-
ici_tensor_parallelism=4

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

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
#!/bin/bash
2+
3+
# This file is documentation for how to get started with Qwen3-Omni-30B-A3B.
4+
5+
# This file runs Step 2 on a v5p-8 TPU VM.
6+
# 2. Run multimodal decoding: text+image, and text+video+audio.
7+
# ---
8+
# Example Usage:
9+
#
10+
# export HF_TOKEN=<your_hf_token>
11+
# export BASE_OUTPUT_PATH=gs://your-gcs-bucket/qwen3-omni-30b-a3b_maxtext_ckpt
12+
# bash tests/end_to_end/tpu/qwen/moe/qwen3-omni-30b-a3b/2_test_qwen3_omni_30b_a3b.sh
13+
# ---
14+
15+
set -ex
16+
17+
export MODEL_NAME="${MODEL_NAME:-qwen3-omni-30b-a3b}"
18+
export TOKENIZER_PATH="${TOKENIZER_PATH:-Qwen/Qwen3-Omni-30B-A3B-Instruct}"
19+
20+
# Base output path where the MaxText checkpoint from Step 1 was written.
21+
export BASE_OUTPUT_PATH="${BASE_OUTPUT_PATH:-gs://your-gcs-bucket/qwen3-omni-30b-a3b_maxtext_ckpt}"
22+
23+
if [ -z "${HF_TOKEN}" ]; then
24+
echo "Error: HF_TOKEN environment variable is not set. Please export your Hugging Face token."
25+
echo "Example: export HF_TOKEN=hf_..."
26+
exit 1
27+
fi
28+
29+
# Strip trailing slash from base path to avoid malformed URIs
30+
BASE_OUTPUT_PATH=${BASE_OUTPUT_PATH%/}
31+
echo "Using BASE_OUTPUT_PATH = ${BASE_OUTPUT_PATH}"
32+
33+
UNSCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/unscanned/0/items
34+
35+
# ---
36+
# Step 2a: Multimodal Decode — text + image
37+
# Uses a test image from the repo assets.
38+
# max_prefill_predict_length accounts for image tokens (~256) + text prompt tokens.
39+
# ---
40+
python3 -m maxtext.inference.decode src/maxtext/configs/base.yml \
41+
model_name=${MODEL_NAME} \
42+
tokenizer_path=${TOKENIZER_PATH} \
43+
tokenizer_type=huggingface \
44+
load_parameters_path=${UNSCANNED_CKPT_PATH} \
45+
hf_access_token=${HF_TOKEN} \
46+
per_device_batch_size=1 \
47+
run_name=qwen3_omni_decode_image \
48+
steps=1 \
49+
async_checkpointing=false \
50+
scan_layers=false \
51+
use_multimodal=true \
52+
prompt='Describe this image in one sentence.' \
53+
image_path='tests/assets/test_image.jpg' \
54+
max_prefill_predict_length=512 \
55+
max_target_length=542 \
56+
add_bos=false \
57+
attention=dot_product \
58+
ici_tensor_parallelism=4
59+
60+
# ---
61+
# Step 2b: Multimodal Decode — text + video + audio
62+
# Passes a video file and enables audio processing from the video track.
63+
# max_prefill_predict_length is set higher to accommodate video frame tokens (~1126)
64+
# plus audio tokens (~77) plus text prompt tokens.
65+
# ---
66+
python3 -m maxtext.inference.decode src/maxtext/configs/base.yml \
67+
model_name=${MODEL_NAME} \
68+
tokenizer_path=${TOKENIZER_PATH} \
69+
tokenizer_type=huggingface \
70+
load_parameters_path=${UNSCANNED_CKPT_PATH} \
71+
hf_access_token=${HF_TOKEN} \
72+
per_device_batch_size=1 \
73+
run_name=qwen3_omni_decode_video \
74+
steps=1 \
75+
async_checkpointing=false \
76+
scan_layers=false \
77+
use_multimodal=true \
78+
use_audio_in_video=true \
79+
prompt='What can you see and hear? Answer in one short sentence.' \
80+
video_path='tests/assets/test_video.mp4' \
81+
max_prefill_predict_length=1240 \
82+
max_target_length=1280 \
83+
add_bos=false \
84+
attention=dot_product \
85+
ici_tensor_parallelism=4

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@ Qwen3 is a family of open-source large language models from the Qwen team at Ali
1717

1818
For more details on Qwen3 architecture, see the [Qwen3 Technical Report](https://arxiv.org/abs/2505.09388).
1919

20+
For more details on Qwen3-Omni architecture, see the [Qwen3-Omni Technical Report](https://arxiv.org/abs/2509.17765).
21+
2022
For more details on Qwen3.5 architecture, see the [Qwen3.5 Blog](https://qwen.ai/blog?id=qwen3.5)
2123

2224
For multimodal functionality (image, video, and audio input), see the [Multimodal Support guide](../../../../../docs/tutorials/posttraining/multimodal.md).

0 commit comments

Comments
 (0)