Installing correct CUDA versions for PyTorch

[cedriclim1]

Problem

Some systems will have different CUDA versions and for the newest release PyTorch=2.11.0 they are straight defaulting to CUDA 13.0: https://github.com/pytorch/pytorch/releases/tag/v2.11.0.

Solution

To allow some flexibility on what CUDA version someone can install on quantem, Claude gave some suggestions using uv sync --group. For example:

[[tool.uv.index]]
name = "pytorch-cu130"
url = "https://download.pytorch.org/whl/cu130"
explicit = true

[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true

[dependency-groups]
cu130 = ["torch>=2.7.0", "torchvision>=0.22.1"]
cu128 = ["torch>=2.7.0", "torchvision>=0.22.1"]
cpu = ["torch>=2.7.0", "torchvision>=0.22.1"]

[tool.uv.sources]
torch = [
  { index = "pytorch-cu130", group = "cu130" },
  { index = "pytorch-cu128", group = "cu128" },
]
torchvision = [
  { index = "pytorch-cu130", group = "cu130" },
  { index = "pytorch-cu128", group = "cu128" },
]

and one would run:

uv sync --group cu130
uv sync --group cu128
uv sync --group cpu  # falls back to PyPI CPU wheels

Maybe we should think about doing this just in case if someone would want more flexibility, but we will default to the latest CUDA version?

[bobleesj]

Just a side comment - for both dev and deployment on microscopy PCs, I no longer uv as this package starts to expand beyond non-python realms that requires complation of cuda and javascript while pulling the latest commits after git pull upstrea dev.

For distribution and installation by HPCs and microscope PCs, conda / mamba has been the go-to solution so far (as many MD, DFT codes do).

Another comment is that, this is the main reason why we should not generally pin versions. Yes, it's good for reproducibility, but the greater goal is to distribute and propagate across machines and OSs.

[bobleesj]

I haven't tested this - sth like:

Prepare env.yml

  name: qt-env                                                                                                            
  channels:                                              
    - pytorch
    - nvidia
    - conda-forge
  dependencies:
    - python=3.13
    - pytorch                                                                                                              
    - torchvision
    - cupy                                                                                                                 
    - numpy                                              
    - scipy
    - h5py
    - matplotlib
    - rsciio                                                                                                               
    - anywidget
    - tqdm                                                                                                                 
    - pip                                                
    - pip:                                                                                                                 
      - -e .     

1. For cu13 HPC:

mamba env create -f env.yml                                                                                      
mamba activate qt-env                                 

2. For cu12 HPC (older):

mamba env create -f env.yml
mamba install -n qt-env -c pytorch -c nvidia pytorch-cuda=12.4 
mamba activate qt-env