Merge branch 'main' into lessw2020/prefill

Author: lessw2020

.gitignore

@@ -6,6 +6,7 @@ __pycache__/
 # C extensions
 *.so
 
+.vscode
 .model-artifacts/
 .venv
 .torchchat
@@ -15,6 +16,7 @@ build/android/*
 et-build/*
 runner-et/cmake-out/*
 runner-aoti/cmake-out/*
+cmake-out/
 
 # pte files
 *.pte
@@ -24,3 +26,7 @@ system_info.txt
 
 # intermediate system file
 .DS_Store
+
+# build artifacts
+checkpoints/
+exportedModels/

install/install_requirements.sh

@@ -52,9 +52,6 @@ PYTORCH_NIGHTLY_VERSION=dev20240814
 # Nightly version for torchvision
 VISION_NIGHTLY_VERSION=dev20240814
 
-# Nightly version for torchao
-AO_NIGHTLY_VERSION=dev20240905
-
 # Nightly version for torchtune
 TUNE_NIGHTLY_VERSION=dev20240910
 
@@ -79,10 +76,6 @@ fi
 REQUIREMENTS_TO_INSTALL=(
   torch=="2.5.0.${PYTORCH_NIGHTLY_VERSION}"
   torchvision=="0.20.0.${VISION_NIGHTLY_VERSION}"
-)
-
-LINUX_REQUIREMENTS_TO_INSTALL=(
-  torchao=="0.5.0.${AO_NIGHTLY_VERSION}"
   torchtune=="0.3.0.${TUNE_NIGHTLY_VERSION}"
 )
 
@@ -94,27 +87,10 @@ LINUX_REQUIREMENTS_TO_INSTALL=(
     "${REQUIREMENTS_TO_INSTALL[@]}"
 )
 
-PLATFORM=$(uname -s)
-
-# Install torchtune and torchao requirements for Linux systems using nightly.
-# For non-Linux systems (e.g., macOS), install torchao from GitHub since nightly
-# build doesn't have macOS build.
-# TODO: Remove this and install nightly build, once it supports macOS
-if [ "$PLATFORM" == "Linux" ];
-then
-  (
-    set -x
-    $PIP_EXECUTABLE install --pre --extra-index-url "${TORCH_NIGHTLY_URL}" --no-cache-dir \
-      "${LINUX_REQUIREMENTS_TO_INSTALL[@]}"
-  )
-else
-  # For torchao need to install from github since nightly build doesn't have macos build.
-  # TODO: Remove this and install nightly build, once it supports macos
-  (
-    set -x
-    $PIP_EXECUTABLE install git+https://github.com/pytorch/ao.git@e11201a62669f582d81cdb33e031a07fb8dfc4f3
-  )
-fi
+(
+  set -x
+  $PIP_EXECUTABLE install torchao=="0.5.0"
+)
 
 if [[ -x "$(command -v nvidia-smi)" ]]; then
   (

torchchat/cli/builder.py

@@ -35,11 +35,8 @@
 from torchchat.utils.measure_time import measure_time
 from torchchat.utils.quantize import quantize_model
 
-# bypass the import issue before torchao is ready on macos
-try:
-    from torchtune.models.convert_weights import meta_to_tune
-except:
-    pass
+from torchtune.models.convert_weights import meta_to_tune
+
 
 
 

torchchat/edge/docs/Android.md

@@ -37,10 +37,10 @@ In your app working directory (for example executorch/examples/demo-apps/android
 copy the AAR to your app libs:
 ```bash
 mkdir -p app/libs
-cp ${TORCHCHAT_ROOT}/android/torchchat/app/libs/executorch.aar ${YOUR_APP_ROOT}/app/libs/executorch.jar
+cp ${TORCHCHAT_ROOT}/android/torchchat/app/libs/executorch.aar ${YOUR_APP_ROOT}/app/libs/executorch.aar
 ```
 
-In your Java app, add the jar file path to your gradle build rule.
+In your Java app, add the aar file path to your gradle build rule.
 ```
 # app/build.grardle.kts
 dependencies {

torchchat/generate.py

@@ -739,7 +739,7 @@ def chat(
         )
         for i in range(num_samples):
             device_sync(device=self.builder_args.device)
-            if i >= 0 and generator_args.chat_mode:
+            if generator_args.chat_mode:
                 prompt = input("User: ")
                 if prompt == "/bye":
                     print("Exiting Chat.\n")
@@ -784,7 +784,6 @@ def chat(
                     )
                     break
 
-            if generator_args.chat_mode and i >= 0:
                 print("Model: ", end="")
 
                 buffer = []
@@ -927,7 +926,8 @@ def callback(x, *, done_generating=False):
                 \nAverage tokens/sec (next tokens): {torch.mean(torch.tensor(aggregate_metrics['next_tokens_per_sec'])).item():.2f} \n\
                 "
         )
-        print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB")
+        if torch.cuda.is_available():
+            print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB")
 
 
 def main(args):

torchchat/model.py

@@ -30,14 +30,9 @@
 
 from torchchat.utils.build_utils import find_multiple, get_precision
 
-# bypass the import issue, if any
-# TODO: remove this once the torchao is ready on macos
-try:
-    from torchtune.models.flamingo import flamingo_decoder, flamingo_vision_encoder
-    from torchtune.modules.model_fusion import DeepFusionModel 
-    from torchtune.models.llama3_1._component_builders import llama3_1 as llama3_1_builder
-except:
-    pass
+from torchtune.models.flamingo import flamingo_decoder, flamingo_vision_encoder
+from torchtune.modules.model_fusion import DeepFusionModel 
+from torchtune.models.llama3_1._component_builders import llama3_1 as llama3_1_builder
 
 config_path = Path(f"{str(Path(__file__).parent)}/model_params")
 

torchchat/utils/scripts/install_utils.sh

@@ -7,6 +7,11 @@
 
 set -ex pipefail
 
+if [ -z "$TORCHCHAT_ROOT" ]; then
+  echo "Defaulting TORCHCHAT_ROOT to $PWD since it is unset."
+  TORCHCHAT_ROOT=$PWD
+fi
+
 install_pip_dependencies() {
   echo "Intalling common pip packages"
   pip3 install wheel "cmake>=3.19" ninja zstd
@@ -20,21 +25,54 @@ function find_cmake_prefix_path() {
   MY_CMAKE_PREFIX_PATH=$path
 }
 
-clone_executorch() {
-  echo "Cloning executorch to ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src"
-  rm -rf ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}
+clone_executorch_internal() {
+  rm -rf ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src
+
   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 $(cat ${TORCHCHAT_ROOT}/install/.pins/et-pin.txt)
-  echo "Install executorch: submodule update"
+  echo "Install ExecuTorch: submodule update"
   git submodule sync
   git submodule update --init
 
   popd
 }
 
+clone_executorch() {
+  echo "Cloning ExecuTorch to ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src"
+
+  # Check if executorch is already cloned and has the correct version
+  if [ -d "${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src/executorch" ]; then
+    pushd ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src/executorch
+
+    # Check if the repo is clean
+    git_status=$(git status --porcelain)
+    if [ -n "$git_status" ]; then
+      echo "ExecuTorch repo is not clean. Removing and recloning."
+      popd
+      clone_executorch_internal
+      return
+    fi
+
+    # Check if the version is the same
+    current_version=$(git rev-parse HEAD)
+    desired_version=$(cat ${TORCHCHAT_ROOT}/install/.pins/et-pin.txt)
+
+    if [ "$current_version" == "$desired_version" ]; then
+      echo "ExecuTorch is already cloned with the correct version. Skipping clone."
+      popd
+      return
+    fi
+
+    echo "ExecuTorch is already cloned but has the wrong version. Removing and recloning."
+    popd
+  fi
+
+  clone_executorch_internal
+}
+
 install_executorch_python_libs() {
   if [ ! -d "${TORCHCHAT_ROOT}/${ET_BUILD_DIR}" ]; then
     echo "Directory ${TORCHCHAT_ROOT}/${ET_BUILD_DIR} does not exist."