Merge branch 'main' into rocm_enablement_staging

Author: petrex

.github/scripts/validate_binaries.sh

@@ -1,2 +1,6 @@
 pip install ${PYTORCH_PIP_PREFIX} torchao --index-url ${PYTORCH_PIP_DOWNLOAD_URL}
+# Intial smoke test, tries importing torchao
 python  ./test/smoke_tests/smoke_tests.py
+# Now we install dev-requirments and try to run the tests
+pip install -r dev-requirements.txt
+pytest test --verbose -s

.github/workflows/build_wheels_linux.yml

@@ -28,7 +28,8 @@ jobs:
       os: linux
       with-cpu: enable
       with-cuda: enable
-      with-rocm: disable
+      with-rocm: enable
+      with-xpu: enable
 
   build:
     needs: generate-matrix

.github/workflows/build_wheels_windows.yml

@@ -0,0 +1,62 @@
+name: Build Windows Wheels
+
+on:
+  pull_request:
+    paths:
+      - build/packaging/**
+      - .github/workflows/build_wheels_windows.yml
+      - setup.py
+  push:
+    branches:
+      - nightly
+      - main
+      - 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]+
+  schedule:
+    - cron: '0 0 * * *'  # Runs at midnight UTC every day
+  workflow_dispatch:
+
+permissions:
+  id-token: write
+  contents: read
+
+jobs:
+  generate-matrix:
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
+    with:
+      package-type: wheel
+      os: windows
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      with-xpu: enable
+      with-cuda: disable
+
+  build:
+    needs: generate-matrix
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - repository: pytorch/ao
+            pre-script: packaging/pre_build_script.sh
+            env-script: packaging/vc_env_helper.bat
+            # post-script: "python packaging/wheel/relocate.py"
+            smoke-test-script: packaging/smoke_test.py
+            package-name: torchao
+    name: ${{ matrix.repository }}
+    uses: pytorch/test-infra/.github/workflows/build_wheels_windows.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 }}
+      env-script: ${{ matrix.env-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/float8_test.yml

@@ -29,7 +29,7 @@ jobs:
             gpu-arch-type: "cuda"
             gpu-arch-version: "12.1"
 
-    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     with:
       timeout: 60
       runner: ${{ matrix.runs-on }}
@@ -38,8 +38,6 @@ jobs:
       script: |
         conda create -n venv python=3.9 -y
         conda activate venv
-        echo "::group::Install newer objcopy that supports --set-section-alignment"
-        yum install -y  devtoolset-10-binutils
         export PATH=/opt/rh/devtoolset-10/root/usr/bin/:$PATH
         python -m pip install --upgrade pip
         pip install ${{ matrix.torch-spec }}

.github/workflows/nightly_smoke_test.yml

@@ -26,7 +26,7 @@ jobs:
             gpu-arch-version: "12.1"
 
 
-    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     with:
       runner: ${{ matrix.runs-on }}
       gpu-arch-type: ${{ matrix.gpu-arch-type }}

.github/workflows/pr-label-check.yml

@@ -0,0 +1,30 @@
+name: PR Label Check
+on:
+  pull_request:
+    types: [opened, labeled, unlabeled, synchronize]
+
+jobs:
+  check-labels:
+    name: Check PR Labels
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check for Topic label
+        run: |
+          # Get the labels using GitHub API
+          LABELS=$(curl -s -H "Authorization: token ${{ secrets.GITHUB_TOKEN }}" \
+          "https://api.github.com/repos/${{ github.repository }}/pulls/${{ github.event.pull_request.number }}" \
+          | jq -r '.labels[].name')
+
+          # Check if there are any labels
+          if [ -z "$LABELS" ]; then
+            echo "::error::This PR requires at least one topic label. Please add a topic from: https://github.com/pytorch/ao/labels?q=topic"
+            exit 1
+          fi
+
+          # Check for Topic label
+          if ! echo "$LABELS" | grep -i "topic:" > /dev/null; then
+            echo "::error::This PR requires at least one label starting with 'topic:'. Available topics can be found at: https://github.com/pytorch/ao/labels?q=topic"
+            exit 1
+          fi
+
+          echo "PR has required topic label"

.github/workflows/regression_test.yml

@@ -18,16 +18,43 @@ env:
   HF_TOKEN: ${{ secrets.HF_TOKEN }}
 
 jobs:
-  test:
+  test-nightly:
     strategy:
       fail-fast: false
       matrix:
         include:
-          - name: CUDA 2.2.2
+          - name: CUDA Nightly
             runs-on: linux.g5.12xlarge.nvidia.gpu
-            torch-spec: 'torch==2.2.2 "numpy<2" '
+            torch-spec: '--pre torch==2.6.0.dev20241101 --index-url https://download.pytorch.org/whl/nightly/cu121'
             gpu-arch-type: "cuda"
             gpu-arch-version: "12.1"
+          - name: CPU Nightly
+            runs-on: linux.4xlarge
+            torch-spec: '--pre torch==2.6.0.dev20241101 --index-url https://download.pytorch.org/whl/nightly/cpu'
+            gpu-arch-type: "cpu"
+            gpu-arch-version: ""
+
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    with:
+      timeout: 120
+      runner: ${{ matrix.runs-on }}
+      gpu-arch-type: ${{ matrix.gpu-arch-type }}
+      gpu-arch-version: ${{ matrix.gpu-arch-version }}
+      script: |
+        conda create -n venv python=3.9 -y
+        conda activate venv
+        python -m pip install --upgrade pip
+        pip install ${{ matrix.torch-spec }}
+        pip install -r dev-requirements.txt
+        pip install .
+        export CONDA=$(dirname $(dirname $(which conda)))
+        export LD_LIBRARY_PATH=$CONDA/lib/:$LD_LIBRARY_PATH
+        pytest test --verbose -s
+  test:
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
           - name: CUDA 2.3
             runs-on: linux.g5.12xlarge.nvidia.gpu
             torch-spec: 'torch==2.3.0'
@@ -38,17 +65,11 @@ jobs:
             torch-spec: 'torch==2.4.0'
             gpu-arch-type: "cuda"
             gpu-arch-version: "12.1"
-          - name: CUDA Nightly
+          - name: CUDA 2.5.1
             runs-on: linux.g5.12xlarge.nvidia.gpu
-            torch-spec: '--pre torch --index-url https://download.pytorch.org/whl/nightly/cu121'
+            torch-spec: 'torch==2.5.1 --index-url https://download.pytorch.org/whl/cu121'
             gpu-arch-type: "cuda"
             gpu-arch-version: "12.1"
-
-          - name: CPU 2.2.2
-            runs-on: linux.4xlarge
-            torch-spec: 'torch==2.2.2 --index-url https://download.pytorch.org/whl/cpu "numpy<2" '
-            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'
@@ -59,9 +80,9 @@ jobs:
             torch-spec: 'torch==2.4.0 --index-url https://download.pytorch.org/whl/cpu'
             gpu-arch-type: "cpu"
             gpu-arch-version: ""
-          - name: CPU Nightly
+          - name: CPU 2.5.1
             runs-on: linux.4xlarge
-            torch-spec: '--pre torch --index-url https://download.pytorch.org/whl/nightly/cpu'
+            torch-spec: 'torch==2.5.1 --index-url https://download.pytorch.org/whl/cpu'
             gpu-arch-type: "cpu"
             gpu-arch-version: ""
 

.github/workflows/ruff_linter.yml

@@ -1,6 +1,12 @@
 name: Code Analysis with Ruff
 
 on:
+  workflow_dispatch:
+    inputs:
+      pr_url:
+        description: 'URL of the PR to fix'
+        required: true
+        type: string
   push:
     branches:
       - main
@@ -13,30 +19,74 @@ on:
 jobs:
   build:
     runs-on: ubuntu-latest
+    permissions:
+      contents: write
+      pull-requests: write
+    
     strategy:
       matrix:
         python-version: ["3.9"]
     steps:
+    - name: Extract PR info
+      if: github.event_name == 'workflow_dispatch'
+      run: |
+        PR_URL=${{ github.event.inputs.pr_url }}
+        PR_NUMBER=$(echo $PR_URL | grep -oE '[0-9]+$')
+        echo "PR_NUMBER=$PR_NUMBER" >> $GITHUB_ENV
+        
     - uses: actions/checkout@v3
+      if: github.event_name == 'workflow_dispatch'
+      with:
+        fetch-depth: 0
+        token: ${{ secrets.GITHUB_TOKEN }}
+    
+    - name: Checkout PR branch
+      if: github.event_name == 'workflow_dispatch'
+      run: |
+        gh pr checkout ${{ env.PR_NUMBER }}
+      env:
+        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        
+    - uses: actions/checkout@v3
+      if: github.event_name != 'workflow_dispatch'
+      with:
+        fetch-depth: 0
+        
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v3
       with:
         python-version: ${{ matrix.python-version }}
+        
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
         pip install ruff==0.6.8
-    - name: Analyzing the code with ruff
+        
+    - name: Regular lint check
+      if: github.event_name != 'workflow_dispatch'
       run: |
         ruff check .
-    - name: Check *all* Python files for F821, F823, and W191
-      run: |
         # --isolated is used to skip the allowlist at all so this applies to all files
         # please be careful when using this large changes means everyone needs to rebase
         ruff check --isolated --select F821,F823,W191
-    - name: Check the allow-listed files for F,I
-      run: |
         ruff check --select F,I
-    - name: Check the allow-listed files for well formatted code
-      run: |
         ruff format --check
+
+    - name: Apply fixes to PR
+      if: github.event_name == 'workflow_dispatch'
+      run: |
+        git config --global user.name 'github-actions[bot]'
+        git config --global user.email 'github-actions[bot]@users.noreply.github.com'
+        
+        # Apply fixes
+        ruff check --select F,I --fix
+        ruff format .
+        
+        # Commit and push if there are changes
+        if [[ -n "$(git status --porcelain)" ]]; then
+          git add .
+          git commit -m "Apply automatic Ruff fixes"
+          git push
+        else
+          echo "No fixes needed!"
+        fi

.pre-commit-config.yaml

@@ -15,6 +15,9 @@ repos:
     hooks:
       # Run the linter.
       - id: ruff
-        args: [--fix]
+        args:
+          - --fix
+          - --select
+          - F,I
       # Run the formatter.
       - id: ruff-format

CITATION.cff

@@ -0,0 +1,9 @@
+cff-version: 1.2.0
+title: "torchao: PyTorch native quantization and sparsity for training and inference"
+message: "If you use this software, please cite it as below."
+type: software
+authors:
+  - given-names: "torchao maintainers and contributors"
+url: "https//github.com/pytorch/torchao"
+license: "BSD-3-Clause"
+date-released: "2024-10-25"