Running Phi-3-vision via ONNX on Jetson Platform
Hi, I'm Jambo a Microsoft Learn Student Ambassador.
This article aims to run the quantized Phi-3-vision model in ONNX format on the Jetson platform and successfully perform inference for image+text dialogue tasks.
Writing Environment:
- Jetpack 6.0 [L4T 36.3.0]
- Compilation Platform: Jetson Orin
- Inference Platform: Jetson Orin Nano
What is Jetson?
The Jetson platform, introduced by NVIDIA, consists of small arm64 devices equipped with powerful GPU computing capabilities. Designed specifically for edge computing and AI applications, Jetson devices run on Linux, enabling complex computing tasks with low power consumption. This makes them ideal for developing embedded AI and machine learning projects.
For other versions of the Phi-3 model, we can use llama.cpp to convert them into GUFF format to run on Jetson, and easily switch between different quantizations. Alternatively, you can conveniently use services like ollama or llamaedge which are based on llama.cpp. More information can be found in the Phi-3CookBook.
However, for the vision version, there is currently no way to convert it into GUFF format (#7444). Additionally, resource-constrained edge devices struggle to run the original model without quantization via transformers. Therefore, we can use ONNX Runtime to run the quantized model in ONNX format.
What is ONNX Runtime?
ONNX Runtime is a high-performance inference engine designed to accelerate and execute AI models in the ONNX (Open Neural Network Exchange) format. The onnxruntime-genai is an API specifically built for LLM (Large Language Model) models, providing a simple way to run models like Llama, Phi, Gemma, and Mistral.
When writing this article, onnxruntime-genai does not have a precompiled version for aarch64 + GPU, so we need to compile it ourselves.
Compiling onnxruntime-genai
Preparation
- Upgrade CMake
sudo apt purge cmake
pip3 install cmake -U
Cloning the onnxruntime-genai Repository
git clone https://github.com/microsoft/onnxruntime-genai
cd onnxruntime-genai
The latest onnxruntime-genai repository cannot be successfully compiled for unknown reasons, so we need to switch to an earlier commit. Below is the latest commit that has been tested and can be successfully compiled:
git checkout 940bc102a317e886f488ad5e120533b96a34ddcd
ONNXRuntime
You can compile ONNXRuntime from the source yourself, but this can be a very time-consuming process for the Jetson platform. Therefore, we will directly use the version compiled by dusty-nv for the Jetson platform. Do not worry about the cu124
in the URL; it runs well on CUDA 12.2.
wget http://jetson.webredirect.org:8000/jp6/cu124/onnxruntime-gpu-1.19.0.tar.gz
mkdir ort
tar -xvf onnxruntime-gpu-1.19.0.tar.gz -C ort
mv ort/include/onnxruntime/onnxruntime_c_api.h ort/include/
rm -rf ort/include/onnxruntime/
Compiling onnxruntime-genai
You should still be in the onnxruntime-genai
directory at this point.
Now we need to prepare to build the Python API. You can use Python >=3.6
for the compilation. JetPack 6.0 comes with Python 3.10 by default, but you can switch to other versions for the compilation. The compiled whl can only be installed on the Python version used during the compilation.
Note: The compilation process will require a significant amount of memory. Therefore, if your Jetson device has limited memory (like the Orin NX), do not use the --parallel parameter.
python3 build.py --use_cuda --cuda_home /usr/local/cuda-12.2 --skip_tests --skip_csharp [--parallel]
The compiled files will be located in the build/Linux/Release/dist/wheel
directory, and we only need the .whl
file. Note that the .whl
file should be around 110 MB.
You can copy the whl file to other Jetson platforms with the same environment (CUDA) for installation.
Note: The generated subdirectory may differ, but we only need the
.whl
file from thebuild
directory.
Installing onnxruntime-genai
If you have multiple CUDA versions, you might need to set the CUDA_PATH
environment variable to ensure it points to the same version used during compilation.
export CUDA_PATH=/usr/local/cuda-12.2
Navigate to the directory where the whl file is located, or copy the whl file to another directory for installation using the following command.
pip3 install *.whl
Running the Phi-3-vision Model
Downloading the Model
Download the Phi-3-vision model for onnx-cuda from huggingface.
pip3 install huggingface-hub[cli]
The FP16 model requires 8 GB of VRAM. If you are running on a device with more resources like the Jetson Orin, you can opt for the FP32 model.
The Int 4 model is a quantized version, requiring only 3 GB of VRAM. This is suitable for more compact devices like the Jetson Orin Nano.
huggingface-cli download microsoft/Phi-3-vision-128k-instruct-onnx-cuda --include cuda-fp16/* --local-dir .
# Or
huggingface-cli download microsoft/Phi-3-vision-128k-instruct-onnx-cuda --include cuda-int4-rtn-block-32/* --local-dir .
Running the Example Script
Download the official example script and an example image.
# Download example script
wget https://raw.githubusercontent.com/microsoft/onnxruntime-genai/main/examples/python/phi3v.py
# Download example image
wget https://onnxruntime.ai/images/table.png
Run the example script.
python3 phi3v.py -m cuda-int4-rtn-block-32
First, input the path to the image, for example, table.png
.
Next, input the prompt text, for example: Convert this image to markdown format
.
```markdown
| Product | Qtr 1 | Qtr 2 | Grand Total |
|---------------------|----------|----------|-------------|
| Chocolade | $744.60 | $162.56 | $907.16 |
| Gummibarchen | $5,079.60| $1,249.20| $6,328.80 |
| Scottish Longbreads | $1,267.50| $1,062.50| $2,330.00 |
| Sir Rodney's Scones | $1,418.00| $756.00 | $2,174.00 |
| Tarte au sucre | $4,728.00| $4,547.92| $9,275.92 |
| Chocolate Biscuits | $943.89 | $349.60 | $1,293.49 |
| Total | $14,181.59| $8,127.78| $22,309.37 |
```
The table lists various products along with their sales figures for Qtr 1, Qtr 2, and the Grand Total. The products include Chocolade, Gummibarchen, Scottish Longbreads, Sir Rodney's Scones, Tarte au sucre, and Chocolate Biscuits. The Grand Total column sums up the sales for each product across the two quarters.
Note: The first round of dialogue during script execution might be slow, but subsequent dialogues will be faster.
We can use Jtop to monitor resource usage:
The above inference is run on Jetson Orin Nano using the Int 4 quantized model. As shown, the Python process occupies 5.4 GB of VRAM for inference, with minimal CPU load and nearly full GPU utilization during inference.
We can modify the example script to use the time
function at key points to measure the inference speed, which is remarkably fast.
All of this is achieved on a device with a power consumption of just 15W.