ARGO: an easy-to-use runtime to improve GNN training performance on multi-core processors

[jasonlin316]

trafficstars

Overview

The GNN training performance on multi-core processors is limited as the current design cannot scale well. We propose a runtime system named ARGO that can improve the scalability of GNN training on multi-core processors. On a CPU platform where the original program can only scale to 16 cores (meaning that no performance improvement is achieved if more than 16 cores are applied), ARGO can further scale the design up to 64 cores, achieving up to 5x speedup compared to the original design without ARGO.

[jasonlin316]

Looks interesting! Do you think we could have an end-to-end example?

Yes. I do have end-to-end example: https://github.com/jasonlin316/ARGO/tree/main/PyG Please see flickr_example_ARGO.py which enables ARGO on the flickr_example.py.

[jasonlin316]

I feel like we should have a complete example because only looking at this class may not be clear enough for most PyG users to adapt it to their use cases IMHO.

Yes, I agree having a complete example is important for the users. @rusty1s Do you want to comment on this? I recall the plan is to add ARGO to the PyG package, and then add an example. Should we finalize this PR first, and then create another PR for the end-to-end example?

[rusty1s]

@jasonlin316 I moved to torch_geometric package directly, moved the doc-string and adjusted input arguments. If you have time, can you follow up with adjusting docs/arguments in the remainder? In addition, how hard would it be to test ARGO in test/nn/models/test_argo.py?

[codecov[bot]]

Codecov Report

Attention: 85 lines in your changes are missing coverage. Please review.

Comparison is base (1b195a0) 89.26% compared to head (15c2080) 89.01%. Report is 2 commits behind head on master.

Files	Patch %	Lines
torch_geometric/nn/models/argo.py	0.00%	85 Missing :warning:

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #8841      +/-   ##
==========================================
- Coverage   89.26%   89.01%   -0.26%     
==========================================
  Files         468      469       +1     
  Lines       29960    30045      +85     
==========================================
  Hits        26744    26744              
- Misses       3216     3301      +85     

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.