diff --git a/pkgs/development/python-modules/bitsandbytes/default.nix b/pkgs/development/python-modules/bitsandbytes/default.nix index 9bfc4a89784fe..1768fc1652a93 100644 --- a/pkgs/development/python-modules/bitsandbytes/default.nix +++ b/pkgs/development/python-modules/bitsandbytes/default.nix @@ -1,31 +1,34 @@ { lib, + torch, + symlinkJoin, buildPythonPackage, fetchFromGitHub, - python, - pythonOlder, + cmake, setuptools, - wheel, - torch, scipy, - symlinkJoin, }: let pname = "bitsandbytes"; - version = "0.43.1"; + version = "0.44.1"; - inherit (torch) cudaCapabilities cudaPackages cudaSupport; - inherit (cudaPackages) backendStdenv cudaVersion; + inherit (torch) cudaPackages cudaSupport; + inherit (cudaPackages) cudaVersion; + + cudaVersionString = lib.replaceStrings [ "." ] [ "" ] (lib.versions.majorMinor cudaVersion); # NOTE: torchvision doesn't use cudnn; torch does! # For this reason it is not included. cuda-common-redist = with cudaPackages; [ - cuda_cccl # - libcublas # cublas_v2.h + (lib.getDev cuda_cccl) # + (lib.getDev libcublas) # cublas_v2.h + (lib.getLib libcublas) libcurand libcusolver # cusolverDn.h - libcusparse # cusparse.h + (lib.getDev libcusparse) # cusparse.h + (lib.getLib libcusparse) + libcusparse.lib ]; cuda-native-redist = symlinkJoin { @@ -33,7 +36,9 @@ let paths = with cudaPackages; [ - cuda_cudart # cuda_runtime.h cuda_runtime_api.h + (lib.getDev cuda_cudart) # cuda_runtime.h cuda_runtime_api.h + (lib.getLib cuda_cudart) + (lib.getStatic cuda_cudart) cuda_nvcc ] ++ cuda-common-redist; @@ -48,47 +53,55 @@ buildPythonPackage { inherit pname version; pyproject = true; - disabled = pythonOlder "3.7"; - src = fetchFromGitHub { owner = "TimDettmers"; repo = "bitsandbytes"; rev = "refs/tags/${version}"; - hash = "sha256-GFbFKPdV96DXPA+PZO4h0zdBclN670fb0PGv4QPHWHU="; + hash = "sha256-yvxD5ymMK5p4Xg7Csx/90mPV3yxUC6QUuF/8BKO2p0k="; }; - postPatch = - '' - substituteInPlace Makefile --replace "/usr/bin/g++" "g++" --replace "lib64" "lib" - substituteInPlace bitsandbytes/cuda_setup/main.py \ - --replace "binary_path = package_dir / self.binary_name" \ - "binary_path = Path('$out/${python.sitePackages}/${pname}')/self.binary_name" - '' - + lib.optionalString torch.cudaSupport '' - substituteInPlace bitsandbytes/cuda_setup/main.py \ - --replace "/usr/local/cuda/lib64" "${cuda-native-redist}/lib" - ''; + # By default, which library is loaded depends on the result of `torch.cuda.is_available()`. + # When `cudaSupport` is enabled, bypass this check and load the cuda library unconditionnally. + # Indeed, in this case, only `libbitsandbytes_cuda124.so` is built. `libbitsandbytes_cpu.so` is not. + # Also, hardcode the path to the previously built library instead of relying on + # `get_cuda_bnb_library_path(cuda_specs)` which relies on `torch.cuda` too. + # + # WARNING: The cuda library is currently named `libbitsandbytes_cudaxxy` for cuda version `xx.y`. + # This upstream convention could change at some point and thus break the following patch. + postPatch = lib.optionalString cudaSupport '' + substituteInPlace bitsandbytes/cextension.py \ + --replace-fail "if cuda_specs:" "if True:" \ + --replace-fail \ + "cuda_binary_path = get_cuda_bnb_library_path(cuda_specs)" \ + "cuda_binary_path = PACKAGE_DIR / 'libbitsandbytes_cuda${cudaVersionString}.so'" + ''; CUDA_HOME = "${cuda-native-redist}"; - preBuild = - if torch.cudaSupport then - with torch.cudaPackages; - let - cudaVersion = lib.concatStrings (lib.splitVersion torch.cudaPackages.cudaMajorMinorVersion); - in - ''make CUDA_VERSION=${cudaVersion} cuda${cudaMajorVersion}x'' - else - ''make CUDA_VERSION=CPU cpuonly''; - nativeBuildInputs = [ + cmake + ]; + + build-system = [ setuptools - wheel - ] ++ lib.optionals torch.cudaSupport [ cuda-native-redist ]; + ]; + + buildInputs = lib.optionals cudaSupport [ cuda-redist ]; + + cmakeFlags = [ + (lib.cmakeFeature "-DCOMPUTE_BACKEND" (if cudaSupport then "cuda" else "cpu")) + ]; + NVCC_PREPEND_FLAGS = lib.optionals cudaSupport [ + "-I${cuda-native-redist}/include" + "-L${cuda-native-redist}/lib" + ]; - buildInputs = lib.optionals torch.cudaSupport [ cuda-redist ]; + preBuild = '' + make -j $NIX_BUILD_CORES + cd .. # leave /build/source/build + ''; - propagatedBuildInputs = [ + dependencies = [ scipy torch ]; @@ -97,11 +110,11 @@ buildPythonPackage { pythonImportsCheck = [ "bitsandbytes" ]; - meta = with lib; { + meta = { description = "8-bit CUDA functions for PyTorch"; homepage = "https://github.com/TimDettmers/bitsandbytes"; changelog = "https://github.com/TimDettmers/bitsandbytes/releases/tag/${version}"; - license = licenses.mit; - maintainers = with maintainers; [ bcdarwin ]; + license = lib.licenses.mit; + maintainers = with lib.maintainers; [ bcdarwin ]; }; }