Bumping ET Pin to 7-19-2024 pytorch#927 (pytorch#931)

* Update et-pin.txt

* Updating MlockConfig to LoadMode

* Bump AO version to pick up AO PR509 fixing infer_schema

* Reverting GH test set up to match user set up

* Fix custom op

* bump pin

---------

Co-authored-by: Mengwei Liu <[email protected]>

Author: Jack-Khuu

.github/workflows/pull.yml

@@ -449,13 +449,12 @@ jobs:
           fi
       - name: Install requirements
         run: |
-          # Have to install ET first because deps of Torchchat might not be the same.
-          export TORCHCHAT_ROOT=$PWD
-          ./scripts/install_et.sh
-
           echo "Intalling pip3 packages"
           ./install_requirements.sh
 
+          export TORCHCHAT_ROOT=$PWD
+          ./scripts/install_et.sh
+
           pip3 list
           python3 -c 'import torch;print(f"torch: {torch.__version__, torch.version.git_version}")'
           python3 -c 'import torchvision;print(f"torchvision: {torchvision.__version__, torchvision.version.git_version}")'

.pins/et-pin.txt

@@ -1 +1 @@
-172574a6be5910a4609e4ed1bef2b6b8475ddb3d
+c7574994ecd775fdaacc0f2de27089526e05b108

build/model_et.py

@@ -8,6 +8,8 @@
 import torch.nn as nn
 from executorch.extension.pybindings import portable_lib as exec_lib
 
+# ET changed the way it's loading the custom ops so it's not included in portable_lib but has to be loaded separately.
+from executorch.examples.models.llama2.custom_ops import sdpa_with_kv_cache # no-qa
 
 class PTEModel(nn.Module):
     def __init__(self, config, path) -> None:

install_requirements.sh

@@ -73,7 +73,7 @@ $PIP_EXECUTABLE install --extra-index-url "${TORCH_NIGHTLY_URL}" \
 
 # 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
-$PIP_EXECUTABLE install git+https://github.com/pytorch/ao.git@ca1b98db60543a1669a32e842762fc008c178376
+$PIP_EXECUTABLE install git+https://github.com/pytorch/ao.git@d36de1b144b73bf753bd082109c2b5d0141abd5b
 if [[ -x "$(command -v nvidia-smi)" ]]; then
   $PYTHON_EXECUTABLE scripts/patch_triton.py
 fi

runner/run.cpp

@@ -156,7 +156,7 @@ void build_transformer(
 #else //__ET_MODEL__
   t->runner = new Module(
       /* path to PTE model */ model_path,
-      /* PTE mmap settings */ Module::MlockConfig::UseMlockIgnoreErrors);
+      /* PTE mmap settings */ Module::LoadMode::MmapUseMlockIgnoreErrors);
 #endif
 }