How did DDP handle x and edge_index input?

[jasperhyp]

Hi! I am a little bit confused on why the example here works. If I understand correctly, DDP will just split inputs to forward on the first dimension. And since the x and edge_index are both inputs, they will essentially be split in the same way, which is problematic because we only want the edge_index to be split.

I actually have an even more complicated case where I am also using an encoder for the node embeddings, and it's on a set of relations (i.e. a knowledge graph). Have been thinking how to use DDP/deepspeed/others for a while and barely made any progress. Would also appreciate any pointers on this end!

[wsad1]

DDP doesn't split the data (see here). Its up to the application using DDP to split the data. In the example you shared, it happens in line 70. After setting up train_idx. NeighborLoader which is a loader built for GNNs, splits the graph into subgraphs it doesn't as you said "split x and edge_index in the same way".

[LukeLIN-web]

DDP doesn't split the edge_index; each process has a subset of train_idx
You can refer https://arxiv.org/pdf/1706.02216.pdf . Especially pay attention to what the mini-batch is in the paper.