diff --git a/pkgs/development/python-modules/bitsandbytes/default.nix b/pkgs/development/python-modules/bitsandbytes/default.nix new file mode 100644 index 000000000000..d63b79a1e20b --- /dev/null +++ b/pkgs/development/python-modules/bitsandbytes/default.nix @@ -0,0 +1,98 @@ +{ lib +, buildPythonPackage +, fetchFromGitHub +, python +, pythonOlder +, pytestCheckHook +, setuptools +, torch +, einops +, lion-pytorch +, scipy +, symlinkJoin +}: + +let + pname = "bitsandbytes"; + version = "0.38.0"; + + inherit (torch) cudaCapabilities cudaPackages cudaSupport; + inherit (cudaPackages) backendStdenv 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 + libcurand + libcusolver # cusolverDn.h + libcusparse # cusparse.h + ]; + + cuda-native-redist = symlinkJoin { + name = "cuda-native-redist-${cudaVersion}"; + paths = with cudaPackages; [ + cuda_cudart # cuda_runtime.h cuda_runtime_api.h + cuda_nvcc + ] ++ cuda-common-redist; + }; + + cuda-redist = symlinkJoin { + name = "cuda-redist-${cudaVersion}"; + paths = cuda-common-redist; + }; + +in +buildPythonPackage { + inherit pname version; + format = "pyproject"; + + disabled = pythonOlder "3.7"; + + src = fetchFromGitHub { + owner = "TimDettmers"; + repo = pname; + rev = "refs/tags/${version}"; + hash = "sha256-gGlbzTDvZNo4MhcYzLvWuB2ec7q+Qt5/LtTbJ0Rc+Kk="; + }; + + postPatch = '' + substituteInPlace Makefile --replace "/usr/bin/g++" "g++" --replace "lib64" "lib" + substituteInPlace bitsandbytes/cuda_setup/main.py \ + --replace "binary_path = package_dir / binary_name" \ + "binary_path = Path('$out/${python.sitePackages}/${pname}')/binary_name" + '' + lib.optionalString torch.cudaSupport '' + substituteInPlace bitsandbytes/cuda_setup/main.py \ + --replace "/usr/local/cuda/lib64" "${cuda-native-redist}/lib" + ''; + + 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 = [ setuptools ] ++ lib.optionals torch.cudaSupport [ cuda-native-redist ]; + buildInputs = lib.optionals torch.cudaSupport [ cuda-redist ]; + + propagatedBuildInputs = [ + torch + ]; + + doCheck = false; # tests require CUDA and also GPU access + nativeCheckInputs = [ pytestCheckHook einops lion-pytorch scipy ]; + + pythonImportsCheck = [ + "bitsandbytes" + ]; + + meta = with lib; { + homepage = "https://github.com/TimDettmers/bitsandbytes"; + description = "8-bit CUDA functions for PyTorch"; + license = licenses.mit; + maintainers = with maintainers; [ bcdarwin ]; + }; +} diff --git a/pkgs/development/python-modules/lion-pytorch/default.nix b/pkgs/development/python-modules/lion-pytorch/default.nix new file mode 100644 index 000000000000..a5a554c1980b --- /dev/null +++ b/pkgs/development/python-modules/lion-pytorch/default.nix @@ -0,0 +1,32 @@ +{ lib +, buildPythonPackage +, pythonOlder +, fetchFromGitHub +, torch +}: + +buildPythonPackage rec { + pname = "lion-pytorch"; + version = "0.0.7"; + format = "setuptools"; + disabled = pythonOlder "3.6"; + + src = fetchFromGitHub { + owner = "lucidrains"; + repo = "lion-pytorch"; + rev = "refs/tags/${version}"; + hash = "sha256-CSb0s3DKv/KpEmCkCR+Y8iwrLdCL9w9Pl6W46cPB420"; + }; + + propagatedBuildInputs = [ torch ]; + + pythonImportsCheck = [ "lion_pytorch" ]; + doCheck = false; # no tests currently + + meta = with lib; { + description = "Optimizer tuned by Google Brain using genetic algorithms"; + homepage = "https://github.com/lucidrains/lion-pytorch"; + license = licenses.mit; + maintainers = with maintainers; [ bcdarwin ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 7d91bd29cb39..7083e9b9f6e3 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -1314,6 +1314,8 @@ self: super: with self; { bitmath = callPackage ../development/python-modules/bitmath { }; + bitsandbytes = callPackage ../development/python-modules/bitsandbytes { }; + bitstring = callPackage ../development/python-modules/bitstring { }; bitstruct = callPackage ../development/python-modules/bitstruct { }; @@ -5643,6 +5645,8 @@ self: super: with self; { linuxfd = callPackage ../development/python-modules/linuxfd { }; + lion-pytorch = callPackage ../development/python-modules/lion-pytorch { }; + liquidctl = callPackage ../development/python-modules/liquidctl { }; lirc = toPythonModule (pkgs.lirc.override {