[llava][18/N] Move token generation loop to a class

Differential Revision: D61047601

Pull Request resolved: #4652

Author: larryliu0820

.buckconfig

@@ -23,3 +23,13 @@
 
 [parser]
   target_platform_detector_spec = target:root//...->prelude//platforms:default target:shim//...->prelude//platforms:default
+
+# Limit the number of files that the buck daemon needs to monitor. If every
+# submodule is cloned recursively, some system can fail to build with "OS file
+# watch limit reached".
+[project]
+  ignore = \
+      .git, \
+      **/.git, \
+      cmake-out, \
+      pip-out

.ci/docker/build.sh

@@ -37,6 +37,10 @@ case "${IMAGE_NAME}" in
     ARM_SDK=yes
     CLANG_VERSION=12
     ;;
+  executorch-ubuntu-22.04-qnn-sdk)
+    QNN_SDK=yes
+    CLANG_VERSION=12
+    ;;
   executorch-ubuntu-22.04-clang12-android)
     LINTRUNNER=""
     CLANG_VERSION=12
@@ -59,6 +63,9 @@ cp ../../requirements-lintrunner.txt ./
 # with a new image hash when the content here is updated
 cp -r ../../examples/arm/ ./arm
 
+# Copy qnn setup script from root to here
+cp -r ../../backends/qualcomm/ ./qualcomm
+
 docker build \
   --no-cache \
   --progress=plain \
@@ -72,6 +79,7 @@ docker build \
   --build-arg "LINTRUNNER=${LINTRUNNER:-}" \
   --build-arg "BUILD_DOCS=${BUILD_DOCS}" \
   --build-arg "ARM_SDK=${ARM_SDK:-}" \
+  --build-arg "QNN_SDK=${QNN_SDK:-}" \
   --build-arg "ANDROID_NDK_VERSION=${ANDROID_NDK_VERSION:-}" \
   -f "${OS}"/Dockerfile \
   "$@" \

.ci/docker/ci_commit_pins/buck2.txt

@@ -1 +1 @@
-2024-02-15
+2024-05-15

.ci/docker/ci_commit_pins/pytorch.txt

@@ -1 +1 @@
-a52b4e22571507abc35c2d47de138497190d2e0a
+b556d31586845fb1e296a975d2b85d9d325205c9

.ci/docker/common/install_android.sh

@@ -11,22 +11,33 @@ set -ex
 [ -n "${ANDROID_NDK_VERSION}" ]
 
 install_prerequiresites() {
-  apt-get update
+  OS=$(grep -oP '(?<=^ID=).+' /etc/os-release | tr -d '"')
+  case "$OS" in
+    amzn)
+      # https://docs.aws.amazon.com/corretto/latest/corretto-17-ug/amazon-linux-install.html
+      yum install -y java-17-amazon-corretto \
+        ca-certificates \
+        ant
+      ;;
+    *)
+      apt-get update
 
-  # NB: Need OpenJDK 17 at the minimum
-  apt-get install -y --no-install-recommends \
-    openjdk-17-jdk \
-    ca-certificates-java \
-    ant
+      # NB: Need OpenJDK 17 at the minimum
+      apt-get install -y --no-install-recommends \
+        openjdk-17-jdk \
+        ca-certificates-java \
+        ant
 
