CI: minimum supported PyTorch check

.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

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

transformer_engine/pytorch/csrc/extensions/ep.cpp

@@ -19,13 +19,18 @@
 #include <cstdio>
 #include <cstdlib>
 #include <string>
-#include <torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp>
-#include <torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp>
 #include <tuple>
 #include <vector>
 
 #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(<torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp>)
+#include <torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp>
+#endif
+
 #ifdef NCCL_HAS_SYMMEM_SUPPORT
 #include <torch/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.hpp>
 #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<py::gil_scoped_release>());
   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);

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.",