SAT: Staleness-Alleviated Training Framework

This repo is the practical realization of our Staleness-Alleviated Training (SAT) framework, which reduces the embedding staleness adaptively, as described in our paper:

Staleness-Alleviated Distributed Graph Neural Network Training via Online Dynamic-Embedding Prediction

SAT is implemented in PyTorch and utilizes the PyTorch Geometric (PyG) library.

A detailed description of SAT can be found in its implementation.

Requirements

• Install PyTorch >= 1.7.0
• Install PyTorch Geometric >= 1.7.0:

pip install torch-scatter -f https://pytorch-geometric.com/whl/torch-${TORCH}+${CUDA}.html
pip install torch-sparse -f https://pytorch-geometric.com/whl/torch-${TORCH}+${CUDA}.html
pip install torch-geometric

where ${TORCH} should be replaced by either 1.7.0 or 1.8.0, and ${CUDA} should be replaced by either cpu, cu92, cu101, cu102, cu110 or cu111, depending on your PyTorch installation.

Installation

python setup.py install

Project Structure

• torch_geometric_autoscale/ contains the source code of SAT
• small_benchmark/ includes experiments to evaluate SAT performance on small-scale graphs

We use Hydra to manage hyperparameter configurations.