Any way to implement muliti-GPU inference of a singular large graph?

[JSLJ23]

Hi pyG team and community,

I am looking for some help to implement an solution that might involve a multi-GPU inference of a single large graph object through the pytorch geometric model.
The context of the problem is basically involving running the MACE Neural Network Force Field which implements pytorch_geometric under the hood and the protein system by which the NNFF will be evaluated on is constructed to as a single mace.data.atomic_data.AtomicData which inherits from torch_geometric.data.Data. The single protein graph has about 5000+ nodes (atoms) and 100k+ edges which is far too large to fit on a single A100 40GB, due to the complexity of the MACE model as well. Would it be possible for this graph to be partitioned across multiple GPUs and for inference to be run on multiple GPUs and consolidated as a single final readout?

Appreciate any advice and help I could get on this.

[wsad1]

Would NeibhorLoader be useful in this case. It samples subgraphs, and the subgraphs could be loaded on your GPU.

[JSLJ23]

Will this work for graph level regression problems? I think the MACE model does a global pooling at the end step to get the total energy of the system so will the NeighborLoader allow for such graph level operations?

[wsad1]

Ah I see, In that case NeighborLoader would not help. I am guessing distributed training added in PyG2.5 might help https://github.com/pyg-team/pytorch_geometric/releases/tag/2.5.0.
cc: @kgajdamo do you think graph level tasks can be solved with distributed training?

[kgajdamo]

Hi, @wsad1 @JSLJ23 ! Distributed training for PyG currently works only on the CPU. If you would like to use Nvidia GPU you could try to use: graphlearn-for-pytorch which supports distributed training. There are also examples for distributed training on GPU in PyG but only for single node multi gpu cases: multi_gpu.
Another my suggestion would be to use Intel GPU, because it has more memory than Nvidia cards.