GraphStorm is a graph machine learning (GML) framework for enterprise use cases. It simplifies the development, training and deployment of GML models for industry-scale graphs by providing scalable training and inference pipelines of Graph Machine Learning (GML) models for extremely large graphs (measured in billons of nodes and edges). GraphStorm provides a collection of built-in GML models and users can train a GML model with a single command without writing any code. To help develop SOTA models, GraphStorm provides a large collection of configurations for customizing model implementations and training pipelines to improve model performance. GraphStorm also provides a programming interface to train any custom GML model in a distributed manner. Users provide their own model implementations and use GraphStorm training pipeline to scale.
GraphStorm is compatible to Python 3.7+. It requires PyTorch 1.13+, DGL 1.0 and transformers 4.3.0+.
GraphStorm can be installed with pip and it can be used to train GNN models in a standalone mode. To run GraphStorm in a distributed environment, we recommend users to using Docker container to reduce envrionment setup efforts. A guideline to setup GraphStorm running environment can be found at here and a full instruction on how to setup distributed training can be found here.
Note: we assume users have setup a GraphStorm standalone environment following the Setup GraphStorm with pip Packages instructions. And users have git cloned the GraphStorm source code into the /graphstorm/ folder to use some complimentatry tools.
Node classification on OGB arxiv graph First, use the below command to download the OGB arxiv data and process it into a DGL graph for the node classification task.
python /graphstorm/tools/gen_ogb_dataset.py --savepath /tmp/ogbn-arxiv-nc/ --retain-original-features true
Second, use the below command to partition this arxiv graph into a distributed graph that GraphStorm can use as its input.
python /graphstorm/tools/partition_graph.py --dataset ogbn-arxiv \
--filepath /tmp/ogbn-arxiv-nc/ \
--num-parts 1 \
--num-trainers-per-machine 4 \
--output /tmp/ogbn_arxiv_nc_train_val_1p_4t
GraphStorm training relies on ssh to launch training jobs. The GraphStorm standalone mode uses ssh services in port 22.
In addition, to run GraphStorm training in a single machine, users need to create a ip_list.txt file that contains one row as below, which will facilitate ssh communication to the machine itself.
127.0.0.1
Users can use the following command to create the simple ip_list.txt file.
touch /tmp/ip_list.txt
echo 127.0.0.1 > /tmp/ip_list.txt
Third, run the below command to train an RGCN model to perform node classification on the partitioned arxiv graph.
python -m graphstorm.run.gs_node_classification \
--workspace /tmp/ogbn-arxiv-nc \
--num-trainers 1 \
--num-servers 1 \
--num-samplers 0 \
--part-config /tmp/ogbn_arxiv_nc_train_val_1p_4t/ogbn-arxiv.json \
--ip-config /tmp/ip_list.txt \
--ssh-port 22 \
--cf /graphstorm/training_scripts/gsgnn_np/arxiv_nc.yaml \
--save-perf-results-path /tmp/ogbn-arxiv-nc/models
Link Prediction on OGB MAG graph First, use the below command to download the OGB MAG data and process it into a DGL graph for the link prediction task. The edge type for prediction is “author,writes,paper”. The command also set 80% of the edges of this type for training and validation (default 10%), and the rest 20% for testing.
python /graphstorm/tools/gen_mag_dataset.py --savepath /tmp/ogbn-mag-lp/ --edge-pct 0.8
Second, use the following command to partition the MAG graph into a distributed format.
python /graphstorm/tools/partition_graph_lp.py --dataset ogbn-mag \
--filepath /tmp/ogbn-mag-lp/ \
--num-parts 1 \
--num-trainers-per-machine 4 \
--target-etypes author,writes,paper \
--output /tmp/ogbn_mag_lp_train_val_1p_4t
Third, run the below command to train an RGCN model to perform link prediction on the partitioned MAG graph.
python -m graphstorm.run.gs_link_prediction \
--workspace /tmp/ogbn-mag-lp/ \
--num-trainers 1 \
--num-servers 1 \
--num-samplers 0 \
--part-config /tmp/ogbn_mag_lp_train_val_1p_4t/ogbn-mag.json \
--ip-config /tmp/ip_list.txt \
--ssh-port 22 \
--cf /graphstorm/training_scripts/gsgnn_lp/mag_lp.yaml \
--node-feat-name paper:feat \
--save-model-path /tmp/ogbn-mag/models \
--save-perf-results-path /tmp/ogbn-mag/models
To learn GraphStorm's full capabilities, please refer to our Documentations and Tutorials.
GraphStorm framework now supports using CPU or NVidia GPU for model training and inference. But it only works with PyTorch-gloo backend. It was only tested on AWS CPU instances or AWS GPU instances equipped with NVidia GPUs including P4, V100, A10 and A100.
Multiple samplers are supported in PyTorch versions <= 1.12 and >= 2.1.0. Please use --num-samplers 0 for other PyTorch versions. More details here.
To use multiple samplers on sagemaker please use PyTorch versions <= 1.12.
This project is licensed under the Apache-2.0 License.