使用Vulkan后端构建和运行ExecuTorch

The ExecuTorch Vulkan Delegate 是用于ExecuTorch的原生GPU委托。

你将在本教程中学习到：

• 如何使用部分GPU委托导出Llama3.2-1B参数模型

• 如何在 Android 上执行部分委托模型

Prerequisites:

• 跟随 设置 ExecuTorch

• 也建议您阅读 ExecuTorch Vulkan Delegate 并遵循该页面中的示例

预备知识

请注意，以下所有步骤都应在 ExecuTorch 仓库的根目录下执行，并假设您已经完成了设置 ExecuTorch 的步骤。

也假设已安装 Android NDK 和 Android SDK，并且设置了以下环境示例。

export ANDROID_NDK=<path_to_ndk>
# Select an appropriate Android ABI for your device
export ANDROID_ABI=arm64-v8a
# All subsequent commands should be performed from ExecuTorch repo root
cd <path_to_executorch_root>
# Make sure adb works
adb --version

将Llama3.2-1B模型降低到Vulkan

注意

生成的模型只会部分委托给Vulkan后端。具体而言，仅二进制算术运算符（aten.add, aten.sub, aten.mul, aten.div）、矩阵乘法运算符（aten.mm, aten.bmm), 以及线性层（aten.linear）会在GPU上通过Vulkan委托执行。模型的其余部分将使用便携式运算符执行。

对LLaMA模型的操作符支持目前正在积极开发中；请查看ExecuTorch仓库的main分支以获取最新功能。

首先，从Llama网站获取consolidated.00.pth、params.json和tokenizer.model文件用于Llama3.2-1B模型。

下载完文件后，可以使用 export_llama 脚本 部分降低 Llama 模型到 Vulkan。

# The files will usually be downloaded to ~/.llama
python -m examples.models.llama.export_llama \
  --disable_dynamic_shape --vulkan -kv --use_sdpa_with_kv_cache -d fp32 \
  --model "llama3_2" \ 
  -c ~/.llama/checkpoints/Llama3.2-1B/consolidated.00.pth \
  -p ~/.llama/checkpoints/Llama3.2-1B/params.json \
  --metadata '{"get_bos_id":128000, "get_eos_ids":[128009, 128001]}'

一个 vulkan_llama2.pte 文件应该已经通过运行脚本创建。

将分词器二进制文件和 vulkan_llama2.pte 推送到您的Android设备上：

adb push ~/.llama/tokenizer.model /data/local/tmp/
adb push vulkan_llama2.pte /data/local/tmp/

在Android上构建并运行LLaMA运行器二进制文件

首先，构建并安装 ExecuTorch 库，然后使用 Android NDK 工具链构建 LLaMA 运行程序二进制文件。

./install_requirements.sh --clean
(mkdir cmake-android-out && \
  cmake . -DCMAKE_INSTALL_PREFIX=cmake-android-out \
    -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake \
    -DANDROID_ABI=$ANDROID_ABI \
    -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
    -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
    -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \
    -DEXECUTORCH_BUILD_VULKAN=ON \
    -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \
    -DEXECUTORCH_BUILD_KERNELS_CUSTOM=ON \
    -DPYTHON_EXECUTABLE=python \
    -Bcmake-android-out && \
  cmake --build cmake-android-out -j16 --target install)

# Build LLaMA Runner library
(rm -rf cmake-android-out/examples/models/llama && \
  cmake examples/models/llama \
    -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake \
    -DANDROID_ABI=$ANDROID_ABI \
    -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \
    -DEXECUTORCH_BUILD_KERNELS_CUSTOM=ON \
    -DCMAKE_INSTALL_PREFIX=cmake-android-out \
    -DPYTHON_EXECUTABLE=python \
    -Bcmake-android-out/examples/models/llama && \
  cmake --build cmake-android-out/examples/models/llama -j16)

最后，在你的 Android 设备上推送并运行 llama runner 二进制文件。请注意，你的设备必须具备足够的 GPU 内存以执行该模型。

adb push cmake-android-out/examples/models/llama/llama_main /data/local/tmp/llama_main

adb shell /data/local/tmp/llama_main \
    --model_path=/data/local/tmp/vulkan_llama2.pte \
    --tokenizer_path=/data/local/tmp/tokenizer.model \
    --prompt "Hello"

请注意，目前由于降低后的图中存在大量委托 blob，模型推理将会非常缓慢，这需要在每个子图之间进行 GPU 的传输。随着越来越多的模型可以被降低到 Vulkan 委托，并且支持诸如量化等技术，性能预计将得到显著提升。