From 1e3f5210e4610d1ead233cbc7259824f7388cfca Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 8 Sep 2026 15:10:30 +0200 Subject: [PATCH 1/3] [CI] Make GitHub Build workflow fit the hosted runner Build a single CUDA arch (sm90) with MAX_JOBS=2, pass build env into the docker-based jobs, add libcuda stub for the sanity import and drop sccache in the CUDA 13 JAX job where it crashes nvcc. Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 48 ++++++++++++++++++++++--------------- 1 file changed, 29 insertions(+), 19 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 285a06ce31..120caa3331 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -37,10 +37,13 @@ jobs: run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation . -v env: NVTE_FRAMEWORK: none - MAX_JOBS: 1 + NVTE_CUDA_ARCHS: "90" + MAX_JOBS: 2 SCCACHE_GHA_ENABLED: "true" - name: 'Sanity check' - run: python3 -c "import transformer_engine" + run: | + ln -sf /usr/local/cuda/lib64/stubs/libcuda.so /usr/lib/x86_64-linux-gnu/libcuda.so.1 + python3 -c "import transformer_engine" working-directory: / pytorch: name: 'PyTorch' @@ -88,12 +91,15 @@ jobs: ' - name: 'Build' - run: docker exec builder bash -c 'pip install --no-build-isolation . -v --no-deps' - env: - NVTE_FRAMEWORK: pytorch - MAX_JOBS: 1 + run: | + docker exec -e NVTE_FRAMEWORK=pytorch -e NVTE_CUDA_ARCHS=90 -e MAX_JOBS=2 builder \ + bash -c 'pip install --no-build-isolation . -v --no-deps' - name: 'Sanity check' - run: docker exec builder bash -c 'python3 tests/pytorch/test_sanity_import.py' + run: | + docker exec builder bash -c '\ + ln -sf /usr/local/cuda/lib64/stubs/libcuda.so /usr/lib/x86_64-linux-gnu/libcuda.so.1 && \ + python3 tests/pytorch/test_sanity_import.py \ + ' jax: name: 'JAX' runs-on: ubuntu-latest @@ -107,17 +113,17 @@ jobs: uses: actions/checkout@v3 with: submodules: recursive - - name: ccache - uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad + # sccache's nvcc wrapper breaks with the CUDA 13.x toolkit in this image - name: 'Build' - run: | - NVTE_CCACHE_BIN=sccache NVTE_USE_CCACHE=1 pip install --no-build-isolation . -v + run: pip install --no-build-isolation . -v env: NVTE_FRAMEWORK: jax - MAX_JOBS: 1 - SCCACHE_GHA_ENABLED: "true" + NVTE_CUDA_ARCHS: "90" + MAX_JOBS: 2 - name: 'Sanity check' - run: python3 tests/jax/test_sanity_import.py + run: | + ln -sf /usr/local/cuda/lib64/stubs/libcuda.so /usr/lib/x86_64-linux-gnu/libcuda.so.1 + python3 tests/jax/test_sanity_import.py all: name: 'All' runs-on: ubuntu-latest @@ -158,9 +164,13 @@ jobs: pip install torch --no-cache-dir --index-url https://download.pytorch.org/whl/cu130 ' - name: 'Build' - run: docker exec builder bash -c 'pip install --no-cache-dir --no-build-isolation . -v --no-deps' - env: - NVTE_FRAMEWORK: all - MAX_JOBS: 1 + run: | + docker exec -e NVTE_FRAMEWORK=all -e NVTE_CUDA_ARCHS=90 -e MAX_JOBS=2 builder \ + bash -c 'pip install --no-cache-dir --no-build-isolation . -v --no-deps' - name: 'Sanity check' - run: docker exec builder bash -c 'python3 tests/pytorch/test_sanity_import.py && python3 tests/jax/test_sanity_import.py' + run: | + docker exec builder bash -c '\ + ln -sf /usr/local/cuda/lib64/stubs/libcuda.so /usr/lib/x86_64-linux-gnu/libcuda.so.1 && \ + python3 tests/pytorch/test_sanity_import.py && \ + python3 tests/jax/test_sanity_import.py \ + ' From 81ba5e8b46d7c4eac9573a328b18f5c6731db53a Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 8 Sep 2026 15:28:34 +0200 Subject: [PATCH 2/3] [CI] Build Core for every CUDA arch in a per-arch matrix; run on main to seed the cache Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 120caa3331..075f91f7bf 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -6,6 +6,8 @@ name: 'Build' on: pull_request: + push: + branches: [main] workflow_dispatch: concurrency: # Group by workflow name + PR number (for PRs) or ref (for branch/tag pushes) @@ -13,8 +15,12 @@ concurrency: cancel-in-progress: true jobs: core: - name: 'Core' + name: 'Core (sm${{ matrix.arch }})' runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + arch: ["70", "80", "89", "90", "100", "120"] container: image: nvcr.io/nvidia/cuda:12.9.2-devel-ubuntu22.04 options: --user root @@ -37,7 +43,7 @@ jobs: run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation . -v env: NVTE_FRAMEWORK: none - NVTE_CUDA_ARCHS: "90" + NVTE_CUDA_ARCHS: ${{ matrix.arch }} MAX_JOBS: 2 SCCACHE_GHA_ENABLED: "true" - name: 'Sanity check' From a666ca5482aaed7227f8f2a90399e85bc4c26d39 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 9 Sep 2026 15:18:25 +0200 Subject: [PATCH 3/3] [CI] Build PyTorch and JAX extensions for sm80 and sm90 to cover both NCCL EP variants Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 075f91f7bf..4c7e8ab516 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -52,8 +52,12 @@ jobs: python3 -c "import transformer_engine" working-directory: / pytorch: - name: 'PyTorch' + name: 'PyTorch (sm${{ matrix.arch }})' runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + arch: ["80", "90"] steps: - name: Move /var/lib/docker/ shell: bash -euxo pipefail {0} @@ -98,7 +102,7 @@ jobs: - name: 'Build' run: | - docker exec -e NVTE_FRAMEWORK=pytorch -e NVTE_CUDA_ARCHS=90 -e MAX_JOBS=2 builder \ + docker exec -e NVTE_FRAMEWORK=pytorch -e NVTE_CUDA_ARCHS=${{ matrix.arch }} -e MAX_JOBS=2 builder \ bash -c 'pip install --no-build-isolation . -v --no-deps' - name: 'Sanity check' run: | @@ -107,8 +111,12 @@ jobs: python3 tests/pytorch/test_sanity_import.py \ ' jax: - name: 'JAX' + name: 'JAX (sm${{ matrix.arch }})' runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + arch: ["80", "90"] container: image: ghcr.io/nvidia/jax:jax-2026-07-21 options: --user root @@ -124,7 +132,7 @@ jobs: run: pip install --no-build-isolation . -v env: NVTE_FRAMEWORK: jax - NVTE_CUDA_ARCHS: "90" + NVTE_CUDA_ARCHS: ${{ matrix.arch }} MAX_JOBS: 2 - name: 'Sanity check' run: |