Update base for Update on "[ET-VK] Enable Partial GPU lowering via Vulkan in stories model export"

## Context

Simple change to add Vulkan Partitioner as a dependency for the llama exporter and runner, and provide a command line flag to invoke the vulkan partitioner during export.

Included a small change to the Vulkan serializer which was needed for everything to work (i.e. enable serializing multiple graph outputs).

Differential Revision: [D54805831](https://our.internmc.facebook.com/intern/diff/D54805831/)

[ghstack-poisoned]

Author: SS-JIA

.ci/docker/ci_commit_pins/pytorch.txt

@@ -1 +1 @@
-f5b99976adcbb01fd71bd0a39ea15bdac6c9e48a
+6ca9ae4f8693639c395544327f7e362441a58c79

.ci/docker/common/install_conda.sh

@@ -44,13 +44,15 @@ install_pip_dependencies() {
 }
 
 fix_conda_ubuntu_libstdcxx() {
+  cat /etc/issue
   # WARNING: This is a HACK from PyTorch core to be able to build PyTorch on 22.04.
-  # The issue still exists with the latest conda 23.10.0-1 at the time of writing
-  # (2023/11/16).
+  # Specifically, ubuntu-20+ all comes lib libstdc++ newer than 3.30+, but anaconda
+  # is stuck with 3.29. So, remove libstdc++6.so.3.29 as installed by
+  # https://anaconda.org/anaconda/libstdcxx-ng/files?version=11.2.0
   #
   # PyTorch sev: https://github.com/pytorch/pytorch/issues/105248
   # Ref: https://github.com/pytorch/pytorch/blob/main/.ci/docker/common/install_conda.sh
-  if grep -e "[12][82].04.[623]" /etc/issue >/dev/null; then
+  if grep -e "2[02].04." /etc/issue >/dev/null; then
     rm "/opt/conda/envs/py_${PYTHON_VERSION}/lib/libstdc++.so.6"
   fi
 }

.github/pytorch-probot.yml

@@ -2,3 +2,5 @@
 ciflow_push_tags:
 - ciflow/nightly
 - ciflow/trunk
+- ciflow/binaries
+- ciflow/binaries/all

.github/workflows/build-wheels-linux.yml

@@ -0,0 +1,57 @@
+# From https://github.com/pytorch/test-infra/wiki/Using-Nova-Reusable-Build-Workflows
+name: Build Linux Wheels
+
+on:
+  pull_request:
+    paths:
+      - build/packaging/**
+      - .github/workflows/build-wheels-linux.yml
+  push:
+    branches:
+      - nightly
+      - release/*
+    tags:
+      # NOTE: Binary build pipelines should only get triggered on release candidate builds
+      # Release candidate tags look like: v1.11.0-rc1
+      - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
+      - ciflow/binaries/*
+  workflow_dispatch:
+
+jobs:
+  generate-matrix:
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
+    with:
+      package-type: wheel
+      os: linux
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      with-cuda: disabled
+      with-rocm: disabled
+
+  build:
+    needs: generate-matrix
+    permissions:
+      id-token: write
+      contents: read
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - repository: pytorch/executorch
+            pre-script: build/packaging/pre_build_script.sh
+            post-script: build/packaging/post_build_script.sh
+            smoke-test-script: build/packaging/smoke_test.py
+            package-name: executorch
+    name: ${{ matrix.repository }}
+    uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
+    with:
+      repository: ${{ matrix.repository }}
+      ref: ""
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      pre-script: ${{ matrix.pre-script }}
+      post-script: ${{ matrix.post-script }}
+      package-name: ${{ matrix.package-name }}
+      smoke-test-script: ${{ matrix.smoke-test-script }}
+      trigger-event: ${{ github.event_name }}

.github/workflows/build-wheels-m1.yml

@@ -0,0 +1,58 @@
+# From https://github.com/pytorch/test-infra/wiki/Using-Nova-Reusable-Build-Workflows
+name: Build M1 Wheels
+
+on:
+  pull_request:
+    paths:
+      - build/packaging/**
+      - .github/workflows/build-wheels-m1.yml
+  push:
+    branches:
+      - nightly
+      - release/*
+    tags:
+      # NOTE: Binary build pipelines should only get triggered on release candidate builds
+      # Release candidate tags look like: v1.11.0-rc1
+      - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
+      - ciflow/binaries/*
+  workflow_dispatch:
+
+jobs:
+  generate-matrix:
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
+    with:
+      package-type: wheel
+      os: macos-arm64
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      with-cuda: disabled
+      with-rocm: disabled
+
+  build:
+    needs: generate-matrix
+    permissions:
+      id-token: write
+      contents: read
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - repository: pytorch/executorch
+            pre-script: build/packaging/pre_build_script.sh
+            post-script: build/packaging/post_build_script.sh
+            smoke-test-script: build/packaging/smoke_test.py
+            package-name: executorch
+    name: ${{ matrix.repository }}
+    uses: pytorch/test-infra/.github/workflows/build_wheels_macos.yml@main
+    with:
+      repository: ${{ matrix.repository }}
+      ref: ""
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      pre-script: ${{ matrix.pre-script }}
+      post-script: ${{ matrix.post-script }}
+      package-name: ${{ matrix.package-name }}
+      runner-type: macos-m1-stable
+      smoke-test-script: ${{ matrix.smoke-test-script }}
+      trigger-event: ${{ github.event_name }}

.gitmodules

@@ -55,3 +55,7 @@
 [submodule "backends/vulkan/third-party/Vulkan-Headers"]
 	path = backends/vulkan/third-party/Vulkan-Headers
 	url = https://github.com/KhronosGroup/Vulkan-Headers
+[submodule "third-party/lm-evaluation-harness"]
+	path = third-party/lm-evaluation-harness
+	url = https://github.com/EleutherAI/lm-evaluation-harness
+	branch = v0.4.1

backends/apple/coreml/quantizer/coreml_quantizer.py

@@ -0,0 +1,7 @@
+# Copyright © 2024 Apple Inc. All rights reserved.
+#
+# Please refer to the license found in the LICENSE file in the root directory of the source tree.
+
+from coremltools.optimize.torch.quantization._coreml_quantizer import (  # noqa: FLAKE8 F401
+    CoreMLQuantizer,
+)

backends/apple/coreml/scripts/generate_test_models.sh

@@ -23,8 +23,8 @@ cd "$EXECUTORCH_ROOT_PATH"
 MODELS=("add" "mul" "mv3")
 for MODEL in "${MODELS[@]}"
 do
-  # TODO: Don't use the script in examples directory. 
-  python3 -m examples.apple.coreml.scripts.export_and_delegate --model_name "$MODEL" --save_processed_bytes
+  # TODO: Don't use the script in examples directory.
+  python3 -m examples.apple.coreml.scripts.export --model_name "$MODEL" --save_processed_bytes
   mv -f "$MODEL""_coreml_all.pte" "$COREML_DIR_PATH/runtime/test/models"
   mv -f "$MODEL""_coreml_all.bin" "$COREML_DIR_PATH/runtime/test/models"
 done

backends/apple/coreml/setup.md

@@ -11,18 +11,18 @@ This is a tutorial for setting up the Core ML backend.
 ```
 cd executorch
 
-./backends/apple/coreml/scripts/install_requirements.sh   
+./backends/apple/coreml/scripts/install_requirements.sh
 
-``` 
+```
 
-3. Run the example script to validate that the **Core ML** backend is set up correctly. 
+3. Run the example script to validate that the **Core ML** backend is set up correctly.
 
 ```
 cd executorch
 
 # Saves add_coreml_all.pte in the current directory if successful.
 
-python3 -m examples.apple.coreml.scripts.export_and_delegate --model_name add 
+python3 -m examples.apple.coreml.scripts.export --model_name add
 
 ```
 
@@ -66,6 +66,6 @@ coreml_backend.xcframework
 - Accelerate.framework
 - CoreML.framework
 - libsqlite3.tbd
-``` 
+```
 
-6. The target could now run a **Core ML** delegated **Program**. 
+6. The target could now run a **Core ML** delegated **Program**.

backends/apple/coreml/test/test_coreml_quantizer.py

@@ -0,0 +1,112 @@
+# Copyright © 2024 Apple Inc. All rights reserved.
+#
+# Please refer to the license found in the LICENSE file in the root directory of the source tree.
+
+from typing import Tuple
+
+import numpy as np
+import pytest
+
+import torch
+
+from coremltools.optimize.torch.quantization.quantization_config import (
+    LinearQuantizerConfig,
+    QuantizationScheme,
+)
+
+from executorch.backends.apple.coreml.quantizer.coreml_quantizer import CoreMLQuantizer
+from torch._export import capture_pre_autograd_graph
+from torch.ao.quantization.quantize_pt2e import (
+    convert_pt2e,
+    prepare_pt2e,
+    prepare_qat_pt2e,
+)
+
+
+class TestCoreMLQuantizer:
+    @staticmethod
+    def quantize_and_compare(
+        model,
+        example_inputs: Tuple[torch.Tensor],
+        quantization_type: str,
+    ) -> None:
+        assert quantization_type in {"PTQ", "QAT"}
+
+        pre_autograd_aten_dialect = capture_pre_autograd_graph(model, example_inputs)
+
+        quantization_config = LinearQuantizerConfig.from_dict(
+            {
+                "global_config": {
+                    "quantization_scheme": QuantizationScheme.symmetric,
+                    "milestones": [0, 0, 10, 10],
+                    "activation_dtype": torch.quint8,
+                    "weight_dtype": torch.qint8,
+                    "weight_per_channel": True,
+                }
+            }
+        )
+        quantizer = CoreMLQuantizer(quantization_config)
+
+        if quantization_type == "PTQ":
+            prepared_graph = prepare_pt2e(pre_autograd_aten_dialect, quantizer)
+        elif quantization_type == "QAT":
+            prepared_graph = prepare_qat_pt2e(pre_autograd_aten_dialect, quantizer)
+
+        prepared_graph(*example_inputs)
+        converted_graph = convert_pt2e(prepared_graph)
+
+        model_output = model(*example_inputs).detach().numpy()
+        quantized_output = converted_graph(*example_inputs).detach().numpy()
+        np.testing.assert_allclose(quantized_output, model_output, rtol=5e-2, atol=5e-2)
+
+    @pytest.mark.parametrize("quantization_type", ("PTQ", "QAT"))
+    def test_conv_relu(self, quantization_type):
+        SHAPE = (1, 3, 256, 256)
+
+        class Model(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.conv = torch.nn.Conv2d(
+                    in_channels=3, out_channels=16, kernel_size=3, padding=1
+                )
+                self.relu = torch.nn.ReLU()
+
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                a = self.conv(x)
+                return self.relu(a)
+
+        model = Model()
+
+        example_inputs = (torch.randn(SHAPE),)
+        self.quantize_and_compare(
+            model,
+            example_inputs,
+            quantization_type,
+        )
+
+    @pytest.mark.parametrize("quantization_type", ("PTQ", "QAT"))
+    def test_linear(self, quantization_type):
+        SHAPE = (1, 5)
+
+        class Model(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.linear = torch.nn.Linear(5, 10)
+
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return self.linear(x)
+
+        model = Model()
+
+        example_inputs = (torch.randn(SHAPE),)
+        self.quantize_and_compare(
+            model,
+            example_inputs,
+            quantization_type,
+        )
+
+
+if __name__ == "__main__":
+    test_runner = TestCoreMLQuantizer()
+    test_runner.test_conv_relu("PTQ")
+    test_runner.test_linear("QAT")