Distributed Graph Neural Networks

"I asked one node for its opinion. To answer, it had to phone three machines, who each phoned three more, and by the time the gossip came back I had aggregated half the internet."

A GNN Layer That Underestimated Its Neighbors

The previous sections of this chapter treated the graph as something to analyze: partition it well (Section 13.2), iterate vertex-centric computations over it (Section 13.3), and run analytics such as PageRank and connected components across the partitions (Section 13.4). We now make the graph something to learn from. A graph neural network (GNN) is a model whose forward pass is itself a message-passing computation over the graph: it produces an embedding for every node by repeatedly mixing each node's representation with those of its neighbors. The same partitioning and vertex-centric machinery returns, but now it must carry gradients backward as well as features forward, and it must coexist with the data-parallel training loop from Chapter 10.

1. Message Passing: A GNN Layer Is Aggregate-Then-Update Intermediate

Almost every GNN in use, GraphSAGE, GCN, GAT, and their descendants, fits a single template called message passing. Each node carries a feature vector, and each layer rewrites that vector in two steps: aggregate a permutation-invariant summary of the neighbors' current vectors, then update the node's own vector from that summary. Let $h_v^{(\ell)}$ be the representation of node $v$ after layer $\ell$, with $h_v^{(0)}$ the input features, and let $\mathcal{N}(v)$ be the neighbors of $v$. One layer is

where $\mathrm{AGG}$ is a symmetric reducer (mean, sum, or max), $W^{(\ell)}$ is a shared learnable weight, and $\sigma$ is a nonlinearity. The weight $W^{(\ell)}$ is identical for every node, exactly the parameter sharing that makes the model generalize across a graph of any size. Crucially, the structure of the computation is dictated not by the model but by the graph: which features get aggregated for $v$ is decided by who $v$'s neighbors are. That is what makes a GNN different from the dense and embedding models of Chapter 11, and what makes it hard to distribute.

This receptive-field view, made precise in the key insight, is the hinge of the whole section. It explains both why GNNs are expressive (a node sees a growing context as layers deepen) and why they are expensive to distribute (that context crosses partition boundaries). Figure 13.5.1 draws the expansion for a single target node and shows where the boundary crossings appear.

2. Neighbor Explosion: Why Full-Neighbor Training Does Not Scale Intermediate

The receptive field grows with depth, and on real graphs it grows catastrophically. If every node had a fixed degree $D$, the $k$-hop neighborhood would hold on the order of $D^k$ nodes, already exponential. Real graphs are worse, because they are not regular: social, web, and citation graphs follow a power-law degree distribution with a few enormously high-degree hubs. A short path through even one hub sweeps a huge fraction of the graph into the receptive field. The practical consequence is blunt: on a graph with hundreds of millions of nodes, the 3-hop neighborhood of a single training node can already touch most of the graph, so computing one node's loss "exactly" with all neighbors means reading nearly the whole feature matrix. This is the neighbor explosion problem.

Naive full-neighbor training, where each layer aggregates over every actual neighbor, is therefore infeasible at scale for two compounding reasons. First, memory: the activations for one mini-batch include every node in the union of the batch's $k$-hop neighborhoods, which can dwarf the batch itself. Second, communication: when the graph is partitioned, every neighbor that lands on another machine is a feature you must pull across the network, and a hub neighbor multiplies those pulls. The standard remedy is to stop insisting on all neighbors and instead aggregate over a bounded random sample of them, which caps the receptive field and the fetch count per layer. That is distributed neighbor sampling, the subject of the very next section.

3. Distributing the GNN: Two Costs That Run at Once Advanced

Put the pieces together. To train a GNN on a graph too big for one machine, we partition the graph structure and the node-feature matrix across $P$ machines (Section 13.2), and we run data-parallel training, where each machine processes a shard of the training nodes and contributes to a shared gradient. A single training step now incurs two distinct communication costs, and a distributed-GNN system is largely an exercise in managing both.

The first cost is feature and message communication: during the forward and backward passes, a machine aggregating a local node's neighbors must fetch the representations of any neighbor that lives on another partition, and push the corresponding gradients back. This cost is graph-shaped; it depends entirely on how many edges cross partition boundaries, which is precisely what the graph-partitioning objective of Section 13.2 tries to minimize. The second cost is the familiar gradient synchronization: because the weight $W^{(\ell)}$ is shared across all nodes and all machines, the per-machine weight gradients must be combined with an all-reduce, exactly the collective that data-parallel training relies on (Section 10.5). The two costs interleave within every step: feature fetches gate the forward pass, the all-reduce gates the optimizer update.

The demonstration below isolates the first, graph-shaped cost on a tiny partitioned graph. It implements one message-passing layer in pure NumPy, runs it for the nodes owned by each of two partitions, and counts how many neighbor features had to be read from the other machine. The point is not the numbers themselves but the mechanism: a node's output literally cannot be computed without reaching across the partition boundary.

