Merge branch 'main' into compression-example-update

.github/workflows/build-container.yml

@@ -53,4 +53,12 @@ jobs:
           build-args: |
             BRANCH=${{github.head_ref}}
           push: true
-          tags: ghcr.io/neuralmagic/sparseml-dev:${{ inputs.name }}
+          tags: ghcr.io/neuralmagic/sparseml-dev:${{ inputs.name }}
+
+      - name: Build Nightly Docker Container
+        if: ${{ inputs.dev == 'false' && inputs.release == 'false'}}
+        uses: docker/build-push-action@v4
+        with:
+          context: ./docker/containers/docker_nightly
+          push: true
+          tags: ghcr.io/neuralmagic/sparseml-nightly:latest, ghcr.io/neuralmagic/sparseml-nightly:${{ steps.date.outputs.date }}

.github/workflows/build-wheel-and-container.yml

@@ -4,15 +4,8 @@ on:
     types: [opened, synchronize, reopened]
     branches:
       - main
-      - 'release/[0-9]+.[0-9]+'
-  push:
-    branches:
-      - 'release/[0-9]+.[0-9]+'
-      - main
-  release:
-    types: [created, published]
   schedule:
-    - cron: '0 0 * * *'
+    - cron: '0 20 * * *'
 
 permissions:
   id-token: write
@@ -23,10 +16,10 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
   cancel-in-progress: true
 
-# if not dev or release, will create a nightly build
+# TODO: do we want to push to nightly everytime we push to main?
+# if not dev or release, will create a nightly build; turning off release for now
 env:
-  PRODUCTION: ${{ github.event_name == 'schedule' || github.event_name == 'release'}}
-  RELEASE: ${{ github.event_name =='release' || startsWith(github.base_ref, 'release/') }}
+  RELEASE: 'false'
   DEV: ${{ github.base_ref == 'main' && github.event_name == 'pull_request'}}
 
 jobs:
@@ -42,8 +35,14 @@ jobs:
           echo "dev=$DEV" >> $GITHUB_OUTPUT
           echo "release=$RELEASE" >> $GITHUB_OUTPUT
 
-  build-wheel-and-push:
+  test-nightly:
     needs: set-outputs
+    if: ${{ needs.set-outputs.outputs.dev  == 'false' && needs.set-outputs.outputs.release  == 'false'}}
+    uses: ./.github/workflows/test-nightly.yml
+
+  build-wheel-and-push:
+    needs: [set-outputs, test-nightly]
+    if: ${{ always() && needs.set-outputs.outputs.dev == 'false' && needs.test-nightly.result == 'success' || always() && needs.set-outputs.outputs.dev  == 'true' && needs.set-outputs.result == 'success' }}
     uses: ./.github/workflows/build-wheel.yml
     with:
       build-label: ubuntu-20.04
@@ -55,22 +54,24 @@ jobs:
       python: '3.10'
     secrets: inherit
 
-  test-wheel-and-push-internal:
-    needs: build-wheel-and-push
-    uses: ./.github/workflows/test-wheel-push-to-internal.yml
+  test-wheel-and-publish:
+    needs: [set-outputs, build-wheel-and-push]
+    if: ${{ always() && !cancelled() && needs.build-wheel-and-push.result == 'success' }}
+    uses: ./.github/workflows/test-wheel-and-publish.yml
     with:
       build-label: ubuntu-20.04
       whl: ${{ needs.build-wheel-and-push.outputs.wheel }}
       python: '3.10'
+      dev: ${{ needs.set-outputs.outputs.dev }}
+      release: ${{ needs.set-outputs.outputs.release  }}
     secrets: inherit
 
-  # TODO: add nightly and release container build steps once wheel build push
-  # to production is automated. Removed until then.
   build-container-and-push:
-    needs: [set-outputs, test-wheel-and-push-internal]
+    needs: [test-wheel-and-publish, set-outputs]
+    if: ${{ always() && !cancelled() && needs.test-wheel-and-publish.result == 'success' }}
     uses: ./.github/workflows/build-container.yml
     with:
-      build-label: k8s-eng-gpu-64G-v100-32G
+      build-label: k8s-eng-gpu-16G-t4-32G
       dev: ${{ needs.set-outputs.outputs.dev }}
       release: ${{ needs.set-outputs.outputs.release  }}
       name: ${{ github.event.number }}

.github/workflows/test-nightly.yml

@@ -1,8 +1,7 @@
 name: Run Nightly Tests
 on:
-  schedule:
-    - cron: '0 20 * * *'
   workflow_dispatch:
+  workflow_call:
 jobs:
   test-nightly-tests:
     runs-on: k8s-mle-gpu-12-vcpu-225GB-ram-2-a6000-48G
@@ -33,6 +32,5 @@ jobs:
         run: |
           pytest tests/sparseml/transformers/obcq -m integration
       - name: Run finetune tests
-        if: always()
         run: |
           pytest tests/sparseml/transformers/finetune -m integration

.github/workflows/test-wheel-push-to-internal.yml → .github/workflows/test-wheel-and-publish.yml

