55# This source code is licensed under the BSD-style license found in the
66# LICENSE file in the root directory of this source tree.
77
8- # Test CUDA model end-to-end, need to run .ci/scripts/export_model_cuda_artifact .sh first
8+ # Test CUDA/Metal model end-to-end, need to run .ci/scripts/export_model_artifact .sh first
99
1010show_help () {
1111 cat << EOF
12- Usage: test_model_cuda_e2e .sh <hf_model> <quant_name> [model_dir]
12+ Usage: test_model_e2e .sh <device> <hf_model> <quant_name> [model_dir]
1313
14- Build and run end-to-end tests for CUDA models.
14+ Build and run end-to-end tests for CUDA/Metal models.
1515
1616Arguments:
17+ device cuda or metal (required)
18+
1719 hf_model HuggingFace model ID (required)
1820 Supported models:
1921 - mistralai/Voxtral-Mini-3B-2507
20- - openai/whisper- small
22+ - openai/whisper series (whisper-{ small, medium, large, large-v2, large-v3, large-v3-turbo})
2123 - google/gemma-3-4b-it
2224
2325 quant_name Quantization type (required)
@@ -27,12 +29,12 @@ Arguments:
2729 - quantized-int4-weight-only
2830
2931 model_dir Directory containing model artifacts (optional, default: current directory)
30- Expected files: model.pte, aoti_cuda_blob.ptd
32+ Expected files: model.pte, aoti_cuda_blob.ptd/aoti_metal_blob.ptd
3133 Tokenizers and test files will be downloaded to this directory
3234
3335Examples:
34- test_model_cuda_e2e .sh "openai/whisper-small" "non-quantized"
35- test_model_cuda_e2e .sh "mistralai/Voxtral-Mini-3B-2507" "quantized-int4-tile-packed" "./model_output"
36+ test_model_e2e .sh metal "openai/whisper-small" "non-quantized"
37+ test_model_e2e .sh cuda "mistralai/Voxtral-Mini-3B-2507" "quantized-int4-tile-packed" "./model_output"
3638EOF
3739}
3840
5557
5658set -eux
5759
58- HF_MODEL=" $1 "
59- QUANT_NAME=" $2 "
60+ DEVICE=" $1 "
61+ HF_MODEL=" $2 "
62+ QUANT_NAME=" $3 "
6063# Download tokenizers, audio, and image files to this directory
61- MODEL_DIR=" ${3 :- .} "
64+ MODEL_DIR=" ${4 :- .} "
6265
6366echo " Testing model: $HF_MODEL (quantization: $QUANT_NAME )"
6467
65- # Make sure model.pte and aoti_cuda_blob .ptd exist
68+ # Make sure model.pte and aoti_${DEVICE}_blob .ptd exist
6669if [ ! -f " $MODEL_DIR /model.pte" ]; then
6770 echo " Error: model.pte not found in $MODEL_DIR "
6871 exit 1
6972fi
70- if [ ! -f " $MODEL_DIR /aoti_cuda_blob .ptd" ]; then
71- echo " Error: aoti_cuda_blob .ptd not found in $MODEL_DIR "
73+ if [ ! -f " $MODEL_DIR /aoti_ ${DEVICE} _blob .ptd" ]; then
74+ echo " Error: aoti_ ${DEVICE} _blob .ptd not found in $MODEL_DIR "
7275 exit 1
7376fi
7477# Locate EXECUTORCH_ROOT from the directory of this script
@@ -91,13 +94,13 @@ case "$HF_MODEL" in
9194 AUDIO_FILE=" poem.wav"
9295 IMAGE_PATH=" "
9396 ;;
94- openai/whisper-small )
95- MODEL_NAME=" whisper "
97+ openai/whisper-* )
98+ MODEL_NAME=" ${HF_MODEL # openai / } "
9699 RUNNER_TARGET=" whisper_runner"
97100 RUNNER_PATH=" whisper"
98101 EXPECTED_OUTPUT=" Mr. Quilter is the apostle of the middle classes"
99102 PREPROCESSOR=" whisper_preprocessor.pte"
100- TOKENIZER_URL=" https://huggingface.co/openai/whisper-small /resolve/main" # @lint-ignore
103+ TOKENIZER_URL=" https://huggingface.co/${HF_MODEL} /resolve/main" # @lint-ignore
101104 TOKENIZER_FILE=" "
102105 AUDIO_URL=" "
103106 AUDIO_FILE=" output.wav"
@@ -117,7 +120,7 @@ case "$HF_MODEL" in
117120 ;;
118121 * )
119122 echo " Error: Unsupported model '$HF_MODEL '"
120- echo " Supported models: mistralai/Voxtral-Mini-3B-2507, openai/whisper- small, google/gemma-3-4b-it"
123+ echo " Supported models: mistralai/Voxtral-Mini-3B-2507, openai/whisper series (whisper-{ small, medium, large, large-v2, large-v3, large-v3-turbo}) , google/gemma-3-4b-it"
121124 exit 1
122125 ;;
123126esac
142145# Download test files
143146if [ " $AUDIO_URL " != " " ]; then
144147 curl -L $AUDIO_URL -o ${MODEL_DIR} /$AUDIO_FILE
145- elif [ " $MODEL_NAME " = " whisper" ]; then
148+ elif [[ " $MODEL_NAME " == * whisper* ] ]; then
146149 conda install -y -c conda-forge " ffmpeg<8"
147150 pip install datasets soundfile torchcodec
148151 python -c " from datasets import load_dataset;import soundfile as sf;sample = load_dataset('distil-whisper/librispeech_long', 'clean', split='validation')[0]['audio'];sf.write('${MODEL_DIR} /$AUDIO_FILE ', sample['array'][:sample['sampling_rate']*30], sample['sampling_rate'])"
@@ -152,14 +155,24 @@ ls -al
152155echo " ::endgroup::"
153156
154157echo " ::group::Build $MODEL_NAME Runner"
158+
159+ if [ " $DEVICE " = " cuda" ]; then
160+ BUILD_BACKEND=" EXECUTORCH_BUILD_CUDA"
161+ elif [ " $DEVICE " = " metal" ]; then
162+ BUILD_BACKEND=" EXECUTORCH_BUILD_METAL"
163+ else
164+ echo " Error: Unsupported device '$DEVICE '. Must be 'cuda' or 'metal'."
165+ exit 1
166+ fi
167+
155168cmake --preset llm \
156- -DEXECUTORCH_BUILD_CUDA =ON \
169+ -D ${BUILD_BACKEND} =ON \
157170 -DCMAKE_INSTALL_PREFIX=cmake-out \
158171 -DCMAKE_BUILD_TYPE=Release \
159172 -Bcmake-out -S.
160173cmake --build cmake-out -j$( nproc) --target install --config Release
161174
162- cmake -DEXECUTORCH_BUILD_CUDA =ON \
175+ cmake -D ${BUILD_BACKEND} =ON \
163176 -DCMAKE_BUILD_TYPE=Release \
164177 -Sexamples/models/$RUNNER_PATH \
165178 -Bcmake-out/examples/models/$RUNNER_PATH /
@@ -168,19 +181,21 @@ echo "::endgroup::"
168181
169182echo " ::group::Run $MODEL_NAME Runner"
170183set +e
171- export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH
184+ if [ " $DEVICE " = " cuda" ]; then
185+ export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH
186+ fi
172187
173188# Build runner command with common arguments
174189RUNNER_BIN=" cmake-out/examples/models/$RUNNER_PATH /$RUNNER_TARGET "
175- RUNNER_ARGS=" --model_path ${MODEL_DIR} /model.pte --data_path ${MODEL_DIR} /aoti_cuda_blob .ptd --temperature 0"
190+ RUNNER_ARGS=" --model_path ${MODEL_DIR} /model.pte --data_path ${MODEL_DIR} /aoti_ ${DEVICE} _blob .ptd --temperature 0"
176191
177192# Add model-specific arguments
178193case " $MODEL_NAME " in
179194 voxtral)
180195 RUNNER_ARGS=" $RUNNER_ARGS --tokenizer_path ${MODEL_DIR} /$TOKENIZER_FILE --audio_path ${MODEL_DIR} /$AUDIO_FILE --processor_path ${MODEL_DIR} /$PREPROCESSOR "
181196 ;;
182- whisper)
183- RUNNER_ARGS=" $RUNNER_ARGS --tokenizer_path ${MODEL_DIR} / --audio_path ${MODEL_DIR} /$AUDIO_FILE --processor_path ${MODEL_DIR} /$PREPROCESSOR "
197+ whisper- * )
198+ RUNNER_ARGS=" $RUNNER_ARGS --tokenizer_path ${MODEL_DIR} / --audio_path ${MODEL_DIR} /$AUDIO_FILE --processor_path ${MODEL_DIR} /$PREPROCESSOR --model_name ${MODEL_NAME} "
184199 ;;
185200 gemma3)
186201 RUNNER_ARGS=" $RUNNER_ARGS --tokenizer_path ${MODEL_DIR} / --image_path $IMAGE_PATH "
0 commit comments