build_wheels_rocm.yml

name: Build AutoGPTQ Wheels with ROCm

on:
  workflow_dispatch:
    inputs:
      version:
        description: 'Version tag of AutoGPTQ to build: v0.5.1'
        default: 'main'
        required: true
        type: string
  workflow_call:
    inputs:
      version:
        description: 'Version tag of AutoGPTQ to build: v0.4.2'
        default: 'main'
        required: true
        type: string

permissions:
  contents: write

jobs:
  build_wheels:
    name: Build wheels for ${{ matrix.os }} and Python ${{ matrix.python }} and RoCm ${{ matrix.rocm }}
    runs-on: ubuntu-20.04
    strategy:
      matrix:
        pyver: ["3.10", "3.11"]
        rocm: ['5.5', '5.6']
    defaults:
      run:
        shell: pwsh
    env:
      ROCM_VERSION: ${{ matrix.rocm }}

    steps:
      - name: Free Disk Space
        uses: jlumbroso/free-disk-space@v1.2.0
        with:
          tool-cache: false
          android: true
          dotnet: true
          haskell: true
          large-packages: false
          swap-storage: false
    
      - uses: actions/checkout@v4
        with:
          repository: 'AutoGPTQ/AutoGPTQ'
          ref: ${{ inputs.version }}
          
      - name: Install ROCm SDK
        shell: bash
        run: |
          [ ! -d /etc/apt/keyrings ] && sudo mkdir --parents --mode=0755 /etc/apt/keyrings
          wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
          echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/$ROCM_VERSION focal main" | sudo tee --append /etc/apt/sources.list.d/rocm.list
          echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
          sudo apt update
          sudo apt install rocm-dev rocsparse-dev rocprim-dev rocthrust-dev rocblas-dev hipblas-dev hipcub-dev hipsparse-dev -y
          echo "/opt/rocm/bin" >> $GITHUB_PATH
          echo "ROCM_PATH=/opt/rocm" >> $GITHUB_ENV
          echo "USE_ROCM=1" >> $GITHUB_ENV
          
      - uses: actions/setup-python@v3
        with:
          python-version: ${{ matrix.pyver }}
          
      - name: Install Dependencies
        run: |
          $packages = 'build wheel safetensors sentencepiece ninja numpy gekko pandas'
          $torver = if ([version]$env:ROCM_VERSION -lt [version]'5.5') {'2.0.1'} else {'2.2.1'}
          $packages += " torch==$torver+rocm$env:ROCM_VERSION torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/rocm$env:ROCM_VERSION"
          
          pip3 install $packages.split(' ')
        
      - name: Build Wheel
        id: build-wheel
        run: |
          if ($(Get-Content 'setup.py' -raw) -match '"version": "(\d+\.(?:\d+\.?(?:dev\d+)?)*)",')
          {
              Write-Output $('::notice file=build-wheels-release.yml,title=Package Version::Detected package version is: {0}' -f $Matches[1])
              Write-Output "PACKAGE_VERSION=$($Matches[1])" >> "$env:GITHUB_OUTPUT"
              $PCKG_VERSION = "$($Matches[1])"
          } else {
              Write-Output '::error file=build-wheels-release.yml::Could not parse version from setup.py! You must upload wheels manually!'
              Write-Output "PACKAGE_VERSION=None" >> "$env:GITHUB_OUTPUT"
              $PCKG_VERSION = $false
          }
          
          $env:PYTORCH_ROCM_ARCH = 'gfx803;gfx900;gfx906:xnack-;gfx908:xnack-;gfx90a:xnack+;gfx90a:xnack-;gfx1010;gfx1012;gfx1030;gfx1100;gfx1101;gfx1102'
          if ([version]$env:ROCM_VERSION -lt [version]'5.5') {$env:PYTORCH_ROCM_ARCH = 'gfx803;gfx900;gfx906:xnack-;gfx908:xnack-;gfx90a:xnack+;gfx90a:xnack-;gfx1010;gfx1012;gfx1030'}
          
          if ([version]$PCKG_VERSION -lt [version]'0.5.0')
          {
              $setup = Get-Content setup.py -raw
              $macroString = "extra_link_args=extra_link_args,`n            define_macros=[('HIPBLAS_USE_HIP_HALF', '1')]"
              New-Item setup.py -itemType File -value $setup.replace('extra_link_args=extra_link_args', $macroString) -force
          }
          
          python setup.py sdist bdist_wheel

      - uses: actions/upload-artifact@v3
        with:
          name: 'linux-rocm-wheels'
          path: ./dist/*.whl
          
      - name: Upload files to a GitHub release
        if: steps.build-wheel.outputs.PACKAGE_VERSION != 'None' && !contains( 'dev', steps.build-wheel.outputs.PACKAGE_VERSION )
        uses: svenstaro/upload-release-action@2.6.1
        continue-on-error: true
        with:
          file: ./dist/*.whl
          tag: ${{ format('v{0}', steps.build-wheel.outputs.PACKAGE_VERSION) }}
          file_glob: true
          overwrite: true
          release_name: ${{ format('v{0}', steps.build-wheel.outputs.PACKAGE_VERSION) }}