import numpy as np

# A small undirected graph: 8 nodes, partitioned across 2 machines.
# Partition 0 holds nodes {0,1,2,3}; partition 1 holds nodes {4,5,6,7}.
edges = [(0,1),(1,2),(2,3),(3,4),(4,5),(5,6),(6,7),(2,5),(1,6),(0,7)]
N, F_in, F_out = 8, 4, 3
part = {0:0,1:0,2:0,3:0, 4:1,5:1,6:1,7:1}   # node -> partition id

# Build adjacency (including self) as neighbor lists.
nbr = {v: {v} for v in range(N)}
for a, b in edges:
    nbr[a].add(b); nbr[b].add(a)

rng = np.random.default_rng(0)
H = rng.standard_normal((N, F_in))           # node feature matrix (lives sharded by partition)
W = rng.standard_normal((F_in, F_out))       # shared layer weight (replicated on every machine)

def message_passing_layer(target_nodes, owner_part):
    """One GNN layer for nodes owned by owner_part: mean-aggregate neighbor
    features, then linear transform. Counts how many neighbor features had to be
    fetched from a REMOTE partition."""
    out = np.zeros((len(target_nodes), F_out))
    remote_fetches = 0
    for i, v in enumerate(target_nodes):
        agg = np.zeros(F_in)
        for u in sorted(nbr[v]):
            agg += H[u]                       # read neighbor feature
            if part[u] != owner_part:         # neighbor lives on another machine
                remote_fetches += 1
        agg /= len(nbr[v])                    # mean aggregation
        out[i] = agg @ W                      # transform
    return out, remote_fetches

p0_nodes = [v for v in range(N) if part[v] == 0]
out0, fetch0 = message_passing_layer(p0_nodes, owner_part=0)
p1_nodes = [v for v in range(N) if part[v] == 1]
out1, fetch1 = message_passing_layer(p1_nodes, owner_part=1)

print("partition 0 owns nodes :", p0_nodes)
print("partition 1 owns nodes :", p1_nodes)
print("remote fetches (part 0):", fetch0)
print("remote fetches (part 1):", fetch1)
print("cross-partition fetches total:", fetch0 + fetch1)
print("output row for node 2  :", np.array2string(out0[2], precision=3))
print("node 2 neighbors       :", sorted(nbr[2]),
      "-> remote:", sorted(u for u in nbr[2] if part[u] != 0))

partition 0 owns nodes : [0, 1, 2, 3]
partition 1 owns nodes : [4, 5, 6, 7]
remote fetches (part 0): 4
remote fetches (part 1): 4
cross-partition fetches total: 8
output row for node 2  : [ 0.158 -0.794  0.3  ]
node 2 neighbors       : [1, 2, 3, 5] -> remote: [5]

Output 13.5.1 makes the dependency concrete: node 2 sits on partition 0, but one of its neighbors (node 5) sits on partition 1, so node 2's layer output cannot be produced without a cross-partition fetch. Scale this from eight nodes to eight hundred million, add a second and third layer so the receptive field reaches two and three hops, and the eight fetches become a torrent that the partitioning of Section 13.2, the sampling of Section 13.6, and the batching strategy of Section 13.7 all exist to tame. The all-reduce on $W$'s gradient, by contrast, is the same size no matter how the graph is cut, because it depends only on the model, not the graph.

import dgl, dgl.nn as dglnn, torch.nn as nn

class GraphSAGE(nn.Module):                 # two message-passing layers
    def __init__(self, in_f, hid, out_f):
        super().__init__()
        self.l1 = dglnn.SAGEConv(in_f, hid, aggregator_type="mean")
        self.l2 = dglnn.SAGEConv(hid, out_f, aggregator_type="mean")
    def forward(self, blocks, x):           # 'blocks' are the sampled k-hop subgraphs
        x = torch.relu(self.l1(blocks[0], x))
        return self.l2(blocks[1], x)

4. Where This Goes Next Beginner

We now have the model (message passing, aggregate-then-update), the obstacle (neighbor explosion on power-law graphs), and the two costs of distributing it (graph-shaped feature fetches plus the model-shaped gradient all-reduce). The rest of the chapter attacks the costs in order. Section 13.6 introduces distributed neighbor sampling, the direct remedy for neighbor explosion, which bounds the receptive field and therefore the cross-partition fetch count per layer. Section 13.7 then weighs mini-batch training, where each step works on a sampled subgraph, against full-graph training, where the entire graph participates in every step, a trade-off between memory, communication, and convergence that only makes sense once you have seen, as in Output 13.5.1, what a single layer over a partitioned graph actually costs.