nathan/nixpkgs
1
0
Fork 0
mirror of https://github.com/NixOS/nixpkgs.git synced 2026-09-17 12:50:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

python3Packages.cupy: 14.0.1 -> 14.1.1 (#551032)

Browse Source
This commit is contained in:
Gaétan Lepage
2026-08-11 16:33:17 +00:00
committed by GitHub
parent b050eb27d9 8b1d9199ec
commit e061588212
4 changed files with 274 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
pkgs/development/python-modules/array-api-compat/default.nix
View File

@@ -52,6 +52,11 @@ buildPythonPackage (finalAttrs: {
pythonImportsCheck = [ "array_api_compat" ];
# Otherwise, cupy will try to write in $HOME (/homeless-shelter)
preCheck = ''
export CUPY_CACHE_DIR=$(mktemp -d)
'';
disabledTests = [
# CUDA (used via cupy) is not available in the testing sandbox
"cupy"

178
pkgs/development/python-modules/cupy/default.nix
View File

@@ -1,85 +1,118 @@
{
lib,
buildPythonPackage,
cudaPackages,
fetchFromGitHub,
fetchpatch,
# build-system
cython,
setuptools,
# nativeBuildInputs
cudaPackages,
symlinkJoin,
addDriverRunpath,
autoAddDriverRunpath,
writableTmpDirAsHomeHook,
# dependencies
numpy,
cuda-pathfinder,
# tests
pytest-mock,
pytestCheckHook,
}:
let
shouldUsePkg = lib.mapNullable (pkg: if pkg.meta.available or true then pkg else null);
# some packages are not available on all platforms
cuda_nvprof = shouldUsePkg (cudaPackages.nvprof or null);
libcutensor = shouldUsePkg (cudaPackages.libcutensor or null);
nccl = shouldUsePkg (cudaPackages.nccl or null);
outpaths = lib.filter (outpath: outpath != null) (
with cudaPackages;
[
cccl # <nv/target>
cuda_cudart
cuda_nvcc # <crt/host_defines.h>
cuda_nvprof
cuda_nvrtc
cuda_nvtx
cuda_profiler_api
libcublas
libcufft
libcurand
libcusolver
libcusparse
libcusparse_lt # cusparseLt.h
]
);
cudatoolkit-joined = symlinkJoin {
name = "cudatoolkit-joined-${cudaPackages.cudaMajorMinorVersion}";
paths =
outpaths ++ lib.concatMap (outpath: lib.map (output: outpath.${output}) outpath.outputs) outpaths;
};
in
buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs: {
pname = "cupy";
version = "14.0.1";
version = "14.1.1";
pyproject = true;
__structuredAttrs = true;
src = fetchFromGitHub {
owner = "cupy";
repo = "cupy";
tag = "v${finalAttrs.version}";
fetchSubmodules = true;
hash = "sha256-TaEJ0BveUCXCRrNq9L49Tfbu0334+cANcVm5qnSOE1Q=";
hash = "sha256-8jreEbQA24V7EAD87z3uFwlr3LPDoGRCvF5vbf2dKvI=";
};
postPatch = ''
substituteInPlace pyproject.toml \
--replace-fail \
"Cython>=3.1,<3.2" \
"Cython"
'';
patches = [
# Let cupy_setup_build.py find static libraries (.a) from buildInputs.
# By default, it looks for them in `$CUDA_PATH/lib{64,}`.
./link-static-libraries.patch
# Fix test collection with pytest>=9
# https://github.com/cupy/cupy/pull/10020
(fetchpatch {
name = "pytest-9-compat.patch";
url = "https://github.com/cupy/cupy/commit/6f52f1541e0e047cdb84286793d6110567bce5c8.patch";
hash = "sha256-D4VjOL0XO8gnfItGC/2MjuWYCJ7eNJYFn8jIyM+EojM=";
})
];
postPatch =
# Inject absolute path to nvcc instead of relying on runtime discovery heuristics
''
substituteInPlace cupy/_environment.py \
--replace-fail \
"nvcc_path = os.environ.get('NVCC', None)" \
"return '${lib.getExe cudaPackages.cuda_nvcc}'"
''
# Inject absolute path to CUDA libraries which get disovered at runtime
+ ''
substituteInPlace cupy/_core/core.pyx \
--replace-fail \
"_cuda_include_dir = find_nvidia_header_directory('cudart')" \
"_cuda_include_dir = '${lib.getInclude cudaPackages.cuda_cudart}/include'"
substituteInPlace cupy/cuda/__init__.py \
--replace-fail \
"from cuda import pathfinder" \
"from ctypes import CDLL" \
--replace-fail \
'pathfinder.load_nvidia_dynamic_lib("cufft")' \
'CDLL("${lib.getLib cudaPackages.libcufft}/lib/libcufft.so")'
substituteInPlace cupy/cuda/cufft.pyx \
--replace-fail \
"from cuda import pathfinder" \
"from ctypes import CDLL" \
--replace-fail \
"loaded_dl = pathfinder.load_nvidia_dynamic_lib('cufft')" \
"loaded_dl = CDLL('${lib.getLib cudaPackages.libcufft}/lib/libcufft.so')" \
--replace-fail \
"handle = loaded_dl._handle_uint" \
"handle = loaded_dl._handle"
substituteInPlace cupy_backends/cuda/api/_runtime_softlink.pxi \
--replace-fail \
"from cuda import pathfinder" \
"from ctypes import CDLL" \
--replace-fail \
"loaded_dl = pathfinder.load_nvidia_dynamic_lib('cudart')" \
"loaded_dl = CDLL('${lib.getLib cudaPackages.cuda_cudart}/lib/libcudart.so')" \
--replace-fail \
"handle = loaded_dl._handle_uint" \
"handle = loaded_dl._handle"
substituteInPlace cupy_backends/cuda/libs/_cnvrtc.pxi \
--replace-fail \
"from cuda import pathfinder" \
"from ctypes import CDLL" \
--replace-fail \
"loaded_dl = pathfinder.load_nvidia_dynamic_lib('nvrtc')" \
"loaded_dl = CDLL('${lib.getLib cudaPackages.cuda_nvrtc}/lib/libnvrtc.so')" \
--replace-fail \
"handle = loaded_dl._handle_uint" \
"handle = loaded_dl._handle"
substituteInPlace cupy/cuda/compiler.py \
--replace-fail \
"cudadevrt = get_cuda_path()" \
"return '${lib.getOutput "static" cudaPackages.cuda_cudart}/lib/libcudadevrt.a'"
'';
env = {
LDFLAGS = toString [
# Fake libcuda.so (the real one is deployed impurely)
"-L${lib.getOutput "stubs" cudaPackages.cuda_cudart}/lib/stubs"
];
# NVCC = "${lib.getExe cudaPackages.cuda_nvcc}"; # FIXME: splicing/buildPackages
CUDA_PATH = "${cudatoolkit-joined}";
};
# See https://docs.cupy.dev/en/v10.2.0/reference/environment.html. Setting both
@@ -92,20 +125,36 @@ buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs
export CUPY_NUM_NVCC_THREADS="$NIX_BUILD_CORES"
'';
enableParallelBuilding = true;
build-system = [
cython
setuptools
];
nativeBuildInputs = [
addDriverRunpath
cudatoolkit-joined
autoAddDriverRunpath
cudaPackages.cuda_nvcc
writableTmpDirAsHomeHook # Needed for pythonImportsCheck
];
buildInputs = [
cudatoolkit-joined
libcutensor
nccl
buildInputs = with cudaPackages; [
cuda_cudart # cuda.h, cuda_runtime.h, libcudart_static.a
cuda_cuxxfilt # nv_decode.h
cuda_nvrtc # nvrtc.h
cuda_nvtx # nvToolsExt.h
cuda_profiler_api # cuda_profiler_api.h
libcublas # cublas_v2.h
libcufft # cufft.h
libcurand # curand.h
libcusolver # cusolverDn.h
libcusparse # cusparse.h
libcusparse_lt # cusparseLt.h
libcutensor # cutensor.h
nccl # nccl.h
(lib.getOutput "static" cuda_cuxxfilt) # libcufilt.a
];
dependencies = [
@@ -113,23 +162,15 @@ buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs
numpy
];
nativeCheckInputs = [
pytest-mock
pytestCheckHook
pythonImportsCheck = [
"cupy"
"cupy_backends"
"cupyx"
];
# Won't work with the GPU, whose drivers won't be accessible from the build
# sandbox
# Won't work with the GPU, whose drivers won't be accessible from the build sandbox
doCheck = false;
postFixup = ''
find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
addDriverRunpath "$lib"
done
'';
enableParallelBuilding = true;
meta = {
description = "NumPy-compatible matrix library accelerated by CUDA";
homepage = "https://cupy.chainer.org/";
@@ -140,5 +181,6 @@ buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs
"x86_64-linux"
];
maintainers = with lib.maintainers; [ GaetanLepage ];
teams = [ lib.teams.cuda ];
};
})

30
pkgs/development/python-modules/cupy/link-static-libraries.patch Normal file
View File

@@ -0,0 +1,30 @@
Instead of looking up static libraries by hand in `$CUDA_PATH/lib{64,}` (which does not work with
the split CUDA redistributables), link them like any other library, so that the compiler wrapper
finds them in the (nix) `buildInputs`.
--- a/install/cupy_builder/cupy_setup_build.py
+++ b/install/cupy_builder/cupy_setup_build.py
@@ -423,10 +423,7 @@
s = copy.deepcopy(settings)
if not no_cuda:
- s['libraries'] = module.libraries
- s['extra_objects'] = [
- _find_static_library(name) for name in module.static_libraries
- ]
+ s['libraries'] = module.static_libraries + module.libraries
compile_args = s.setdefault('extra_compile_args', [])
link_args = s.setdefault('extra_link_args', [])
@@ -490,9 +487,8 @@
if (not no_cuda and not use_hip
and PLATFORM_LINUX
and f == 'cupy.cuda.function'):
- s_file['extra_objects'] = (
- s_file.get('extra_objects', [])
- + [_find_static_library('cufilt')])
+ s_file['libraries'] = (
+ ['cufilt'] + s_file.get('libraries', []))
sources = module_extension_sources(f, use_cython, no_cuda)
extension = setuptools.Extension(name, sources, **s_file)

171
pkgs/development/python-modules/deepspeed/default.nix
View File

@@ -2,45 +2,56 @@
lib,
buildPythonPackage,
fetchFromGitHub,
symlinkJoin,
cudaPackages,
# build-system
setuptools,
torch,
# dependencies
einops,
hjson,
msgpack,
ninja,
numpy,
cupy,
cutlass,
packaging,
psutil,
py-cpuinfo,
pydantic,
torch,
tqdm,
nix-update-script,
cudaPackages,
symlinkJoin,
# cuda-only:
nvidia-ml-py,
# cupy,
# cutlass,
# tests
accelerate,
addBinToPathHook,
pytest-xdist,
pytestCheckHook,
tabulate,
transformers,
writableTmpDirAsHomeHook,
}:
let
cudaVersion = cudaPackages.cudaMajorMinorVersion;
inherit (torch) cudaCapabilities cudaSupport;
cuda-native-redist = symlinkJoin {
name = "cuda-native-redist-${cudaVersion}";
name = "cuda-native-redist";
paths = with cudaPackages; [
(lib.getDev cuda_cudart)
(lib.getLib cuda_cudart)
(lib.getStatic cuda_cudart)
cuda_nvcc
(lib.getBin cuda_nvcc)
];
};
in
buildPythonPackage (finalAttrs: {
buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {
pname = "deepspeed";
version = "0.19.2";
version = "0.19.5";
pyproject = true;
__structuredAttrs = true;
@@ -48,14 +59,49 @@ buildPythonPackage (finalAttrs: {
owner = "deepspeedai";
repo = "DeepSpeed";
tag = "v${finalAttrs.version}";
hash = "sha256-Nw1rw65hdqhARFR7W+XmmRT/pLkCi5nTF+6R9L3bLyo=";
fetchSubmodules = true;
hash = "sha256-OIy8AehiwKhbdWrWslwc3yimAH4QYobhgg/inkHDx/U=";
};
postPatch = ''
substituteInPlace deepspeed/ops/op_builder/builder.py \
--replace-fail \
'import distutils' \
'import setuptools._distutils'
''
# `setup.py` appends `+{git_hash}` (here, `+unknown`) to the version unless `build.txt` exists,
# in which case its content is appended instead.
+ ''
touch build.txt
''
+ lib.optionalString cudaSupport (
# Hardcode CUDA_HOME to nix store path for JIT op compilation at runtime
''
substituteInPlace deepspeed/ops/op_builder/builder.py \
--replace-fail \
"cuda_home = torch.utils.cpp_extension.CUDA_HOME" \
"cuda_home = '${cuda-native-redist}'"
''
# Hardcode CUTLASS_PATH to nix store path
+ ''
substituteInPlace deepspeed/ops/op_builder/evoformer_attn.py \
--replace-fail \
'self.cutlass_path = os.environ.get("CUTLASS_PATH")' \
'self.cutlass_path = "${lib.getInclude cudaPackages.cutlass}/include/cutlass"'
''
);
build-system = [
setuptools
torch
];
# `deepspeed.ops.transformer.inference.triton` creates its autotune cache directory
# (`$TRITON_HOME`, defaulting to `$HOME/.triton`) at import time.
preBuild = ''
export TRITON_HOME=$(mktemp -d)
'';
dependencies = [
einops
hjson
@@ -66,53 +112,94 @@ buildPythonPackage (finalAttrs: {
psutil
py-cpuinfo
pydantic
setuptools
torch
tqdm
]
++ lib.optionals cudaSupport [
cutlass
cupy
nvidia-ml-py
];
postPatch = ''
substituteInPlace deepspeed/ops/op_builder/builder.py \
--replace-fail 'import distutils' 'import setuptools._distutils'
''
+ lib.optionalString cudaSupport ''
# Hardcode CUDA_HOME to nix store path for JIT op compilation at runtime
substituteInPlace deepspeed/ops/op_builder/builder.py \
--replace-fail \
"cuda_home = torch.utils.cpp_extension.CUDA_HOME" \
"cuda_home = '${cuda-native-redist}'"
# Hardcode CUTLASS_PATH to nix store path
substituteInPlace deepspeed/ops/op_builder/evoformer_attn.py \
--replace-fail \
"self.cutlass_path = os.environ.get(\"CUTLASS_PATH\")" \
"self.cutlass_path = '${cutlass}'"
'';
env = lib.optionalAttrs cudaSupport {
TORCH_CUDA_ARCH_LIST = lib.concatStringsSep ";" cudaCapabilities;
};
preConfigure = ''
# setuptools writes to ~/.cache during builds
export HOME=$TMPDIR
preCheck = ''
rm -rf deepspeed
''
# Tests JIT-compile CPU ops, and the op builder pairs the ISA it detects (`-D__AVX512__`) with
# `-march=native`, which the cc wrapper strips. Let it through, otherwise the intrinsics fail to
# compile: "inlining failed in call to 'always_inline'".
+ ''
export NIX_ENFORCE_NO_NATIVE=0
'';
pythonImportsCheck = [
"deepspeed"
pythonImportsCheck = [ "deepspeed" ];
nativeCheckInputs = [
accelerate
addBinToPathHook # Tests run the deepspeed CLI
ninja
pytest-xdist
pytestCheckHook
tabulate
transformers
writableTmpDirAsHomeHook
];
passthru.updateScript = nix-update-script { };
pytestFlags = [
# Tests memory consumption grows significantly with the number of parallel processes
# -> Limit the number of parallel jobs to prevent OOMing
"--maxprocesses=16"
];
enabledTestPaths = [
# The test suite seems to be very resource intensive. We stick to unit tests for now.
"tests/unit/"
];
disabledTestPaths = [
# Require unpackaged `mup`
"tests/unit/runtime/test_mup_optimizers.py"
];
disabledTests = [
# AssertionError: assert 2 == 1
# + where 2 = len(['[pinned-memory checkpoint] crossed 32.000 GB: 33.000 GB pinned across 1 ...
"test_checkpoint_emits_info"
# Fixture handling is incompatible with pytest >= 8.4, upstream pins pytest < 8.4:
# https://github.com/deepspeedai/DeepSpeed/pull/7327
"TestDistributedFixture"
"TestZeROElasticCheckpoint"
"TestZeROPPLoadCheckpoint"
"TestZeROUniversalCheckpointDP"
# Fail since we patch the cutlass/cuda path logic
"test_include_paths_accepts_cutlass_include_dir_directly"
"test_include_paths_finds_cutlass_from_cmake_prefix_path"
"test_include_paths_finds_cutlass_from_compiler_include_path"
"test_include_paths_finds_python_package_candidate_without_env"
"test_include_paths_reports_missing_cutlass"
"test_include_paths_uses_cutlass_path_env"
"test_installed_cuda_version_reports_a_runtime_only_cuda_home_as_missing"
# Too long
"TestWarmupCosineLR"
];
# Some python tests are skipped unless a GPU is visible
passthru.gpuCheck = finalAttrs.finalPackage.overrideAttrs {
requiredSystemFeatures = [ "cuda" ];
};
meta = {
description = "Deep learning optimization library that makes distributed training and inference easy, efficient, and effective.";
description = "Deep learning optimization library that makes distributed training and inference easy, efficient, and effective";
homepage = "https://www.deepspeed.ai/";
downloadPage = "https://github.com/deepspeedai/deepspeed";
changelog = "https://github.com/deepspeedai/DeepSpeed/releases/tag/${finalAttrs.src.tag}";
license = lib.licenses.asl20;
mainProgram = "deepspeed";
maintainers = with lib.maintainers; [ jlesquembre ];
teams = [ lib.teams.cuda ];
};
})