-  # Cleanup package manager
-  apt-get autoclean && apt-get clean
-  rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
+      # Cleanup package manager
+      apt-get autoclean && apt-get clean
+      rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
+    ;;
+  esac
 }
 
 install_ndk() {
   NDK_INSTALLATION_DIR=/opt/ndk
-  mkdir -p "${NDK_INSTALLATION_DIR}"
+  rm -rf "${NDK_INSTALLATION_DIR}" && mkdir -p "${NDK_INSTALLATION_DIR}"
 
   pushd /tmp
   # The NDK installation is cached on ossci-android S3 bucket
@@ -54,7 +65,7 @@ install_cmdtools() {
 
 install_sdk() {
   SDK_INSTALLATION_DIR=/opt/android/sdk
-  mkdir -p "${SDK_INSTALLATION_DIR}"
+  rm -rf "${SDK_INSTALLATION_DIR}" && mkdir -p "${SDK_INSTALLATION_DIR}"
 
   # These are the tools needed to build Android apps
   yes | /opt/cmdline-tools/bin/sdkmanager --sdk_root="${SDK_INSTALLATION_DIR}" --install "platforms;android-34"

.ci/docker/common/install_base.sh

@@ -23,7 +23,8 @@ install_ubuntu() {
     unzip \
     gdb \
     rsync \
-    libssl-dev
+    libssl-dev \
+    zip
 
   # Cleanup package manager
   apt-get autoclean && apt-get clean

.ci/docker/requirements-ci.txt

@@ -1,13 +1,17 @@
 mpmath==1.3.0
-numpy==1.25.2
+numpy==1.21.3; python_version == '3.10'
+numpy==1.23.2; python_version == '3.11'
+numpy; python_version >= '3.12'
 PyYAML==6.0.1
 ruamel.yaml==0.17.32
 sympy==1.12
 timm==0.6.13
 tomli==2.0.1
 torchsr==1.0.4
-transformers==4.36.0
+transformers==4.38.0
 zstd==1.5.5.1
+pandas==2.0.3; python_version == '3.10'
+pandas; python_version >= '3.11'
 pytest==7.2.0
 pytest-cov==4.1.0
 expecttest==0.1.6

.ci/docker/ubuntu/Dockerfile

@@ -82,5 +82,10 @@ COPY --chown=ci-user:ci-user ./arm /opt/arm
 # Set up ARM SDK if needed
 RUN if [ -n "${ARM_SDK}" ]; then git config --global user.email "[email protected]"; git config --global user.name "OSS CI"; bash /opt/arm/setup.sh --i-agree-to-the-contained-eula /opt/arm-sdk; chown -R ci-user:ci-user /opt/arm-sdk; fi
 
+ARG QNN_SDK
+COPY --chown=ci-user:ci-user ./qualcomm /opt/qualcomm
+# Set up QNN SDK if needed
+RUN if [ -n "${QNN_SDK}" ]; then git config --global user.email "[email protected]"; git config --global user.name "OSS CI"; fi
+
 USER ci-user
 CMD ["bash"]

.ci/scripts/build-qnn-sdk.sh

@@ -0,0 +1,19 @@
+#!/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
+
+build_qnn_backend() {
+  echo "Start building qnn backend."
+  export ANDROID_NDK_ROOT=/opt/ndk
+  export QNN_SDK_ROOT=/tmp/qnn/2.23.0.240531
+  export EXECUTORCH_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
+
+  bash backends/qualcomm/scripts/build.sh --skip_aarch64 --job_number 2 --release
+}
+
+build_qnn_backend

.ci/scripts/build_llama_android.sh

@@ -0,0 +1,56 @@
+#!/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 -exu
+
+# shellcheck source=/dev/null
+source "$(dirname "${BASH_SOURCE[0]}")/utils.sh"
+
+install_executorch_and_backend_lib() {
+  echo "Installing executorch and xnnpack backend"
+  rm -rf cmake-android-out && mkdir cmake-android-out
+  ANDROID_NDK=/opt/ndk
+  BUCK2=buck2
+  ANDROID_ABI=arm64-v8a
+  cmake -DBUCK2="${BUCK2}" \
+    -DCMAKE_TOOLCHAIN_FILE="${ANDROID_NDK}/build/cmake/android.toolchain.cmake" \
+    -DANDROID_ABI="${ANDROID_ABI}" \
+    -DANDROID_PLATFORM=android-23 \
+    -DCMAKE_INSTALL_PREFIX=cmake-android-out \
+    -DCMAKE_BUILD_TYPE=Release \
+    -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
+    -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
+    -DEXECUTORCH_BUILD_XNNPACK=ON \
+    -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \
+    -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \
+    -DEXECUTORCH_BUILD_KERNELS_CUSTOM=ON \
+    -DXNNPACK_ENABLE_ARM_BF16=OFF \
+    -Bcmake-android-out .
+
+  cmake --build cmake-android-out -j4 --target install --config Release
+}
+
+build_llama_runner() {
+    echo "Building llama runner for Android..."
+    ANDROID_ABI=arm64-v8a
+    cmake -DBUCK2="${BUCK2}" \
+    -DCMAKE_TOOLCHAIN_FILE="$ANDROID_NDK"/build/cmake/android.toolchain.cmake  \
+    -DANDROID_ABI="${ANDROID_ABI}" \
+    -DANDROID_PLATFORM=android-23 \
+    -DCMAKE_INSTALL_PREFIX=cmake-android-out \
+    -DCMAKE_BUILD_TYPE=Release -DPYTHON_EXECUTABLE=python \
+    -DEXECUTORCH_BUILD_XNNPACK=ON \
+    -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \
+    -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \
+    -DEXECUTORCH_BUILD_KERNELS_CUSTOM=ON \
+    -Bcmake-android-out/examples/models/llama2 examples/models/llama2
+
+    cmake --build cmake-android-out/examples/models/llama2 -j4 --config Release
+}
+install_flatc_from_source
+install_executorch_and_backend_lib
+build_llama_runner

.ci/scripts/gather_test_models.py

@@ -13,7 +13,6 @@
 from examples.models import MODEL_NAME_TO_MODEL
 from examples.xnnpack import MODEL_NAME_TO_OPTIONS
 
-
 DEFAULT_RUNNERS = {
     "linux": "linux.2xlarge",
     "macos": "macos-m1-stable",
@@ -24,9 +23,11 @@
         "w2l": "linux.12xlarge",
         "ic4": "linux.12xlarge",
         "resnet50": "linux.12xlarge",
+        "llava": "linux.12xlarge",
         # This one causes timeout on smaller runner, the root cause is unclear (T161064121)
         "dl3": "linux.12xlarge",
         "emformer_join": "linux.12xlarge",
+        "emformer_predict": "linux.12xlarge",
     }
 }
 
@@ -35,9 +36,11 @@
     # Just some examples on how custom timeout can be set
     "linux": {
         "mobilebert": 90,
+        "emformer_predict": 360,
     },
     "macos": {
         "mobilebert": 90,
+        "emformer_predict": 360,
     },
 }
 
@@ -56,7 +59,7 @@ def parse_args() -> Any:
         "-e",
         "--event",
         type=str,
-        choices=["pull_request", "push"],
+        choices=["pull_request", "push", "schedule"],
         required=True,
         help="GitHub CI Event. See https://docs.github.com/en/actions/using-workflows/workflow-syntax-for-github-actions#on",
     )
@@ -83,7 +86,21 @@ def model_should_run_on_event(model: str, event: str) -> bool:
     We put higher priority and fast models to pull request and rest to push.
     """
     if event == "pull_request":
-        return model in ["add", "ic3", "mv2", "mv3", "resnet18", "vit"]
+        return model in ["mv3", "vit"]
+    elif event == "push":
+        # 'emformer_predict' is running super slow. Only run it periodically
+        return model not in ["emformer_predict"]
+    else:
+        return True
+
+
+def model_should_run_on_target_os(model: str, target_os: str) -> bool:
+    """
+    A helper function to decide whether a model should be tested on a target os (linux/macos).
+    For example, a big model can be disabled in macos due to the limited macos resources.
+    """
+    if target_os == "macos":
+        return model not in ["llava"]
     return True
 
 
@@ -119,6 +136,9 @@ def export_models_for_ci() -> dict[str, dict]:
         if not model_should_run_on_event(name, event):
             continue
 
+        if not model_should_run_on_target_os(name, target_os):
+            continue
+
         if backend == "xnnpack":
             if name not in MODEL_NAME_TO_OPTIONS:
                 continue

.ci/scripts/setup-macos.sh

@@ -28,15 +28,18 @@ install_buck() {
   fi
 
   pushd .ci/docker
-
-  # TODO(huydo): This is a one-off copy of buck2 2024-02-15 to unblock Jon and
+  # TODO(huydo): This is a one-off copy of buck2 2024-05-15 to unblock Jon and
   # re-enable ShipIt. It’s not ideal that upgrading buck2 will require a manual
   # update the cached binary on S3 bucket too. Let me figure out if there is a
   # way to correctly implement the previous setup of installing a new version of
   # buck2 only when it’s needed. AFAIK, the complicated part was that buck2
-  # --version doesn't say anything w.r.t its release version, i.e. 2024-02-15.
+  # --version doesn't say anything w.r.t its release version, i.e. 2024-05-15.
   # See D53878006 for more details.
-  BUCK2=buck2-aarch64-apple-darwin.zst
+  #
+  # If you need to upgrade buck2 version on S3, please reach out to Dev Infra
+  # team for help.
+  BUCK2_VERSION=$(cat ci_commit_pins/buck2.txt)
+  BUCK2=buck2-aarch64-apple-darwin-${BUCK2_VERSION}.zst
   curl -s "https://ossci-macos.s3.amazonaws.com/${BUCK2}" -o "${BUCK2}"
 
   zstd -d "${BUCK2}" -o buck2
@@ -74,16 +77,18 @@ install_sccache() {
 
   export PATH="${SCCACHE_PATH}:${PATH}"
 
-  # Create temp directory for sccache shims
-  TMP_DIR=$(mktemp -d)
-  trap 'rm -rfv ${TMP_DIR}' EXIT
+  # Create temp directory for sccache shims if TMP_DIR doesn't exist
+  if [ -z "${TMP_DIR:-}" ]; then
+    TMP_DIR=$(mktemp -d)
+    trap 'rm -rfv ${TMP_DIR}' EXIT
+    export PATH="${TMP_DIR}:$PATH"
+  fi
 
   write_sccache_stub "${TMP_DIR}/c++"
   write_sccache_stub "${TMP_DIR}/cc"
   write_sccache_stub "${TMP_DIR}/clang++"
   write_sccache_stub "${TMP_DIR}/clang"
 
-  export PATH="${TMP_DIR}:$PATH"
   sccache --zero-stats || true
 }
 
@@ -104,6 +109,12 @@ print_cmake_info() {
   codesign -f -s - "${CMAKE_EXEC}" || true
 }
 
+setup_macos_env_variables() {
+  CMAKE_PREFIX_PATH=$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')
+  export CMAKE_PREFIX_PATH
+}
+
+setup_macos_env_variables
 # NB: we need buck2 in all cases because cmake build also depends on calling
 # buck2 atm
 install_buck

.ci/scripts/setup-qnn-deps.sh

@@ -0,0 +1,29 @@
+#!/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_qnn() {
+  echo "Start installing qnn."
+  QNN_INSTALLATION_DIR=/tmp/qnn
+  mkdir -p "${QNN_INSTALLATION_DIR}"
+
+  curl -Lo /tmp/v2.23.0.24.06.24.zip "https://softwarecenter.qualcomm.com/api/download/software/qualcomm_neural_processing_sdk/v2.23.0.24.06.24.zip"
+  echo "Finishing downloading qnn sdk."
+  unzip -qo /tmp/v2.23.0.24.06.24.zip -d /tmp
+  echo "Finishing unzip qnn sdk."
+
+
+  # Print the content for manual verification
+  ls -lah "/tmp/qairt"
+  mv "/tmp/qairt"/* "${QNN_INSTALLATION_DIR}"
+  echo "Finishing installing qnn '${QNN_INSTALLATION_DIR}' ."
+
+  ls -lah "${QNN_INSTALLATION_DIR}"
+}
+
+install_qnn