Everything that makes a GNN better — deeper aggregation, LSTM aggregators, bigger fanout — multiplies the neighborhood it must hold in GPU memory, and the obvious fix of shrinking the batch silently changes what the model learns; Betty instead splits each batch into micro-batches that accumulate gradients (mathematically identical training), partitioning along a redundancy-embedded graph so shared neighbors aren't cloned into every piece, with an analytic memory model choosing how many pieces — turning OOM configurations into single-GPU runs with up to 48% less peak memory than generic partitioners.
The argument in one glance
-
Problem
Every GNN improvement is a memory bomb
Graph neural networks learn a node's representation by recursively aggregating its neighbors' features — so the data one training example needs grows with every hop. All the levers that improve accuracy pull the same fuse: deeper aggregation (papers push from 1 to 112 layers) grows neighborhoods exponentially, larger fanout admits more neighbors, and advanced aggregators like LSTM multiply intermediate state. On a 24 GB GPU, an LSTM aggregator on a medium dataset already exceeds memory; a 4-layer model OOMs outright. And unlike ordinary DNNs, the memory goes into the INPUT — node feature vectors are 55% of the footprint.
-
What people did
Shrink the batch, or buy more GPUs
Two standard escapes. Scale out: DistDGL partitions the graph across machines — but training one relational GCN on a 7-billion-edge graph takes 32 NVIDIA T4 GPUs, a hardware bill many model scientists cannot pay. Or shrink the mini-batch so each iteration touches a smaller subgraph — an appealing zero-hardware fix.
-
Why it fell short
Batch size is a hyperparameter, and GNN batches don't cut cleanly
Shrinking the mini-batch changes the EFFECTIVE batch size — a statistical property, not an implementation detail: training curves fluctuate, accuracy degrades after a few epochs, and learning-rate schedules must be retuned. The DNN world's answer is gradient accumulation over micro-batches — but that assumes samples are independent. A GNN batch is a multi-level bipartite graph whose output nodes SHARE input neighbors: cut it naively and every shared neighbor is cloned into each piece that needs it. Partitioning with Metis inflates nodes 23.5%, memory 27.6%, and time 47.6%; and the power-law degree distribution leaves pieces up to 31.9% apart in memory, so the biggest one still OOMs.
-
This paper's idea
Partition along the shared-neighbor structure, sized by a memory model
Betty builds a redundancy-embedded graph (REG) whose edge weights count the neighbors two output nodes SHARE — so a standard min-cut partitioner (Metis) applied to REG groups co-neighbor nodes together and minimizes duplication. An analytic memory model (within 8% of actual, even for LSTM) then estimates each micro-batch's footprint without running it, incrementing the partition count until the largest piece fits. Gradients accumulate across micro-batches, so training is mathematically identical to the full batch: every OOM configuration now runs on one 24 GB GPU, with up to 48.3% less peak memory than generic partitioners.
Why this problem is worth solving
A a neural network that operates on
a graph — social networks, citation graphs, product co-purchases —
where each node carries a feature vector and each layer updates a
node’s representation by aggregating its neighbors’ representations
(the “message passing” pattern)
the cap on how many neighbors are sampled per node
per layer — fanout 10 over 3 layers means one output node can pull
in up to 10³ input nodes
This matters because the community keeps discovering that the expensive knobs are the good ones: deeper aggregation captures longer-range structure, larger fanout preserves more neighborhood information than aggressive sampling, and heavyweight aggregators (LSTM, attention) beat simple means on accuracy. The paper’s motivating measurements show all four knobs hitting the 24 GB wall of an RTX 6000 — LSTM aggregators, a 4th layer, hidden size 512, fanout beyond 20, each an out-of-memory error on the 2.4-million-node ogbn-products dataset. The escape hatches are unpalatable: cut sampling rates (loses accuracy), or rent a fleet (32 T4 GPUs for one big model). Meanwhile the memory analysis reveals the enemy is not the model — parameters, activations, and optimizer state are the usual DNN suspects, but here input node features alone are 55% of the footprint. That points at a purely systems-level fix: keep the training math identical, and manage which slice of the input graph is resident when.
What people did before, and where it broke
The DNN world solved “batch too big for memory” long ago with
split a batch into K micro-batches,
run forward+backward on each sequentially, sum the gradients, and
apply one weight update at the end — the update is mathematically
identical to processing the whole batch at once, so convergence and
hyperparameters are untouched
For GNNs, the 1:1 assumption collapses. A GNN batch is a
the actual data structure of a GNN
mini-batch — layer by layer, a bipartite graph from source nodes
(neighbors whose features flow in) to destination nodes (nodes
being updated); stacked over k layers, the batch’s labeled output
nodes sit at the top of a k-level hierarchy of these bipartite
blocks
- Generic graph partitioners optimize the wrong objective. Metis, range, and random partitioning balance node counts or cut edges — none of them models duplication. Measured: partitioning a batch 8 ways with Metis carries 23.5% more redundant nodes than Betty, costing 27.6% more memory and 47.6% more time. At 64 micro-batches on ogbn-products, range/random/Metis each balloon to ~118 million extra input-node loads; the duplicated nodes are re-fetched from CPU and re-aggregated on GPU every epoch.
- Distributed-GNN partitioners solve a different problem. DistDGL, AliGraph, NeuGraph, and ROC partition the GRAPH across GPUs to minimize communication or balance load — they spend more hardware rather than less memory, and none models redundancy.
- Load imbalance sneaks in through the degree distribution. Real graphs are power-law: DGL’s in-degree bucketing groups same-degree nodes for dense aggregation, but the long tail lands in one overflow bucket (“bucketing explosion”). After a clean partition, one micro-batch can still hold ~19% more heavy nodes than its sibling — and peak GPU memory is set by the WORST micro-batch, which may still OOM while the average looks fine. Worse, memory is nonlinear in the partition count, so you can’t predict how many pieces are enough without trying.
How Betty works
Betty's pipeline (a Python module on DGL)
-
Sample the full batch
Standard DGL sampling produces the batch's multi-level bipartite structure: labeled output nodes plus their (fanout-bounded) sampled neighborhoods, layer by layer. Nothing changes for the user — same model, same hyperparameters.
-
Build the redundancy-embedded graph new
On the CPU, from the last layer's bipartite adjacency matrix A, compute C = A-transpose times A. Entry (i, j) counts exactly how many neighbors output nodes i and j SHARE. Drop non-output nodes and self-loops; what remains is a graph over output nodes whose edge weights are "how much would it cost to separate these two."
-
Min-cut the REG with Metis new
A K-way min-cut partition of the REG is exactly a minimum-redundancy split of the batch: cutting a heavy edge means separating two nodes that share many neighbors, which duplicates all of them. Betty reuses battle-tested Metis — the transformation, not the partitioner, is the novelty.
-
Estimate memory, re-partition until it fits new
An analytic model sums eight components per micro-batch — input features, labels, bipartite block edges, hidden-layer outputs, aggregator intermediates (for LSTM, in-degree times bucket size times hidden size times a PyTorch-specific constant), gradients, and Adam's two optimizer states — and takes the peak. If the largest micro-batch exceeds GPU capacity, try K+1 partitions and re-estimate; no trial-and-error training runs. Estimation error stays under 8% even for LSTM, the hardest case.
-
Generate micro-batch bipartites new
Each output-node group is expanded back into its own hierarchical bipartite (DGL blocks) with an index dictionary mapping micro-batch-local node IDs back to the raw graph — a bookkeeping layer DGL itself lacks.
-
Train with gradient accumulation
Micro-batches stream from CPU memory to the GPU one at a time; each runs forward and backward with the same weights; intermediates are freed after each backward; gradients accumulate; one optimizer step per full batch. The update equals full-batch training exactly — no hyperparameter changes, no accuracy loss.
The idea
Turn “redundancy” into something a partitioner can see. The core trick is a change of representation. Redundancy lives in the bipartite structure — how many source nodes two output nodes share — but partitioners only understand edge weights. The redundancy-embedded graph makes the invisible explicit: one matrix multiplication, C = AᵀA, computes for every pair of output nodes exactly the number of shared neighbors (entry c-ij sums over all nodes k that point to both i and j). Now the min-cut objective IS the redundancy objective: any partitioner that minimizes cut weight on the REG is minimizing cloned nodes in the micro-batches. Betty gets to reuse Metis — decades of engineering — pointed at a problem Metis was never designed for.
Never run out of memory by never guessing. REG partitioning fixes duplication but not the power-law imbalance, and micro-batch memory is nonlinear in K — so Betty refuses to guess. It prices each candidate micro-batch with an analytic model built from the things it can count before touching the GPU: nodes, edges, feature dimensions, per-in-degree bucket sizes, and known per-framework constants (an LSTM aggregator materializes 18 intermediate tensors per node in PyTorch). Peak memory is the maximum of the aggregator phase and the gradient phase, since aggregator intermediates die when backpropagation starts. The search loop is then trivial — K-way, estimate, K+1 if the biggest piece is over budget — and cheap, because estimation never executes training. A wrong guess costs a re-estimate, not an OOM crash after minutes of setup.
Keep the math sacred. Everything above is packaging; the training semantics are untouched. Gradients from the K micro-batches sum to exactly the full-batch gradient, so accuracy, convergence curves, and hyperparameters carry over unchanged — the paper shows convergence curves for 1, 2, 4, and 8 micro-batches lying on top of each other. This is the property mini-batch shrinking cannot offer, and it makes Betty a transparent systems knob rather than a modeling decision.
Partitioning one batch, step by step
A 196K-node ogbn-products batch through Betty
Output nodes = the batch's labeled nodes (destinations of the last layer). Source nodes = neighbors whose features must be loaded. REG = redundancy-embedded graph over output nodes. K = number of micro-batches. Numbers from the paper's ogbn-products experiments (2.4M nodes, 62M edges, 24 GB RTX 6000).
Step 1 of 7
The batch that will not fit
Full-batch training on ogbn-products means 196,615 output nodes — and through 3 layers of fanout (25, 35, 40), millions of input feature vectors. With a Mean aggregator it squeaks by; switch to LSTM, add a layer, or widen the hidden size, and the GPU's 24 GB is gone before the first update.
input features: 55% of total memory the knobs that improve accuracy all inflate this
Why it works
- Redundancy is a pairwise property of output nodes (shared neighbors), and AᵀA computes exactly that pairwise count → the REG’s min-cut objective coincides with minimum duplication → a generic, decades-optimized partitioner (Metis) solves a problem it was never designed for, at the cost of one sparse matrix multiply on the CPU.
- Gradient accumulation is exact, not approximate → the effective batch size — a hyperparameter that governs convergence — never changes → model scientists get memory relief without retuning anything, which is precisely what shrinking the mini-batch cannot offer.
- Peak memory is set by the LARGEST micro-batch, and micro-batch memory is analytically countable (nodes, edges, buckets, framework constants) → estimating before executing turns partition-count selection from trial-and-error OOM crashes into a fast arithmetic loop → under-8% estimation error keeps the chosen K safe in practice.
- Input features dominate GNN memory (55%) and live per-node → any reduction in duplicated nodes cuts memory AND CPU-GPU transfer AND aggregation compute together → Betty’s computation efficiency (nodes processed per second) stays flat as K grows, while redundancy-blind partitions degrade.
- CPU memory is large and idle during GPU training → holding the full batch host-side and streaming micro-batches converts the memory-capacity problem into a bandwidth cost → a single commodity GPU substitutes for a multi-GPU cluster on workloads that are memory-bound, not compute-bound.
What it costs
- Sequential micro-batches slow the epoch. Memory savings buy time: more micro-batches mean more redundancy (at K=64, +440% redundant nodes vs K=1) and lower GPU utilization. Epoch time roughly doubles from K=1 to K=16 on ogbn-products. The paper’s guidance: K of 4–8 balances memory against speed.
- Partitioning overhead. REG construction (a sparse AᵀA) and Metis run on the CPU per batch; the authors list optimizing this as future work.
- The estimator can still miss. Under-8% error is small, but a micro-batch running near capacity can theoretically OOM; folding the error margin into the re-partitioning loop is left as future work.
- Single-GPU by design. Betty trades time for memory on one device; it does not (yet) compose with multi-GPU data parallelism — the natural extension the authors name.
- Bookkeeping. Micro-batch-local vs full-batch vs raw-graph node indices require an extra mapping dictionary; a small but real implementation layer on top of DGL.
The evaluation
Experiments run on one 24 GB RTX 6000 with GraphSAGE and GAT over five graphs from tiny to large: Cora, Pubmed, Reddit (115M edges), ogbn-arxiv, and ogbn-products (2.4M nodes, 62M edges). The headline is qualitative: every OOM configuration from the motivation section runs — LSTM aggregators, 5 layers, hidden 512, fanout 800 — on the same GPU. Accuracy tables match full-batch DGL (the equivalence is mathematical, and the measured convergence curves for 2/4/8 micro-batches lie on the full-batch curve).
Ways to fit a too-big GNN batch on one GPU
| Property | Shrink the mini-batch | Micro-batch via generic partitioners (range / random / Metis) | Betty (REG + memory-aware) |
|---|---|---|---|
| Effect on training semantics | |||
| Redundant input-node loads (ogbn-products, 64 pieces) | |||
| Peak GPU memory | Reduced, but at the accuracy cost above; 0.50 GB at 64 batches | ||
| Training time | |||
| Knowing how many pieces you need | Trial and error |
Check yourself
Check your understanding
1. Gradient accumulation works perfectly for image classifiers: split the batch K ways, accumulate, update. Predict what happens when you apply the same K-way split to a GNN batch's output nodes.
2. Betty's redundancy-embedded graph sets the edge weight between output nodes i and j to entry (i, j) of AᵀA. Why does handing THIS graph to a standard min-cut partitioner like Metis minimize redundancy?
3. After REG partitioning, why does Betty still need a memory model and a re-partitioning loop — and why can't it just pick the piece count K once, up front?
4. Why does the paper reject the simplest fix — just train with smaller mini-batches — when it demonstrably also reduces memory?