[CI] Add workflow verifying build against minimum supported PyTorch and JAX

[pggPL]

Description

Adds a GitHub Actions workflow (minimum_versions.yml) that builds TE against pinned minimum framework versions and runs the sanity imports. The workflow versions are set in a single place (MIN_TORCH_VERSION: 2.8.0, MIN_JAX_VERSION: 0.5.3). This guards the minimum framework versions selected for CI; the older package requirement declarations are left unchanged for now.

Context — the declared lower bounds are older than the CI-supported floors (verified empirically and against pytorch git history):

• build_tools/pytorch.py currently declares torch>=2.1 and remains unchanged in this PR. Current main does not compile against torch 2.1 (std::optional vs c10::optional in generator APIs — alias exists since 2.2; usable c10::Stream pybind caster since 2.4) and does not import below 2.6 (torch.distributed.fsdp._fully_shard moved there in 2.6, torch.distributed.tensor.DTensor export since 2.5). The minimum PyTorch version supported and enforced by this workflow is 2.8.0.
• pyproject.toml currently declares jax>=0.5.0 and remains unchanged in this PR, but transformer_engine/jax/sharding.py imports jax.sharding.get_abstract_mesh, which only exists since jax 0.5.3. The build compiles against 0.5.0 but the import fails; the minimum JAX version supported and enforced by this workflow is 0.5.3.

The declared package requirements can be reconciled separately.

The PyTorch job builds NVTE_CUDA_ARCHS=90 so the NCCL EP path is exercised; the JAX job builds sm70. Both use MAX_JOBS=2 (the 7GB public runner OOMs at full parallelism) and sccache backed by the GitHub Actions cache (~10 min warm). The sm90 sanity import uses the CUDA toolkit's libcuda.so stub since the runner has no GPU driver (the EP-enabled core library links the driver via NCCL).

Building sm90 against an older torch surfaced two gaps in the NCCL EP build gating, fixed here:

1. The torch extension enabled NVTE_WITH_NCCL_EP based only on the env var (default on), while the common library gates NCCL EP on having an arch >= 90 in the target list. A build targeting only archs < 90 compiled the EP extension against a core library with no nvte_ep_* symbols. The extension now uses the same nccl_ep_enabled() gate as the common CMake side.
2. ep.cpp includes torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp / NCCLSymmetricMemory.hpp and uses get_window()/get_offset(), all of which only exist since torch 2.11 — i.e. a default source build on Hopper+ hard-failed for any older torch. The extension now auto-skips NCCL EP (with a log message) when the installed torch lacks the symm-mem headers; an explicit NVTE_WITH_NCCL_EP=1 raises a clear error instead. Builds against torch >= 2.11 are unaffected.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Add .github/workflows/minimum_versions.yml with two jobs, on pull_request / workflow_dispatch:
  • "Minimum supported PyTorch": build with torch==2.8.0 on sm90 (NCCL EP path exercised) + sanity import
  • "Minimum supported JAX": build with jax==0.5.3 (CPU jaxlib) on sm70 + sanity import
• build_tools/pytorch.py: gate NVTE_WITH_NCCL_EP/USE_NCCL defines on nccl_ep_enabled() (matching the common CMake gate) and on the installed torch providing the c10d symm-mem headers (torch >= 2.11); auto-skip with a message, error only when EP was explicitly requested

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

🤖 Generated with Claude Code

[greptile-apps]

Greptile Summary

The PR adds CI jobs that build and sanity-import Transformer Engine against the selected minimum PyTorch and JAX versions. It also aligns NCCL EP architecture gating and makes zero-copy support conditional on the installed PyTorch headers.

• Adds containerized minimum-version PyTorch and JAX workflow jobs.
• Reuses the common NCCL EP architecture gate for the PyTorch extension.
• Compiles EP without zero-copy on older PyTorch versions and reports unsupported zero-copy requests explicitly.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains.

No blocking failure remains.

Important Files Changed

Filename	Overview
.github/workflows/minimum_versions.yml	Adds minimum-framework build and sanity-import jobs for PyTorch 2.8.0 and JAX 0.5.3.
build_tools/pytorch.py	Aligns the PyTorch extension's NCCL EP define with the common architecture-aware build gate.
transformer_engine/pytorch/csrc/extensions/ep.cpp	Makes PyTorch symmetric-memory headers optional and exposes whether zero-copy support was compiled.
transformer_engine/pytorch/ep.py	Rejects zero-copy bootstrap requests when the native extension lacks symmetric-memory support.

_{Reviews (6): Last reviewed commit: "Compile out EP zero-copy path when torch..." | Re-trigger Greptile}

[pggPL]

Added @phu0ngng as a reviewer since EP raised pytorch version to 2.11, please look at the pytorch build part if this is ok.
After dealing with that we are able to support 2.6, but we plan to support >= 2.8 for TE 3.0 if I remember correctly.

[pggPL]

Also note we do not support versions we declare in build for JAX 0.5.3 is needed instead of 0.5.0

[phu0ngng]

LGTM. Thanks!

[fheinecke]

Is this tracked in documentation anywhere today?