ZEROGNN: GPU-Resident Sampling in GNN Training

• ZEROGNN is a system design for sampling-based GNN training that replaces CPU-mediated control with fully GPU-resident execution.
• It employs DRMB, DLM, and MFD to dynamically manage runtime metadata and preserve CUDA Graph replayability despite iteration variability.
• Experimental results show average end-to-end speedups of 5.28× and up to 17.69× gains in sampling performance, along with significant memory savings.

ZEROGNN is a systems design for sampling-based graph neural network training that makes a dynamically shaped mini-batch pipeline execute as a fully GPU-resident, replayable workflow by removing the CPU from the metadata-driven control loop (Gong et al., 28 May 2026). It targets a specific bottleneck in neighbor-sampled GNN training: runtime-generated metadata such as sampled vertex counts, edge counts, and frontier sizes determine tensor allocation sizes and kernel launch parameters, so conventional systems repeatedly export scalars from GPU to CPU, synchronize, and relaunch subsequent work from the host. ZEROGNN replaces that host-mediated loop with device-resident metadata, device-mediated dynamic execution inside a fixed launch structure, and a statistically tight execution envelope that restores CUDA Graph replayability under iteration-to-iteration variation (Gong et al., 28 May 2026).

1. Problem setting and system model

ZEROGNN is defined around standard mini-batch GNN training with neighbor sampling. For an $N$-layer GNN, a mini-batch $V_s^1$ is expanded hop by hop into source and destination sets $V_s^i, V_d^i$ and per-hop edge sets whose cardinalities vary across iterations (Gong et al., 28 May 2026). Those runtime sizes determine relabeling tables, feature and label gathering footprints, temporary scratch space, and launch dimensions. The pipeline considered in the paper comprises subgraph sampling, ID translation or relabeling, feature and label copy, and then GNN computation on the sampled subgraph (Gong et al., 28 May 2026).

The central systems claim is that the dominant bottleneck in such workloads is often not GPU arithmetic itself but host-device orchestration overhead caused by metadata-dependent control (Gong et al., 28 May 2026). In the conventional path, neighbor sampling kernels produce counts on the GPU; the counts are copied back to the CPU as scalars; the host decides allocation sizes and launch parameters; and the host then relaunches downstream kernels. Because this pattern recurs within each hop and across hops, it introduces persistent synchronization, CPU logic on the critical path, and GPU idle time (Gong et al., 28 May 2026).

The paper formalizes the dependence of launch geometry on runtime work size $N \in \{ |V|, |E| \}$ with

$$\text{grid} \approx \left\lceil \frac{N}{T} \right\rceil,$$

where $T$ is the per-block work quota (Gong et al., 28 May 2026). It also defines the mechanism-level metric of GPU Execution Fraction as

$$\text{GPU Execution Fraction} = \frac{\text{GPU Time}}{\text{End-to-End Training Runtime}},$$

while noting that the printed formula in the paper is typographically malformed (Gong et al., 28 May 2026). A low value indicates runtime dominated by orchestration and idle gaps rather than active GPU execution.

A concrete characterization reported for DGL on Reddit at batch size 128 is that only 45% of end-to-end runtime is GPU execution, implying that 55% is effectively overhead or idle time (Gong et al., 28 May 2026). The runtime breakdown on Reddit with GraphSAGE is given as 26% sampling, 66% training, and 8% feature or label copy, which the paper uses to motivate optimization of sampling and training as the dominant stages (Gong et al., 28 May 2026).

2. Host-mediated dependency barriers

The paper describes the conventional dependency pattern as a host-mediated dependency barrier, abbreviated HMDB (Gong et al., 28 May 2026). Its per-step structure is “Produce $\rightarrow$ Export $\rightarrow$ Consume $\rightarrow$ Relaunch”: GPU kernels produce metadata; the metadata is exported to the CPU; the CPU consumes it to provision memory and compute launch parameters; and the CPU relaunches follow-up GPU work (Gong et al., 28 May 2026). This loop occurs repeatedly because the frontier of hop $V_s^1$0 depends on the output of hop $V_s^1$1 (Gong et al., 28 May 2026).

