[ci] Add workflows to validate TE QA test cases

.github/workflows/build.yml

@@ -8,30 +8,6 @@ on:
   pull_request:
   workflow_dispatch:
 jobs:
-  core:
-    name: 'Core'
-    runs-on: ubuntu-latest
-    container:
-      image: nvcr.io/nvidia/cuda:12.1.0-devel-ubuntu22.04
-      options: --user root
-    steps:
-      - name: 'Dependencies'
-        run: |
-          apt-get update
-          apt-get install -y git python3.9 pip cudnn9-cuda-12
-          pip install cmake==3.21.0 pybind11[global] ninja nvidia-mathdx==25.1.1
-      - name: 'Checkout'
-        uses: actions/checkout@v3
-        with:
-          submodules: recursive
-      - name: 'Build'
-        run: pip install --no-build-isolation . -v
-        env:
-          NVTE_FRAMEWORK: none
-          MAX_JOBS: 1
-      - name: 'Sanity check'
-        run: python3 -c "import transformer_engine"
-        working-directory: /
   pytorch:
     name: 'PyTorch'
     runs-on: ubuntu-latest
@@ -55,43 +31,3 @@ jobs:
           MAX_JOBS: 1
       - name: 'Sanity check'
         run: python3 tests/pytorch/test_sanity_import.py
-  jax:
-    name: 'JAX'
-    runs-on: ubuntu-latest
-    container:
-      image: ghcr.io/nvidia/jax:jax
-      options: --user root
-    steps:
-      - name: 'Dependencies'
-        run: pip install pybind11[global] nvidia-mathdx==25.1.1
-      - name: 'Checkout'
-        uses: actions/checkout@v3
-        with:
-          submodules: recursive
-      - name: 'Build'
-        run: pip install --no-build-isolation . -v
-        env:
-          NVTE_FRAMEWORK: jax
-          MAX_JOBS: 1
-      - name: 'Sanity check'
-        run: python3 tests/jax/test_sanity_import.py
-  all:
-    name: 'All'
-    runs-on: ubuntu-latest
-    container:
-      image: ghcr.io/nvidia/jax:jax
-      options: --user root
-    steps:
-      - name: 'Dependencies'
-        run: pip install torch pybind11[global] einops onnxscript nvidia-mathdx==25.1.1
-      - name: 'Checkout'
-        uses: actions/checkout@v3
-        with:
-          submodules: recursive
-      - name: 'Build'
-        run: pip install --no-build-isolation . -v --no-deps
-        env:
-          NVTE_FRAMEWORK: all
-          MAX_JOBS: 1
-      - name: 'Sanity check'
-        run: python3 tests/pytorch/test_sanity_import.py && python3 tests/jax/test_sanity_import.py

.github/workflows/qa-format.yml

@@ -0,0 +1,32 @@
+name: format_check
+
+on:
+  pull_request:
+    branches: [ "main" ]
+    types: [opened, synchronize, reopened]
+
+jobs:
+  format:
+    runs-on: ubuntu-22.04
+    env:
+      PRID: ${{ github.event.pull_request.number }}
+      BRANCH: main
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event.pull_request.base.ref }}
+
+      - name: Merge PR to sub-branch
+        run: |
+          git fetch origin pull/${PRID}/merge
+          git checkout -b test FETCH_HEAD
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+
+      - name: Run pre-commit
+        run: bash ./qa/format.sh