Benefits of transferring data to GPU before passing it to NeighborLoader

[liadgerstman]

Hi, I understand that neighbor sampling using NeighborLoader is supposed to mainly run on CPU and then we need to transfer the sampled batches from the CPU to the GPU for the training part.
So:

1. Why is it claimed in reddit.py that the nodes' features/labels are sent to the GPU 'for faster access during sampling'? Does it mean that the sampling is being done on the GPU in this case? I am asking this because I could not see the GPU kernels doing sampling when profiling this case.
   . If the sampling is indeed being processed on a CPU- why aren't the batches being sent to the GPU as well?

2. I wanted to clarify that the sampling process itself only starts when the next() is called when we iterate over NeighborLoader iterator and not we construct NeighborLoader. Hence, the sampling is not done here:
   
   but here:
   

Thank you for the clarification

[akihironitta]

Hey @liadgerstman! I haven't benchmarked this, but I guess there should be some speedup. AFAIK, this NeighborLoader (1) samples node indices first on CPU, and (2) packs node features and labels (data.{x,y}) of sampled node indices into cuda tensors directly instead of CPU tensors. If you don't transfer the data.{x,y} to the device beforehand, you will have to transfer the sampled data next(train_loader_iter).{x,y} to the device in each iteration, which may create performance difference.

[liadgerstman]

Hey @akihironitta, I profiled reddit.py and it looks like the gpu is idle during NeighborLoader execution - only CPU is involved (nodes' features/labels are sent to the gpu before it starts). That is why I asked the questions:
Is the sampling process actually being executed when NeighborLoader is called here? Or is it done lazily by the time the next() function is called while you iterate over the train_loader in the for loop.
In either ways, the gpu is not involved in the sampling process (according to the profiling) and that makes the comment in the reddit.py quiet confusing (The comment that says that nodes' features/labels are sent to the GPU 'for faster access during sampling).
I am interested also in understanding if gpu sampling is even possible using PYG's NeighborLoader- I tried to pass dataset[0] to the gpu before calling to NeighborLoader but it fails when you call NeighborLoader. When you only transfer x, y, edge_index to the gpu it doesn't fail but it looks like the sampling is still done on the cpu.

[liadgerstman]

Hi @akihironitta and @rusty1s thanks for your answers!
So, to sum things up,

• The comment : # features/labels are sent to GPU so you don't have to in each iteration is now much better 👍
• Regarding @akihironitta 's profiling (thank you for looking at that!), we can see that data transfers being executed in between the iterations slow down the whole training process, and we can see that nothing is being executed on the GPU except for some overhead functionality probably. By the way, why not transferring the edge_index to the GPU as well before sampling? (I guess that it is because we want to avoid device transfer of unnecessary edge_index elements?)
• The neighbor sampling is being done lazily when next() is called.

Thank you again for the great answers

[akihironitta]

The comment : # features/labels are sent to GPU so you don't have to in each iteration is now much better 👍

I missed this until @rusty1s commented above, but the current comment is still true (and clear enough IMO) as index_select runs faster on GPU as you can see in the profiling result. Do you have any suggestion on how to update the comment? Feel free to send a PR if any :)

why not transferring the edge_index to the GPU as well before sampling?

I think there won't be much difference because sampling node indices, which uses edge_index, will run on CPU anyway as sampling on GPU isn't supported yet (pyg-team/pyg-lib#23).

[liadgerstman]

@akihironitta, do you suggest that the comment is still true because the index_select is being done on the GPU? Is that the only thing being done on the GPU in the context of sampling?
I was thinking that since the sampling itself (except for the final index_select) is currently not executed on the GPU , it is a bit confusing saying that transferring the data to the GPU is because of 'fast access during sampling'. I would suggest giving the context of index_select in the comment. BTW, now that we have this discussion here I guess it is no longer an issue.

[akihironitta]

do you suggest that the comment is still true because the index_select is being done on the GPU? Is that the only thing being done on the GPU in the context of sampling?

Yes, and Yes from what I see from the profiler. You can also check it on your own by running my code snippet above.

Thanks for sharing your thoughts. Feel free to send a PR if you have any suggestions around the comment :)