The bottleneck is particularly acute in modern deep learning stacks because tensor allocation remains host-centric. The paper states that framework allocators such as PyTorch tensor creation expect allocation sizes as CPU-resident scalars (Gong et al., 28 May 2026). Thus even when the computation itself is already on the GPU, the control plane remains partially anchored on the CPU. The resulting host-device orchestration overhead, abbreviated HDOO, remains roughly constant even when sampled subgraphs become smaller, so it dominates more strongly when per-iteration GPU kernels are short (Gong et al., 28 May 2026).

ZEROGNN is positioned against two prior families of mitigation. CUDA Graphs reduce per-iteration launch overhead only when execution structure, memory addresses, and launch topology are stable enough to be captured once and replayed many times; ordinary sampling-based GNN training violates those assumptions because later work depends on metadata produced by earlier GPU kernels (Gong et al., 28 May 2026). GPU dynamic parallelism can move some launch decisions onto the device, but the paper argues that a pilot-kernel-plus-nested-launch strategy incurs substantial overhead and scheduling cost, so it may solve functionality without solving efficiency (Gong et al., 28 May 2026).

This framing is significant because it shifts the optimization target. Rather than treating the workload as primarily limited by kernel quality or language-level overhead, ZEROGNN identifies metadata-driven host involvement as the root systems problem in this regime (Gong et al., 28 May 2026). A plausible implication is that similar benefits may arise in other dynamic workloads whose execution and allocation decisions are driven by device-generated metadata, although that broader generalization is not directly benchmarked in the paper.

3. Core architecture: DRMB, DLM, and MFD

ZEROGNN consists of three tightly connected components: Device-Resident Metadata Buffer (DRMB), Device-Side Launch Mediation (DLM), and Metadata-Free Dispatcher (MFD) (Gong et al., 28 May 2026). Together they produce a fixed host launch structure while resolving dynamic behavior entirely on the GPU.

Component	Function	Key effect
DRMB	Stores runtime metadata in preallocated GPU memory	Removes mid-iteration CPU scalar export
DLM	Uses conservative overprovisioning inside a fixed launch skeleton	Preserves replayable launch structure
MFD	Preallocates buffers under a conservative yet tight execution envelope	Restores stable addresses for replay

DRMB is the foundation. Instead of materializing runtime metadata such as sampled $V_s^1$2, $V_s^1$3, and frontier sizes as CPU scalars, ZEROGNN stores them in preallocated GPU memory and passes device pointers downstream (Gong et al., 28 May 2026). Kernels dereference those pointers directly on-device. Because the number of sampling hops is fixed by model depth, the number of metadata slots is fixed as well, so the buffers can be allocated once at initialization and reused at stable addresses every iteration (Gong et al., 28 May 2026). This matters for CUDA Graph compatibility because capture cannot include CPU-GPU synchronizations in the middle of the graph and replay expects stable pointer identities (Gong et al., 28 May 2026).

DLM addresses dynamic execution inside a fixed launch skeleton. The paper distinguishes two forms of launch dynamism: dynamic grid-size allocation, where the number of blocks depends on runtime $V_s^1$4 or $V_s^1$5, and dynamic numbers of kernel calls, such as recursive or multi-round scan or prefix-sum routines whose execution depth depends on input size (Gong et al., 28 May 2026). Rather than computing exact launch counts on the host each iteration or using nested device-side launches, DLM relies on conservative overprovisioning: the host issues a launch with an upper-bound grid size fixed at capture time, each kernel reads the true runtime size from DRMB, and excess threads or blocks exit via ordinary boundary checks (Gong et al., 28 May 2026).

The paper argues that this overprovisioning is efficient because overallocated blocks that fail the bounds check incur negligible cost; experimentally, a state-of-the-art SpMM kernel maintains near-constant runtime even when the grid is overallocated by +20% to +180% on Reddit and OGBN-Products (Gong et al., 28 May 2026). In effect, DLM transforms launch-time dynamism into in-kernel specialization (Gong et al., 28 May 2026). For control-flow dynamism, the host-side launch chronology remains fixed while device-visible state determines whether additional rounds perform useful work (Gong et al., 28 May 2026).

