Move PT pin and AO commit hash to benefit from SDPA MPS implementation (#964)

* Move PT pin and AO commit hash to benefit from SDPA MPS implementation in PT

* Move PT nightly to 2024-07-28

* Shuffle weights prior to _convert_weight_to_int4pack call in gguf_loader

* Uninstall torchao on M1 before re-installing it

* Update ET pin

* Uninstall torchao on M1 before re-installing it

* Update ET pin

Author: manuelcandales

.github/workflows/pull.yml

@@ -627,11 +627,14 @@ jobs:
       runner: macos-m1-stable  # neeps MPS, was macos-m1-stable
       script: |
         set -x
-        # NS: Remove previous installation  of torch first
-        # as this script does not isntall anything into conda env but rather as system dep
+        # NS/MC: Remove previous installation of torch and torchao first
+        # as this script does not install anything into conda env but rather as system dep
         pip3 uninstall -y torch || true
         set -eou pipefail
 
+        pip3 uninstall -y torchao || true
+        set -eou pipefail
+
         echo "::group::Print machine info"
         uname -a
         sysctl machdep.cpu.brand_string
@@ -736,10 +739,12 @@ jobs:
       runner: macos-m1-stable  # needs MPS, was macos-m1-stable
       script: |
         set -x
-        # NS: Remove previous installation  of torch first
-        # as this script does not isntall anything into conda env but rather as system dep
+        # NS/MC: Remove previous installation of torch and torchao first
+        # as this script does not install anything into conda env but rather as system dep
         pip3 uninstall -y torch || true
+        set -eou pipefail
 
+        pip3 uninstall -y torchao || true
         set -eou pipefail
 
         echo "::group::Print machine info"

.pins/et-pin.txt

@@ -1 +1 @@
-c7574994ecd775fdaacc0f2de27089526e05b108
+91298923a0076c1b41059efb6dad2876426e4b03

build/gguf_loader.py

@@ -170,8 +170,9 @@ def load_model_and_state_dict(
             if load_state_dict:
                 q, s, z = Q4_0.unpack(t)
                 scales_and_zeros = pack_scales_and_zeros(s, z)
+                q_uint8 = (q[::, ::2] << 4 | q[::, 1::2]).to(torch.uint8)
                 weight_int4pack = torch.ops.aten._convert_weight_to_int4pack(
-                    q, inner_k_tiles
+                    q_uint8, inner_k_tiles
                 )
                 state_dict[f"{fqn}.weight"] = weight_int4pack
                 state_dict[f"{fqn}.scales_and_zeros"] = scales_and_zeros

install_requirements.sh

@@ -47,7 +47,7 @@ fi
 # NOTE: If a newly-fetched version of the executorch repo changes the value of
 # NIGHTLY_VERSION, you should re-run this script to install the necessary
 # package versions.
-NIGHTLY_VERSION=dev20240710
+NIGHTLY_VERSION=dev20240728
 
 # Uninstall triton, as nightly will depend on pytorch-triton, which is one and the same
 (
@@ -82,7 +82,7 @@ REQUIREMENTS_TO_INSTALL=(
 # TODO: Remove this and install nightly build, once it supports macos
 (
   set -x
-  $PIP_EXECUTABLE install git+https://github.com/pytorch/ao.git@d36de1b144b73bf753bd082109c2b5d0141abd5b
+  $PIP_EXECUTABLE install git+https://github.com/pytorch/ao.git@d477c0e59b458b5617dcb3e999290a87df3070d8
 )
 if [[ -x "$(command -v nvidia-smi)" ]]; then
   (

quantization/qops.py

@@ -401,8 +401,9 @@ def _prepare_weight_and_scales_and_zeros(
         weight_int32, scales_and_zeros = group_quantize_tensor(
             weight_bf16, n_bit=4, groupsize=groupsize
         )
+        weight_uint8 = (weight_int32[::, ::2] << 4 | weight_int32[::, 1::2]).to(torch.uint8)
         weight_int4pack = torch.ops.aten._convert_weight_to_int4pack(
-            weight_int32, inner_k_tiles
+            weight_uint8, inner_k_tiles
         )
         return weight_int4pack, scales_and_zeros