[Paper Review] Sequential Aggregation and Rematerialization: Distributed Full-batch Training of Graph Neural Networks on Large Graphs

Hesham Mostafa|arXiv (Cornell University)|Nov 11, 2021

Advanced Graph Neural Networks4 citations

TL;DR

This paper introduces Sequential Aggregation and Rematerialization (SAR), a distributed full-batch training method for Graph Neural Networks (GNNs) that drastically reduces memory usage by sequentially reconstructing and freeing computational graph components during backpropagation. SAR enables training on large graphs—such as ogbn-papers100M and ogbn-products—by ensuring per-worker memory scales linearly with the number of workers, achieving up to 50% memory reduction on 128 workers and significant speedups with optimized attention kernels.

ABSTRACT

We present the Sequential Aggregation and Rematerialization (SAR) scheme for distributed full-batch training of Graph Neural Networks (GNNs) on large graphs. Large-scale training of GNNs has recently been dominated by sampling-based methods and methods based on non-learnable message passing. SAR on the other hand is a distributed technique that can train any GNN type directly on an entire large graph. The key innovation in SAR is the distributed sequential rematerialization scheme which sequentially re-constructs then frees pieces of the prohibitively large GNN computational graph during the backward pass. This results in excellent memory scaling behavior where the memory consumption per worker goes down linearly with the number of workers, even for densely connected graphs. Using SAR, we report the largest applications of full-batch GNN training to-date, and demonstrate large memory savings as the number of workers increases. We also present a general technique based on kernel fusion and attention-matrix rematerialization to optimize both the runtime and memory efficiency of attention-based models. We show that, coupled with SAR, our optimized attention kernels lead to significant speedups and memory savings in attention-based GNNs.We made the SAR GNN training library publicy available: \\url{https://github.com/IntelLabs/SAR}.

Motivation & Objective

• To address the memory bottleneck in full-batch training of GNNs on large, densely connected graphs where the computational graph becomes prohibitively large to store.
• To enable scalable, distributed full-batch training of any GNN architecture without relying on sampling or non-learnable message passing.
• To minimize communication overhead in distributed training by avoiding redundant materialization of intermediate tensors, especially for attention-based models.
• To provide a practical, memory-efficient baseline for evaluating sampling-based and non-learnable message-passing GNN methods.
• To optimize performance for attention-based GNNs like GAT by eliminating storage of attention coefficients and computing them on-the-fly.

Proposed method

• SAR uses a distributed domain-parallel training scheme where the input graph is partitioned across N workers, and each worker processes only its assigned subgraph.
• Instead of materializing the entire GNN computational graph during the forward pass, SAR delays graph construction until the backward pass, reconstructing it piecewise.
• During backpropagation, SAR sequentially rematerializes and frees computational graph components, ensuring at most two partitions are stored per worker at any time.
• The method achieves memory scaling of O(2/N) per worker, independent of graph density, enabling training on arbitrarily large graphs by adding more workers.
• For attention-based models, SAR integrates kernel fusion and on-the-fly attention coefficient computation to avoid storing large attention matrices.
• The approach is generalizable to any domain-parallel training setup where outputs depend on cross-worker inputs, including spatially-parallel CNNs.

Experimental results

Research questions

• RQ1 Can full-batch GNN training be scaled to large graphs without relying on sampling or non-learnable message passing?
• RQ2 How can memory usage in distributed GNN training be reduced to enable training on graphs too large to fit in memory?
• RQ3 What is the communication and memory overhead of rematerializing computational graphs during backpropagation in a distributed setting?
• RQ4 Can rematerialization be optimized to avoid materializing expensive intermediate tensors like attention coefficients in GAT models?
• RQ5 How does SAR compare in memory efficiency and training speed to sampling-based and non-learnable message-passing GNN methods on large-scale benchmarks?

Key findings

• SAR reduces peak memory consumption per worker by up to 50% when training a GraphSage model on 128 workers on ogbn-papers100M, with memory scaling linearly as 2/N.
• On ogbn-products, SAR achieves a 2.5x speedup over DGL’s GAT implementation when using optimized attention kernels, due to reduced memory pressure and on-the-fly coefficient computation.
• For GraphSage, SAR achieves comparable runtime to domain-parallel training with lower communication volume, while cutting memory usage by half at 128 workers.
• The optimized attention kernel (FAK) used with SAR reduces memory usage by avoiding storage of attention coefficients, improving forward pass speed without degrading backward pass performance.
• SAR enables the largest reported full-batch GNN training to date on ogbn-papers100M and ogbn-products, demonstrating feasibility on graphs with over 100 million nodes.
• The method is communication-avoiding for many GNN variants, making memory savings effectively free in terms of runtime overhead.