MFD solves memory provisioning. A naive strategy is to allocate for the theoretical maximum sampled subgraph implied by batch size and fan-outs. If $V_s^1$6 is the mini-batch size, $V_s^1$7 the number of hops, and $V_s^1$8 the fan-out at hop $V_s^1$9, the paper gives the worst-case vertex bound

$V_s^i, V_d^i$0

approximately $V_s^i, V_d^i$1 when $V_s^i, V_d^i$2 (Gong et al., 28 May 2026). The paper names this pessimistic baseline MaxSG or max reserved memory (Gong et al., 28 May 2026). Because real graphs have degree limits, repeated sampling of popular neighbors, and deduplication across hops, this bound is far looser than practical sampled sizes (Gong et al., 28 May 2026).

4. Execution envelope, replayability, and overflow handling

MFD replaces MaxSG with what the paper calls a conservative yet tight execution envelope $V_s^i, V_d^i$3 (Gong et al., 28 May 2026). The key empirical and probabilistic observation is that the deduplicated sampled size is highly stable across iterations despite stochastic neighbor sampling (Gong et al., 28 May 2026). The deduplicated sampled vertex count is formalized as

$V_s^i, V_d^i$4

where $V_s^i, V_d^i$5 is an indicator equal to 1 if vertex $V_s^i, V_d^i$6 is sampled at least once and 0 otherwise (Gong et al., 28 May 2026). Under sampling with replacement,

$V_s^i, V_d^i$7

with $V_s^i, V_d^i$8 modeled as proportional to degree and $V_s^i, V_d^i$9 the total number of draw opportunities; the source text notes that the PDF is truncated around this formula but that the intended exponent is clearly $N \in \{ |V|, |E| \}$0 (Gong et al., 28 May 2026). Since the $N \in \{ |V|, |E| \}$1 are non-identical Bernoulli variables, $N \in \{ |V|, |E| \}$2 is Poisson-binomial, and for large graphs the paper approximates it as normal to obtain concentration bounds (Gong et al., 28 May 2026).

The central formal statement is Lemma 1: $N \in \{ |V|, |E| \}$3 where $N \in \{ |V|, |E| \}$4 is the coefficient of variation of $N \in \{ |V|, |E| \}$5, and $N \in \{ |V|, |E| \}$6 is a Gaussian quantile adjusted for confidence level $N \in \{ |V|, |E| \}$7 over $N \in \{ |V|, |E| \}$8 repeated iterations (Gong et al., 28 May 2026). The result states that the normalized fluctuation band is narrow when the coefficient of variation is small, and MFD uses this concentration behavior to define a fixed per-workload envelope for memory and launch resources (Gong et al., 28 May 2026).

With $N \in \{ |V|, |E| \}$9, all dynamic buffers are allocated once during initialization and reused at stable addresses (Gong et al., 28 May 2026). Runtime kernels still consume true sizes from DRMB, so excess capacity is harmless and out-of-range threads or blocks exit (Gong et al., 28 May 2026). In this way, MFD restores the two CUDA Graph prerequisites that dynamic workloads usually violate: stable addresses and stable launch topology (Gong et al., 28 May 2026).

End-to-end execution proceeds in three phases. During warm-up, ZEROGNN runs iterations without CUDA Graph replay in order to initialize pools and materialize buffers under the chosen envelope (Gong et al., 28 May 2026). During capture, it records a steady-state iteration into a CUDA Graph object (Gong et al., 28 May 2026). During replay, each training iteration becomes essentially a single CPU graph launch plus GPU-side execution (Gong et al., 28 May 2026). To satisfy stable I/O pointer requirements, minibatch features and labels are copied into statically allocated buffers before replay (Gong et al., 28 May 2026).

