Merge branch 'main' into main

.github/workflows/doc_build.yml

@@ -43,16 +43,9 @@ jobs:
           python -m pip install -e .
           cd docs
           python -m pip install -r requirements.txt
-      - name: Get the torchtune version
-        run: |
-          # Get the github.ref_name and save into the
-          # REF_NAME variable. This will be passed in
-          # conf.py to display the version in the
-          # site dropdown
-          REF_NAME=${{ github.ref_name }}
-          TORCHAO_VERSION_DOCS="${REF_NAME}"
-          echo "$TORCHAO_VERSION_DOCS"
       - name: Build docs
+        env:
+          TORCHAO_VERSION_DOCS: ${{ github.ref }}
         run: |
           cd docs
           make html
@@ -61,19 +54,41 @@ jobs:
           name: Doc-Build
           path: docs/build/html/
 
+  doc-preview:
+    runs-on: [self-hosted, linux.2xlarge]
+    needs: build_docs
+    if: ${{ github.event_name == 'pull_request' }}
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+      - name: Download artifact
+        uses: actions/download-artifact@v3
+        with:
+          name: Doc-Build
+          path: docs
+      - name: Upload docs preview
+        uses: seemethere/upload-artifact-s3@v5
+        if: ${{ github.event_name == 'pull_request' }}
+        with:
+          retention-days: 14
+          s3-bucket: doc-previews
+          if-no-files-found: error
+          path: docs
+          s3-prefix: pytorch/ao/${{ github.event.pull_request.number }}
+
   upload:
     runs-on: ubuntu-latest
     permissions:
       # Grant write permission here so that the doc can be pushed to gh-pages branch
       contents: write
     needs: build_docs
-    if: github.repository == 'pytorch-labs/ao' && github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/heads/release/') || startsWith(github.ref, 'refs/tags/v') || github.event_name == 'workflow_dispatch')
+    if: github.repository == 'pytorch/ao' && github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/v') || github.event_name == 'workflow_dispatch')
     steps:
       - name: Checkout
         uses: actions/checkout@v3
         with:
           ref: gh-pages
-          persist-credentials: false
+          persist-credentials: true
       - name: Download artifact
         uses: actions/download-artifact@v3
         with:
@@ -89,23 +104,20 @@ jobs:
       - name: Move and commit changes
         run: |
           set -euo pipefail
-          REF_TYPE=${{ github.ref_type }}
-          REF_NAME=${{ github.ref_name }}
-          if [[ "${REF_TYPE}" == branch ]]; then
-            TARGET_FOLDER="${REF_NAME}"
-          elif [[ "${REF_TYPE}" == tag ]]; then
-            case "${REF_NAME}" in
-              *-rc*)
-                echo "Aborting upload since this is an RC tag: ${REF_NAME}"
-                exit 0
-                ;;
-              *)
-                TARGET_FOLDER=$(echo "${REF_NAME}" | sed 's/v\([0-9]\+\)\.\([0-9]\+\)\.[0-9]\+/\1.\2/')
-                ;;
-            esac
-          fi
+          # Get github.ref for the output doc folder. By default "main"
+          # If matches a tag like refs/tags/v1.12.0-rc3 or
+          # refs/tags/v1.12.0 convert to 1.12
+          GITHUB_REF=${{ github.ref }}
 
+          # Convert refs/tags/v1.12.0rc3 into 1.12.
+          # Adopted from https://github.com/pytorch/pytorch/blob/main/.github/workflows/_docs.yml#L150C11-L155C13
+          if [[ "${GITHUB_REF}" =~ ^refs/tags/v([0-9]+\.[0-9]+)\.* ]]; then
+            TARGET_FOLDER="${BASH_REMATCH[1]}"
+          else
+            TARGET_FOLDER="main"
+          fi
           echo "Target Folder: ${TARGET_FOLDER}"
