[Feat] Improve CI

- Enable c8g instance type to enhance testing
- Add builds as artifacts to re-use in testing stage
- Filter out fp16 and bf16 pytorch tests when
  ONEDNN_DEFAULT_FPMATH_MODE="BF16" to mitigate accuracy failures
- Add TORCH_RELEASE_ID as part of build wheel to match PEP440 standard

Author: Radu2k

.github/workflows/pytorch.yml

@@ -58,13 +58,14 @@ concurrency:
 permissions: read-all
 
 jobs:
-  build-and-test:
+  build-image:
     strategy:
       matrix:
-        config: [
-          { name: c7g, label: ah-ubuntu_22_04-c7g_8x-100 }
-        ]
-
+        config:
+          [
+            { name: c7g, label: ah-ubuntu_22_04-c7g_8x-100 },
+            { name: c8g, label: ah-ubuntu_22_04-c8g_8x }
+           ]
     runs-on: ${{ matrix.config.label }}
     steps:
       - name: Checkout Tool-Solutions
@@ -81,7 +82,45 @@ jobs:
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
 
-      - name: Run smoke tests
+      - name: Save image as a artifact
+        run: docker save toolsolutions-pytorch:latest -o toolsolutions-pytorch-image-${{ matrix.config.name }}.tar
+
+      - name: Upload build artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: toolsolutions-pytorch-image-${{ matrix.config.name }}
+          path: toolsolutions-pytorch-image-${{ matrix.config.name }}.tar
+          compression-level: 9
+          retention-days: 1
+
+  test:
+    needs: build-image
+    strategy:
+      matrix:
+        config:
+          [
+            { name: c7g, label: ah-ubuntu_22_04-c7g_8x-100 },
+            { name: c8g, label: ah-ubuntu_22_04-c8g_8x }
+           ]
+        onednn_fpmath_mode: [FP32, BF16]
+    runs-on: ${{ matrix.config.label }}
+    env:
+      ONEDNN_DEFAULT_FPMATH_MODE: ${{ matrix.onednn_fpmath_mode }}
+
+    steps:
+      - name: Download image artifact
+        uses: actions/download-artifact@v5
+        with:
+          name: toolsolutions-pytorch-image-${{ matrix.config.name }}
+          path: .
+
+      - name: Set up Docker
+        uses: docker/setup-docker-action@v4
+
+      - name: Load Docker image
+        run: docker load -i toolsolutions-pytorch-image-${{ matrix.config.name }}.tar
+
+      - name: Run smoke tests for
         run: docker run --rm toolsolutions-pytorch:latest ./test-examples.sh
 
       - name: Run unit tests

ML-Frameworks/pytorch-aarch64/build-wheel.sh

@@ -108,7 +108,7 @@ docker exec -t $TORCH_BUILD_CONTAINER bash -c "rm -rf $PYTORCH_ROOT/dist"
 # pytorch/.ci/aarch64_linux/aarch64_wheel_ci_build.py for this logic.
 build_date=$(cd $PYTORCH_HOST_DIR && git log --pretty=format:%cs -1 | tr -d '-')
 version=$(cat $PYTORCH_HOST_DIR/version.txt| tr -d [:space:] )
-OVERRIDE_PACKAGE_VERSION=${version%??}.dev$build_date
+OVERRIDE_PACKAGE_VERSION="${version%??}.dev${build_date}${TORCH_RELEASE_ID:+"+$TORCH_RELEASE_ID"}"
 # We unset MAX_JOBS from 12 (written to /tmp/env by populate_binary_env.sh) to
 # let any downstream functions to decide. Currently nproc when ninja is used,
 # but ideally we would let ninja make the right choice, but nproc is good enough

ML-Frameworks/pytorch-aarch64/examples/run_unit_tests.sh

@@ -17,5 +17,14 @@
 # limitations under the License.
 # *******************************************************************************
 
-OMP_NUM_THREADS=16 python pytorch/test/test_mkldnn.py
-OMP_NUM_THREADS=16 python pytorch/test/test_transformers.py
+if [[ "${ONEDNN_DEFAULT_FPMATH_MODE:-}" == "BF16" ]]; then
+    OMP_NUM_THREADS=16 python -m unittest pytorch/test/test_mkldnn.py -k lower_precision -k bf16 -k bfloat16 -k float16
+else
+    OMP_NUM_THREADS=16 python -m unittest pytorch/test/test_mkldnn.py
+fi
+
+if [[ "${ONEDNN_DEFAULT_FPMATH_MODE:-}" == "BF16" ]]; then
+    OMP_NUM_THREADS=16 python -m unittest pytorch/test/test_transformers.py -k bfloat16 -k float16
+else
+    OMP_NUM_THREADS=16 python -m unittest pytorch/test/test_transformers.py
+fi