Home · ASPLOS 2024 · Machine Learning Systems

Hector: An Efficient Programming and Compilation Framework for Implementing Relational Graph Neural Networks in GPU Architectures

Kun Wu, Mert Hidayetoglu, Xiang Song, Sitao Huang, Da Zheng, Israt Nisa and 1 other

Conference
ASPLOS 2024
Published
2024
Pages
528-544
Citations
3

Neural networks that reason over graphs with many kinds of nodes and edges — like a citation network of authors, papers, and venues — are accurate but painfully slow and memory-hungry on GPUs, partly because today's frameworks bloat memory with giant temporary copies and force engineers to hand-write a new kernel for every combination of edge type and data format; this work from Illinois and AWS builds Hector, a two-level compiler that separates the model, the data layout, and the low-level tuning, generates all the GPU code from just two templates, and skips the redundant copies, running up to about 44 times faster in training while never running out of memory.

The argument in one glance

  1. Problem

    Relational graph networks are accurate but slow and memory-hungry on GPUs

    Relational graph neural networks handle graphs with many node and edge types by giving each type its own weights. But they are hard to run fast on GPUs: the message-passing computation is memory-intensive and graph-dependent, the interface-to-kernel gap forces giant temporary tensors, and the explosion of typed kernels is not covered by pre-built libraries.

  2. What people did

    Use existing frameworks or hand-write kernels

    Run relational graph networks with PyTorch-based graph frameworks and their operators, or have engineers hand-write optimized GPU kernels for each typed operation.

  3. Why it fell short

    Frameworks bloat memory; hand-written kernels cost months

    PyTorch-based frameworks materialize huge temporary tensors — a copy of the right weight matrix for every edge — wasting memory and hitting out-of-memory errors, and fall back on slow implementations that launch a tiny kernel per relation type. Hand-writing optimized kernels is enormous effort, taking a full-time engineer over a month for a single typed layer.

  4. This paper's idea

    A two-level IR that decouples model, layout, and tuning (Hector)

    Hector is a two-level intermediate representation and code generator for relational graph networks on GPUs. A higher inter-operator level expresses model semantics and data-layout choices in a decoupled way; a lower intra-operator level specializes templates into CUDA kernels. It generates all code from just two templates — a matrix-multiply template and a node-and-edge traversal template — operating on tensor views to skip redundant copies, and adds compact materialization and operator reordering. It runs up to 9.9 times faster in inference and 43.7 times in training, with no out-of-memory errors.

Why this problem is worth solving

Graph neural networks are widely used, and the most valuable graphs are exactly the ones that are hardest to compute on.

A graph neural network is a neural network that operates on a graph: it propagates and transforms the features attached to nodes using the graph’s connectivity, through a procedure called message passing. In each layer, an edgewise operation produces a message along every edge, and then a nodewise aggregation gathers the messages arriving at each destination node to update its features. Graph neural networks power applications like fraud detection and recommendation, and open-source frameworks such as the Deep Graph Library and PyTorch Geometric provide optimized operations for them. But most of those optimizations assume a homogeneous graph, where all nodes and edges are the same kind. Real-world graphs are usually heterogeneous — they contain multiple types of nodes and edges. A citation graph, for example, might have authors and articles as different node types, connected by different relations such as citing or writing. To capture this structure, relational graph neural networks (models like RGCN, RGAT, and HGT) give each node type and each edge type its own dedicated weights and data paths. This makes them versatile and accurate, and increasingly popular — but also far more complex to execute efficiently.

The complexity translates into three concrete obstacles on a GPU.

The first obstacle is that the computation is memory-intensive and depends on the input graph. To use the GPU’s fast matrix-multiply kernels for message passing, the system must first traverse the graph and scatter and gather tensor data into the right shapes, and the resulting memory-access patterns are dictated by the graph’s structure and data layout, not just the model. Because the connectivity drives the computation, no single optimization strategy fits every model-and-graph combination, and the type dimension of relational networks enlarges the design space further. The second obstacle is a gap between the high-level programming interface and the low-level kernel calls, which causes extra data movement. The clearest example is an edgewise typed linear layer, which multiplies a vector on each edge by the weight matrix belonging to that edge’s type. To do this with standard operators, PyTorch-based systems materialize a huge temporary three-dimensional tensor holding a copy of the correct weight matrix for every single edge — an enormous, redundant duplication that wastes memory. The third obstacle is that pre-built kernel libraries are inadequate here. The number of kernels needed multiplies with the choices of adjacency-matrix storage format, and relational networks’ typed nodes and edges multiply it again, so developers must either hand-write optimized kernels — one typed layer reportedly took a full-time engineer over a month — or accept slow fallbacks, such as a loop that launches a tiny separate GPU kernel for each relation type, badly underutilizing the hardware.

