From ee0cca1e5dedcdeeb0d64fc58e36143485fa50a9 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 30 Jan 2026 19:04:32 +0000 Subject: [PATCH 01/22] fix Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index d80564274e..0ea4d6b0a9 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -78,10 +78,7 @@ 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 MAX_JOBS=1 -e NVTE_FRAMEWORK=pytorch 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' jax: @@ -148,9 +145,6 @@ 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 MAX_JOBS=1 -e NVTE_FRAMEWORK=all 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' From df6a81b62b7f5a3c90beb29a6ef9f7ee5d7073c4 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 30 Jan 2026 21:42:44 +0000 Subject: [PATCH 02/22] fix Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 33 +++++++++++++++++++++++++++------ 1 file changed, 27 insertions(+), 6 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 0ea4d6b0a9..fab933383e 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -39,6 +39,15 @@ jobs: name: 'PyTorch' runs-on: ubuntu-latest steps: + - name: Free up disk space + run: | + sudo rm -rf /usr/local/share/boost + sudo rm -rf "$AGENT_TOOLSDIRECTORY" + sudo rm -rf /opt/hostedtoolcache + sudo rm -rf /usr/share/swift + sudo rm -rf /usr/local/.ghcup + df -h + - name: Move /var/lib/docker/ shell: bash -euxo pipefail {0} run: sudo mv /var/lib/docker/ "${GITHUB_WORKSPACE}/docker" @@ -46,9 +55,9 @@ jobs: - name: Maximize build space uses: easimon/maximize-build-space@c28619d8999a147d5e09c1199f84ff6af6ad5794 with: - root-reserve-mb: 5120 + root-reserve-mb: 4096 temp-reserve-mb: 32 - swap-size-mb: 10240 + swap-size-mb: 4096 remove-dotnet: 'true' remove-android: 'true' remove-haskell: 'true' @@ -74,7 +83,9 @@ jobs: apt-get update && \ apt-get install -y git python3.9 pip cudnn9-cuda-12 && \ pip install cmake torch ninja pydantic importlib-metadata>=1.0 packaging pybind11 numpy einops onnxscript && \ - apt-get clean \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* && \ + pip cache purge \ ' - name: 'Build' @@ -109,6 +120,15 @@ jobs: name: 'All' runs-on: ubuntu-latest steps: + - name: Free up disk space + run: | + sudo rm -rf /usr/local/share/boost + sudo rm -rf "$AGENT_TOOLSDIRECTORY" + sudo rm -rf /opt/hostedtoolcache + sudo rm -rf /usr/share/swift + sudo rm -rf /usr/local/.ghcup + df -h + - name: Move /var/lib/docker/ shell: bash -euxo pipefail {0} run: sudo mv /var/lib/docker/ "${GITHUB_WORKSPACE}/docker" @@ -116,9 +136,9 @@ jobs: - name: Maximize build space uses: easimon/maximize-build-space@c28619d8999a147d5e09c1199f84ff6af6ad5794 with: - root-reserve-mb: 5120 + root-reserve-mb: 4096 temp-reserve-mb: 32 - swap-size-mb: 10240 + swap-size-mb: 4096 remove-dotnet: 'true' remove-android: 'true' remove-haskell: 'true' @@ -142,7 +162,8 @@ jobs: run: | docker exec builder bash -c '\ pip install cmake==3.21.0 pybind11[global] einops onnxscript && \ - pip install torch --no-cache-dir --index-url https://download.pytorch.org/whl/cu130 + pip install torch --no-cache-dir --index-url https://download.pytorch.org/whl/cu130 && \ + pip cache purge ' - name: 'Build' run: docker exec -e MAX_JOBS=1 -e NVTE_FRAMEWORK=all builder bash -c 'pip install --no-cache-dir --no-build-isolation . -v --no-deps' From daa3bf3a55fdea866d7fcbc2ad38957da09cbb9a Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 30 Jan 2026 23:06:57 +0000 Subject: [PATCH 03/22] fix Signed-off-by: Pawel Gadzinski --- .github/workflows/deploy_nightly_docs.yml | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/.github/workflows/deploy_nightly_docs.yml b/.github/workflows/deploy_nightly_docs.yml index b4e015d2da..a8e5ee5ba2 100644 --- a/.github/workflows/deploy_nightly_docs.yml +++ b/.github/workflows/deploy_nightly_docs.yml @@ -7,6 +7,7 @@ name: Deploy nightly docs on: push: branches: [ "main" ] + workflow_dispatch: jobs: build: uses: ./.github/workflows/docs.yml @@ -21,9 +22,8 @@ jobs: name: "te_docs" path: "html" - name: Prepare for pages - uses: actions/upload-pages-artifact@v1.0.7 + uses: actions/upload-pages-artifact@v3 with: - name: github-pages path: "html" deploy: needs: prepare @@ -36,4 +36,5 @@ jobs: runs-on: ubuntu-latest steps: - name: Deploy - uses: actions/deploy-pages@v2.0.0 + id: deployment + uses: actions/deploy-pages@v4 From d47399ac69572c82c17da062f2390dc18a0bd4bb Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Sat, 31 Jan 2026 01:08:44 +0000 Subject: [PATCH 04/22] fix Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index fab933383e..51a1ea3148 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -75,16 +75,13 @@ jobs: - name: Start named container run: | - docker run -v $(pwd):$(pwd) -w $(pwd) --name builder -d nvcr.io/nvidia/cuda:12.8.0-devel-ubuntu22.04 sleep infinity + docker run -v $(pwd):$(pwd) -w $(pwd) --name builder -d ghcr.io/nvidia/jax:jax sleep infinity - name: 'Dependencies' run: | docker exec builder bash -c '\ - apt-get update && \ - apt-get install -y git python3.9 pip cudnn9-cuda-12 && \ - pip install cmake torch ninja pydantic importlib-metadata>=1.0 packaging pybind11 numpy einops onnxscript && \ - apt-get clean && \ - rm -rf /var/lib/apt/lists/* && \ + pip install cmake==3.21.0 pybind11[global] ninja pydantic importlib-metadata>=1.0 packaging numpy einops onnxscript && \ + pip install torch --no-cache-dir && \ pip cache purge \ ' From d2091f2f4a42e7ccff99951c382663e6c1154c93 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 3 Feb 2026 04:20:34 +0000 Subject: [PATCH 05/22] ccache size limit Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 51a1ea3148..70917f6d78 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -32,6 +32,7 @@ jobs: NVTE_FRAMEWORK: none MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" + SCCACHE_CACHE_SIZE: "1G" - name: 'Sanity check' run: python3 -c "import transformer_engine" working-directory: / @@ -111,6 +112,7 @@ jobs: NVTE_FRAMEWORK: jax MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" + SCCACHE_CACHE_SIZE: "1G" - name: 'Sanity check' run: python3 tests/jax/test_sanity_import.py all: From 4dc93233aecceded3c87de5bbe0f2d08796ca878 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 3 Feb 2026 04:25:33 +0000 Subject: [PATCH 06/22] ccache Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 119 +++++++++++------------------------- 1 file changed, 34 insertions(+), 85 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 70917f6d78..7041884bf4 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -39,57 +39,30 @@ jobs: pytorch: name: 'PyTorch' runs-on: ubuntu-latest + container: + image: ghcr.io/nvidia/jax:jax + options: --user root steps: - - name: Free up disk space + - name: 'Dependencies' run: | - sudo rm -rf /usr/local/share/boost - sudo rm -rf "$AGENT_TOOLSDIRECTORY" - sudo rm -rf /opt/hostedtoolcache - sudo rm -rf /usr/share/swift - sudo rm -rf /usr/local/.ghcup - df -h - - - name: Move /var/lib/docker/ - shell: bash -euxo pipefail {0} - run: sudo mv /var/lib/docker/ "${GITHUB_WORKSPACE}/docker" - - - name: Maximize build space - uses: easimon/maximize-build-space@c28619d8999a147d5e09c1199f84ff6af6ad5794 - with: - root-reserve-mb: 4096 - temp-reserve-mb: 32 - swap-size-mb: 4096 - remove-dotnet: 'true' - remove-android: 'true' - remove-haskell: 'true' - remove-codeql: 'true' - build-mount-path: '/var/lib/docker/' - - - name: Restore /var/lib/docker/ - shell: bash -euxo pipefail {0} - run: sudo sh -c "mv ${GITHUB_WORKSPACE}/docker/* /var/lib/docker" - + pip install cmake==3.21.0 pybind11[global] ninja pydantic importlib-metadata>=1.0 packaging numpy einops onnxscript + pip install torch --no-cache-dir + pip cache purge - name: 'Checkout' uses: actions/checkout@v3 with: submodules: recursive - - - name: Start named container - run: | - docker run -v $(pwd):$(pwd) -w $(pwd) --name builder -d ghcr.io/nvidia/jax:jax sleep infinity - - - name: 'Dependencies' - run: | - docker exec builder bash -c '\ - pip install cmake==3.21.0 pybind11[global] ninja pydantic importlib-metadata>=1.0 packaging numpy einops onnxscript && \ - pip install torch --no-cache-dir && \ - pip cache purge \ - ' - + - name: ccache + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad - name: 'Build' - run: docker exec -e MAX_JOBS=1 -e NVTE_FRAMEWORK=pytorch builder bash -c 'pip install --no-build-isolation . -v --no-deps' + run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation . -v --no-deps + env: + NVTE_FRAMEWORK: pytorch + MAX_JOBS: 1 + SCCACHE_GHA_ENABLED: "true" + SCCACHE_CACHE_SIZE: "1G" - name: 'Sanity check' - run: docker exec builder bash -c 'python3 tests/pytorch/test_sanity_import.py' + run: python3 tests/pytorch/test_sanity_import.py jax: name: 'JAX' runs-on: ubuntu-latest @@ -118,53 +91,29 @@ jobs: all: name: 'All' runs-on: ubuntu-latest + container: + image: ghcr.io/nvidia/jax:jax + options: --user root steps: - - name: Free up disk space + - name: 'Dependencies' run: | - sudo rm -rf /usr/local/share/boost - sudo rm -rf "$AGENT_TOOLSDIRECTORY" - sudo rm -rf /opt/hostedtoolcache - sudo rm -rf /usr/share/swift - sudo rm -rf /usr/local/.ghcup - df -h - - - name: Move /var/lib/docker/ - shell: bash -euxo pipefail {0} - run: sudo mv /var/lib/docker/ "${GITHUB_WORKSPACE}/docker" - - - name: Maximize build space - uses: easimon/maximize-build-space@c28619d8999a147d5e09c1199f84ff6af6ad5794 - with: - root-reserve-mb: 4096 - temp-reserve-mb: 32 - swap-size-mb: 4096 - remove-dotnet: 'true' - remove-android: 'true' - remove-haskell: 'true' - remove-codeql: 'true' - build-mount-path: '/var/lib/docker/' - - - name: Restore /var/lib/docker/ - shell: bash -euxo pipefail {0} - run: sudo sh -c "mv ${GITHUB_WORKSPACE}/docker/* /var/lib/docker" - + pip install cmake==3.21.0 pybind11[global] einops onnxscript + pip install torch --no-cache-dir + pip cache purge - name: 'Checkout' uses: actions/checkout@v3 with: submodules: recursive - - - name: Start named container - run: | - docker run -v $(pwd):$(pwd) -w $(pwd) --name builder -d ghcr.io/nvidia/jax:jax sleep infinity - - - name: 'Dependencies' - run: | - docker exec builder bash -c '\ - pip install cmake==3.21.0 pybind11[global] einops onnxscript && \ - pip install torch --no-cache-dir --index-url https://download.pytorch.org/whl/cu130 && \ - pip cache purge - ' + - name: ccache + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad - name: 'Build' - run: docker exec -e MAX_JOBS=1 -e NVTE_FRAMEWORK=all builder bash -c 'pip install --no-cache-dir --no-build-isolation . -v --no-deps' + run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation . -v --no-deps + env: + NVTE_FRAMEWORK: all + MAX_JOBS: 1 + SCCACHE_GHA_ENABLED: "true" + SCCACHE_CACHE_SIZE: "1G" - name: 'Sanity check' - run: docker exec builder bash -c 'python3 tests/pytorch/test_sanity_import.py && python3 tests/jax/test_sanity_import.py' + run: | + python3 tests/pytorch/test_sanity_import.py + python3 tests/jax/test_sanity_import.py From 4171efee373d05ffe5f4cd01abc34cab13ac5e29 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 3 Feb 2026 04:31:11 +0000 Subject: [PATCH 07/22] ccache Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 7041884bf4..e2636f4fcb 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -32,7 +32,7 @@ jobs: NVTE_FRAMEWORK: none MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - SCCACHE_CACHE_SIZE: "1G" + SCCACHE_CACHE_SIZE: "5G" - name: 'Sanity check' run: python3 -c "import transformer_engine" working-directory: / @@ -45,9 +45,8 @@ jobs: steps: - name: 'Dependencies' run: | - pip install cmake==3.21.0 pybind11[global] ninja pydantic importlib-metadata>=1.0 packaging numpy einops onnxscript - pip install torch --no-cache-dir - pip cache purge + pip install --no-cache-dir cmake==3.21.0 pybind11[global] ninja pydantic importlib-metadata>=1.0 packaging numpy einops onnxscript + pip install --no-cache-dir torch - name: 'Checkout' uses: actions/checkout@v3 with: @@ -60,7 +59,7 @@ jobs: NVTE_FRAMEWORK: pytorch MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - SCCACHE_CACHE_SIZE: "1G" + SCCACHE_CACHE_SIZE: "5G" - name: 'Sanity check' run: python3 tests/pytorch/test_sanity_import.py jax: @@ -85,7 +84,7 @@ jobs: NVTE_FRAMEWORK: jax MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - SCCACHE_CACHE_SIZE: "1G" + SCCACHE_CACHE_SIZE: "5G" - name: 'Sanity check' run: python3 tests/jax/test_sanity_import.py all: @@ -97,9 +96,8 @@ jobs: steps: - name: 'Dependencies' run: | - pip install cmake==3.21.0 pybind11[global] einops onnxscript - pip install torch --no-cache-dir - pip cache purge + pip install --no-cache-dir cmake==3.21.0 pybind11[global] einops onnxscript + pip install --no-cache-dir torch - name: 'Checkout' uses: actions/checkout@v3 with: @@ -112,7 +110,7 @@ jobs: NVTE_FRAMEWORK: all MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - SCCACHE_CACHE_SIZE: "1G" + SCCACHE_CACHE_SIZE: "5G" - name: 'Sanity check' run: | python3 tests/pytorch/test_sanity_import.py From b44ec7439c09f33bb8e61407be8dcbc10a1ef634 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 3 Feb 2026 04:39:21 +0000 Subject: [PATCH 08/22] ccache Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 4 ---- 1 file changed, 4 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index e2636f4fcb..43c010023f 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -30,7 +30,6 @@ jobs: run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation . -v env: NVTE_FRAMEWORK: none - MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" SCCACHE_CACHE_SIZE: "5G" - name: 'Sanity check' @@ -57,7 +56,6 @@ jobs: run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation . -v --no-deps env: NVTE_FRAMEWORK: pytorch - MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" SCCACHE_CACHE_SIZE: "5G" - name: 'Sanity check' @@ -82,7 +80,6 @@ jobs: NVTE_CCACHE_BIN=sccache NVTE_USE_CCACHE=1 pip install --no-build-isolation . -v env: NVTE_FRAMEWORK: jax - MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" SCCACHE_CACHE_SIZE: "5G" - name: 'Sanity check' @@ -108,7 +105,6 @@ jobs: run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation . -v --no-deps env: NVTE_FRAMEWORK: all - MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" SCCACHE_CACHE_SIZE: "5G" - name: 'Sanity check' From 23ea443f2bac8843ad49ffc881e9c0c912112b35 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 3 Feb 2026 04:43:27 +0000 Subject: [PATCH 09/22] ccache Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 43c010023f..6a1cbde7d4 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -10,7 +10,7 @@ on: jobs: core: name: 'Core' - runs-on: ubuntu-latest + runs-on: ubuntu-latest-8-cores container: image: nvcr.io/nvidia/cuda:12.1.0-devel-ubuntu22.04 options: --user root @@ -37,7 +37,7 @@ jobs: working-directory: / pytorch: name: 'PyTorch' - runs-on: ubuntu-latest + runs-on: ubuntu-latest-8-cores container: image: ghcr.io/nvidia/jax:jax options: --user root @@ -62,7 +62,7 @@ jobs: run: python3 tests/pytorch/test_sanity_import.py jax: name: 'JAX' - runs-on: ubuntu-latest + runs-on: ubuntu-latest-8-cores container: image: ghcr.io/nvidia/jax:jax options: --user root @@ -86,7 +86,7 @@ jobs: run: python3 tests/jax/test_sanity_import.py all: name: 'All' - runs-on: ubuntu-latest + runs-on: ubuntu-latest-8-cores container: image: ghcr.io/nvidia/jax:jax options: --user root From a0a528f021cfe6e228300219c948610ef20eebc5 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 3 Feb 2026 04:47:23 +0000 Subject: [PATCH 10/22] ccache Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 6a1cbde7d4..405e2c40aa 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -10,7 +10,7 @@ on: jobs: core: name: 'Core' - runs-on: ubuntu-latest-8-cores + runs-on: ubuntu-latest-4-cores container: image: nvcr.io/nvidia/cuda:12.1.0-devel-ubuntu22.04 options: --user root @@ -37,7 +37,7 @@ jobs: working-directory: / pytorch: name: 'PyTorch' - runs-on: ubuntu-latest-8-cores + runs-on: ubuntu-latest-4-cores container: image: ghcr.io/nvidia/jax:jax options: --user root @@ -62,7 +62,7 @@ jobs: run: python3 tests/pytorch/test_sanity_import.py jax: name: 'JAX' - runs-on: ubuntu-latest-8-cores + runs-on: ubuntu-latest-4-cores container: image: ghcr.io/nvidia/jax:jax options: --user root @@ -86,7 +86,7 @@ jobs: run: python3 tests/jax/test_sanity_import.py all: name: 'All' - runs-on: ubuntu-latest-8-cores + runs-on: ubuntu-latest-4-cores container: image: ghcr.io/nvidia/jax:jax options: --user root From 95c333f50ece0abd699fd93cc9acbb658f5b91fa Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 3 Feb 2026 04:51:00 +0000 Subject: [PATCH 11/22] ccache Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 405e2c40aa..980509d263 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -10,7 +10,7 @@ on: jobs: core: name: 'Core' - runs-on: ubuntu-latest-4-cores + runs-on: ubuntu-latest container: image: nvcr.io/nvidia/cuda:12.1.0-devel-ubuntu22.04 options: --user root @@ -37,7 +37,7 @@ jobs: working-directory: / pytorch: name: 'PyTorch' - runs-on: ubuntu-latest-4-cores + runs-on: ubuntu-latest container: image: ghcr.io/nvidia/jax:jax options: --user root @@ -45,7 +45,7 @@ jobs: - name: 'Dependencies' run: | pip install --no-cache-dir cmake==3.21.0 pybind11[global] ninja pydantic importlib-metadata>=1.0 packaging numpy einops onnxscript - pip install --no-cache-dir torch + pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cu130 - name: 'Checkout' uses: actions/checkout@v3 with: @@ -62,7 +62,7 @@ jobs: run: python3 tests/pytorch/test_sanity_import.py jax: name: 'JAX' - runs-on: ubuntu-latest-4-cores + runs-on: ubuntu-latest container: image: ghcr.io/nvidia/jax:jax options: --user root @@ -86,7 +86,7 @@ jobs: run: python3 tests/jax/test_sanity_import.py all: name: 'All' - runs-on: ubuntu-latest-4-cores + runs-on: ubuntu-latest container: image: ghcr.io/nvidia/jax:jax options: --user root @@ -94,7 +94,7 @@ jobs: - name: 'Dependencies' run: | pip install --no-cache-dir cmake==3.21.0 pybind11[global] einops onnxscript - pip install --no-cache-dir torch + pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cu130 - name: 'Checkout' uses: actions/checkout@v3 with: From cb3fa26fda34650cb392bbac757fce0adbb1110c Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 3 Feb 2026 04:55:58 +0000 Subject: [PATCH 12/22] ccache Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 980509d263..b064c17fff 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -30,6 +30,7 @@ jobs: run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation . -v env: NVTE_FRAMEWORK: none + MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" SCCACHE_CACHE_SIZE: "5G" - name: 'Sanity check' @@ -56,6 +57,7 @@ jobs: run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation . -v --no-deps env: NVTE_FRAMEWORK: pytorch + MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" SCCACHE_CACHE_SIZE: "5G" - name: 'Sanity check' @@ -80,6 +82,7 @@ jobs: NVTE_CCACHE_BIN=sccache NVTE_USE_CCACHE=1 pip install --no-build-isolation . -v env: NVTE_FRAMEWORK: jax + MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" SCCACHE_CACHE_SIZE: "5G" - name: 'Sanity check' From cc3c5b1925e32f2de898a547fb58efb0f83a5cdb Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 3 Feb 2026 04:56:00 +0000 Subject: [PATCH 13/22] ccache Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index b064c17fff..7387dfbe5b 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -108,6 +108,7 @@ jobs: run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation . -v --no-deps env: NVTE_FRAMEWORK: all + MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" SCCACHE_CACHE_SIZE: "5G" - name: 'Sanity check' From 52f6cb20919bc697de54b8d9feda2d8e3c70fa23 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 3 Feb 2026 07:11:31 +0000 Subject: [PATCH 14/22] ccache Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 7387dfbe5b..e98cde121d 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -32,7 +32,7 @@ jobs: NVTE_FRAMEWORK: none MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - SCCACHE_CACHE_SIZE: "5G" + SCCACHE_CACHE_SIZE: "2G" - name: 'Sanity check' run: python3 -c "import transformer_engine" working-directory: / @@ -59,7 +59,7 @@ jobs: NVTE_FRAMEWORK: pytorch MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - SCCACHE_CACHE_SIZE: "5G" + SCCACHE_CACHE_SIZE: "2G" - name: 'Sanity check' run: python3 tests/pytorch/test_sanity_import.py jax: @@ -84,7 +84,7 @@ jobs: NVTE_FRAMEWORK: jax MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - SCCACHE_CACHE_SIZE: "5G" + SCCACHE_CACHE_SIZE: "2G" - name: 'Sanity check' run: python3 tests/jax/test_sanity_import.py all: @@ -110,7 +110,7 @@ jobs: NVTE_FRAMEWORK: all MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - SCCACHE_CACHE_SIZE: "5G" + SCCACHE_CACHE_SIZE: "2G" - name: 'Sanity check' run: | python3 tests/pytorch/test_sanity_import.py From 4ebbe220675291476198b9adeb73e979cbec5946 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 3 Feb 2026 07:38:09 +0000 Subject: [PATCH 15/22] ccache Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 32 ++++++++++++++++++-------------- 1 file changed, 18 insertions(+), 14 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index e98cde121d..c3033b8fcb 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -32,7 +32,6 @@ jobs: NVTE_FRAMEWORK: none MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - SCCACHE_CACHE_SIZE: "2G" - name: 'Sanity check' run: python3 -c "import transformer_engine" working-directory: / @@ -40,13 +39,15 @@ jobs: name: 'PyTorch' runs-on: ubuntu-latest container: - image: ghcr.io/nvidia/jax:jax + image: nvcr.io/nvidia/cuda:12.1.0-devel-ubuntu22.04 options: --user root steps: - name: 'Dependencies' run: | - pip install --no-cache-dir cmake==3.21.0 pybind11[global] ninja pydantic importlib-metadata>=1.0 packaging numpy einops onnxscript - pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cu130 + apt-get update + apt-get install -y git python3.9 pip cudnn9-cuda-12 + pip install cmake==3.21.0 pybind11[global] ninja pydantic importlib-metadata>=1.0 packaging numpy einops onnxscript + pip install torch --index-url https://download.pytorch.org/whl/cu121 - name: 'Checkout' uses: actions/checkout@v3 with: @@ -59,18 +60,21 @@ jobs: NVTE_FRAMEWORK: pytorch MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - SCCACHE_CACHE_SIZE: "2G" - 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 + image: nvcr.io/nvidia/cuda:12.1.0-devel-ubuntu22.04 options: --user root steps: - name: 'Dependencies' - run: pip install cmake==3.21.0 pybind11[global] + run: | + apt-get update + apt-get install -y git python3.9 pip cudnn9-cuda-12 + pip install cmake==3.21.0 pybind11[global] ninja + pip install jax[cuda12] - name: 'Checkout' uses: actions/checkout@v3 with: @@ -78,26 +82,27 @@ jobs: - name: ccache uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad - name: 'Build' - run: | - NVTE_CCACHE_BIN=sccache NVTE_USE_CCACHE=1 pip install --no-build-isolation . -v + run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation . -v env: NVTE_FRAMEWORK: jax MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - SCCACHE_CACHE_SIZE: "2G" - 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 + image: nvcr.io/nvidia/cuda:12.1.0-devel-ubuntu22.04 options: --user root steps: - name: 'Dependencies' run: | - pip install --no-cache-dir cmake==3.21.0 pybind11[global] einops onnxscript - pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cu130 + apt-get update + apt-get install -y git python3.9 pip cudnn9-cuda-12 + pip install cmake==3.21.0 pybind11[global] ninja pydantic importlib-metadata>=1.0 packaging numpy einops onnxscript + pip install torch --index-url https://download.pytorch.org/whl/cu121 + pip install jax[cuda12] - name: 'Checkout' uses: actions/checkout@v3 with: @@ -110,7 +115,6 @@ jobs: NVTE_FRAMEWORK: all MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - SCCACHE_CACHE_SIZE: "2G" - name: 'Sanity check' run: | python3 tests/pytorch/test_sanity_import.py From 89d2985ce14f94260b947b2099c07ba19dcab45b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 3 Feb 2026 10:09:01 +0000 Subject: [PATCH 16/22] ccache Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index c3033b8fcb..8badd9b840 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -12,7 +12,7 @@ jobs: name: 'Core' runs-on: ubuntu-latest container: - image: nvcr.io/nvidia/cuda:12.1.0-devel-ubuntu22.04 + image: nvcr.io/nvidia/cuda:13.0.0-devel-ubuntu22.04 options: --user root steps: - name: 'Dependencies' @@ -20,6 +20,7 @@ jobs: apt-get update apt-get install -y git python3.9 pip cudnn9-cuda-12 pip install cmake==3.21.0 pybind11[global] ninja + git config --global --add safe.directory '*' - name: 'Checkout' uses: actions/checkout@v3 with: @@ -39,7 +40,7 @@ jobs: name: 'PyTorch' runs-on: ubuntu-latest container: - image: nvcr.io/nvidia/cuda:12.1.0-devel-ubuntu22.04 + image: nvcr.io/nvidia/cuda:13.0.0-devel-ubuntu22.04 options: --user root steps: - name: 'Dependencies' @@ -47,7 +48,8 @@ jobs: apt-get update apt-get install -y git python3.9 pip cudnn9-cuda-12 pip install cmake==3.21.0 pybind11[global] ninja pydantic importlib-metadata>=1.0 packaging numpy einops onnxscript - pip install torch --index-url https://download.pytorch.org/whl/cu121 + pip install torch --index-url https://download.pytorch.org/whl/cu130 + git config --global --add safe.directory '*' - name: 'Checkout' uses: actions/checkout@v3 with: @@ -66,15 +68,16 @@ jobs: name: 'JAX' runs-on: ubuntu-latest container: - image: nvcr.io/nvidia/cuda:12.1.0-devel-ubuntu22.04 + image: nvcr.io/nvidia/cuda:13.0.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 - pip install jax[cuda12] + pip install cmake==3.21.0 pybind11[global] ninja packaging + pip install jax + git config --global --add safe.directory '*' - name: 'Checkout' uses: actions/checkout@v3 with: @@ -93,7 +96,7 @@ jobs: name: 'All' runs-on: ubuntu-latest container: - image: nvcr.io/nvidia/cuda:12.1.0-devel-ubuntu22.04 + image: nvcr.io/nvidia/cuda:13.0.0-devel-ubuntu22.04 options: --user root steps: - name: 'Dependencies' @@ -101,8 +104,9 @@ jobs: apt-get update apt-get install -y git python3.9 pip cudnn9-cuda-12 pip install cmake==3.21.0 pybind11[global] ninja pydantic importlib-metadata>=1.0 packaging numpy einops onnxscript - pip install torch --index-url https://download.pytorch.org/whl/cu121 + pip install torch --index-url https://download.pytorch.org/whl/cu130 pip install jax[cuda12] + git config --global --add safe.directory '*' - name: 'Checkout' uses: actions/checkout@v3 with: From b5f554ea3d84af9011128d7450fce7cee53391fa Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 3 Feb 2026 13:20:06 +0000 Subject: [PATCH 17/22] ccache Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 8badd9b840..31d73418df 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -76,7 +76,7 @@ jobs: apt-get update apt-get install -y git python3.9 pip cudnn9-cuda-12 pip install cmake==3.21.0 pybind11[global] ninja packaging - pip install jax + pip install jax[cuda13] flax[cuda13] git config --global --add safe.directory '*' - name: 'Checkout' uses: actions/checkout@v3 @@ -105,7 +105,7 @@ jobs: apt-get install -y git python3.9 pip cudnn9-cuda-12 pip install cmake==3.21.0 pybind11[global] ninja pydantic importlib-metadata>=1.0 packaging numpy einops onnxscript pip install torch --index-url https://download.pytorch.org/whl/cu130 - pip install jax[cuda12] + pip install jax[cuda13] flax[cuda13] git config --global --add safe.directory '*' - name: 'Checkout' uses: actions/checkout@v3 From 30cd354f4d69fe84796c8955407eee1d31f4649c Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 3 Feb 2026 13:47:43 +0000 Subject: [PATCH 18/22] ccache Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 31d73418df..b2b77827b8 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -15,6 +15,10 @@ jobs: image: nvcr.io/nvidia/cuda:13.0.0-devel-ubuntu22.04 options: --user root steps: + - name: 'Set up swap' + uses: pierotofy/set-swap-space@master + with: + swap-size-gb: 10 - name: 'Dependencies' run: | apt-get update @@ -43,6 +47,10 @@ jobs: image: nvcr.io/nvidia/cuda:13.0.0-devel-ubuntu22.04 options: --user root steps: + - name: 'Set up swap' + uses: pierotofy/set-swap-space@master + with: + swap-size-gb: 10 - name: 'Dependencies' run: | apt-get update @@ -71,6 +79,10 @@ jobs: image: nvcr.io/nvidia/cuda:13.0.0-devel-ubuntu22.04 options: --user root steps: + - name: 'Set up swap' + uses: pierotofy/set-swap-space@master + with: + swap-size-gb: 10 - name: 'Dependencies' run: | apt-get update @@ -99,6 +111,10 @@ jobs: image: nvcr.io/nvidia/cuda:13.0.0-devel-ubuntu22.04 options: --user root steps: + - name: 'Set up swap' + uses: pierotofy/set-swap-space@master + with: + swap-size-gb: 10 - name: 'Dependencies' run: | apt-get update From ba96b4277b6751d1522e1551bf8c6ce8252791e3 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 3 Feb 2026 13:58:51 +0000 Subject: [PATCH 19/22] ccache Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 20 ++++---------------- 1 file changed, 4 insertions(+), 16 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index b2b77827b8..05281ef7e8 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -15,10 +15,6 @@ jobs: image: nvcr.io/nvidia/cuda:13.0.0-devel-ubuntu22.04 options: --user root steps: - - name: 'Set up swap' - uses: pierotofy/set-swap-space@master - with: - swap-size-gb: 10 - name: 'Dependencies' run: | apt-get update @@ -37,6 +33,7 @@ jobs: NVTE_FRAMEWORK: none MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" + NVTE_CUDA_ARCHS: "120" - name: 'Sanity check' run: python3 -c "import transformer_engine" working-directory: / @@ -47,10 +44,6 @@ jobs: image: nvcr.io/nvidia/cuda:13.0.0-devel-ubuntu22.04 options: --user root steps: - - name: 'Set up swap' - uses: pierotofy/set-swap-space@master - with: - swap-size-gb: 10 - name: 'Dependencies' run: | apt-get update @@ -70,6 +63,7 @@ jobs: NVTE_FRAMEWORK: pytorch MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" + NVTE_CUDA_ARCHS: "120" - name: 'Sanity check' run: python3 tests/pytorch/test_sanity_import.py jax: @@ -79,10 +73,6 @@ jobs: image: nvcr.io/nvidia/cuda:13.0.0-devel-ubuntu22.04 options: --user root steps: - - name: 'Set up swap' - uses: pierotofy/set-swap-space@master - with: - swap-size-gb: 10 - name: 'Dependencies' run: | apt-get update @@ -102,6 +92,7 @@ jobs: NVTE_FRAMEWORK: jax MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" + NVTE_CUDA_ARCHS: "120" - name: 'Sanity check' run: python3 tests/jax/test_sanity_import.py all: @@ -111,10 +102,6 @@ jobs: image: nvcr.io/nvidia/cuda:13.0.0-devel-ubuntu22.04 options: --user root steps: - - name: 'Set up swap' - uses: pierotofy/set-swap-space@master - with: - swap-size-gb: 10 - name: 'Dependencies' run: | apt-get update @@ -135,6 +122,7 @@ jobs: NVTE_FRAMEWORK: all MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" + NVTE_CUDA_ARCHS: "120" - name: 'Sanity check' run: | python3 tests/pytorch/test_sanity_import.py From 7ad602fc7d1e9da9c521379ac7f62238ac489a4d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 3 Feb 2026 14:05:27 +0000 Subject: [PATCH 20/22] ccache Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 05281ef7e8..fc3e801727 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -33,7 +33,7 @@ jobs: NVTE_FRAMEWORK: none MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - NVTE_CUDA_ARCHS: "120" + NVTE_CUDA_ARCHS: "90" - name: 'Sanity check' run: python3 -c "import transformer_engine" working-directory: / @@ -63,7 +63,7 @@ jobs: NVTE_FRAMEWORK: pytorch MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - NVTE_CUDA_ARCHS: "120" + NVTE_CUDA_ARCHS: "90" - name: 'Sanity check' run: python3 tests/pytorch/test_sanity_import.py jax: @@ -92,7 +92,7 @@ jobs: NVTE_FRAMEWORK: jax MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - NVTE_CUDA_ARCHS: "120" + NVTE_CUDA_ARCHS: "90" - name: 'Sanity check' run: python3 tests/jax/test_sanity_import.py all: @@ -122,7 +122,7 @@ jobs: NVTE_FRAMEWORK: all MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - NVTE_CUDA_ARCHS: "120" + NVTE_CUDA_ARCHS: "90" - name: 'Sanity check' run: | python3 tests/pytorch/test_sanity_import.py From d3bfeeba5cc7553786f065e35427b16d1a867fa3 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 4 Feb 2026 06:49:31 +0000 Subject: [PATCH 21/22] added blackwell Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index fc3e801727..6ff0a57a02 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -33,7 +33,7 @@ jobs: NVTE_FRAMEWORK: none MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - NVTE_CUDA_ARCHS: "90" + NVTE_CUDA_ARCHS: "90;100" - name: 'Sanity check' run: python3 -c "import transformer_engine" working-directory: / @@ -63,7 +63,7 @@ jobs: NVTE_FRAMEWORK: pytorch MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - NVTE_CUDA_ARCHS: "90" + NVTE_CUDA_ARCHS: "90;100" - name: 'Sanity check' run: python3 tests/pytorch/test_sanity_import.py jax: @@ -92,7 +92,7 @@ jobs: NVTE_FRAMEWORK: jax MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - NVTE_CUDA_ARCHS: "90" + NVTE_CUDA_ARCHS: "90;100" - name: 'Sanity check' run: python3 tests/jax/test_sanity_import.py all: @@ -122,7 +122,7 @@ jobs: NVTE_FRAMEWORK: all MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - NVTE_CUDA_ARCHS: "90" + NVTE_CUDA_ARCHS: "90;100" - name: 'Sanity check' run: | python3 tests/pytorch/test_sanity_import.py From 03488d929b10785892ab56fbf2fd638c10836a96 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 4 Feb 2026 07:46:16 +0000 Subject: [PATCH 22/22] added blackwell Signed-off-by: Pawel Gadzinski --- .github/workflows/build.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 6ff0a57a02..f0ad600691 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -33,7 +33,7 @@ jobs: NVTE_FRAMEWORK: none MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - NVTE_CUDA_ARCHS: "90;100" + NVTE_CUDA_ARCHS: "100" - name: 'Sanity check' run: python3 -c "import transformer_engine" working-directory: / @@ -63,7 +63,7 @@ jobs: NVTE_FRAMEWORK: pytorch MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - NVTE_CUDA_ARCHS: "90;100" + NVTE_CUDA_ARCHS: "100" - name: 'Sanity check' run: python3 tests/pytorch/test_sanity_import.py jax: @@ -92,7 +92,7 @@ jobs: NVTE_FRAMEWORK: jax MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - NVTE_CUDA_ARCHS: "90;100" + NVTE_CUDA_ARCHS: "100" - name: 'Sanity check' run: python3 tests/jax/test_sanity_import.py all: @@ -122,7 +122,7 @@ jobs: NVTE_FRAMEWORK: all MAX_JOBS: 1 SCCACHE_GHA_ENABLED: "true" - NVTE_CUDA_ARCHS: "90;100" + NVTE_CUDA_ARCHS: "100" - name: 'Sanity check' run: | python3 tests/pytorch/test_sanity_import.py