search

qitianwu / SGFormer

The official implementation for NeurIPS2023 paper "SGFormer: Simplifying and Empowering Transformers for Large-Graph Representations"
213 stars 18 forks source link

Running out of memory on ogbn-arxiv #8

Closed ryandeng1 closed 7 months ago

ryandeng1 commented 9 months ago

Hi,

I am currently trying to run the code in the script for ogbn-arxiv, but I am running out of memory. Right now, it attempts to do a matrix multiply between two tensors which are each 169k by 256, and it allocated 100GB of VRam to do so.

I am wondering what settings I need to change in the script to get it to run?

Specifically, https://github.com/qitianwu/SGFormer/blob/466ff764f46e95a7bd3cd2985be08e205008776b/large/ours.py#L152 is where the error occurs.

Am I supposed to use main-batch.py instead of main.py?

Thanks!

qitianwu commented 9 months ago

Thanks for letting us know about this issue. This is indeed a mistake in the original code that adopts the quadratic attention and comments out the linear attention. Now the bug has been fixed with the correct version of the linear attention we actually used for SGFormer. The code should run smoothly on the large datasets without OOM (for arxiv, we use main.py and for other large datasets, we use main-batch.py)