gador/nixpkgs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

python3Packages.transformer-engine: 2.14 -> 2.15 (#519564)

Browse Source
This commit is contained in:
Connor Baker
2026-05-13 22:00:02 +00:00
committed by GitHub
parent 1a11bb4e94 5a9a35ff48
commit 0f60144c1d
2 changed files with 51 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkgs/development/python-modules/transformer-engine/cuda-libs-paths.patch
+3 -58
View File
@@ -1,8 +1,8 @@
diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py
index 40933f17..f2eb337c 100644
index 40933f17..e59f43bb 100644
--- a/transformer_engine/common/__init__.py
+++ b/transformer_engine/common/__init__.py
@@ -239,119 +239,7 @@ def _get_sys_extension() -> str:
@@ -239,64 +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."""
@@ -64,66 +64,11 @@ index 40933f17..f2eb337c 100644
- raise RuntimeError(f"{lib_name} shared object not found.")
-
- return path_found, ctypes_handles
-
-
-@functools.lru_cache(maxsize=None)
-def _load_cuda_library_from_system(lib_name: str):
- """
- Attempts to load shared object file installed via system/cuda-toolkit.
-
- `lib_name`: Name of library to load without extension or `lib` prefix.
- """
-
- # Where to look for the shared lib in decreasing order of preference.
- paths = (
- os.environ.get(f"{lib_name.upper()}_HOME"),
- os.environ.get(f"{lib_name.upper()}_PATH"),
- os.environ.get("CUDA_HOME"),
- os.environ.get("CUDA_PATH"),
- "/usr/local/cuda",
- )
-
- for path in paths:
- if path is None:
- continue
- libs = glob.glob(f"{path}/**/lib{lib_name}{_get_sys_extension()}*", recursive=True)
- libs = [lib for lib in libs if "stub" not in lib]
- libs.sort(reverse=True, key=os.path.basename)
- if libs:
- return True, ctypes.CDLL(libs[0], mode=ctypes.RTLD_GLOBAL)
-
- # Search in LD_LIBRARY_PATH.
- try:
- _lib_handle = ctypes.CDLL(f"lib{lib_name}{_get_sys_extension()}", mode=ctypes.RTLD_GLOBAL)
- return True, _lib_handle
- except OSError:
- return False, None
-
-
-@functools.lru_cache(maxsize=None)
-def _load_cuda_library(lib_name: str):
- """
- Load given shared library.
- Prioritize loading from system/toolkit
- before checking python packages.
- """
-
- # Attempt to locate library in system.
- found, handle = _load_cuda_library_from_system(lib_name)
- if found:
- return True, handle
-
- # Attempt to locate library in Python dist-packages.
- found, handle = _load_cuda_library_from_python(lib_name)
- if found:
- return False, handle
-
- raise RuntimeError(f"{lib_name} shared object not found.")
+ return "@cudart_include_dir@"
@functools.lru_cache(maxsize=None)
@@ -366,18 +254,9 @@ if "NVTE_PROJECT_BUILDING" not in os.environ or bool(int(os.getenv("NVTE_RELEASE
@@ -366,18 +309,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.
pkgs/development/python-modules/transformer-engine/default.nix
+48 -49
View File
@@ -4,7 +4,6 @@
buildPythonPackage,
fetchFromGitHub,
replaceVars,
fetchpatch,
python,
cudaPackages,
@@ -34,6 +33,9 @@
onnx,
onnxscript,
# passthru
transformer-engine,
cudaSupport ? config.cudaSupport,
cudaCapabilities ?
if withPytorch then torch.cudaCapabilities else cudaPackages.flags.cudaCapabilities,
@@ -41,6 +43,7 @@
withPytorch ? true,
withJax ? true,
withNvshmem ? false,
withCusolvermp ? false,
}:
let
@@ -77,7 +80,7 @@ let
in
buildPythonPackage.override { stdenv = backendStdenv; } (finalAttrs: {
pname = "transformer-engine";
version = "2.14";
version = "2.15";
pyproject = true;
__structuredAttrs = true;
@@ -87,42 +90,18 @@ buildPythonPackage.override { stdenv = backendStdenv; } (finalAttrs: {
tag = "v${finalAttrs.version}";
# Their CMakeLists.txt does not easily let us inject dependencies
fetchSubmodules = true;
hash = "sha256-yxcUn75blB5ssEqGXZFDUrBv2/WM8yzumJCt5olV5Po=";
hash = "sha256-0a6etDttNoTL1pe2OZb1CcS0/AtozeAG8NFz2Hkppn8=";
};
patches =
optionals cudaSupport [
(replaceVars ./cuda-libs-paths.patch {
libcudnn_so = "${getLib cudaPackages.cudnn}/lib/libcudnn.so";
libnvrtc_so = "${getLib cudaPackages.cuda_nvrtc}/lib/libnvrtc.so";
libcurand_so = "${getLib cudaPackages.libcurand}/lib/libcurand.so";
patches = optionals cudaSupport [
(replaceVars ./cuda-libs-paths.patch {
libcudnn_so = "${getLib cudaPackages.cudnn}/lib/libcudnn.so";
libnvrtc_so = "${getLib cudaPackages.cuda_nvrtc}/lib/libnvrtc.so";
libcurand_so = "${getLib cudaPackages.libcurand}/lib/libcurand.so";
cudart_include_dir = "${getInclude cudaPackages.cuda_cudart}/include";
})
# https://github.com/NVIDIA/TransformerEngine/pull/2832
(fetchpatch {
name = "fix-cuda-arch-cmake-logic";
url = "https://github.com/NVIDIA/TransformerEngine/commit/fca261ecd09c318d22e7eeebda79632eed8cb9e4.patch";
hash = "sha256-nph01cIfmjN7RUFZif5ORoz29CEHWwetiHMEZVnnOyY=";
})
]
++ optionals withNvshmem [
# https://github.com/NVIDIA/TransformerEngine/pull/2815
(fetchpatch {
name = "fix-nvshmem-build";
url = "https://github.com/NVIDIA/TransformerEngine/commit/e83c09742166dfef3f871cfa1407605feafb3afe.patch";
hash = "sha256-5pf0Dg1XL7oAQjR1JZcdgbeaGj9qw9G5+i9Ac0iff64=";
})
]
++ optionals (withMpi && withJax) [
# https://github.com/NVIDIA/TransformerEngine/pull/2835
(fetchpatch {
name = "fix-jax-extension-build-with-mpi";
url = "https://github.com/NVIDIA/TransformerEngine/commit/2dd31bb849e83cce51c7d169db883862063d3a95.patch";
hash = "sha256-QSRMetseYPGGZCgGkS9rIj9nJdazCD4hv2IgPc+ClSM=";
})
];
cudart_include_dir = "${getInclude cudaPackages.cuda_cudart}/include";
})
];
postPatch =
# Patch build-system requirements:
@@ -171,6 +150,9 @@ buildPythonPackage.override { stdenv = backendStdenv; } (finalAttrs: {
NVTE_ENABLE_NVSHMEM = if withNvshmem then 1 else 0;
NVSHMEM_HOME = optionalString withNvshmem cudaPackages.libnvshmem.outPath;
NVTE_WITH_CUSOLVERMP = if withCusolvermp then 1 else 0;
CUSOLVERMP_HOME = optionalString withCusolvermp (getLib cudaPackages.libcusolvermp).outPath;
};
build-system = [
@@ -216,6 +198,9 @@ buildPythonPackage.override { stdenv = backendStdenv; } (finalAttrs: {
]
++ optionals withMpi [
mpi # mpi.h
]
++ optionals withCusolvermp [
cudaPackages.libcusolvermp
];
runtimeDependencies = optionals withNvshmem [
@@ -245,24 +230,38 @@ buildPythonPackage.override { stdenv = backendStdenv; } (finalAttrs: {
torch
];
# When built with nvshmem support `dlopen`ing libtransformer_engine.so `dlopen`s
# libnvidia-ml.so.1 which is provided by the GPU driver at run time:
# OSError: libnvidia-ml.so.1: cannot open shared object file: No such file or directory
pythonImportsCheck = optionals (!withNvshmem) (
[
"transformer_engine"
]
++ optionals withJax [
"transformer_engine_jax"
]
++ optionals withPytorch [
"transformer_engine_torch"
]
);
dontUsePythonImportsCheck =
# When built with cusolvermp support `dlopen`ing libtransformer_engine.so `dlopen`s
# libcuda.so.1 which is provided by the GPU driver at run time:
# OSError: libcuda.so.1: cannot open shared object file: No such file or directory
withCusolvermp
# When built with nvshmem support `dlopen`ing libtransformer_engine.so `dlopen`s
# libnvidia-ml.so.1 which is provided by the GPU driver at run time:
# OSError: libnvidia-ml.so.1: cannot open shared object file: No such file or directory
|| withNvshmem;
pythonImportsCheck = [
"transformer_engine"
]
++ optionals withJax [
"transformer_engine_jax"
]
++ optionals withPytorch [
"transformer_engine_torch"
];
# Almost all tests require GPU access
doCheck = false;
passthru.tests = {
withMpi = transformer-engine.override { withMpi = true; };
withPytorch = transformer-engine.override { withPytorch = true; };
withJax = transformer-engine.override { withJax = true; };
withNvshmem = transformer-engine.override { withNvshmem = true; };
withCusolvermp = transformer-engine.override { withCusolvermp = true; };
};
meta = {
description = "Library for accelerating Transformer models on NVIDIA GPUs";
homepage = "https://github.com/NVIDIA/TransformerEngine";