Set torch-backend for gpu drivers

[Jacob-Chmura]

Purpose

The current pyproject.toml only specifies torch >= 2.5.1 which defaults to cpu wheels. Nobody will care to run our library in CPU only mode, so the default should be to provide gpu support.

Solution

The UV documentation suggests using the following (experimental) feature https://docs.astral.sh/uv/reference/settings/#pip_torch-backend, which:

When set, uv will ignore the configured index URLs for packages in the PyTorch ecosystem, and will instead use the defined backend. For example, when set to cpu, uv will use the CPU-only PyTorch index; when set to cu126, uv will use the PyTorch index for CUDA 12.6.
The auto mode will attempt to detect the appropriate PyTorch index based on the currently installed CUDA drivers.

I think this is reasonable

• If user has no gpu drivers configured, it should use cpu
• Otherwise, it should use the cuda wheel that matches your drivers

Alternatives

• Explicitly force torch == 2.5.1+cu121 so that all users have the same cuda-compatible version (needs care if their drivers are out of date, but completely reproducible with our code)
• Keep install cpu only, but add documentation that requires user to issue uv pip install torch==2.5.1+cu121 --extra-index-url https://download.pytorch.org/whl/cu121. Most generic but requires an extra command on the behalf of the user.
• Get people to just use our docker image and avoid these setup

Note: Pypi only has cpu torch wheels, so when we publish there, we'll have to deal with this problem as well.

Relevant Prs

Fix #80

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

📢 Thoughts on this report? Let us know!

[ntgbaoo]

I am using a machine with Windows OS for development. uv pip install torch==2.5.1+cu121 --extra-index-url https://download.pytorch.org/whl/cu121 doesn't resolve the issue of CUDA not available, although I installed CUDA driver already. I tried:
uv pip install torch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 --index-url https://download.pytorch.org/whl/cu121
It did resolve the issue.

I would personally prefer Docker image setup. Since dealing with these dependencies, especially for Windows users, is sometimes painful.

[Jacob-Chmura]

I tried: uv pip install torch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 --index-url https://download.pytorch.org/whl/cu121 It did resolve the issue.

Can you see if uv pip install torch --torch-backend=auto works on windows? As is described here, torch wheels are hosted differently on different platforms and with different drivers. So I want to find the simplest setup that works for most people.

I would personally prefer Docker image setup. Since dealing with these dependencies, especially for Windows users, is sometimes painful.

I agree. Unfortunetely, I'm not sure that most people in academia will feel the same way. They would probably rather just uv pip install into a python env that they don't really care about then setup docker. But I could be wrong, and we can discuss this.

[shenyangHuang]

I think the proposed solution is reasonable

• If user has no gpu drivers configured, it should use cpu
• Otherwise, it should use the cuda wheel that matches your drivers

This is how PyG resolves this, I think on our website if we have separate section showing how you can install a CPU vs. GPU version it might be good.
the TLDR is that for users, it is very easy if they can clearly see for CPU version what is the command to run for install and for GPU version, what is the command to run to install.

Also, there would be researchers who don't have GPU compute but would like to run simple methods like EdgeBank and get into TG research. Plus another use case is if someone just want to use TGM to process and checkout their own local TG dataset, they also don't need the GPU as no ML workload is performed.

[ntgbaoo]

Can you see if uv pip install torch --torch-backend=auto works on windows?

@Jacob-Chmura It doesn't work on Windows. Seems like you need to specify torchvision and torchaudio as follows:
uv pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124 to make it works.
This may be related: astral-sh/uv#7202

@shenyangHuang's point also makes sense

[Jacob-Chmura]

Seems like you need to specify torchvision and torchaudio as follows: uv pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124 to make it works.

Understood. I will investigate this when I have some time tomorrow. It seems to me there should be an easy way to get cuda wheels on windows without needing to pip install these auxilary packages.

Also, based on our discussion in the meeting today, we will not be officially supporting windows for the time being. Therefore, I am happy to make small changes that make it more widely compatible to other platforms, but for now, we will target linux and keep those stable. Perhaps look into windows subsystem for linux: https://github.com/microsoft/WSL