Add FP8 runners + tweak building FP8 image

.github/workflows/build_docker_images.yml

@@ -102,9 +102,15 @@ jobs:
         id: date
         run: |
           echo "date=$(date '+%Y-%m-%d')" >> $GITHUB_ENV
+          # Get the previous month
+          echo "base_year=$(date -d 'last month' '+%y')" >> $GITHUB_ENV
+          echo "base_month=$(date -d 'last month' '+%m')" >> $GITHUB_ENV
       - name: Build and Push GPU
         uses: docker/build-push-action@v4
         with:
           file: benchmarks/fp8/transformer_engine/Dockerfile
           push: true
-          tags: huggingface/accelerate:gpu-fp8-transformerengine-nightly-${{ env.date }}
+          tags: huggingface/accelerate:gpu-fp8-transformerengine-nightly-${{ env.date }}
+          build-args: |
+            BASE_YEAR=${{ env.base_year }}
+            BASE_MONTH=${{ env.base_month }}

.github/workflows/fp8_runner.yml

@@ -0,0 +1,37 @@
+name: Test FP8 Runner
+
+on:
+  workflow_dispatch:
+
+env:
+  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+jobs:
+  set-prev-day:
+    runs-on: ubuntu-latest
+    outputs:
+      prev-day: ${{ steps.set-prev-day.outputs.prev-day }}
+    steps:
+      - name: Set PREV_DAY
+        id: set-prev-day
+        run: |
+          PREV_DAY=$(date -d "yesterday" '+%Y-%m-%d')
+          echo "prev-day=$PREV_DAY" >> $GITHUB_OUTPUT
+  run-fp8-tests:
+    needs: set-prev-day
+    runs-on:
+      group: aws-g6e-12xlarge
+    container:
+      image: huggingface/accelerate:gpu-fp8-transformerengine-nightly-${{ needs.set-prev-day.outputs.prev-day }}
+      options: --gpus all --shm-size "16gb"
+    steps:
+      - uses: actions/checkout@v3
+      - name: Install the library
+        run: |
+            pip install -e .[test_prod,test_fp8]
+      - name: Show installed libraries
+        run: |
+          pip freeze
+      - name: Run TE FP8 tests
+        run: |
+          python -m pytest -s -v ./tests/test_fp8.py
+

benchmarks/fp8/transformer_engine/Dockerfile

@@ -1,4 +1,7 @@
-FROM nvcr.io/nvidia/pytorch:24.07-py3
+ARG BASE_YEAR=25
+ARG BASE_MONTH=03
+
+FROM nvcr.io/nvidia/pytorch:${BASE_YEAR}.${BASE_MONTH}-py3
 
 RUN pip install transformers evaluate datasets
 RUN git clone https://github.com/huggingface/accelerate.git

setup.py

@@ -40,6 +40,7 @@
 extras["deepspeed"] = ["deepspeed"]
 extras["rich"] = ["rich"]
 
+extras["test_fp8"] = ["torchao"] # note: TE for now needs to be done via pulling down the docker image directly
 extras["test_trackers"] = ["wandb", "comet-ml", "tensorboard", "dvclive", "mlflow", "matplotlib"]
 extras["dev"] = extras["quality"] + extras["testing"] + extras["rich"]