What people did before, and where it broke

  • Use PyTorch-based graph frameworks. Run relational graph networks with the operators provided by libraries like the Deep Graph Library. Where it broke. They materialize huge temporary tensors (a weight-matrix copy per edge), wasting memory and hitting out-of-memory errors, and fall back on slow per-relation kernel launches.
  • Hand-write optimized GPU kernels. Have engineers craft fast CUDA kernels for each typed operation. Where it broke. The effort is enormous — over a month of a full-time engineer for a single typed layer — and it multiplies with every combination of type and storage format.
  • Rely on pre-built kernel libraries. Reuse libraries like cuBLAS and cuSPARSE for the heavy operations. Where it broke. They do not cover the explosion of kernels that typed edges and nodes crossed with storage formats create, so they leave much of the work unoptimized.
  • The gap. No framework let a developer express a relational graph network compactly and get fast, memory-efficient GPU code, decoupling the model, the data layout, and the low-level tuning. Providing exactly that is what Hector does.

How Hector works

From framework operators and hand-written kernels to a two-level generating IR

  1. new

  2. new

  3. new

  4. new

  5. new

The idea

Separate the model, the data layout, and the low-level tuning. Hector’s foundation is a two-level intermediate representation that decouples three things existing systems tangle together: what the model computes, how its tensors are laid out in memory, and how each operator is optimized into a kernel — so a change to any one does not force rewriting the others, which is what collapses the programming effort from months of hand-tuned kernels to a few lines of model code. Building a compiler as layered representations that separate concerns and generate efficient GPU kernels, so the high-level model and the low-level optimizations can each evolve independently, is the same design the deep-learning-training compiler work covered elsewhere on this site follows, splitting its stack so that an expert-optimized core and the flexible, model-specific parts are handled at different levels rather than fused into one rigid whole.

Skip the redundant copies instead of materializing them. A typed linear layer naively builds a giant temporary tensor holding a copy of the right weight matrix for every edge; Hector instead has its operators work on views of the small set of per-type weights and uses compact materialization to keep intermediate tensors small, eliminating the redundant duplication that wasted memory and caused out-of-memory failures. Recognizing that a system is spending huge effort moving and storing data it does not truly need, and arranging to avoid that movement, is the same principle the warehouse-scale memory-traffic work covered elsewhere on this site is built on, cutting the needless transfers and materializations so that scarce memory and bandwidth are spent only on data that matters.

Generate the whole kernel explosion from a couple of templates. Because the number of kernels needed explodes with edge types and storage formats, hand-writing each is untenable; Hector instead generates all of them by specializing just two generalized templates — a matrix-multiply template and a node-and-edge traversal template — so a combinatorial space of kernels comes from a small, reusable base. Handling a large space of cases by filling in a small number of general templates, especially a traversal template, rather than writing every case by hand, is the same strategy the traversal-synthesis work covered elsewhere on this site takes, expressing many concrete traversals as instantiations of templates so that a single template captures what would otherwise be a sprawl of bespoke code.

Following a typed linear layer through Hector

Applying per-edge-type weights without a giant temporary tensor

A trace of Hector compiling an edgewise typed linear layer, following the paper's two-level IR and compaction.

Step 1 of 5

The layer must apply each edge type's weight

A relational graph layer must multiply, on each edge, the vector by the weight matrix belonging to that edge's type — a typed linear transformation, and the memory-hungry heart of the model.

There are only a few weight matrices, one per edge type.

Why it works

  • Performance is dominated by memory movement. Because message passing must traverse the graph and scatter and gather tensor data, with patterns set by the specific graph → the computation is memory-intensive and cannot be optimized one-size-fits-all → so a system that minimizes data movement per model and graph is what unlocks speed.
  • Decoupling concerns tames the design space. Because model semantics, data layout, and operator optimizations are entangled in existing systems → a small change forces rewriting hand-tuned kernels → so a two-level IR that separates them lets each be optimized independently, cutting programming effort enormously.
  • Views beat huge temporaries. Because a typed linear layer naively materializes a giant per-edge weight tensor → operating on views of the small per-type weights removes the redundant copies → so both memory footprint and access counts drop sharply.
  • Two templates cover the kernel explosion. Because the number of kernels needed explodes with types and storage formats → generating them all from a matrix-multiply template and a traversal template avoids hand-writing each → so a month of engineering per layer becomes automatic code generation.
  • Compaction removes the memory wall. Because the largest cost and the out-of-memory failures come from bloated intermediate tensors → compact materialization keeps them small → so Hector fits larger graphs, raises GPU utilization, and never runs out of memory.

