SC '22 | HGL: Accelerating Heterogeneous GNN Training with Holistic Representation and Optimization

[jasperzhong]

https://dl.acm.org/doi/abs/10.5555/3571885.3571980

https://github.com/ytgui/HGL-proto/tree/master

[jasperzhong]

HetG可以理解为是a collection of graphs，每一个relation是一个graph. 传统的做法是将HomoGNN模型如GCN用于每个HetG decomposed出来的homo subgraph上. 比如DGL的RGAT可以分解成多个GAT.

# homogeneous model
conv = nn.GATConv(...)
# heterogeneous model
convs = nn.HeteroGraphConv({
  'paper_cite_paper': nn.GATConv(...),
  'author_write_paper': nn.GATConv(...),
}, aggregate=’sum’)

现有的系统把HomoGNN layers当作HetGNN的building blocks. 有两个问题

1. HetGNN需要更多的memory优化. 因为每个relation都有一个subgraph. 如果relation很多，如KG的场景，就会导致fragmented memory.
2. HetGNN需要更多的计算上的并行. 因为HetGNN可以分解为多个HomoGNN，计算量与relation数量是成正比的，因此有更多并行的机会. 而现有的系统就是sequentially执行的，没有利用到并行.

于是他们提出了HGL，将HomoGNN和HetGNN用Holistic Intermediate Representation (HIR)来表示，用HIR来表示HomoGNN/HetGNN的计算图DAG后，又用tailored optimizations: graph stitching, operator fusion和operator bundling来进行优化. HIR只包括了forward computation.

HIR这个描述有点抽象. 看代码可能更直接一些.

[jasperzhong]

HGL中DGLGraph如何表示？调用dgl.DGLGraph.adj_external()，对每个etype转成csr. HomoGraph只有一个CSR，HeteGraph每个relation有一个CSR.

tracer很有意思，继承了torch tensor，然后实现了__torch_function__().

https://pytorch.org/docs/stable/notes/extending.html#subclassing-torch-tensor

这种写法是每次调用pytorch tensor operation的时候，都会执行这一段，有点像一个wrapper. 这个tracer在这里面记录了function name, args, kwargs等等信息，然后用做一次forward pass，得到一个DAG. 然后walk这个DAG，转成他们自己的IR，这个IR除了包括一些常见的op (add/sub/mul etc)，还包括tensor, graph, OpVertFunc, OpEdgeFunc. 这个IR和普通DNN的IR最重要区别是多了几个GNN的概念: OpVertFunc/OpEdgeFunc的输入是一个HomoGraph，HeteGraph在IR里面是拆成了多个HomoGraph.

这个IR怎么执行，就是walk一边这个DAG，遇到什么op就执行对应的kernel（可能是pytorch的，可能是他自己写的fused kernel）

[jasperzhong]

有了HIR后就做了三件事:

1. operator fusion:

• vertical fusion: 一个把SDDMM, leakyReLU和softmax fuse了起来.

另一个省略掉copy_u步骤，比如DGL的GCN的message passing先有一个copy_u('h', 'm')，然后aggregate_sum('m', 'v')，这里fuse起来就是跳过copy_u这一步，直接aggregate_sum.

• horizontal fusion: 把cross relation的并行起来. 这里做的还挺奇怪的，是把g-SDDMM和g-SPMM fuse起来. 最后实验是RGAT用了这个优化，但效果不显著.

后面ablation study显示，vertical fusion非常有用，horizontal fusion效果不显著.

2. graph stitching: 后面的ablation study显示，graph stitching是主要的优化收益来源. 他们是把有相同ending type的relation subgraphs给stitch在一起，有相同ending type的话，adjacency matrix有一个维度是一样的. 例如, $A1 \in R{N\times N_1}, A2 \in R{N\times N_2}$. 但有相同ending type的subgraphs数量也可能有很多，他们这里没有直接把这些subgraphs全部fuse在一起，而是把subgraphs分成几个bin，每一个bin的fuse在一起. 他们设置了一个bin capacity (the maximum number of edges)为全部subgraphs的number of edges之和的一半.

3. 

这里不太明白为什么不直接全部fuse，而是分成多个bin fuse？为了给后面留一点优化空间？

6. operator bundling: 这个说的是把多个有相同输入的GEMM kernels合并成一个kernel. 我觉得这里主要说的是multi-head attention，多个head都有相同的输入h，只是权重不一样. 背后调用的是cublasSgemmStridedBatched这个API.

(为啥Table VI显示R-GCN也能用这个优化？）

[jasperzhong]

实验

HomoGNN也是有一些收益的，不是很大. HeteGNN收益非常显著，但都是在很小KG图上跑的.

这个ablation study也很重要. 可以看到主要收益来自于graph stitching和vertical fusion.

---

换成ogbn-mag效果还能有多少？

[jasperzhong]

结果是容易OOM....