teddykoker / torchsort

Fast, differentiable sorting and ranking in PyTorch
https://pypi.org/project/torchsort/
Apache License 2.0
765 stars 33 forks source link

Update Pre-built Wheels #78

Closed batuhanfaik closed 8 months ago

batuhanfaik commented 8 months ago

Problem

Building torchsort using conda as the package manager on AWS EC2 instances with CUDA toolkit 12.1+ is quite cumbersome.

Possible Solution

It would be fantastic if we could update the pre-built wheels to support newer versions of CUDA.

Wheels for the following versions should be added:

python>=3.10
torch>=2.1
cuda>=12.0

I specifically require the following wheel file: torchsort-0.1.9+pt21cu121-cp311-cp311-linux_x86_64.whl

Why?

As noted in the README, installing gxx_linux-64 should not be required when using conda. Steps provided in #52 are also obsolete for CUDA v12+, since cudatoolkit-dev goes up to CUDA v11.7 (other methods mentioned in the same issue are also obsolete). One can obtain the latest nvcc from cuda-nvcc if necessary. My attempts using gxx_linux-64==9.5.0 and cuda-nvcc==12.1.1 were successful. Below is my environment yaml file if anybody would like to reproduce.

name: torchsort_p311
channels:
    - conda-forge
    - pytorch
    - nvidia
    - nvidia/label/cuda-12.1.1
dependencies:
    - python>=3.11
    - pytorch>=2
    - pytorch-cuda==12.1
    # I don't like this a single bit
    - cuda-nvcc
    - gxx_linux-64>=9.5,<10
    - pip
    - pip:
        - torchsort
teddykoker commented 8 months ago

Adding these pre-built wheels should be as simple as updating the verison matrix in the release workflow. The cuda install script will also need to be updated to include cuda 12 .deb file location. I should be able to get around to doing this shortly.

teddykoker commented 8 months ago

@batuhanfaik I added binaries for torch 2.1 and cuda 12.1 to the releases page. You should be able to install with:

pip install https://github.com/teddykoker/torchsort/releases/download/v0.1.9/torchsort-0.1.9+pt21cu121-cp311-cp311-linux_x86_64.whl

Let me know if you are having any troubles.