What it costs (trade-offs and scope)

  • Backward propagation stays latency-bound. Profiling shows that higher memory efficiency lets Hector fit larger input and raise throughput in forward propagation, but backward propagation does not benefit from larger input because it is bound by the latency of atomic updates and outer products. The optimizations help forward more than backward.
  • The best options vary by model and dataset. Compact materialization and linear operator reordering give the largest gains in different combinations for different models and datasets, so getting the most out of Hector means selecting among options rather than one fixed setting.
  • It is specialized to relational graph networks on GPUs. Hector’s IR and templates capture the specific structure of relational graph neural networks — typed nodes and edges, message passing — and target GPUs. Its benefits are tied to that domain, not to arbitrary computation.
  • Two templates bound expressiveness. Generating everything from a matrix-multiply template and a traversal template is powerful and compact, but a computation that does not fit those two templates is outside what the code generator can express.
  • The speedups are against specific baselines and models. The up-to 9.9 times inference and 43.7 times training speedups, and the further up-to 3.8 times gains, are measured on RGCN, RGAT, and HGT with particular datasets against particular state-of-the-art systems; other models, datasets, and baselines would shift the numbers.
  • It is a framework to adopt. Realizing the benefit means expressing models in Hector’s representation and using its code generator, a real toolchain rather than a drop-in change to an existing framework, though the artifact is publicly available.

The evaluation

Three ways to run relational graph neural networks on GPUs

PropertyPyTorch-based frameworkHand-written kernelsHector
Avoids huge temporary tensors
Programming effortLow but slow
Covers the typed-kernel explosionOnly what is written
Out-of-memory on large graphsCommonVariesNone with compaction
SpeedBaselineFast if hand-tunedUp to 9.9x inference, 43.7x training
Three ways to run relational graph neural networks on GPUs, as compared in the paper. PyTorch-based frameworks are easy to program but materialize huge temporary tensors (a weight-matrix copy per edge), wasting memory, hitting out-of-memory errors, and falling back on slow per-relation kernel launches. Hand-writing optimized kernels is fast but takes enormous effort — over a month of a full-time engineer for a single typed layer — and does not scale to the explosion of typed kernels. Hector, a two-level intermediate representation and code generator, decouples model semantics, data layout, and operator-specific optimizations, generates all kernels from a matrix-multiply template and a node-and-edge traversal template operating on tensor views, and adds compact materialization and linear operator reordering. On RGCN, RGAT, and HGT with Deep Graph Library and Open Graph Benchmark datasets, it achieves up to 9.9 times inference and 43.7 times training speedup over state-of-the-art systems, with the two optimizations adding up to a further 3.8 times, and it incurs no out-of-memory errors with compaction enabled. As an indicator of effort, 51 lines of model code generate about 8000 lines of CUDA and C++.Underlined values explain themselves when clicked.

Check yourself

Check your understanding

  1. 1. What makes relational graph neural networks harder to run than ordinary graph neural networks?

  2. 2. Why does a typed linear layer waste so much memory in existing frameworks?

  3. 3. What is the core design of Hector, and why does it help?

  4. 4. Hector speeds up training more than the profiling suggests it could speed up backward propagation. What does this reveal?

Sources

The paper is “Hector: An Efficient Programming and Compilation Framework for Implementing Relational Graph Neural Networks in GPU Architectures” by Kun Wu, Mert Hidayetoglu, Xiang Song, Sitao Huang, Da Zheng, Israt Nisa, and Wen-Mei Hwu, published at ASPLOS 2024, DOI 10.1145/3620666.3651322. The authors are at the University of Illinois Urbana-Champaign and Amazon Web Services.

This explainer is based on the full text of the paper, obtained from the authors’ arXiv version (arXiv:2301.06284). The problem framing (graph neural networks and message passing, the shift from homogeneous to heterogeneous graphs with multiple node and edge types, relational graph neural networks such as RGCN, RGAT, and HGT giving each type its own weights, and the three challenges — memory-intensive graph-dependent computation, the interface-to-kernel gap causing extra data movement via a huge temporary per-edge weight tensor, and the inadequacy of pre-built kernel libraries for the explosion of typed kernels, illustrated by the month-long effort to implement a typed linear layer and the slow per-relation kernel launches), the Hector design (the two-level intermediate representation with an inter-operator level for decoupled model semantics and data layout and an intra-operator level for template specialization to CUDA; the decoupling of model semantics, data layout, and operator-specific optimizations; the generation from a matrix-multiply template and a node-and-edge traversal template; operating on tensor views to avoid redundant copies; and the two optimizations of compact materialization and linear operator reordering), and the results (up to 9.9 times inference and 43.7 times training speedup over state-of-the-art systems on RGCN, RGAT, and HGT with Deep Graph Library and Open Graph Benchmark datasets, a further up to 3.8 times inference and 2.7 times training from the two optimizations, no out-of-memory errors with compaction, the 51 lines of model code generating about 8000 lines of CUDA and C++, and the profiling finding that forward propagation benefits from memory efficiency while backward propagation is latency-bound by atomic updates and outer products) are all drawn directly from the paper. The artifacts are publicly available. Background explanations of graph neural networks, message passing, matrix multiplication on GPUs, intermediate representations, and tensor materialization are standard and provided to make the paper’s contribution understandable.

Related papers on Machine Learning Systems

See all 212 →