summaryrefslogtreecommitdiff
path: root/Omni/Bild/Deps/bitsandbytes.nix
diff options
context:
space:
mode:
authorBen Sima <ben@bsima.me>2024-11-15 14:55:37 -0500
committerBen Sima <ben@bsima.me>2024-12-21 10:06:49 -0500
commit6513755670892983db88a6633b8c1ea6019c03d1 (patch)
tree44e9eccdb7a3a74ab7e96a8fee7572dd6a78dc73 /Omni/Bild/Deps/bitsandbytes.nix
parentae7b7e0186b5f2e0dcd4d5fac0a71fa264caedc2 (diff)
Re-namespace some stuff to Omni
I was getting confused about what is a product and what is internal infrastructure; I think it is good to keep those things separate. So I moved a bunch of stuff to an Omni namespace, actually most stuff went there. Only things that are explicitly external products are still in the Biz namespace.
Diffstat (limited to 'Omni/Bild/Deps/bitsandbytes.nix')
-rw-r--r--Omni/Bild/Deps/bitsandbytes.nix86
1 files changed, 86 insertions, 0 deletions
diff --git a/Omni/Bild/Deps/bitsandbytes.nix b/Omni/Bild/Deps/bitsandbytes.nix
new file mode 100644
index 0000000..eb32aac
--- /dev/null
+++ b/Omni/Bild/Deps/bitsandbytes.nix
@@ -0,0 +1,86 @@
+{ lib, buildPythonPackage, fetchFromGitHub, python, pythonOlder, pytestCheckHook
+, setuptools, torch, einops, lion-pytorch, scipy, symlinkJoin }:
+
+let
+ pname = "bitsandbytes";
+ version = "0.38.0";
+
+ inherit (torch) cudaPackages cudaSupport;
+ inherit (cudaPackages) cudaVersion;
+
+ # NOTE: torchvision doesn't use cudnn; torch does!
+ # For this reason it is not included.
+ cuda-common-redist = with cudaPackages; [
+ cuda_cccl # <thrust/*>
+ 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 ];
+ };
+}