Python using GPUs#

Many Python packages are published in both a CPU and a GPU flavour. Getting the GPU one is mostly a matter of making sure your installer actually resolves it – which, for the conda-based tools, does not happen by itself on Levante.

The easiest start is a combination that is known to work: the specifications at Ready-made environments are built and tested on Levante every day, and the ones that can make use of a GPU node are marked in the GPU column.

Be aware that GPU-enabled environments are large. A PyTorch environment with CUDA is roughly 2 GB, see quota and file system hints.

Using uv#

There is nothing to configure. PyPI ships the CUDA libraries as ordinary packages, so they are installed along with everything else:

uv init --python 3.13 my_project
cd my_project
uv add torch

Using pixi or micromamba#

conda-forge decides between the CPU and the GPU build of a package by looking for a __cuda virtual package, which describes the CUDA version supported by the driver of the machine that is resolving the environment. Levante’s login nodes have no GPU, so __cuda is absent there and the solver quietly settles for the CPU build:

$ micromamba create -n my_env --dry-run pytorch
...
+ pytorch  2.13.0  cpu_mkl_py314_hf0508d8_102

Nothing warns you about this. The environment installs and imports fine, and the GPU simply stays idle at runtime.

The fix is to tell the solver which CUDA version to assume. Use the version from the CUDA Version field of nvidia-smi on a GPU node.

With micromamba this is an environment variable, and it has to be set again for every create and every install into that environment:

CONDA_OVERRIDE_CUDA=13 micromamba create -n my_env pytorch

With pixi you record it once, in the platforms entry of your pixi.toml, replacing the plain platforms = ["linux-64"] that pixi init wrote:

[workspace]
platforms = [{ platform = "linux-64", cuda = "13" }]

Because this lives in the manifest it travels with your project – including to machines where it is not true. On Levante, declare what the GPU nodes actually offer. Only lower it if the same project has to build on a machine with an older driver, and be aware of the trade-off: a build for an older CUDA version runs fine here, since a CUDA 12 build works with a driver that supports 13 and both target the A100 equally well, but it also ships older CUDA math libraries (cuBLAS, cuDNN, NCCL) and may leave some performance on the table. Benchmark before settling for an old version in production.

If the project is also used on machines without a GPU, keep the plain entry for those platforms and add cuda only where you need it:

[workspace]
platforms = ["osx-arm64", { platform = "linux-64", cuda = "13" }]

With that in place the solver picks the CUDA build, which you can recognise by the build string and by the CUDA packages that come with it:

+ libtorch  2.13.0  cuda130_mkl_h1ca3d63_302
+ pytorch   2.13.0  cuda130_mkl_py314_h05291b0_302

Note

Creating the environment on a GPU node instead is not a workaround. micromamba would indeed detect __cuda there, but pixi resolves against the requirements declared in your manifest and ignores what the machine it runs on happens to offer – you get the same CPU build. Occupying a GPU to install packages would be wasteful in any case. Declare the version.

Ready-made specifications#

The specifications from Ready-made environments bring their packages and channels, but they cannot bring the CUDA declaration with them: an environment.yaml has no place to put it, and neither has the pixi.toml that pixi init --import generates from one. For a specification marked GPU you have to add it yourself:

# uv -- nothing to add
wget -nc $BASE/pyproject.toml
uv sync

# micromamba -- set the override on the create command
wget -nc $BASE/environment.yaml
CONDA_OVERRIDE_CUDA=13 micromamba create -f environment.yaml

# pixi -- import first, then edit platforms in the generated pixi.toml
wget -nc $BASE/environment.yaml
pixi init --import environment.yaml
# set platforms = [{ platform = "linux-64", cuda = "13" }], then
pixi install

Checking that it worked#

GPU support can only be tested on a GPU node. Request one, then ask the package itself:

salloc -A <your_project> -p gpu --gres=gpu:1 -t 10

uv run python -c "import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name(0))"
# True NVIDIA A100-SXM4-40GB

If this prints False while nvidia-smi on the same node shows a GPU, you are running a CPU build. Check the build string with micromamba list pytorch or pixi list pytorch.

Getting help#

As a quick fallback, the pytorch module provides a system wide environment with GPU support for some packages.

Installing GPU-enabled packages can get involved – it depends on the package, its version, its dependencies and the CUDA version. If you struggle with your GPU-enabled Python environment, get in touch with our User Support.

On this page