diff --git a/.github/workflows/minimum_versions.yml b/.github/workflows/minimum_versions.yml new file mode 100644 index 0000000000..0f46427ea4 --- /dev/null +++ b/.github/workflows/minimum_versions.yml @@ -0,0 +1,85 @@ +# 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.10 python3-pip cudnn9-cuda-12 libnccl-dev libnccl2 + python3.10 --version + python3.10 -m pip install torch==${MIN_TORCH_VERSION} + python3.10 -m 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 python3.10 -m 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.10 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.10 python3-pip cudnn9-cuda-12 + python3.10 --version + python3.10 -m pip install jax==${MIN_JAX_VERSION} "flax>=0.7.1" + python3.10 -m 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 python3.10 -m 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.10 tests/jax/test_sanity_import.py diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index 2bb238c522..98331ccbd8 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 @@ -89,7 +90,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(): 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 diff --git a/transformer_engine/pytorch/csrc/extensions/ep.cpp b/transformer_engine/pytorch/csrc/extensions/ep.cpp index cef489dbbb..216979f1f1 100644 --- a/transformer_engine/pytorch/csrc/extensions/ep.cpp +++ b/transformer_engine/pytorch/csrc/extensions/ep.cpp @@ -19,13 +19,18 @@ #include #include #include -#include -#include #include #include #include "transformer_engine/comm_window.h" +// torch's NCCL symm-mem headers (zero-copy path) exist only since torch 2.11; +// without them NCCL_HAS_SYMMEM_SUPPORT stays undefined and zero-copy is compiled out. +#if __has_include() +#include +#include +#endif + #ifdef NCCL_HAS_SYMMEM_SUPPORT #include #endif @@ -152,6 +157,14 @@ size_t check_mxfp8_scale_pair(const at::Tensor& send_scale, const at::Tensor& re bool ep_get_zero_copy() { return g_zero_copy_enabled.load(std::memory_order_relaxed); } +bool ep_zero_copy_supported() { +#ifdef NCCL_HAS_SYMMEM_SUPPORT + return true; +#else + return false; +#endif +} + // ── Bootstrap ──────────────────────────────────────────────────────────────── // Borrows torch's NCCL host comm (from ``ProcessGroupNCCL._comm_ptr()``). // ``group_name`` is captured for the symm-mem window resolver. @@ -498,6 +511,8 @@ void register_ep_bindings(pybind11::module_& m) { m.def("ep_finalize", &ep_finalize, "Tear down the EP backend. Idempotent.", py::call_guard()); m.def("ep_get_zero_copy", &ep_get_zero_copy, "Return the current EP zero-copy toggle state."); + m.def("ep_zero_copy_supported", &ep_zero_copy_supported, + "Return True when the extension was built with NCCL symm-mem (zero-copy) support."); m.def("ep_handle_mem_size", &ep_handle_mem_size, "Return the handle_mem byte size for the given layer config.", py::arg("top_k"), py::arg("dispatch_output_per_expert_alignment") = 0); diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 2799bfdf5d..c520cf0289 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -145,6 +145,11 @@ def ep_bootstrap( raise ValueError("ep_bootstrap: drop_on_overflow requires recv_capacity_per_rank") _check_nccl_runtime_version() if zero_copy: + if not tex.ep_zero_copy_supported(): + raise RuntimeError( + "ep_bootstrap: zero_copy=True requires the Transformer Engine torch extension " + "built with NCCL symm-mem support (torch >= 2.11 with USE_NCCL)." + ) warnings.warn( "ep_bootstrap(zero_copy=True) is experimental; the symm-mem IO path " "and its alias contracts on EpBuffer slots are subject to change.",