The paper also includes an overflow-safe fallback because $$\text{grid} \approx \left\lceil \frac{N}{T} \right\rceil,$$0 is a statistical guarantee rather than an absolute one (Gong et al., 28 May 2026). If overflow is detected, ZEROGNN replays a cached safe graph associated with a more conservative envelope rather than performing host-driven reprovisioning (Gong et al., 28 May 2026). The source description notes that the paper’s discussion of “chaining an identical graph twice” is awkwardly phrased, but the intended operational point is that correctness is preserved without ad hoc host reconfiguration (Gong et al., 28 May 2026).

5. Experimental evaluation

The experimental platform is a single node with two NVIDIA A100 GPUs, 80 GB each, running CUDA 11.7, with GraphSAGE as the model following GraphSAGE or Cluster-GCN settings (Gong et al., 28 May 2026). Datasets include Cora, Hollywood, LiveJournal, OGBN-Products, Reddit, Orkut, and OGBN-papers100M (Gong et al., 28 May 2026). The paper reports the following dataset statistics (Gong et al., 28 May 2026).

Dataset	Vertices	Edges
Cora	2,708	10,858
Hollywood	1,069,127	112,613,308
LiveJournal	4,847,571	137,987,546
OGBN-Products	2,449,029	123,718,280
Reddit	232,965	229,231,784
Orkut	3,072,627	234,370,166
OGBN-papers100M	111,059,956	1,615,685,872

The same table in the paper also lists feature dimensions and class counts: Cora has 1,433 features and 7 classes; Hollywood 150 features and 7 classes; LiveJournal 150 features and 7 classes; OGBN-Products 100 features and 47 classes; Reddit 602 features and 41 classes; Orkut 150 features and 7 classes; and OGBN-papers100M 128 features and 172 classes (Gong et al., 28 May 2026). Baselines are DGL, GraphPy, MariusGNN in some experiments, and an internal CUDA dynamic parallelism baseline denoted CU-DPI (Gong et al., 28 May 2026).

The main quantitative results are strong. Across datasets, ZEROGNN achieves average end-to-end speedups of $$\text{grid} \approx \left\lceil \frac{N}{T} \right\rceil,$$1 over DGL, $$\text{grid} \approx \left\lceil \frac{N}{T} \right\rceil,$$2 over GraphPy, and $$\text{grid} \approx \left\lceil \frac{N}{T} \right\rceil,$$3 over CU-DPI (Gong et al., 28 May 2026). For sampling-only runtime, the reported gains are $$\text{grid} \approx \left\lceil \frac{N}{T} \right\rceil,$$4 over DGL in the abstract and $$\text{grid} \approx \left\lceil \frac{N}{T} \right\rceil,$$5 in the body, $$\text{grid} \approx \left\lceil \frac{N}{T} \right\rceil,$$6 over GraphPy, and $$\text{grid} \approx \left\lceil \frac{N}{T} \right\rceil,$$7 over CU-DPI (Gong et al., 28 May 2026). In both the sampling stage and end-to-end training, ZEROGNN sustains approximately 100% GPU execution fraction, whereas DGL and GraphPy remain substantially lower (Gong et al., 28 May 2026).

Memory efficiency is another central result. Relative to DGL, ZEROGNN uses up to $$\text{grid} \approx \left\lceil \frac{N}{T} \right\rceil,$$8 less memory, and relative to MaxSG it delivers about $$\text{grid} \approx \left\lceil \frac{N}{T} \right\rceil,$$9 memory savings on average (Gong et al., 28 May 2026). Its memory consumption is reported as comparable to the “optimal” GraphPy approach that allocates precisely according to runtime metadata, which the paper presents as empirical validation of the execution-envelope strategy (Gong et al., 28 May 2026).

The paper also reports controlled sensitivity studies. On Reddit, across batch sizes, ZEROGNN averages $T$0 speedup over DGL and $T$1 over GraphPy; at batch size 64 it reaches $T$2 over DGL, and at batch size 4096 it still provides $T$3 (Gong et al., 28 May 2026). Across sampling depths, speedup decreases as depth increases because GPU work grows while orchestration overhead stays relatively constant: $T$4 at 2 layers and $T$5 at 5 layers (Gong et al., 28 May 2026). For a simulated OGBN-papers100M setup in which topology remains on GPU while the full feature-table bottleneck is abstracted away, ZEROGNN beats GraphPy by $T$6 to $T$7, with an average of $T$8, while DGL runs out of memory (Gong et al., 28 May 2026).

