python3Packages.transformer-engine: 2.12 -> 2.13 (#508557)

This commit is contained in:
Gaétan Lepage
2026-05-11 11:29:01 +00:00
committed by GitHub
2 changed files with 16 additions and 13 deletions
@@ -1,8 +1,8 @@
diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py
index 02388d2e..f2eb337c 100644
index 40933f17..f2eb337c 100644
--- a/transformer_engine/common/__init__.py
+++ b/transformer_engine/common/__init__.py
@@ -239,117 +239,7 @@ def _get_sys_extension() -> str:
@@ -239,119 +239,7 @@ def _get_sys_extension() -> str:
def _nvidia_cudart_include_dir() -> str:
"""Returns the include directory for cuda_runtime.h if exists in python environment."""
@@ -12,11 +12,13 @@ index 02388d2e..f2eb337c 100644
- return ""
-
- # Installing some nvidia-* packages, like nvshmem, create nvidia name, so "import nvidia"
- # above doesn't through. However, they don't set "__file__" attribute.
- if nvidia.__file__ is None:
- return ""
- # above doesn't throw. However, they don't set "__file__" attribute.
- if nvidia.__file__ is not None:
- nvidia_root = Path(nvidia.__file__).parent
- else:
- nvidia_root = Path(nvidia.__path__[0]) # namespace package
-
- include_dir = Path(nvidia.__file__).parent / "cuda_runtime"
- include_dir = nvidia_root / "cuda_runtime"
- return str(include_dir) if include_dir.exists() else ""
-
-
@@ -121,7 +123,7 @@ index 02388d2e..f2eb337c 100644
@functools.lru_cache(maxsize=None)
@@ -364,18 +254,9 @@ if "NVTE_PROJECT_BUILDING" not in os.environ or bool(int(os.getenv("NVTE_RELEASE
@@ -366,18 +254,9 @@ if "NVTE_PROJECT_BUILDING" not in os.environ or bool(int(os.getenv("NVTE_RELEASE
# `_load_cuda_library` is used for packages that must be loaded
# during runtime. Both system and pypi packages are searched
# and an error is thrown if not found.
@@ -77,8 +77,9 @@ let
in
buildPythonPackage.override { stdenv = backendStdenv; } (finalAttrs: {
pname = "transformer-engine";
version = "2.12";
version = "2.14";
pyproject = true;
__structuredAttrs = true;
src = fetchFromGitHub {
owner = "NVIDIA";
@@ -86,7 +87,7 @@ buildPythonPackage.override { stdenv = backendStdenv; } (finalAttrs: {
tag = "v${finalAttrs.version}";
# Their CMakeLists.txt does not easily let us inject dependencies
fetchSubmodules = true;
hash = "sha256-/e11kacSYPKdjVEKAo3x/CarzKhO3tiTsMjYWLzHbls=";
hash = "sha256-yxcUn75blB5ssEqGXZFDUrBv2/WM8yzumJCt5olV5Po=";
};
patches =
@@ -102,8 +103,8 @@ buildPythonPackage.override { stdenv = backendStdenv; } (finalAttrs: {
# https://github.com/NVIDIA/TransformerEngine/pull/2832
(fetchpatch {
name = "fix-cuda-arch-cmake-logic";
url = "https://github.com/GaetanLepage/TransformerEngine/commit/a3cf63e0d03dd9af1d494854949387f1ae677bf0.patch";
hash = "sha256-g2aIF0fROsExEjuNiyI62/rrCOXYyOjyQIOn6rCrUyI=";
url = "https://github.com/NVIDIA/TransformerEngine/commit/fca261ecd09c318d22e7eeebda79632eed8cb9e4.patch";
hash = "sha256-nph01cIfmjN7RUFZif5ORoz29CEHWwetiHMEZVnnOyY=";
})
]
++ optionals withNvshmem [
@@ -118,8 +119,8 @@ buildPythonPackage.override { stdenv = backendStdenv; } (finalAttrs: {
# https://github.com/NVIDIA/TransformerEngine/pull/2835
(fetchpatch {
name = "fix-jax-extension-build-with-mpi";
url = "https://github.com/GaetanLepage/TransformerEngine/commit/f68cd3cab34972a899ad0069e2c4ee806e8bc6fb.patch";
hash = "sha256-u0ljg1FwY0QjR+ETswpzWV+Sbv00JHI5CSrNQ/9zsuA=";
url = "https://github.com/NVIDIA/TransformerEngine/commit/2dd31bb849e83cce51c7d169db883862063d3a95.patch";
hash = "sha256-QSRMetseYPGGZCgGkS9rIj9nJdazCD4hv2IgPc+ClSM=";
})
];