+
           mkdir -p "${TARGET_FOLDER}"
           rm -rf "${TARGET_FOLDER}"/*
           mv docs/* "${TARGET_FOLDER}"
@@ -114,4 +126,4 @@ jobs:
           git config user.email '[email protected]'
           git add "${TARGET_FOLDER}" || true
           git commit -m "auto-generating sphinx docs" || true
-          git push -f      
+          git push -f

.github/workflows/regression_test.yml

@@ -8,49 +8,59 @@ on:
     branches:
       - main
 
+concurrency:
+  group: regression_test-${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}
+  cancel-in-progress: true
+
+env:
+  HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
+
 jobs:
   test:
     strategy:
+      fail-fast: false
       matrix:
         include:
           - name: CUDA 2.2.2
-            runs-on: 4-core-ubuntu-gpu-t4
+            runs-on: linux.g5.12xlarge.nvidia.gpu
             torch-spec: 'torch==2.2.2'
-          - name: CUDA 2.3 RC
-            runs-on: 4-core-ubuntu-gpu-t4
-            torch-spec: 'torch==2.3.0 --index-url https://download.pytorch.org/whl/test/cu121'
+            gpu-arch-type: "cuda"
+            gpu-arch-version: "12.1"
+          - name: CUDA 2.3
+            runs-on: linux.g5.12xlarge.nvidia.gpu
+            torch-spec: 'torch==2.3.0'
+            gpu-arch-type: "cuda"
+            gpu-arch-version: "12.1"
           - name: CUDA Nightly
-            runs-on: 4-core-ubuntu-gpu-t4
+            runs-on: linux.g5.12xlarge.nvidia.gpu
             torch-spec: '--pre torch --index-url https://download.pytorch.org/whl/nightly/cu121'
+            gpu-arch-type: "cuda"
+            gpu-arch-version: "12.1"
           - name: CPU 2.2.2
-            runs-on: 32-core-ubuntu
+            runs-on: linux.4xlarge
             torch-spec: 'torch==2.2.2 --index-url https://download.pytorch.org/whl/cpu'
-          - name: CPU 2.3 RC
-            runs-on: 32-core-ubuntu
-            torch-spec: 'torch==2.3.0 --index-url https://download.pytorch.org/whl/test/cpu'
+            gpu-arch-type: "cpu"
+            gpu-arch-version: ""
+          - name: CPU 2.3
+            runs-on: linux.4xlarge
+            torch-spec: 'torch==2.3.0 --index-url https://download.pytorch.org/whl/cpu'
+            gpu-arch-type: "cpu"
+            gpu-arch-version: ""
           - name: Nightly CPU
-            runs-on: 32-core-ubuntu
+            runs-on: linux.4xlarge
             torch-spec: '--pre torch --index-url https://download.pytorch.org/whl/nightly/cpu'
-    runs-on: ${{ matrix.runs-on }}
-    steps:
-    - uses: actions/checkout@v2
-
-    - name: Set up Python
-      uses: actions/setup-python@v2
-      with:
-        python-version: '3.9'
+            gpu-arch-type: "cpu"
+            gpu-arch-version: ""
 
-    - name: Install dependencies
-      run: |
+    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    with:
+      runner: ${{ matrix.runs-on }}
+      gpu-arch-type: ${{ matrix.gpu-arch-type }}
+      gpu-arch-version: ${{ matrix.gpu-arch-version }}
+      script: |
         python -m pip install --upgrade pip
         pip install ${{ matrix.torch-spec }}
         pip install -r requirements.txt
         pip install -r dev-requirements.txt
-
-    - name: Install package
-      run: |
-        pip install .
-
-    - name: Run tests
-      run: |
+        python setup.py install
         pytest test --verbose -s

.gitignore

@@ -261,7 +261,7 @@ docs/dev
 docs/build
 docs/source/tutorials/*
 docs/source/gen_modules/*
-docs/source/sg_execution_times.rst
+docs/source/sg_execution_times
 
 # LevelDB files
 *.sst

README.md

@@ -1,20 +1,28 @@
 # torchao: PyTorch Architecture Optimization
 
-**Note: This repository is currently under heavy development - if you have suggestions on the API or use-cases you'd like to be covered, please open an github issue**
+**Note: This repository is currently under heavy development - if you have suggestions on the API or use-cases you'd like to be covered, please open an GitHub issue**
 
 ## Introduction
-
 torchao is a PyTorch native library for optimizing your models using lower precision dtypes, techniques like quantization and sparsity and performant kernels.
 
-## Our Goals
-torchao embodies PyTorch’s design philosophy [details](https://pytorch.org/docs/stable/community/design.html), especially "usability over everything else". Our vision for this repository is the following:
+## Get Started
+To try out our APIs, you can check out API examples in [quantization](./torchao/quantization) (including `autoquant`), [sparsity](./torchao/sparsity), [dtypes](./torchao/dtypes).
 
-* Composability: Native solutions for optimization techniques that compose with both `torch.compile` and `FSDP` 
-    * For example, for QLoRA for new dtypes support
-* Interoperability: Work with the rest of the PyTorch ecosystem such as torchtune, gpt-fast and ExecuTorch
-* Transparent Benchmarks: Regularly run performance benchmarking of our APIs across a suite of Torchbench models and across hardware backends
-* Heterogeneous Hardware: Efficient kernels that can run on CPU/GPU based server (w/ torch.compile) and mobile backends (w/ ExecuTorch).
-* Infrastructure Support: Release packaging solution for kernels and a CI/CD setup that runs these kernels on different backends. 
+## Installation
+**Note: this library makes liberal use of several new features in pytorch, it's recommended to use it with the current nightly or latest stable version of PyTorch.**
+
+1. From PyPI:
+```Shell
+pip install torchao
+```
+
+2. From Source:
+
+```Shell
+git clone https://github.com/pytorch-labs/ao
+cd ao
+pip install -e .
+```
 
 ## Key Features
 The library provides
@@ -28,6 +36,17 @@ The library provides
 3. [Sparsity algorithms](./torchao/sparsity) such as Wanda that help improve accuracy of sparse networks
 4. Integration with other PyTorch native libraries like [torchtune](https://github.com/pytorch/torchtune) and [ExecuTorch](https://github.com/pytorch/executorch)
 
+
+## Our Goals
+torchao embodies PyTorch’s design philosophy [details](https://pytorch.org/docs/stable/community/design.html), especially "usability over everything else". Our vision for this repository is the following:
+
+* Composability: Native solutions for optimization techniques that compose with both `torch.compile` and `FSDP` 
+    * For example, for QLoRA for new dtypes support
+* Interoperability: Work with the rest of the PyTorch ecosystem such as torchtune, gpt-fast and ExecuTorch
+* Transparent Benchmarks: Regularly run performance benchmarking of our APIs across a suite of Torchbench models and across hardware backends
+* Heterogeneous Hardware: Efficient kernels that can run on CPU/GPU based server (w/ torch.compile) and mobile backends (w/ ExecuTorch).
+* Infrastructure Support: Release packaging solution for kernels and a CI/CD setup that runs these kernels on different backends. 
+
 ## Interoperability with PyTorch Libraries
 
 torchao has been integrated with other repositories to ease usage
@@ -36,33 +55,12 @@ torchao has been integrated with other repositories to ease usage
 * [Executorch](https://github.com/pytorch/executorch/tree/main/examples/models/llama2#quantization) is integrated with GPTQ for both 8da4w (int8 dynamic activation, with int4 weight) and int4 weight only quantization.
 
 ## Success stories
-Our kernels have has been used to achieve SOTA inference performance on
+Our kernels have been used to achieve SOTA inference performance on
 
 1. Image segmentation models with [sam-fast](pytorch.org/blog/accelerating-generative-ai)
 2. Language models with [gpt-fast](pytorch.org/blog/accelerating-generative-ai-2)
 3. Diffusion models with [sd-fast](pytorch.org/blog/accelerating-generative-ai-3)
 
-
-## Installation
-
-**Note: this library makes liberal use of several new features in pytorch, its recommended to use it with the current pytorch nightly if you want full feature coverage. If not, the subclass APIs may not work, though the module swap api's will still work.**
-
-1. From PyPI:
-```Shell
-pip install torchao
-```
-
-2. From Source:
-
-```Shell
-git clone https://github.com/pytorch-labs/ao
-cd ao
-pip install -e .
-```
-
-## Get Started
-To try out our APIs, you can check out API examples in [quantization](./torchao/quantization) (including `autoquant`), [sparsity](./torchao/sparsity), [dtypes](./torchao/dtypes).
-
 ## License
 
 `torchao` is released under the [BSD 3](https://github.com/pytorch-labs/ao/blob/main/LICENSE) license.