For multi-GPU strong scaling, the paper argues that host overhead becomes an even stronger limiter because per-GPU subgraphs shrink while per-worker orchestration overhead does not (Gong et al., 28 May 2026). On a 2-GPU machine, ZEROGNN achieves up to $T$9 speedup over GraphPy depending on batch size, and its own 1-to-2 GPU strong scaling is $$\text{GPU Execution Fraction} = \frac{\text{GPU Time}}{\text{End-to-End Training Runtime}},$$0 to $$\text{GPU Execution Fraction} = \frac{\text{GPU Time}}{\text{End-to-End Training Runtime}},$$1, close to ideal and stable across batch sizes (Gong et al., 28 May 2026).

6. Significance, limitations, and nomenclature

The main contributions identified in the paper are threefold. First, it isolates metadata-driven host involvement, rather than generic framework overhead alone, as the root systems problem in sampling-based GNN training (Gong et al., 28 May 2026). Second, it proposes a concrete architecture—DRMB, DLM, and MFD—that removes the CPU from the metadata-driven control loop while preserving correctness under iteration-to-iteration variation (Gong et al., 28 May 2026). Third, it shows that this restructuring restores CUDA Graph replayability for a class of workloads previously considered too dynamic for effective capture and replay (Gong et al., 28 May 2026).

The limitations are also explicit. The formal development in the main text is incomplete in places: some equations are typographically truncated, no explicit pseudocode is provided, and the safe-graph overflow discussion is underexplained (Gong et al., 28 May 2026). The approach relies on the stability of deduplicated sampled sizes; if a workload exhibited much higher variance or adversarially bursty subgraph sizes, the envelope would plausibly need to be looser, reducing memory efficiency (Gong et al., 28 May 2026). ZEROGNN also does not optimize GPU kernels themselves, so relative speedups shrink when GPU computation becomes dominant at larger batch sizes or deeper models (Gong et al., 28 May 2026).

The term “ZEROGNN” should be distinguished from several similarly named but conceptually different lines of work. “ZeroG: Investigating Cross-dataset Zero-shot Transferability in Graphs” studies cross-dataset zero-shot node classification with a language-model-centered graph framework for disjoint graphs and disjoint label spaces $$\text{GPU Execution Fraction} = \frac{\text{GPU Time}}{\text{End-to-End Training Runtime}},$$2 and $$\text{GPU Execution Fraction} = \frac{\text{GPU Time}}{\text{End-to-End Training Runtime}},$$3 (Li et al., 2024). “A Review on Zeroing Neural Networks” surveys zeroing neural networks for time-varying optimization and control, where the defining mechanism is an error-zeroing dynamics such as $$\text{GPU Execution Fraction} = \frac{\text{GPU Time}}{\text{End-to-End Training Runtime}},$$4, rather than graph learning or GNN systems (Jiang et al., 1 Jul 2025). “Design of a SiPM-on-Tile ZDC for the future EIC and its Performance with Graph Neural Networks” concerns a zero-degree calorimeter reconstructed with graph neural networks, a detector-and-reconstruction problem unrelated to metadata-driven sampling pipelines in GNN training (Milton et al., 2024).

Within systems research on GNN training, however, ZEROGNN denotes a specific thesis: highly dynamic sampling-based GNN pipelines need not remain host-mediated. Once runtime metadata remains on-device, dynamic work is mediated within fixed launches, and memory is provisioned through a conservative yet tight execution envelope, the pipeline can become almost purely GPU-executed and replayable under CUDA Graphs (Gong et al., 28 May 2026). This suggests a broader methodological lesson: in dynamic deep learning workloads, the control plane can be as important a bottleneck as the kernels themselves.