Enhance dependency management in pyproject.toml and CI workflows by adding transformer-engine extras for CUDA 12 and 13. Update EXTRAS_TAG in CI configurations to include transformer-engine-cu12, ensuring compatibility with the latest dependencies. This change addresses potential issues with dynamic CUDA dependency resolution in transformer-engine-torch. (#1683)

Author: peterdsharpe

.github/regen-ci-deps-lock.sh

@@ -43,7 +43,7 @@ set -euo pipefail
 
 # These MUST match the workflow `env:` values in
 # .github/workflows/github-{pr,nightly-uv}.yml.  Bump in lockstep.
-EXTRAS_TAG="${EXTRAS_TAG:-cu12,natten-cu12,utils-extras,mesh-extras,nn-extras,model-extras,datapipes-extras,uq-extras,gnns,sym,perf}"
+EXTRAS_TAG="${EXTRAS_TAG:-cu12,natten-cu12,utils-extras,mesh-extras,nn-extras,model-extras,datapipes-extras,uq-extras,gnns,sym,transformer-engine-cu12}"
 UV_VERSION="${UV_VERSION:-0.11.7}"
 # Matches the `--find-links` URL committed to .github/ci-requirements.txt.
 # Bump the torch-X.Y.Z+cu128 segment in lockstep with the locked torch version.

.github/workflows/github-nightly-uv.yml

@@ -86,7 +86,7 @@ env:
   # with no extras home (moto, scikit-image, pyg_lib, earth2grid, ...) are
   # installed by the `Install CI-only test dependencies` step inside the
   # setup-uv-env composite action.
-  EXTRAS_TAG: "cu12,natten-cu12,utils-extras,mesh-extras,nn-extras,model-extras,datapipes-extras,uq-extras,gnns,sym,perf"
+  EXTRAS_TAG: "cu12,natten-cu12,utils-extras,mesh-extras,nn-extras,model-extras,datapipes-extras,uq-extras,gnns,sym,transformer-engine-cu12"
 
   # ---- Cache key prefixes ------------------------------------------------
   # Inlined literally because GitHub Actions does not allow env-to-env

.github/workflows/github-pr.yml

@@ -63,7 +63,7 @@ env:
   # with no extras home (moto, scikit-image, pyg_lib, earth2grid, ...) are
   # installed by the `Install CI-only test dependencies` step inside the
   # setup-uv-env composite action.
-  EXTRAS_TAG: "cu12,natten-cu12,utils-extras,mesh-extras,nn-extras,model-extras,datapipes-extras,uq-extras,gnns,sym,perf"
+  EXTRAS_TAG: "cu12,natten-cu12,utils-extras,mesh-extras,nn-extras,model-extras,datapipes-extras,uq-extras,gnns,sym,transformer-engine-cu12"
 
   # ---- Cache key prefixes (shared with nightly) --------------------------
   # The `-fullextras` suffix is bumped relative to the previous prefix so

pyproject.toml

@@ -76,6 +76,10 @@ conflicts = [
         { extra = "natten-cu12" },
         { extra = "natten-cu13" },
     ],
+    [
+        { extra = "transformer-engine-cu12" },
+        { extra = "transformer-engine-cu13" },
+    ],
 ]
 
 [tool.uv.extra-build-dependencies]
@@ -152,6 +156,28 @@ natten = [
     { index = "natten-cu130-whl", extra = "natten-cu13" },
 ]
 
+# transformer-engine-torch is an sdist whose setup.py dynamically pins
+# transformer_engine_cu{12,13} based on torch.version.cuda in the build
+# environment.  Without this override, uv lock freezes whichever CUDA
+# variant happened to be in the build env, which broke CI when PyTorch
+# 2.12 flipped the PyPI default to cu130 (the cu13 binary then got
+# pulled into a CUDA 12.8 runtime, producing missing libcublas.so.13
+# errors at import time).  We provide the CUDA-agnostic subset of the
+# dependencies here; the matching cuXX binary is pulled by the explicit
+# transformer-engine-cu12 / transformer-engine-cu13 extras below.
+# Static deps mirror TE 2.15's build_tools/pytorch.py::install_requirements().
+[[tool.uv.dependency-metadata]]
+name = "transformer-engine-torch"
+requires-dist = [
+    "torch>=2.1",
+    "einops",
+    "onnxscript",
+    "onnx",
+    "packaging",
+    "pydantic",
+    "nvdlfw-inspect",
+]
+
 #####################################################################
 # Flags Controlling the local build of physicsnemo
 #####################################################################
@@ -248,6 +274,21 @@ natten-cu12 = [
 natten-cu13 = [
     "natten>=0.21.5",
 ]
+# Transformer Engine extras (mutually exclusive via [tool.uv] conflicts).
+# Pair with the matching CUDA backend extra, e.g.:
+#   pip install nvidia-physicsnemo[cu12,transformer-engine-cu12]
+#   pip install nvidia-physicsnemo[cu13,transformer-engine-cu13]
+# Uses the metapackage's `core_cu12` / `core_cu13` extras (introduced in
+# transformer-engine 2.14.0) to explicitly pin the CUDA backend binary,
+# so the resolution doesn't depend on torch.version.cuda in the build
+# env.  See [[tool.uv.dependency-metadata]] above for the matching
+# override on transformer-engine-torch's dynamic CUDA dep.
+transformer-engine-cu12 = [
+    "transformer_engine[pytorch,core_cu12]>=2.14.0",
+]
+transformer-engine-cu13 = [
+    "transformer_engine[pytorch,core_cu13]>=2.14.0",
+]
 utils-extras = [
     "wandb",
     "mlflow>=3.12.0",
@@ -295,10 +336,6 @@ sym = [
     "sympy>=1.12",
 ]
 
-perf = [
-    "transformer_engine[pytorch]",
-]
-
 
 #####################################################################
 # Linting configuration