Add support for pytorch GPU under Docker #662
Conversation
Dockerfile.release
Outdated
| fi && \ | ||
| python3 -m pip install --no-cache-dir dgl==2.5.0 -f https://data.dgl.ai/wheels/torch-2.6/repo.html --break-system-packages; \ | ||
| fi && \ | ||
| python3 -m pip install --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/cu124 --break-system-packages && \ |
There was a problem hiding this comment.
We should do a GPU specific image because this (I think) will make the image quite a bit bigger, as it will install everything CUDA-related along with it.
An easy-ish way to do this might be to create a new dockerfile, with the base image being pytorch/pytorch, which I think is Ubuntu based. I think there is also rocm/pytorch too. The downside to this is that it appears to only be built for amd64...
There was a problem hiding this comment.
Yep, I was also thinking about that, so most likely we need:
- mage with pytorch for CPU ARM
- mage with pytorch for CPU x86
- one mage with pytorch 2.8 cuda 12.6 x86
- one mage with pytorch 2.8 rocm X.Y x86
NOTE: DGL doesn't get built for 2.8, we managed to build for 2.6 -> we might have to keep it. There we can also get rid of the DGL and substitute it with e.g. pytorch-geometric.
|
The issue is that during Dockerfile build, CPU version gets installed, when memgraph starts, it loads the CPU version, and then it's not possible to reload the loaded module using the
CALL mg.load_all().Should we only install the GPU version -> https://stackoverflow.com/questions/58511598/can-both-the-gpu-and-cpu-versions-of-pytorch-be-installed-in-the-same-conda-envi? -> NO, because of the size, take a look at the comment below