GATv2 Run out of Memory on Heterogeneous Graph

[abdalgader-a-zz]

Hi there 👋,

I ran a heterogeneous model with 2-layers (256 -> 256, 256 - > 128) of GATv2. The number of model's parameters are 1.7M. It rise an OOM error.

Subgraph is on SparseTensor format with 46,635 nodes with 128 dimension and 12M edges.

GPU Spec: Tesla P100 SXM2 16GB

PS: The program run normally with SAGEConv.

Any help or suggestions ?

Thanks!

[wsad1]

It's kind of surprising that you are getting an OOM. Could you maybe post some snippet of your code.

[rusty1s]

GATv2Conv will materialize the edge representation, e.g., it will compute 12M messages. One option to mitigate OOMs would be to reduce the number of edges in your graph, e.g., by edge dropout.

[abdalgader-a-zz]

Thanks @wsad1 for your replying. It ran OOM during the forward step, code is quite similar to this with using NeighborLoader for sampling subgraph. Maybe the graph is the key point, as I'm using undirected OGB-MAG with extra 3 edge types.

Thanks @rusty1s for your illustration. Would it possible to illustrate why it's not the case with SAGEConv? Is this due to the attention implementation on GATv2Conv?

[rusty1s]

Yes, it is due to attention mechanism. We learn multiple edge weights here, which does not let us use sparse-matrix multiplication. There is an open PR that tries to eliminate this, see #2380.