@@ -1,4 +1,4 @@
-name: Test Wheel and Push to Internal PyPi
+name: Test Wheel and Publish
 on:
   workflow_call:
     inputs:
@@ -11,9 +11,15 @@ on:
         required: true
       python:
         type: string
+      dev:
+        type: string
+        required: true 
+      release:
+        type: string
+        required: true
 
 jobs:
-  test-wheel-and-push-internal:
+  test-wheel-and-publish:
     runs-on: ${{ inputs.build-label }}
     steps:
     - uses: actions/setup-python@v4
@@ -36,24 +42,37 @@ jobs:
         filename: ${{ inputs.whl }}
         dst: dist_s3
 
-    - name: Set Env
-      run: |
-        pip3 install virtualenv
-        virtualenv venv
-        source venv/bin/activate
-  
     - name: Fetch name of whl
       run: |
           echo "FILENAME=$(echo dist_s3/*.whl)" >> $GITHUB_ENV
 
     - name: Install whl
       run: |
-          pip3 install $FILENAME[dev]
+          pip3 install $FILENAME[dev,onnxruntime,torch,torchvision,transformers]
 
     - name: Checkout code
       uses: actions/checkout@v3
 
     - name: Remove src files and run tests
       run: |
+            pwd
             rm -rf src
-            make test
+            make test
+
+    - name: Make directory for wheel
+      run: |
+          mkdir dist_s3
+          
+    - name: Pull from s3
+      uses: neuralmagic/nm-actions/actions/s3_pull@main
+      with:
+        filename: ${{ inputs.whl }}
+        dst: dist_s3
+
+    - name: Publish Nightly Wheel
+      if: ${{ inputs.DEV == 'false' && inputs.RELEASE == 'false'}}
+      uses: neuralmagic/nm-actions/actions/publish-whl@main
+      with:
+        username: ${{ secrets.PYPI_PUBLIC_USER }}
+        password: ${{ secrets.PYPI_PUBLIC_AUTH }}
+        whl: ./$FILENAME

src/sparseml/modifiers/quantization/gptq/utils/gptq_wrapper.py

@@ -103,6 +103,14 @@ def fasterprune(
             W = W.t()
         W = W.float()
 
+        sparsity = tensor_sparsity(W)
+        preserve_zeros = sparsity >= SPARSITY_THRESHOLD
+        W_nz_mask = (
+            (~torch.isclose(W, torch.zeros(1, device=W.device).float())).float()
+            if preserve_zeros
+            else None
+        )
+
         tick = time.time()
 
         dead = torch.diag(self.H) == 0
@@ -119,17 +127,6 @@ def fasterprune(
         self.H = torch.linalg.cholesky(self.H, upper=True)
         Hinv = self.H
 
-        sparsity = tensor_sparsity(W)
-        mask = (
-            torch.where(
-                W == 0,
-                torch.tensor(1, dtype=torch.bool),
-                torch.tensor(0, dtype=torch.bool),
-            )
-            if sparsity >= SPARSITY_THRESHOLD
-            else None
-        )
-
         # See section 3.4 of https://arxiv.org/abs/2203.07259
         for i1 in range(0, self.columns, blocksize):
             i2 = min(i1 + blocksize, self.columns)
@@ -141,21 +138,13 @@ def fasterprune(
             Losses1 = torch.zeros_like(W1)
             Hinv1 = Hinv[i1:i2, i1:i2]
 
-            if sparsity >= SPARSITY_THRESHOLD:
-                tmp = (
-                    (~mask[:, i1:i2])
-                    * W1**2
-                    / (torch.diag(Hinv1).reshape((1, -1))) ** 2
-                )
-                thresh = torch.sort(tmp.flatten())[0][int(tmp.numel() * sparsity)]
-                mask1 = tmp <= thresh
+            if preserve_zeros:
+                W1_nz_mask = W_nz_mask[:, i1:i2]
 
             for i in range(count):
                 w = W1[:, i]
                 d = Hinv1[i, i]
                 q = w.clone()
-                if sparsity >= SPARSITY_THRESHOLD:
-                    q[mask1[:, i]] = 0
 
                 if hasattr(self.layer, "weight_fake_quant"):
                     scale = self.layer.weight_fake_quant.scale
@@ -216,13 +205,21 @@ def fasterprune(
                 Losses1[:, i] = (w - q) ** 2 / d**2
 
                 err1 = (w - q) / d
-                W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
+                w1_err = err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
+                if preserve_zeros:
+                    W1[:, i:] -= w1_err * W1_nz_mask[:, i:]
+                else:
+                    W1[:, i:] -= w1_err
                 Err1[:, i] = err1
 
             W[:, i1:i2] = Q1
             Losses += torch.sum(Losses1, 1) / 2
 
-            W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])
+            w_err = Err1.matmul(Hinv[i1:i2, i2:])
+            if preserve_zeros:
+                W[:, i2:] -= w_err * W_nz_mask[:, i2:]
+            else:
+                W[:, i2:] -= w_err
 
         _LOGGER.info("time %.2f" % (time.time() - tick))
         _LOGGER.info("error %.2f" % torch.sum(Losses).item())