diff --git a/.github/workflows/minimum_versions.yml b/.github/workflows/minimum_versions.yml new file mode 100644 index 0000000000..ea50b5ee05 --- /dev/null +++ b/.github/workflows/minimum_versions.yml @@ -0,0 +1,83 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# Verify that TE builds and imports against the minimum supported framework versions. +name: 'Minimum supported versions' +on: + pull_request: + workflow_dispatch: +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true +env: + MIN_TORCH_VERSION: "2.8.0" + MIN_JAX_VERSION: "0.5.3" +jobs: + min-pytorch: + name: 'Minimum supported PyTorch' + runs-on: ubuntu-latest + container: + image: nvcr.io/nvidia/cuda:12.8.0-devel-ubuntu22.04 + options: --user root + steps: + - name: 'Dependencies' + run: | + apt-get update + apt-get install -y --allow-change-held-packages git python3.9 pip cudnn9-cuda-12 libnccl-dev libnccl2 + pip install torch==${MIN_TORCH_VERSION} + pip install cmake pybind11[global] ninja pydantic importlib-metadata>=1.0 packaging einops onnxscript "nvidia-cudnn-frontend>=1.25.0" + - name: 'Checkout' + uses: actions/checkout@v3 + with: + submodules: recursive + - name: ccache + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad + - name: 'Build' + run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation --no-deps . -v + env: + NVTE_FRAMEWORK: pytorch + # Single arch to keep memory/build size down; sm90 so the NCCL EP + # path (and its torch-version gate) is exercised + NVTE_CUDA_ARCHS: "90" + # Prefer apt cudnn9 headers over the pip nvidia-cudnn-cu12 copy + CUDNN_INCLUDE_PATH: /usr/include/x86_64-linux-gnu + CUDNN_LIBRARY_PATH: /usr/lib/x86_64-linux-gnu + # Full parallelism OOMs the 7GB runner (exit 137) + MAX_JOBS: 2 + SCCACHE_GHA_ENABLED: "true" + - name: 'Sanity check' + # No GPU driver on the runner; the sm90 build links libcuda via NCCL EP + run: | + ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/lib/x86_64-linux-gnu/libcuda.so.1 + python3 tests/pytorch/test_sanity_import.py + min-jax: + name: 'Minimum supported JAX' + runs-on: ubuntu-latest + container: + image: nvcr.io/nvidia/cuda:12.8.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 jax==${MIN_JAX_VERSION} "flax>=0.7.1" + pip install cmake pybind11[global] ninja packaging pydantic "nvidia-cudnn-frontend>=1.25.0" + - name: 'Checkout' + uses: actions/checkout@v3 + with: + submodules: recursive + - name: ccache + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad + - name: 'Build' + run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation --no-deps . -v + env: + NVTE_FRAMEWORK: jax + NVTE_CUDA_ARCHS: "70" + CUDNN_INCLUDE_PATH: /usr/include/x86_64-linux-gnu + CUDNN_LIBRARY_PATH: /usr/lib/x86_64-linux-gnu + MAX_JOBS: 2 + SCCACHE_GHA_ENABLED: "true" + - name: 'Sanity check' + run: python3 tests/jax/test_sanity_import.py diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index 2bb238c522..4a660e3711 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -15,6 +15,7 @@ cuda_version, get_cuda_include_dirs, debug_build_enabled, + nccl_ep_enabled, setup_mpi_flags, ) from typing import List @@ -46,6 +47,24 @@ def test_requirements() -> List[str]: ] +def nccl_ep_supported_by_torch() -> bool: + """Return True when the installed torch can build the NCCL EP extension. + + ep.cpp needs torch's c10d symm-mem headers (torch >= 2.11). Mirrors the + auto-skip/raise policy of nccl_ep_enabled() on the torch-version axis. + """ + from torch.utils.cpp_extension import include_paths + + header = Path("torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp") + if any((Path(p) / header).exists() for p in include_paths()): + return True + message = f"Installed torch does not provide {header}; NCCL EP requires torch >= 2.11." + if bool(int(os.getenv("NVTE_WITH_NCCL_EP", "0"))): + raise RuntimeError(f"NVTE_WITH_NCCL_EP=1 was set but: {message}") + print(f"[NCCL EP] {message} Skipping NCCL EP in the torch extension.") + return False + + def setup_pytorch_extension( csrc_source_files, csrc_header_files, @@ -89,7 +108,7 @@ def setup_pytorch_extension( # Mirror the NCCL EP gate from setup.py / common CMake. When disabled, the # ep.cpp source no-ops at the #ifdef boundary; without the define it would # produce undefined references to nvte_ep_*. - if bool(int(os.getenv("NVTE_WITH_NCCL_EP", "1"))): + if nccl_ep_enabled() and nccl_ep_supported_by_torch(): cxx_flags.append("-DNVTE_WITH_NCCL_EP") # PyTorch's symm-mem headers gate the NCCL_HAS_SYMMEM_* feature macros on # USE_NCCL. The EP extension shares the symm-mem NCCL comm with torch, so