Skip to content

Commit

Permalink
Enable android runner with custom sdpa op (pytorch#363)
Browse files Browse the repository at this point in the history
Summary:
This diff
- refactors install_et.sh into a bunch of utils
- Uses those utils in build_android.sh to minimize duplication.
- Makes sure taht we are building with custom sdpa op

Test Plan:

Model export
python export.py --quant '{"linear:a8w4dq" : {"groupsize": 256}}'
--checkpoint-path /home/kimishpatel/models/llama2/stories/stories110M.pt
--params-path /home/kimishpatel/models/llama2/stories/params.json
--output-pte-path /tmp/stories110m_a8w4dq.pte
python utils/tokenizer.py --tokenizer-model=/tmp/tokenizer.model

linux:
./scripts/install_et.sh
rm -rf build/cmake-out/
cmake -S ./runner-et -B build/cmake-out -G Ninja
cmake --build ./build/cmake-out

./build/cmake-out/runner_et /tmp/stories110m_a8w4dq.pte -z /tmp/tokenizer.bin -t 0 -n 120

android:
./runner-et/build_android.sh
adb push ./build/cmake-out-android/runner_et /data/local/tmp/
adb push /tmp/stories110m_a8w4dq.pte /data/local/tmp/
adb push /tmp/tokenizer.bin /data/local/tmp/
adb shell "cd /data/local/tmp && ./runner_et ./stories110m_a8w4dq.pte -z
./tokenizer.bin -t 0 -n 120"

Will add build commands to ci in the next PR

Reviewers:

Subscribers:

Tasks:

Tags:
  • Loading branch information
kimishpatel authored and malfet committed Jul 17, 2024
1 parent ee280db commit 1f7f01e
Show file tree
Hide file tree
Showing 5 changed files with 158 additions and 77 deletions.
1 change: 1 addition & 0 deletions export_et_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,7 @@ def forward(self, x, freqs_cis, mask, input_pos=None):


def replace_attention_with_custom_sdpa_attention(module: nn.Module):
from executorch.examples.models.llama2.custom_ops import sdpa_with_kv_cache # noqa
for name, child in module.named_children():
if isinstance(child, Attention):
setattr(module, name, CustomSDPAAttention(child))
Expand Down
33 changes: 30 additions & 3 deletions runner-et/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ target_link_libraries(
run PRIVATE
executorch
extension_module
${TORCHCHAT_ROOT}/et-build/src/executorch/${CMAKE_OUT_DIR}/lib/libextension_data_loader.a # This one gets installed in build directory by ExecuTorch
${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src/executorch/${CMAKE_OUT_DIR}/extension/data_loader/libextension_data_loader.a # This one does not get installed by ExecuTorch
optimized_kernels
portable_kernels
cpublas
Expand All @@ -58,5 +58,32 @@ target_link_options_shared_lib(xnnpack_backend)
target_link_options_shared_lib(XNNPACK)
target_link_options_shared_lib(pthreadpool)
target_link_options_shared_lib(cpuinfo)
target_link_options_shared_lib(executorch)
target_link_libraries(run PRIVATE "$<LINK_LIBRARY:WHOLE_ARCHIVE,${TORCHCHAT_ROOT}/et-build/src/executorch/${CMAKE_OUT_DIR}/examples/models/llama2/custom_ops/libcustom_ops.a>")
# Not clear why linking executorch as whole-archive outside android/apple is leading
# to double registration. Most likely because of linkage issues.
# Will figure this out later. Until then use this.
if(ANDROID OR APPLE)
target_link_options_shared_lib(executorch)
endif()
target_link_libraries(runner_et PRIVATE
"$<LINK_LIBRARY:WHOLE_ARCHIVE,${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src/executorch/${CMAKE_OUT_DIR}/examples/models/llama2/custom_ops/libcustom_ops.a>")

# This one is needed for cpuinfo where it uses android specific log lib
if(ANDROID)
target_link_libraries(runner_et PRIVATE log)
endif()

# Adding target_link_options_shared_lib as commented out below leads to this:
#
# CMake Error at Utils.cmake:22 (target_link_options):
# Cannot specify link options for target
# "/Users/scroy/etorch/torchchat/et-build/src/executorch/${CMAKE_OUT_DIR}/examples/models/llama2/custom_ops/libcustom_ops_lib.a"
# which is not built by this project.
# Call Stack (most recent call first):
# Utils.cmake:30 (macos_kernel_link_options)
# CMakeLists.txt:41 (target_link_options_shared_lib)
#
#target_link_options_shared_lib("${TORCHCHAT_ROOT}/et-build/src/executorch/${CMAKE_OUT_DIR}/examples/models/llama2/custom_ops/libcustom_ops_lib.a") # This one does not get installed by ExecuTorch

# This works on mac, but appears to run into issues on linux
# It is needed to solve:
# E 00:00:00.055965 executorch:method.cpp:536] Missing operator: [8] llama::sdpa_with_kv_cache.out
41 changes: 11 additions & 30 deletions runner-et/build_android.sh
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@

set -ex

source "$(dirname "${BASH_SOURCE[0]}")/../scripts/install_utils.sh"

if ["${ANDROID_NDK}" == ""]; then
echo "Please set ANDROID_NDK enviornment variable."
echo "For example it can be /Users/guest/Desktop/android-ndk-r26."
Expand All @@ -17,45 +19,24 @@ else
fi

export CMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake
export DANDROID_ABI=arm64-v8a
export DANDROID_PLATFORM=android-23
export ANDROID_ABI=arm64-v8a
export ANDROID_PLATFORM=android-23
export ET_BUILD_DIR="et-build-android"
export CMAKE_OUT_DIR="cmake-out-android"
export EXECUTORCH_BUILD_CUSTOM_OPS_AOT="OFF"
export EXECUTORCH_BUILD_CUSTOM="ON"
export CMAKE_OUT_DIR="cmake-out-android"
# export DCMAKE_INSTALL_PREFIX=cmake-out-android
#

install_executorch() {
echo "Cloning executorch to ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src"
ET_BUILD_DIR="${TORCHCHAT_ROOT}/${ET_BUILD_DIR}"
rm -rf ${ET_BUILD_DIR}
mkdir -p ${ET_BUILD_DIR}/src
pushd ${ET_BUILD_DIR}/src
git clone https://github.com/pytorch/executorch.git
cd executorch
git checkout viable/strict
echo "Install executorch: submodule update"
git submodule sync
git submodule update --init

echo "Applying fixes"
cp ${TORCHCHAT_ROOT}/scripts/fixes_et/module.cpp ${ET_BUILD_DIR}/src/executorch/extension/module/module.cpp # ET uses non-standard C++ that does not compile in GCC
cp ${TORCHCHAT_ROOT}/scripts/fixes_et/managed_tensor.h ${ET_BUILD_DIR}/src/executorch/extension/runner_util/managed_tensor.h # ET is missing headers for vector/memory. This causes downstream issues when building runner-et.

CMAKE_OUT_DIR="cmake-out-android"
echo "Building and installing C++ libraries"
echo "Inside: ${PWD}"
mkdir ${CMAKE_OUT_DIR}
cmake -DCMAKE_BUILD_TYPE=Release -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake -DANDROID_ABI=arm64-v8a -DANDROID_PLATFORM=android-23 -DCMAKE_INSTALL_PREFIX=cmake-out-android -DEXECUTORCH_ENABLE_LOGGING=ON -DEXECUTORCH_LOG_LEVEL=Info -DEXECUTORCH_BUILD_OPTIMIZED=ON -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON -DEXECUTORCH_BUILD_XNNPACK=ON -S . -B ${CMAKE_OUT_DIR} -G Ninja
cmake --build ${CMAKE_OUT_DIR}
cmake --install ${CMAKE_OUT_DIR} --prefix ${ET_BUILD_DIR}/install
popd
}

build_runner_et() {
rm -rf build/cmake-out-android
echo "ET BUILD DIR IS ${ET_BUILD_DIR}"
cmake -DCMAKE_BUILD_TYPE=Release -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake -DANDROID_ABI=arm64-v8a -DANDROID_PLATFORM=android-23 -S ./runner-et -B build/cmake-out-android -G Ninja
cmake --build build/cmake-out-android/ -j16 --config Release
}

# install_executorch
find_cmake_prefix_path
clone_executorch
install_executorch
build_runner_et
53 changes: 9 additions & 44 deletions scripts/install_et.sh
Original file line number Diff line number Diff line change
Expand Up @@ -5,55 +5,20 @@
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

set -exu
set -ex

install_pip_dependencies() {
echo "Intalling common pip packages"

pip3 install wheel
pip3 install "cmake>=3.19"
pip3 install ninja
pip3 install zstd
pushd ${TORCHCHAT_ROOT}
pip3 install -r ./requirements.txt
popd
}

install_executorch() {
echo "Cloning executorch to ${TORCHCHAT_ROOT}/et-build/src"
rm -rf ${TORCHCHAT_ROOT}/et-build
mkdir -p ${TORCHCHAT_ROOT}/et-build/src
pushd ${TORCHCHAT_ROOT}/et-build/src
git clone https://github.com/pytorch/executorch.git
cd executorch
git checkout viable/strict
echo "Install executorch: submodule update"
git submodule sync
git submodule update --init

echo "Building and installing python libraries"
if [ "${ENABLE_ET_PYBIND}" = false ]; then
echo "Not installing pybind"
bash ./install_requirements.sh
else
echo "Installing pybind"
bash ./install_requirements.sh --pybind xnnpack
fi
pip3 list

echo "Building and installing C++ libraries"
echo "Inside: ${PWD}"
mkdir cmake-out
cmake -DCMAKE_BUILD_TYPE=Release -DEXECUTORCH_ENABLE_LOGGING=ON -DEXECUTORCH_LOG_LEVEL=Info -DEXECUTORCH_BUILD_CUSTOM=ON -DEXECUTORCH_BUILD_OPTIMIZED=ON -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON -DEXECUTORCH_BUILD_XNNPACK=ON -S . -B cmake-out -G Ninja
cmake --build cmake-out
cmake --install cmake-out --prefix ${TORCHCHAT_ROOT}/et-build/install
popd
}
source "$(dirname "${BASH_SOURCE[0]}")/install_utils.sh"

if [ "${ET_BUILD_DIR}" == "" ]; then
ET_BUILD_DIR="et-build"
fi

ENABLE_ET_PYBIND="${1:-true}"

pushd ${TORCHCHAT_ROOT}
find_cmake_prefix_path
install_pip_dependencies
install_executorch $ENABLE_ET_PYBIND
clone_executorch
install_executorch_python_libs $ENABLE_ET_PYBIND
install_executorch
popd
107 changes: 107 additions & 0 deletions scripts/install_utils.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
#!/bin/bash
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

set -ex

install_pip_dependencies() {
echo "Intalling common pip packages"

pip3 install wheel
pip3 install "cmake>=3.19"
pip3 install ninja
pip3 install zstd
pushd ${TORCHCHAT_ROOT}
pip3 install -r ./requirements.txt
popd
}

function find_cmake_prefix_path() {
path=`python -c "from distutils.sysconfig import get_python_lib;print(get_python_lib())"`
MY_CMAKE_PREFIX_PATH=$path
}

clone_executorch() {
echo "Cloning executorch to ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src"
rm -rf ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}
mkdir -p ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src
pushd ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src
git clone https://github.com/pytorch/executorch.git
cd executorch
git checkout viable/strict
echo "Install executorch: submodule update"
git submodule sync
git submodule update --init

echo "Applying fixes"
cp ${TORCHCHAT_ROOT}/scripts/fixes_et/module.cpp ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src/executorch/extension/module/module.cpp # ET uses non-standard C++ that does not compile in GCC
cp ${TORCHCHAT_ROOT}/scripts/fixes_et/managed_tensor.h ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src/executorch/extension/runner_util/managed_tensor.h # ET is missing headers for vector/memory. This causes downstream issues when building runner-et.
popd
}

install_executorch_python_libs() {
if [ ! -d "${TORCHCHAT_ROOT}/${ET_BUILD_DIR}" ]; then
echo "Directory ${TORCHCHAT_ROOT}/${ET_BUILD_DIR} does not exist."
echo "Make sur eyou run clone_executorch"
exit 1
fi
pushd ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src
cd executorch

echo "Building and installing python libraries"
echo "Building and installing python libraries"
if [ "${ENABLE_ET_PYBIND}" = false ]; then
echo "Not installing pybind"
bash ./install_requirements.sh
else
echo "Installing pybind"
bash ./install_requirements.sh --pybind xnnpack
fi
pip3 list
popd
}

install_executorch() {
# AOT lib has to be build for model export
# So by default it is built, and you can explicitly opt-out
EXECUTORCH_BUILD_CUSTOM_OPS_AOT_VAR=OFF
if [ "${EXECUTORCH_BUILD_CUSTOM_OPS_AOT}" == "" ]; then
EXECUTORCH_BUILD_CUSTOM_OPS_AOT_VAR=ON
fi

# but for runner not
EXECUTORCH_BUILD_CUSTOM_VAR=OFF
if [ ! ["${EXECUTORCH_BUILD_CUSTOM}" == ""] ]; then
EXECUTORCH_BUILD_CUSTOM_VAR=ON
fi
echo "${EXECUTORCH_BUILD_CUSTOM_OPS_AOT_VAR}"
echo "${EXECUTORCH_BUILD_CUSTOM_VAR}"
if [ ! -d "${TORCHCHAT_ROOT}/${ET_BUILD_DIR}" ]; then
echo "Directory ${TORCHCHAT_ROOT}/${ET_BUILD_DIR} does not exist."
echo "Make sure you run clone_executorch"
exit 1
fi
pushd ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src
cd executorch

if [ "${CMAKE_OUT_DIR}" == "" ]; then
CMAKE_OUT_DIR="cmake-out"
fi

CROSS_COMPILE_ARGS=""
if [ "${CMAKE_OUT_DIR}" == "cmake-out-android" ]; then
CROSS_COMPILE_ARGS="-DCMAKE_TOOLCHAIN_FILE=${CMAKE_TOOLCHAIN_FILE} -DANDROID_ABI=${ANDROID_ABI} -DANDROID_PLATFORM=${ANDROID_PLATFORM}"
fi

echo "Building and installing C++ libraries"
echo "Inside: ${PWD}"
rm -rf ${CMAKE_OUT_DIR}
mkdir ${CMAKE_OUT_DIR}
cmake -DCMAKE_PREFIX_PATH=${MY_CMAKE_PREFIX_PATH} -DCMAKE_BUILD_TYPE=Release -DEXECUTORCH_ENABLE_LOGGING=ON -DEXECUTORCH_LOG_LEVEL=Info -DEXECUTORCH_BUILD_CUSTOM_OPS_AOT=${EXECUTORCH_BUILD_CUSTOM_OPS_AOT_VAR} -DEXECUTORCH_BUILD_CUSTOM=${EXECUTORCH_BUILD_CUSTOM_VAR} -DEXECUTORCH_BUILD_OPTIMIZED=ON -DEXECUTORCH_BUILD_OPTIMIZED=ON -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON -DEXECUTORCH_BUILD_XNNPACK=ON ${CROSS_COMPILE_ARGS} -S . -B ${CMAKE_OUT_DIR} -G Ninja
cmake --build ${CMAKE_OUT_DIR}
cmake --install ${CMAKE_OUT_DIR} --prefix ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/install
popd
}

0 comments on commit 1f7f01e

Please sign in to comment.