#  : Automatically Pushing the Envelope of Multi-GPU System
for Billion-Scale GNN Training

Jie Sun1, Li Su2, Zuocheng Shi1, Wenting Shen2, Zeke Wang1
Lei Wang2, Jie Zhang1, Yong Li2, Wenyuan Yu2, Jingren Zhou2, Fei Wu1,3
1 Collaborative Innovation Center of Artificial Intelligence, Zhejiang
University, China
2 Alibaba Group
3 Shanghai Institute for Advanced Study of Zhejiang University, China

###### Abstract

Graph neural network(GNN) has been widely applied in real-world applications,
such as product recommendation in e-commerce platforms and risk control in
financial management systems. Several cache-based GNN systems have been built
to accelerate GNN training in a single machine with multiple GPUs. However,
these systems fail to train billion-scale graphs efficiently, which is a
common challenge in the industry. In this work, we propose , a system that
automatically pushes the envelope of multi-GPU systems for accelerating
billion-scale GNN training. First, we design a hierarchical graph partitioning
mechanism that significantly improves the multi-GPU cache performance. Second,
we build a unified multi-GPU cache that helps to minimize the PCIe traffic
incurred by caching both graph topology and features with the highest hotness.
Third, we develop an automatic cache management mechanism that adapts the
multi-GPU cache plan according to the hardware specifications and various
graphs to maximize the overall training throughput. Evaluations on various GNN
models and multiple datasets show that supports training billion-scale GNNs in
a single machine and significantly outperforms the state-of-the-art cache-
based systems on small graphs.

##  1 Introduction

Graph neural networks (GNNs), such as [22, 16, 39, 10, 49, 8], are a class of
deep learning algorithms that learn the low-dimensional embedding using the
structure and attribute information of graphs. The learned embedding can be
further used in machine-learning tasks including node classification and link
prediction. GNNs have been successfully applied in many real-world
applications, such as recommendation systems in e-commerce platforms, fraud
detection and risk control in financial management, and molecular property
prediction in drug development [47, 25, 48, 36, 13]. Systems such as DGL [41],
PyG [31], and Graph-Learn [54] are proposed to ease the development and
training of GNN models.

It is common to apply GNNs over large-scale graphs in industrial scenarios.
For example, in Alibaba’s Taobao recommendation system, the user behavior
graph contains more than one billion vertices and tens of billions of edges
[54]. In addition, as graphs are often skewed, it is infeasible to aggregate
all neighboring vertices when training a specific vertex. Sampling-based mini-
batch training, such as GraphSAGE [16], is proposed to extend GNN training to
very large graphs. In the sampling-based GNN training, there are two key steps
of data preparations before training a batch: (1) sampling the multi-hop sub-
graph for each vertex in the batch, and (2) extracting the features of
vertices in sampled sub-graphs. Systems such as DGL [41] and PyG [31] store
the graph data in the CPU memory, prepare the training data of mini-batches
using CPUs, and utilize GPUs for model training. As this approach requires
transferring the sampled sub-graphs and high-dimension feature data to the GPU
for every batch, the end-to-end training throughput is severely limited by the
CPU-GPU data transferring bandwidth [23, 46]. In addition, the throughput of
graph sampling using CPU is often insufficient to keep up with the throughput
of GPU training, especially in multi-GPU machines.

Several cache-based approaches have been proposed to speed up GNN training
[23, 33, 46, 29]. As it is the feature data that accounts for a majority of
the CPU-GPU data transferring, caching the features of frequently accessed
vertices in GPU can significantly reduce the amount of transferred data. To
improve the throughput of graph sampling, GPU-based sampling has also been
adopted in GNN systems [41, 46, 33].

We identify that existing approaches face severe limitations or performance
issues in multi-GPU training, particularly when the graph is large. First, the
multi-GPU cache scalability of existing cache-based systems is poor. Some
cache-based GNN systems [46, 33] shuffle the training set across all GPUs and
replicate an identical feature cache across all GPUs or NVLink
cliques111NVLink clique denotes a group of GPUs where each pair of GPUs are
connected with NVLink. to facilitate data parallel training. The cache
capacity is constrained by the memory of a single GPU or NVLink clique (an
NVLink clique only consists of two GPUs in some multi-GPU architectures),
resulting in poor cache performance when scaling up the number of GPUs (see
the experiment in Figure 2). PaGraph [23] partitions the graph using a self-
reliant algorithm and caches nodes with the highest in-degree for different
partitions in different GPUs, trying to make use of data locality inside each
partition. As partitions in PaGraph include the complete L-hop neighbors of
their training vertices, there is a significant overlap between the caches of
different partitions, resulting in the same duplication issue as the
aforementioned cache-based GNN systems. Second, when adopting GPU-based graph
sampling, existing systems manage the graph topology in a very coarse-grained
manner: the topology has to be completely stored in a single GPU [46, 41, 33]
or in the CPU memory [41, 33]. The former approach puts a hard limit on the
graph scale, and further squeezes the cache capacity for features. The latter
storing the topology in the CPU and accessing it from GPU would result in very
low utilization of the PCIe bandwidth, as the data access of graph sampling is
usually random and fine-grained.

This paper presents , a GNN system that fully explores the hardware
capabilities of modern multi-GPU servers for training large-scale graphs in a
single machine. proposes two key designs to fully exploit the memory space of
multi-GPUs for feature and topology cache. First, to avoid cache replication,
we propose NVLink-aware hierarchical graph partitioning technique that helps
scale the cache on multi-GPU memory efficiently according to the specific
hardware structure. first partitions the graph with minimal edge-cut and
assigns each partition exclusively to an NVLink clique, and then uses hash
partition to further map the training vertices to GPUs inside each NVLink
clique. Second, we propose a hotness-aware unified cache that manages both the
feature and topology cache in a vertex-centric data structure. We enable an
NVLink-enhanced cache space for the unified cache and prioritize the topology
and features with the highest hotness to fill the cache, so as to improve the
multi-GPU memory utilization.

The above designs pose a new challenge to . Given a fixed size of GPU memory,
it is hard to manually decide the optimal fractions of topology and feature
cache such that the overall training throughput is maximized. To solve the
challenge, we propose an automatic cache management mechanism. Specifically,
we build a cost model in the mechanism to evaluate the key factor to the
overall throughput, i.e., PCIe traffic, of both graph sampling and feature
extraction in the training phase, which is used to guide the allocations of
cache spaces for graph topology and feature. Overall, the three key designs in
enable automatic caching optimization and full utilization of hardware
capability of various modern GPU servers. Experiments show that can outperform
state-of-the-art cache-based GNN systems up to $4.32\times$.

In summary, the contributions of this paper include:

  1. 1.

We propose an NVLink-aware hierarchical graph partitioning technique that
helps minimize cache replication between NVLink cliques and extends the
threshold of cache capacity beyond the limit of an NVLink clique.

  2. 2.

We propose a hotness-aware unified cache to store topology and features with
the highest hotness in GPU memory, so as to improve the GPU memory
utilization.

  3. 3.

We present an automatic cache management mechanism that searches for the
optimal cache plan without requiring extra knowledge of hardware
specifications and GNN performance details from users.

  4. 4.

We implement that fully explores the hardware capabilities of multi-GPU
systems targeting billion-scale GNN training, not supported by existing cache-
based GNN systems, in a single server.

##  2 Preliminaries

In this section, we introduce the basic concept of GNN and the workflow of
mini-batch GNN training.

###  2.1 Graph Neural Networks

Given a graph $G=(V,E)$, where each vertex is associated with a vector of data
as its features $X_{v},v\in V$, Graph Neural Networks(GNNs) learn a low-
dimensional embedding for each target vertex by stacking multiple GNNs layers
$L$. For each layer $l,l\in L$, vertex $v$ updates its activation by
aggregating features or hidden activations of its neighbors $N(v),v\in V$:

$\begin{split}a_{v}^{l}&=AGGREGATE^{l}({h_{u}^{l-1}|u\in N(v)})\\\
h_{v}^{l}&=UPDATE^{l}(a_{v}^{l},h_{v}^{l-1})\end{split}$ (1)

###  2.2 Mini-batch GNNs Training

Figure 1: The workflow of 2-hop GraphSAGE training.

Mini-batch training is a practical solution for scaling GNN training on very
large graphs. Neighbor sampling is used to generate mini-batches, allowing
sampling-based GNN models to handle unseen vertices. For example,
GraphSAGE[16] samples multiple hops of neighbors for training as shown in
Figure 1. The workflow of GraphSAGE training follows a vertex-centric
computation paradigm including the following steps: 1, selecting a mini-batch
of training vertices from the training set. 2, uniformly sampling the multiple
hops of fixed-size neighbors for each training vertex. 3, extracting the
features of the sub-graph consisting of the training vertices and their
neighbors to generate the mini-batch training data. Finally, performing
AGGREGATE and UPDATE according to Equations 1, as well as executing the
forward and backward propagation to update the model parameters.

##  3 Observation and Motivation

When training large-scale graphs whose size exceeds the capacity of GPU memory
on a multi-GPU server, the major performance bottleneck becomes the data
movement from CPU to GPUs under the constraint of PCIe bandwidth. To this end,
existing works [41, 46, 33] intend to relieve the PCIe bandwidth bottleneck by
caching the hottest graph features on GPU memory. Though these cache-based
approaches significantly reduce PCIe traffic, we still identify two issues of
these existing cache-based GNN systems when training large-scale graphs: 1)
poor multi-GPU cache scalability, and 2) coarse-grained GPU memory management
for graph topology. In the following, we discuss each issue and the
corresponding observation that motivates the design of .

###  3.1 Multi-GPU Cache Scalability

As feature extraction occupies most of the data transferring from CPU to GPU,
cache-based systems like GNNLab [46] maintain a global feature cache for
vertices which are more frequently accessed via a pre-sampling phase. As
training vertices are globally shuffled among all training GPUs, GNNLab
replicates this cache across all GPUs involved in model training. Since a
single GPU’s memory space is quite limited, the fraction of cached features
would inevitably become lower when the graph size grows, resulting in a lower
cache hit ratio even on multi-GPU servers. To increase the cache capacity, the
cache mechanism in Quiver [33] leverages high-speed NVLinks to support inter-
GPU cache between NVLink-connected GPUs. Different from GNNLab, Quiver
replicates feature cache between NVLink cliques and averagely hashes the
features among GPUs in the same NVLink clique. However, this mechanism could
still lead to poor cache scalability, especially when the NVLink clique is
relatively small. E.g., the Siton server used in Table 1 has $4$ NVLink
cliques, each of which contains only $2$ GPUs. Figure 2 illustrates that, in
systems like Quiver, the PCIe transactions incurred by CPU-GPU data
transferring stop decreasing when the number of GPUs is larger than the size
of NVLink clique. This result shows that the cache performance in the above
GNN systems cannot scale well with the increasing number of GPUs in modern
servers.

To solve the scalability issue incurred by cache replication, PaGraph [23]
partitions the graph in a self-reliance approach and maintains an independent
cache for each partition using an in-degree-based metric on different GPUs. To
train an L-layer GNN model, PaGraph extends every partition with redundant
vertices and edges to include all the L-hop neighbor vertices for each train
vertex in this partition. Each GPU only trains its own partition and
synchronizes its local gradients periodically to update the model. However,
the inclusion of the L-hop neighbor vertices leads to heavily duplicated cache
contents on all GPUs. Figure 2 shows that the PaGraph exhibits a similar cache
performance as GNNLab which adopts the cache replication mechanism. We further
implement a PaGraph-plus design to alleviate the cache duplication issue in
PaGraph. Specifically, we replace the graph partitioning algorithm in PaGraph
with the XtraPulp [34] algorithm that minimizes edge-cuts between partitions
and adopts a pre-sampling-based hotness metric to select the vertex features
to be cached. Although PaGraph-plus achieves higher cache hit rates compared
to PaGraph, the cache hit rates on different GPUs are very unbalanced as
different partitions have various graph distributions. Figure 3 illustrates
the load imbalance issue of PaGraph-plus by measuring the cache hit rates of
eight GPUs. We observe that the hit rate varies by up to 17%.

To sum up, for systems that globally shuffle the training vertices among GPUs
in every iteration, such as GNNLab and Quiver, cache replication cannot be
completely eliminated as each GPU may randomly access any vertex in the entire
graph. Whereas the high-speed NVLinks between GPUs can be used to reduce the
replication factor and expand the cache capacity. For systems that locally
shuffle training vertices in each partition to produce mini-batches for
different GPUs, such as PaGraph, the cache replication problem could be
alleviated only when the model layer is small (e.g., less than $2$). PaGraph-
plus can further reduce cache duplication but faces another issue of
unbalanced cache hit rates among GPUs.

(a) 2 GPUs per NVLink clique

(b) 4 GPUs per NVLink clique

Figure 2: Comparing the cache scalability of cache-based GNN systems using the
Products [17] dataset and 2-hop GraphSAGE [16] model in terms of normalized
CPU-GPU PCIe transactions. The cache ratio is set to 5% $|V|$ on every GPU.
The tested platforms are Siton (a) and DGX-V100 (b) servers, as shown in Table
1. Figure 3: Cache hit rates of different systems in a server with 8 GPUs. The
cache ratio is set to 5% $|V|$ on every GPU. The graph sampling follows the
2-hop GraphSAGE [16] model’s setting using the Products [17] dataset. “NVx"
means utilizing NVLink clique with x GPUs.

Observation O1: Graph partitioning can be suitably guided by hardware
structure.  Different from Quiver, GNNLab, PaGraph, and PaGraph-plus do not
take advantage of the NVLink between GPUs, which is a common capability in
modern multi-GPU servers. As GPUs inside the same NVLink clique can access
each other’s memory via the low-latency high-throughput NVLink, an NVLink
clique can hold the entire cache of a partition, which can be randomly sliced
and averagely allocated among GPUs inside a clique. This hardware-coherent
design can balance the cache hit ratios between intra-clique GPUs. As the
number of partitions is reduced to the number of NVLink cliques, it is more
likely that the partitions follow a similar distribution (see the cache hit
rate distribution of in Figure 3). Inspired by O1, we propose an NVLink-aware
hierarchical partitioning to preserve multi-GPU cache scalability in (Section
4.1).

(a) PCIe throughput vs. payload size

(b)

Figure 4: (a) The PCIe 3.0 throughput under different payload sizes of PCIe
requests. (b) The PCIe traffic reduction rate for Paper100M with the growth of
the cache capacity. The cache is on a single GPU and selected after pre-
sampling.

###  3.2 Coarse-grained GPU Memory Management for Graph Topology

In multi-GPU servers, the throughput of CPU-based graph sampling may not catch
up with the throughput of GPU-based training. To improve the end-to-end
training throughput, recent GNN systems [41, 46, 33] adopt GPUs to accelerate
graph sampling. We observe that all these systems apply a very coarse-grained
memory management mechanism for graph topology. In particular, they store the
entire graph topology either in CPU memory or in a single GPU, depending on
the size of graph topology: the graph topology is stored in CPU memory when it
is too large or exceeds the capacity of a single GPU. The approach of storing
the entire graph topology in a single GPU sets a hard limit on the scale of
the graph. For example, a V100 GPU with 16GB memory can store at most $4$
billion edges [16] without considering any other memory usage of feature cache
and model training. When storing the graph topology in CPU memory, GPUs can
directly access the graph topology via a unified virtual memory address (UVA
[27]) technique. While the data access pattern of graph sampling is usually
random and fine-grained. E.g., Figure 4a shows that the PCIe throughput of
graph sampling is much lower than feature extraction. A large number of
sampling PCIe transactions with small payload sizes will increase the CPU-GPU
PCIe contention and lead to low bandwidth utilization.

Observation O2: The access of graph topology is skewed as graph features.
Existing cache-based GNN systems [46, 23, 33] only maintain feature cache in
GPU to reduce the CPU-GPU communication costs. However, we observe that the
performance gain of the per-unit feature cache decreases once the cache
capacity exceeds a threshold (see Figure 4b). We observe that the access of
graph topology during graph sampling is also skewed as the access of features.
Instead of allocating all the available GPU memory (except for the reservation
for model training) for feature cache, it is reasonable to cache a subset of
graph topology, i.e., edges of vertices that are frequently accessed during
sampling, in the GPU memory to accelerate GPU sampling. Figure 4b shows that a
relatively small topology cache can obviously reduce the number of PCIe
transactions incurred by GPU sampling. Motivated by O2, we propose a hotness-
aware unified cache in . Specifically, caches both graph topology and graph
features with the goal of minimizing CPU-GPU communication overhead (see
Section 4.2). Under the capacity limit of GPU memory, it is difficult to
manually decide the optimal fractions of topology and feature cache. solves
this challenge with an automatic cache management mechanism, which can
generate the optimal cache plan without requiring knowledge of hardware
specifications from users.

##  4 Design of

In order to address the aforementioned performance issues of existing cache-
based GNN systems, we propose , a cache-optimal GNN system that can push the
envelope of the multi-GPU system automatically for billion-scale GNN training.
The overall design of is presented in Figure 5. We propose an NVLink-aware
hierarchical partitioning technique (Section 4.1) in that facilitates scaling
up the cache capacity and reducing cache duplication in multi-GPU servers. To
utilize GPU cache for both graph sampling and feature extraction, we present a
hotness-aware unified cache (Section 4.2) that maintains both the topology and
feature caches to optimize the overhead of PCIe traffic. We also develop an
automatic cache management mechanism (Section 4.3) to automatically decide the
memory allocations for both topology and feature caches.

Figure 5: Design overview of . consists of three main contributions C1, C2,
and C3.

###  4.1 NVLink-aware Hierarchical Partitioning

Motivated by observation O1, we propose a simple yet effective graph
partitioning mechanism, referred to as hierarchical partitioning, to
facilitate cache scalability in . Different from conventional graph
partitioning algorithms which partition all edges/vertices of a graph into
multiple tablets, hierarchical partitioning in aims to divide the training
vertices/edges into multiple disjoint tablets. The inputs of hierarchical
partitioning are an NVLink topology matrix $M_{T}$ of the underlying multi-GPU
server and a graph $G$. The output is an assignment plan disseminating
training vertices/edges among GPUs. Specifically, the process of hierarchical
partitioning mainly consists of four steps:

S1: NVLink Clique Detection.  With the topology matrix $M_{T}$ of the server,
employs a MaxCliqueDyn algorithm [44] to identify the NVLink clique sets in
$M_{T}$, and outputs the number of NVLink cliques $K_{c}$ and the number of
GPUs in each clique $K_{g}$.

S2: Inter-clique Graph Partitioning. To reduce the cache duplication between
NVLink cliques, uses an edge-cut minimizing partitioning algorithm, e.g.,
METIS [21] and XtraPulp [34], to split the input graph $G$ into $K_{c}$
partitions, i.e., $P_{1}$, $P_{2}$, …, $P_{Kc}$, such that nodes are balanced
among partitions and inter-partition edge-cuts are minimized. The training
vertex set in $P_{i}$ is denoted as $VP_{i}$. As the training vertices are
randomly selected from $G$, the training vertex sets of different partitions
are almost of the equal size. The number of partitions is equal to the number
of detected NVLink cliques, and each NVLink clique hosts the cache for a
dedicated partition. This way, can reduce the cache duplication between NVLink
cliques and take advantage of cache locality within an NVLink clique.

S3: Intra-clique Training Vertex Partitioning. As GPUs within an NVLink clique
can access each other’s memory via low-latency high-throughput NVLink
interconnect, hierarchical partitioning further hashes the training vertex set
of each partition into $K_{g}$ tablets, where $K_{g}$ is the GPU number in a
clique. E.g., $VP_{i}$ is split into $VP_{i}[1]$ and $VP_{i}[2]$ if $K_{g}$
equals 2. Each tablet is exclusively mapped to a GPU in the corresponding
NVLink clique. We explain how to generate the cache for each training vertex
tablet in Section 4.2.

S4: Training Vertex Assignment. Finally, assigns training vertices of each
tablet to a corresponding GPU as the batch seeds, which will then be shuffled
locally to generate mini-batches for graph sampling and training.

As such, provides better cache scalability and load balancing compared to
existing systems. Figure 2 shows the cache performance of improves with the
increase of GPUs almost linearly. Figure 3 illustrates that has smaller
fluctuations in the cache hit rates on multi-GPU servers with NVLink cliques
of various sizes.

###  4.2 Hotness-aware Unified Cache

Motivated by the observation O2, we propose a hotness-aware unified cache to
cache both graph topology and graph features. In this Section, we introduce
the detailed mechanism of the unified cache.

####  4.2.1 Cache Structure

The unified cache consists of two parts: the topology cache and the feature
cache. In particular, the topology cache maintains out-edge neighbor IDs for
each selected hot vertex in the format of a compressed sparse row (CSR). As
for the feature cache, stores the feature vectors of selected hot vertices in
the format of a 2D array, where each row is the feature vector of a selected
hot vertex. Note that, the selected vertices in the topology and feature
caches could be different.

####  4.2.2 Cache Construction

The construction of the unified cache is divided into three steps: (1) pre-
sampling, (2) cache candidate selection, and (3) cache initialization. All the
GPUs/NVLink cliques perform these steps concurrently to construct their own
unified cache.

S1: Pre-sampling. Similar to GNNLab [46], adopts a pre-sampling phase222During
pre-sampling, graph topology is stored in the CPU memory. to estimate the
hotness metrics of graph topology and feature data during the training phase.
Once the process of hierarchical partitioning is completed, the training
vertex tablet assigned to each GPU is determined, which is used as the input
for pre-sampling. The output of pre-sampling includes two hotness matrices:
topology hotness matrix $H_{T}$ and feature hotness matrix $H_{F}$. Each
matrix’s row represents the GPU IDs within an NVLink clique, the column
represents the vertex IDs, and the element $H_{ij}$ of either matrix
represents the hotness of the $j$-th vertex in the $i$-th GPU. During the pre-
sampling, each GPU conducts a local shuffle on its own training vertex tablet
to generate seeds for mini-batches, performs graph sampling for each mini-
batch, and updates the corresponding row in $H_{T}$ and $H_{F}$. Figure 6
shows a pre-sampling example. For $H_{T}$, whenever an edge is traversed
during sampling, the hotness of its source vertex is incremented by $1$. For
$H_{F}$, the hotness for each vertex that appears in the sample results of the
mini-batch is incremented by 1. Additionally, uses Intel® Performance Counter
Monitor (PCM) [18] to collect the summation of PCIe transactions number,
$N_{TSUM}$, generated by all GPUs in an NVLink clique during pre-
sampling.333$N_{TSUM}$ is further used by cost model’s evaluation.

Figure 6: Update the hotness matrices of graph topology and features by pre-
sampling. For simplicity, we only show the result for GPU 1.

S2: Cache Candidate Selection.  The objective of cache candidate selection is
to select and disseminate the hot topology sub-structures and features among
GPUs within the same NVLink clique based on pre-sampled hotness matrices. Thus
this phase is conducted in the unit of NVLink clique, and each clique requires
one GPU to perform the computation. The detailed process of cache candidate
selection is presented in Algorithm 1. In brief, this algorithm computes the
global topology/feature hotness of all vertices, i.e., $A_{T}$ and $A_{F}$, in
the NVLink clique by conducting a column-wise sum on $H_{T}$ and $H_{F}$,
respectively (Line 1). $A_{T}$ and $A_{F}$ are then sorted in descending order
to generate $Q_{T}$ and $Q_{F}$(Line 2). Next, We iterate $Q_{T}$ and $Q_{F}$
in order and assign every visited vertex to the GPU with the highest local
hotness in $H_{T}$ and $H_{F}$. For each GPU, we maintain two queues, i.e.,
$G_{T}$, $G_{F}$, whose order represents the priority of vertices to be
included in this GPU cache. The outputs of Algorithm 1 are further used by the
cost model (see Section 4.3) to generate the physical cache plan.

S3: Cache Initialization and Fill-up.  ’s cache management automatically
decides the cache ratio for topology and feature so that the overall
throughput is maximized (see Section 4.3). Guided by this mechanism, allocates
memory for both the topology and feature cache ($TC$ and $FC$) of each GPU,
and fetches the corresponding topology and feature data from CPU memory to
fill up each GPU cache according to the corresponding cache orders in $G_{T}$
and $G_{F}$.

Algorithm 1 Complete Sharing with Local Preference (CSLP)

Input :

$K_{g}$: number of GPUs per NVLink clique

$H_{F}$: feature hotness matrix

$H_{T}$: topology hotness matrix

Output :

${A_{F}}$: accumulated vertex-wise feature hotness vector

${A_{T}}$: accumulated vertex-wise topology hotness vector

${Q_{T}}$: vertex ID queue representing clique-level topology order,

${Q_{F}}$: vertex ID queue representing clique-level feature order

${G_{T}}$: vertex ID queue representing GPU-level topology order

${G_{F}}$: vertex ID queue representing GPU-level feature order

/* Step 1: Accumulate each vertex’s hotness from $K_{g}$ GPUs. */

${A_{F}}$ = $H_{F}.columnWiseSum()$; ${A_{T}}$ = $H_{T}.columnWiseSum()$;

/* Step 2: Sort vertices in ${A_{F}}$ and ${A_{T}}$ */

${Q_{F}}$ <\- SortbyKeyDescend(${A_{F}}$); ${Q_{T}}$ <\-
SortbyKeyDescend(${A_{T}}$);

/* Step 3: Assign each vertex to the GPU with the highest local hotness. */

for  _$v\\_id$ in ${Q_{T}}$_ do

2 ${gpu\\_id}$ = $max(H_{T}[1:K_{g}][v\\_id]).index$;
$G_{T}[gpu\\_id].push(v\\_id)$;

3 end for

4for  _$v\\_id$ in ${Q_{F}}$_ do

5 ${gpu\\_id}$ = $max(H_{F}[1:K_{g}][v\\_id]).index$;
$G_{F}[gpu\\_id].push(v\\_id)$;

6 end for

###  4.3 Automatic Cache Management

The design of the unified cache poses a new challenge: how to properly specify
the cache size for graph topology and features under the constraint of GPU
memory such that the overall training throughput is maximized.

The general idea is to predict the overall throughput under different cache
plans and search for the best cache plan that maximizes overall throughput. We
define the cache plan as a cache memory management setting ($B$, $\alpha$) at
the NVLink clique granularity, where $B$ is the multi-GPU cache memory size in
an NVLink clique and $\alpha$ is the memory ratio for topology cache. $B$ is
identical among NVLink cliques and is by default set as the total multi-GPU
memory minus the size of GPU memory reserved for GNN models and intermediate
buffers in an NVLink clique. We need three steps to determine the optimal
cache memory management setting ($B$, $\alpha$), as discussed in Sections
4.3.1, 4.3.2, and 4.3.3.

####  4.3.1 Estimating Overall Throughput

The key goal of this Section is to build the relationship between the overall
throughput and a cache plan. We build the relationship by estimating a key
factor: the total PCIe traffic $N_{total}$, due to two reasons. First, the
PCIe traffic is the major bottleneck of the overall system throughput, and
lower PCIe traffic leads to higher overall system throughput. Second, varying
cache plans major results in the variance of PCIe traffic.444Though NVLink
traffic is also influenced by the cache plan, we neglect it since NVLink has a
much higher bandwidth than PCIe.  Because each NVLink clique maintains caches
for its own partition, we independently select the optimal cache plan for each
NVLink clique so as to minimize the PCIe traffic of each NVLink clique. Thus,
the overall system’s PCIe traffic is minimized.

####  4.3.2  Cost Model to Estimate $N_{total}$

The key goal of this Section is to present a cost model to estimate
$N_{total}$ under a specific cache plan ($B$, $\alpha$). First, given a
specific cache plan ($B$, $\alpha$), we can calculate the topology cache size
$m_{T}$ and the feature cache size $m_{F}$. Second, we find which vertices’
topology/features should be stored in the topology/feature cache. Third, we
estimate the PCIe traffic for graph sampling ($N_{T}$) and for feature
extraction ($N_{F}$) with the current topology/feature cache utilization. At
last, we estimate $N_{total}$ by adding up $N_{T}$ and $N_{F}$, as shown in
Equation 2.

$\displaystyle N_{total}$ $\displaystyle=N_{T}+N_{F}$ (2)

To estimate $N_{T}$ and $N_{F}$, we need to collect other information apart
from a given cache plan: the hotness vectors $A_{T}$ and $A_{F}$, the
summation of PCIe transaction number $N_{TSUM}$ incurred by graph sampling,
and the order queues of topology/feature cache candidates, $Q_{T}$ and
$Q_{F}$.

Estimating $N_{T}$. We estimate $N_{T}$ when the memory size of a topology
cache under one specific cache plan ($B$, $\alpha$) is $m_{T}$, where
$m_{T}=B\times\alpha$. The estimation consists of three steps.

First, with $m_{T}$, we decide which vertices’ topology should be cached. We
define ${V}$ as the set of all vertices in the graph. And we define
${V_{Tcache}}$ as the set of all vertices whose topology is cached under
current topology cache size $m_{T}$. To get ${V_{Tcache}}$, we increase
vertices and their topology into the cache with the growth of occupied
topology cache memory by the order $Q_{T}$. Until the overall occupied
topology cache memory size reaches $m_{T}$, we record ${V_{Tcache}}$. Equation
3 illustrates the relation between $m_{T}$ and ${V_{Tcache}}$, where $nc(v)$
means the neighbor count of the vertex $v$. Here we assume the data types are
Uint64 and Uint32 for the row and the column indices of the compressed sparse
row format (CSR), respectively. We use $s_{uint64}$ and $s_{uint32}$ to denote
the number of bytes to store a single Uint64 and Uint32 data accordingly.

$\sum_{v\in{V_{Tcache}}}(nc(v)\times s_{uint32}+s_{uint64})=m_{T}$ (3)

Second, once we get ${V_{Tcache}}$, we can calculate the ratio of the PCIe
transaction reduced by the topology cache by Equation 4. Let $a_{T}(v)$ mean
the topology hotness of a specific vertex $v$ ($a_{T}(v)\in A_{T}$).

$R_{T}=\frac{\sum_{v\in{V_{Tcache}}}a_{T}(v)}{\sum_{v\in{V}}a_{T}(v)})$ (4)

Third, we get $N_{T}$ by multiplying the entire PCIe transaction $N_{TSUM}$
with the ratio of PCIe transactions that can not be reduced by the topology
cache. We can get $N_{T}$ by Equation 5.

$N_{T}=N_{TSUM}\times(1-R_{T})$ (5)

Estimating $N_{F}$. We explain how to calculate $N_{F}$ when the feature cache
memory size is $m_{F}$, where $m_{F}=B\times(1-\alpha)$. There are also three
steps in estimation.

First, given $m_{F}$, we decide which vertices’ features should be cached. We
define ${V_{Fcache}}$ as the set of vertices whose feature data is cached.
Then we increase the vertices with their feature into cache by the order
$Q_{F}$ until the occupied feature cache memory size reaches $m_{F}$, as
defined in Equation 6. $D$ represents the dimension of a feature vector and
the feature data is the Float32 type each of which needs $s_{float32}$ bytes
to store.

$\sum_{v\in{V_{Fcache}}}D\times s_{float32}=m_{F}$ (6)

Second, as shown in Equation 7, we calculate the total number of features
$U_{F}$ that still needs transferring through PCIe with a feature cache.

$U_{F}=\sum_{v\in{V}}(a_{F}(v))-\sum_{v\in{V_{Fcache}}}(a_{F}(v))$ (7)

Third, we get $N_{F}$ by multiplying the transaction number needed by
transferring one vertex’s feature with the total number of features to be
transferred, $U_{F}$, as shown in Equation 8. Here $CLS$ means the transferred
cache line size. $CLS$ might be different for various CPUs and GPUs. We can
get the $CLS$ from PCM. E.g., $CLS$ equals 64 in our machine settings. And
$a_{F}(v)$ means the hotness of a specific vertex $v$ ($a_{F}(v)\in A_{F}$).

$N_{F}=(\lceil\frac{D\times s_{float32}}{CLS}\rceil)\times U_{F}$ (8)

####  4.3.3 Searching for Optimal Cache Plan in Parallel

The key goal of this Section is to efficiently determine the optimal cache
plan for each clique. As discussed in Section 4.3.1, we search for the optimal
cache plan independently with one GPU for each NVLink clique. In each NVLink
clique, we first need to traverse $\alpha$ from 0 to 1 by an interval
$\Delta\alpha$ 555$\Delta\alpha$ is set to be $0.01$ by default. to generate
the candidate cache plans, and the calculate $N_{total}$ accordingly. Then we
need to search $N_{total}$ sequences and find the smallest one with the
dedicated $\alpha$. To minimize overhead, the process is well parallelized,
including four steps:

First, we generate all the candidate cache plans in parallel and get sequences
of $m_{T}$ and $m_{F}$ in each setting.

Second, we get the boundaries of cached vertices set ${V_{Tcache}}$ and
${V_{Fcache}}$ using Equations 3 and 6, where the boundaries are the largest
cached vertices’ indexes in $Q_{T}$ and $Q_{F}$. To do so, we get the topology
and feature memory size of every single vertex in parallel and store them in
two arrays, $S_{Tsingle}$ and $S_{Fsingle}$, following the vertices order,
$Q_{T}$ and $Q_{F}$. Next, we calculate the cumulative sum of $S_{Tsingle}$
and $S_{Fsingle}$ by a parallel inclusive scan and get $S_{Tsum}$ and
$S_{Fsum}$. Then for each cache plan with $m_{T}$ and $m_{F}$, we use a
parallel binary search towards $S_{Tsum}$ and $S_{Fsum}$ to get the boundary
indexes of vertices, respectively.

Third, we get the $R_{T}$ and $U_{F}$ according to Equations 4 and 7. To do
so, we calculate the cumulative sum of $A_{T}$ and $A_{F}$ by a parallel
inclusive scan and get $A_{Tsum}$ and $A_{Fsum}$. Then for each cache plan, we
lookup $A_{Tsum}$ and $A_{Fsum}$ with the boundary indexes of vertices set
${V_{Tcache}}$, ${V_{Fcache}}$, and get $\sum_{v\in{V_{Tcache}}}a_{T}(v)$ and
$\sum_{v\in{V_{Fcache}}}a_{F}(v)$, respectively. Similarly, after lookup the
largest indexes in $A_{Tsum}$ and $A_{Fsum}$, we get $\sum_{v\in{V}}a_{F}(v)$
and $\sum_{v\in{V}}a_{F}(v)$. As such, we can get the corresponding $R_{T}$
and $U_{F}$.

At last, we calculate $N_{T}$ and $N_{F}$ for each cache plan according to
Equation 5 and 8. Then we get $N_{total}$ by Equation 2 and search in parallel
for the smallest $N_{total}$ with the corresponding $\alpha$.

After getting the optimal cache plans ($B$, $\alpha$), can automatically
allocate the cache space and fill up the cache.

##  5 Implementation of

Figure 7: An example of fine-grained GNN training pipeline for 2-hop GraphSAGE
model.

mainly consists of two components, which are the sampling server and the
training backend. The sampling server is implemented from scratch and the
training backend is built on top of Pytorch [31]. The sampling server is
responsible for generating sampled results, and the training backend takes the
sampled results as input to train the GNN models.

In , every GPU executes the graph sampling, feature extraction, and model
training stages, and all these stages are scheduled in a fine-grained pipeline
to fully utilize the GPU computation cycles. Figure 7 illustrates how the
training process is pipelined for a 2-hop GraphSAGE [16] model. In order to
improve the overall throughput, we design an inter-batch pipeline overlapping
the tasks of the sampling server and the training backend for different
batches. E.g., the training of batch $B_{i}$ can be overlapped with the
sampling and feature extraction of batch $B_{i+1}$. To further improve the
throughput of sampling and feature extraction, we design an intra-batch
pipeline inside the sampling server. Specifically, we break down the workloads
of the sampling server into four types, each of which corresponds to a type of
operator: (1) Batch generator shuffles the local training vertices to generate
seeds for mini batches; (2) Neighbor sampler executes the L-hop neighbor
sampling; (4) Feature extractor extracts the feature of the batch seeds and
vertices in the sampled results; (4) Graph constructor is used to generating
the subgraph based on the sampled results. For the same batch, graph sampling
and graph construction can be overlapped with feature extraction.

##  6 Evaluation

###  6.1 Experimental Setting

(a) DGX-V100, GraphSAGE

(b) DGX-V100, GraphSAGE

(c) DGX-A100, GraphSAGE

(d) DGX-A100, GraphSAGE

(e) DGX-V100, GCN

(f) DGX-V100, GCN

(g) DGX-A100, GCN

(h) DGX-A100, GCN

Figure 8: Overall performance of comparing with state-of-the-art systems.
“$\times$” denotes OOM (out of memory).

Experimental Platform. The experiments are conducted using three different GPU
servers: DGX-V100, Siton, and DGX-A100, as shown in Table 1. For DGX-A100, we
set the upper limit of GPU memory to 40 GB.

Table 1: GPU Server Statistics. Server | DGX-V100 | Siton | DGX-A100
---|---|---|---
GPU Type | 16GB-V100x8 | 40GB-A100x8 | 80GB-A100x8
NVLink Topo. |  $K_{c}=2$, $K_{g}=4$ |  $K_{c}=4$, $K_{g}=2$ |  $K_{c}=1$, $K_{g}=8$
PCIe Gen. | 3.0x16 | 4.0x16 | 4.0x16
PCIe Topo. |  4 switches, 2 GPUs/switch |  2 switches, 4 GPUs/switch |  4 switches, 2 GPUs/switch
CPU Mem. | 384GB | 1TB | 1TB
CPU Core Num. |  96 |  104 |  128
Sockets, NUMA Num. | 2, 1 | 2, 2 | 2, 1

GNN Models. We use two sampling-based GNN models: GraphSAGE [16] and GCN [22],
which both adopt a 2-hop random neighbor sampling. The sampling fan-outs are
25 and 10. The dimension of the hidden layers in both models is set to 256.
Similar to existing work [46], the batch size is set to 8000. Unless
explicitly explained, node classification is used as the GNN task.

Datasets.  We conduct our experiments on multiple real-world graph datasets
with various scales. Table 2 shows the dataset characteristics. The Products
(PR) and Paper100M (PA) are available in Open Graph Benchmark [17]. The Com-
Friendster (CO) graph is an online gaming network [45]. And the Uk-Union
(UKS), UK-2014 (UKL), and Clue-web (CL) are from WebGraph [5, 4, 2, 3]. As CO,
UKS, UKL, and CL have no feature, we manually generate the features with the
dimension specified as $128$ or $256$. Following PR’s setting, we choose 10%
of vertices from each graph as training vertices.

Table 2: Dataset Statistics. Dataset | PR | PA | CO | UKS | UKL | CL
---|---|---|---|---|---|---
Vertices | 2.4M | 111M | 65M | 133M | 0.79B | 1B
Edges | 120M | 1.6B | 1.8B | 5.5B | 47.2B | 42.5B
Topology Storage | 640M | 6.4GB | 7.2GB | 22GB | 189GB | 170GB
Feature Size | 100 | 128 | 256 | 256 | 128 | 128
Feature Storage | 960M | 56GB | 65GB | 136GB | 400GB | 512GB

Baselines.  We use DGL [41], PaGraph [23] and GNNLab [46] as the baseline
systems. The DGL version is v0.9.1, which supports accessing graph topology
and features via the UVA technique. We don’t compare with Quiver [33] in the
overall performance experiment as its open-sourced version cannot support
training on servers with 8 GPUs. Instead, we implement a Quiver-like multi-GPU
cache mechanism in for comparison in Section 6.3.

###  6.2 End-to-end Performance

We compare the end-to-end performance of with baseline systems on the DGX-V100
and DGX-A100 servers. On the DGX-V100 server, we evaluate PR, PA, CO, and UKS
graphs whose graph topology and features can fit into 384 GB CPU memory. On
the DGX-A100 server, we evaluate all six graphs. As PaGraph and GNNLab are
implemented using CUDA 10 which cannot support A100 GPU, we exclude them from
the experiments using DGX-A100.

Baseline Configuration. For all the baselines, we manually adjust their
configurations to achieve optimal performance. DGL uses the UVA mode, where
sampling is performed in GPU, and the topology and features are all stored in
CPU memory. The number of worker threads in PaGraph is set to be $64$ to
maximize the CPU sampling throughput. For GNNLab, we adjust the numbers of
sampling and training GPUs such that the overall throughput is maximized. In
contrast, relies on its automatic cache management mechanism to generate the
unified cache plan.

Evaluation Metrics.  We record the average epoch time for all systems. We also
use PCM [18] to measure the maximum PCIe counter value across different
sockets and report the normalized values based on the result of DGL for all
systems.

Figure 9: Effect of graph partition strategies (NoPart: no partitioning; Edge-
cut: partitioning minimizing edge-cut; Hierarchical: hierarchical
partitioning) to multi-GPU cache in terms of cache hit rate, with different
NVLink infrastructures. (noNV: disable NVLinks; NV2: $K_{c}$ = 4 and $K_{g}$ =
2; NV4: $K_{c}$ = 2 and $K_{g}$ = 4; NV8: $K_{c}$ = 1 and $K_{g}$ = 8; ).

.

Support training on large graphs.  As shown in Figures 8a, 8e, 8c and 8g,
outperform all the baseline systems in every setting. Specifically, achieves
3.78-5.69$\times$ speedup for GraphSAGE (3.5-5.19$\times$ for GCN) on DGX-V100
and 2.89-4.77$\times$ speedup for GraphSAGE (2.34-4.45$\times$ for GCN) on
DGX-A100 over DGL(UVA). Figures 8b, 8f, 8d and 8h show that, compared with the
baselines, can sufficiently utilize the multi-GPU cache to minimize PCIe
traffic incurred by CPU-GPU data transferring. GNNLab runs out of GPU memory
for UKS on DGX-V100 as the size of graph topology exceeds the capacity of
single GPU memory. PaGraph runs out of the CPU memory for most graphs except
for PR on DGX-V100, as the memory management in PaGraph incurs extra memory
overheads, including duplicated multi-hop neighbors in CPU memory and
redundant intermediate buffers generated during computation.

Speedup over SOTA system on small graphs.  achieves 1.39-4.18$\times$ speedup
for GraphSAGE (1.29-4.32$\times$ speedup for GCN) over GNNLab on the small
graphs (PR, PA, CO). The performance gain mainly comes from two aspects.
First, Figure 8b and 8f show that significantly reduces the PCIe traffic for
PA and CO, as it has a scalable multi-GPU cache design compared with GNNLab.
The reduction of PCIe traffic relieves the CPU-GPU communication bottleneck
such that the overall performance is improved. Second, can use all GPUs for
model training, while GNNLab needs to allocate several GPUs for sampling
exclusively due to its factored design. In , the graph sampling is overlapped
by model training due to the fine-grained pipeline (see Section 5). E.g., when
training GraphSAGE using the PR dataset, all the topology and feature data can
be stored in GPU memory in both and GNNLab. However, can use 8 GPUs for
training while GNNLab only uses 4 GPUs for training (see Figures 8a).

###  6.3 Effect of Hierarchical Partitioning

In this experiment, we examine the effect of hierarchical partitioning in . We
report the cache hit rates under different partition strategies in all three
GPU servers: DGX-V100 (NV4: $K_{c}$ = 2 and $K_{g}$ = 4), Siton (NV2: $K_{c}$
= 4 and $K_{g}$ = 2) and DGX-A100 (NV8: $K_{c}$ = 1 and $K_{g}$ = 8).

Figure 10: Data transferring in feature extraction of PA dataset on DGX-V100
(NV4). The rows and columns of each matrix denote the destination and source
of data transferring. The right-most (red) column records the data
transferring volume from CPU to GPU via PCIe. The middle (green) columns
represent the GPU-GPU data transferring volume. We normalize the recorded
values based on the CPU-GPU data transferring volumes in GNNLab.

####  6.3.1 Cache Performance

Baselines. For a fair comparison, we implement the cache designs of GNNLab,
PaGraph-plus (described in Section 3.1), and Quiver-plus in and compare their
cache hit rates. Specifically, GNNLab maintains a globally replicated cache
among all GPUs without using NVLinks (noPart+noNV). Quiver-plus enables NVLink
and maintains replicated cache among NVLink cliques (noPart+NV2 / noPart+NV4 /
noPart+NV8). PaGraph-plus takes the XtraPulp [34] partitioning which minimizes
across-partition edge-cuts and disables NVLinks (Edge-cut+noNV). uses
hierarchical partitioning (inter-NVLink-clique partitioning: XtraPulp) and
enables NVLink (Hierarchical+NV2 / Hierarchical+NV4 / Hierarchical+NV8). We
use the pre-sampling hotness metric for all these cache designs. The in-
degree-based hotness metric in the original PaGraph and Quiver design are
replaced with the pre-sampling hotness metric in Pagraph-plus and Quiver-plus,
which has a better performance on cache hit rates [46].

The datasets used in this experiment are PR, CO, UKL, and CL. We vary the
cache ratio from 1.25% $|V|$ to 10% $|V|$ for PR and CO. For UKL and CL whose
sizes are relatively large, the cache ratio varies from 1.25% $|V|$ to 5%
$|V|$. Figure 9 shows that, for almost all the experiment settings, has the
highest cache hit rate. Specifically, obviously outperforms Quiver-plus in the
cases of NV2 and NV4, since can reduce the inter-NVLink-clique cache
duplication and achieves higher multi-GPU memory utilization compared with
Quiver-plus. For the case of NV8, as all GPUs are in the same NVLink clique,
the inter-clique graph partitioning in can be skipped, and hierarchical
partitioning turns into hash partitioning among all the GPUs, which is
identical to Quiver-plus in the case of NV8. outperforms PaGraph-plus because
it has much less cache duplication. Specifically, PaGraph-plus’s cache
mechanism may replicate vertices with high global hotness on multiple GPUs.
Compared with GNNLab, has higher cache hit rates as it can scale up the cache
capacity with the increase of GPUs, while GNNLab replicates the same feature
cache across all GPUs. These results demonstrate that can effectively adapt
the cache plan to optimize the cache performance for multi-GPU servers with
various NVLink topologies.

####  6.3.2 Data Transferring in Feature Extraction

In this experiment, we demonstrate the GPU-GPU and CPU-GPU data transferring
volume during feature extraction using the PA dataset. Specifically, we
perform the graph sampling and feature extraction stages using the PA graph on
DGX-V100 (NV4) and record the data transferring volumes of feature extraction
on each GPU in the format of a traffic matrix. We use GNNLab, PaGraph-plus,
and Quiver-plus as the baselines, and set the feature cache ratio on each GPU
to 2.5% $|V|$. The results are presented in Figure 10. We can see that ’s data
transferring volume from CPU to GPU is the smallest, indicating the best cache
performance among the compared systems. As it is the GPU with the largest CPU-
GPU data transferring volume that dominates the overall performance, although
’s CPU-GPU volumes on some GPUs are higher than PaGraph-plus, can still
outperform PaGraph-plus because its largest CPU-GPU volume is lower than that
of PaGraph-plus.

####  6.3.3 Model Convergence

Compared with global shuffling (randomly generating batch seeds from the
vertex set of the entire graph), recent studies [28, 23] show that local
shuffling (generating batch seeds within partitions) brings negligible impact
on the rate of model convergence. adopts local shuffling, and we conduct an
experiment on the Siton server (NV2) to compare its convergence speed with
global shuffling on both GraphSAGE and GCN using the PR dataset. The results
in Figure 11 show that the local shuffling of could catch up with the
convergence speed of global shuffling.

(a) GraphSAGE

(b) GCN

Figure 11: Comparing local shuffling and global shuffling on model convergence
(NoPart: no partitioning; Hierarchical: hierarchical partitioning).

###  6.4 Effect of Unified Cache

Different from existing cache-based systems, ’s unified cache also takes graph
topology into account. In this experiment, we demonstrate the benefits of
topology cache.

We compare the training epoch time of unified cache in with two baselines: (1)
storing all topology in the CPU (denoted as TopoCPU) and (2) replicating the
entire topology in every single GPU (denoted as TopoGPU). For a fair
comparison, we implement both TopoCPU and TopoGPU in and use the same GPU
memory volume for the three settings. Among the three settings, TopoCPU has
the most GPU memory available for the feature cache, and the TopoGPU has the
least GPU memory for the feature cache or even runs out of GPU memory. We
evaluate PA, CO, and UKS on DGX-V100 and evaluate UKL and CL on DGX-A100.

As shown in Figure 12, the unified cache outperforms the other two baselines
for all graphs. This result demonstrates that, when the size of the feature
cache exceeds a threshold, the increase of cache hit rate slows down. In this
case, caching some hot topology data in GPU memory will save the system from
severe PCIe contention incurred by graph sampling and benefit the overall GNN
training throughput.

Figure 12: The impact of topology cache. “$\times$” means OOM (out of memory).

###  6.5 Evaluation of Cost Model

proposes the cost model to guide allocating GPU memory for both graph topology
and feature cache. In this experiment, we evaluate the effectiveness of this
mechanism. Specifically, we compare the predicted PCIe traffic with the
experimental per-epoch execution time of graph sampling and feature
extraction. In the experiment using the PA dataset, the GPU memory allocated
for the cache is $10$ GB. And in the experiment using the UKS dataset, the GPU
memory allocated for the cache is $8$ GB. When varying the size of the
topology cache, the size of the feature cache is adjusted accordingly. Figure
13 shows that our cost model can precisely predict the trend of per-epoch
execution time without manual interference.

(a) PA, Single GPU

(b) UKS, DGX-V100 (NV4)

Figure 13: Evaluation of cost model. The left y-axis means the PCIe
transaction number predicted by the cost model. The right y-axis represents
the experimental per-epoch graph sampling and feature extraction time.

###  6.6 Partitioning Cost

In this experiment, we study the partitioning cost in . We run our experiment
on the UKL dataset that has the largest number of edges among all the
datasets, resulting in the highest cost of edge-cut partitioning. We also
present the results of the PA data (medium size) to show the partitioning
costs of different graph scales. We partition PA on DGX-V100 and UKL on Siton
using the XtraPulp algorithm. For node classification, we set the training set
to be 10% of the total edges for both graphs. For link prediction, we set the
training set to be 80% of total edges. When the graph is too large to be
partitioned in memory, like UKL, we randomly sample a fraction of edges (25%
for UKL) and keep all vertices in the graph such that the subgraph can be
partitioned in memory. This technique can obviously speedup graph partitioning
and preserves a low edge-cut ratio.

Table 3 shows the preprocessing cost of ’s hierarchical partitioning. We
observe that the partitioning cost is tolerable, because 1) we only partition
the graph once but can use the partitioning results for multiple GNN training
jobs, and 2) the GNN task like link prediction needs multiple epochs to
converge while a single epoch often costs a long time to finish.

Table 3: Evaluation of Partitioning Cost. Dataset | PA (DGX-V100) | UKL (Siton)
---|---|---
Graph Partition(min) | 7.2 | 75
Data Loading From Disk To Memory(min) | 0.32 | 3.5
Node Classification Epoch(s) | 1.98 | 15.6
Link Prediction Epoch(min) | 49.8 | 402

##  7 Related Work

To our knowledge, is the first work that automatically pushes the envelope of
multi-GPU systems for billion-scale GNN training. In the following, we
contrast and existing works in the following aspects.

GNN Frameworks.  Several GNN systems [20, 41, 11, 54, 23, 46, 33, 26, 50, 42,
12, 37, 52] have emerged in recent years. Most of these GNN systems are built
on top of deep learning frameworks like Pytorch [31], TensorFlow [1] and MXNet
[9].

GPU Sampling.  NextDoor [19] and C-SAW [30] focus on accelerating GPU sampling
kernel. DGL [41] also supports GPU sampling in its recent release. Quiver [33]
can support GPU sampling with the entire topology either stored in the single
GPU or in the CPU memory. GNNLab [46] adopts a factored design where each GPU
is dedicated to graph sampling or model training exclusively. In contrast,
uses all GPUs for end-to-end GNN acceleration.

Graph Partitioning.  Graph partitioning such as [21, 34, 35, 38, 15, 14, 32,
6], has been widely adopted in GNN systems. DGL [41] adopts METIS [21] to
partition the graph. PaGraph [23] adopts a self-reliant partitioning strategy
with the goal of achieving balanced training vertex allocation across GPUs and
improving data locality on every GPU. DGCL [7] adopts a partitioning algorithm
to partition the graph’s physical edges and features and store them among
distributed machines. In contrast, adopts hierarchical partitioning to
automatically partition graphs to each GPU in a single multi-GPU server
accordingly to GPU interconnections.

GPU Feature Cache. PaGraph [23], BGL [24], GNNLab [46], Quiver [33] and [29]
explore feature caching on GPU to accelerate GNN training. PaGraph [23] and
Quiver [33] use the in-degree of vertexes as the hotness metric. BGL [24]
applies a FIFO dynamic cache policy and selects training vertices in a BFS
order for a higher cache hit rate, but hinders model convergence and incurs
cache replacement overheads. [29] uses a weighted reverse PageRank algorithm
as a hotness metric. GNNLab [46] uses vertices’ access frequencies in the pre-
sampling epoch as a hotness metric. In contrast, automatically caches both
features and topology with the highest hotness. And statically partitions the
graph with minimal edge-cut to preserve intra-partition data locality. Figures
9 and 11 show that Legion can achieve a high cache hit rate even with small
cache ratios without compromising the model convergence rate.

Large Graph Systems. SSD-based GNN systems [40] and distributed GNN systems
[24, 51, 53, 12] also aim at large-graph training and propose distinct
approaches to solve I/O problems at various levels. MariusGNN [40] minimizes
I/O between SSD and CPU by including valid graph data in a single swap as much
as possible. Systems like BGL [24], DistDGLv2 [53], and P3 [12] optimize
network I/O between distributed machines, whose network performance can be
improved when introducing GPU-centric SmartNIC [43]. In contrast, focuses on
utilizing GPU caches to minimize PCIe traffic from CPU memory to multiple
GPUs, which is orthogonal to the above systems.

##  8 Conclusion

We present , a system that automatically pushes the envelope of multi-GPU
systems for billion-scale GNN training. has three key innovations. First, we
propose an NVLink-aware hierarchical partitioning technique that helps
minimize cache replication and extends the threshold of cache capacity beyond
the limit of a single GPU or NVLink clique. Second, we propose a novel
hotness-aware unified cache mechanism that helps accelerate both graph
sampling and feature extraction. Third, we present an automatic cache
management mechanism enabling optimal cache planning without requiring extra
knowledge of hardware specifications and GNN performance details from users.
Experiments show Legion outperforms SOTA cache-based GNN systems up to 4.32×
and supports training on billion-scale graphs. And is open-sourced at
https://github.com/RC4ML/Legion.

Acknowledgements.  We thank our shepherd Anand Iyer and anonymous reviewers
for their detailed feedback. The work is supported by the following grants:
the Program of Zhejiang Province Science and Technology (2022C01044), a
research grant from Alibaba Group through the Alibaba Innovative Research
(AIR) Program, the Fundamental Research Funds for the Central Universities
226-2022-00151, Key Laboratory for Corneal Diseases Research of Zhejiang
Province, Starry Night Science Fund of Zhejiang University Shanghai Institute
for Advanced Study (SN-ZJU-SIAS-0010). Zeke Wang and Fei Wu are the
corresponding authors.

## References

  * [1] Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, Manjunath Kudlur, Josh Levenberga, Sherry Moore Rajat Monga, Derek G. Murray, Benoit Steiner, Paul Tucker, Vijay Vasudevan, Martin Wicke Pete Warden, Yuan Yu, and Xiaoqiang Zheng.  Tensorflow: a system for large-scale machine learning.  In OSDI, 2016.
  * [2] Paolo Boldi, Bruno Codenotti, Massimo Santini, and Sebastiano Vigna.  Ubicrawler: A scalable fully distributed web crawler.  Software: Practice & Experience, 2004.
  * [3] Paolo Boldi, Andrea Marino, Massimo Santini, and Sebastiano Vigna.  Bubing: Massive crawling for the masses.  In WWW, 2014.
  * [4] Paolo Boldi, Marco Rosa, Massimo Santini, and Sebastiano Vigna.  Layered label propagation: A multiresolution coordinate-free ordering for compressing social networks.  In WWW, 2011.
  * [5] Paolo Boldi and Sebastiano Vigna.  The web graph framework: Compression techniques.  In WWW, 2004.
  * [6] Erik G Boman, Karen D Devine, and Sivasankaran Rajamanickam.  Scalable matrix computations on large scale-free graphs using 2d graph partitioning.  In SC, 2013.
  * [7] Zhenkun Cai, Xiao Yan, Yidi Wu, Kaihao Ma, James Cheng, and Fan Yu.  DGCL: An efficient communication library for distributed GNN training.  In Eurosys, 2021.
  * [8] Jie Chen, Tengfei Ma, and Cao Xiao.  Fastgcn: fast learning with graph convolutional networks via importance sampling.  arXiv preprint arXiv:1801.10247, 2018.
  * [9] Tianqi Chen, Mu Li, Yutian Li, Min Lin, Naiyan Wang, Minjie Wang, Tianjun Xiao, Bing Xu, Chiyuan Zhang, and Zheng Zhang.  Mxnet: A flexible and efficient machine learning library for heterogeneous distributed systems.  arXiv preprint arXiv:1512.01274, 2015.
  * [10] Wei-Lin Chiang, Xuanqing Liu, Si Si, Yang Li, Samy Bengio, and Cho-Jui Hsieh.  Cluster-gcn: An efficient algorithm for training deep and large graph convolutional networks.  In SIGKDD, 2019.
  * [11] Matthias Fey and Jan Eric Lenssen.  Fast graph representation learning with pytorch geometric.  arXiv preprint arXiv:1903.02428, 2019.
  * [12] Swapnil Gandhi and Anand Padmanabha Iyer.  P3: Distributed deep graph learning at scale.  In OSDI, 2021.
  * [13] Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl.  Neural message passing for quantum chemistry.  In ICML, 2017.
  * [14] Joseph E Gonzalez, Yucheng Low, Haijie Gu, Danny Bickson, and Carlos Guestrin.  Powergraph: Distributed graph-parallel computation on natural graphs.  In OSDI, 2012.
  * [15] Joseph E Gonzalez, Reynold S Xin, Ankur Dave, Daniel Crankshaw, Michael J Franklin, and Ion Stoica.  Graphx: Graph processing in a distributed dataflow framework.  In OSDI, 2014.
  * [16] Will Hamilton, Zhitao Ying, and Jure Leskovec.  Inductive representation learning on large graphs.  NeurIPS, 2017.
  * [17] Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec.  Open graph benchmark: Datasets for machine learning on graphs.  NIPS, 2020.
  * [18] Intel.  PCM.  https://github.com/intel/pcm, 2022.
  * [19] Abhinav Jangda, Sandeep Polisetty, Arjun Guha, and Marco Serafini.  Accelerating graph sampling for graph machine learning using gpus.  In Eurosys, 2021.
  * [20] Zhihao Jia, Sina Lin, Mingyu Gao, Matei Zaharia, and Alex Aiken.  Improving the accuracy, scalability, and performance of graph neural networks with roc.  MLSys, 2020.
  * [21] George Karypis and Vipin Kumar.  Metis: A software package for partitioning unstructured graphs, partitioning meshes, and computing fill-reducing orderings of sparse matrices.  1997\.
  * [22] Thomas N Kipf and Max Welling.  Semi-supervised classification with graph convolutional networks.  arXiv preprint arXiv:1609.02907, 2016.
  * [23] Zhiqi Lin, Cheng Li, Youshan Miao, Yunxin Liu, and Yinlong Xu.  Pagraph: Scaling gnn training on large graphs via computation-aware caching.  In SoCC, 2020.
  * [24] Tianfeng Liu, Yangrui Chen, Dan Li, Chuan Wu, Yibo Zhu, Jun He, Yanghua Peng, Hongzheng Chen, Hongzhi Chen, and Chuanxiong Guo.  Bgl: Gpu-efficient gnn training by optimizing graph data i/o and preprocessing.  arXiv preprint arXiv:2112.08541, 2021.
  * [25] Yang Liu, Xiang Ao, Zidi Qin, Jianfeng Chi, Jinghua Feng, Hao Yang, and Qing He.  Pick and choose: a gnn-based imbalanced learning approach for fraud detection.  In WWW, 2021.
  * [26] Lingxiao Ma, Zhi Yang, Youshan Miao, Jilong Xue, Ming Wu, Lidong Zhou, and Yafei Dai.  Neugraph: Parallel deep neural network computation on large graphs.  In USENIX ATC, 2019.
  * [27] Mark Harris.  Unified Memory for CUDA Beginners.  https://developer.nvidia.com/blog/unified-memory-cuda-beginners/, 2017.
  * [28] Qi Meng, Wei Chen, Yue Wang, Zhi-Ming Ma, and Tie-Yan Liu.  Convergence analysis of distributed stochastic gradient descent with shuffling.  Neurocomputing, 2019.
  * [29] Seung Won Min, Kun Wu, Mert Hidayetoglu, Jinjun Xiong, Xiang Song, and Wen-mei Hwu.  Graph neural network training and data tiering.  In SIGKDD, 2022.
  * [30] Santosh Pandey, Lingda Li, Adolfy Hoisie, Xiaoye S Li, and Hang Liu.  C-saw: A framework for graph sampling and random walk on gpus.  In SC, 2020.
  * [31] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  Pytorch: An imperative style, high-performance deep learning library.  NeurIPS, 2019.
  * [32] Fabio Petroni, Leonardo Querzoni, Khuzaima Daudjee, Shahin Kamali, and Giorgio Iacoboni.  Hdrf: Stream-based partitioning for power-law graphs.  In CIKM, 2015.
  * [33] QuiverTeam.  Quiver.  https://github.com/quiver-team/torch-quiver, 2021.
  * [34] George M Slota, Sivasankaran Rajamanickam, Karen Devine, and Kamesh Madduri.  Partitioning trillion-edge graphs in minutes.  In IPDPS, 2017.
  * [35] Isabelle Stanton and Gabriel Kliot.  Streaming graph partitioning for large distributed graphs.  In SIGKDD, 2012.
  * [36] Chang Su, Yu Hou, and Fei Wang.  Gnn-based biomedical knowledge graph mining in drug development.  In Graph Neural Networks: Foundations, Frontiers, and Applications. 2022.
  * [37] John Thorpe, Yifan Qiao, Jonathan Eyolfson, Shen Teng, Guanzhou Hu, Zhihao Jia, Keval Vora, Ravi Netravali, Miryung Kim, and Guoqing Harry Xu.  Dorylus: Affordable, scalable, and accurate nn training with distributed cpu servers and serverless threads.  In OSDI, 2021.
  * [38] Charalampos Tsourakakis, Christos Gkantsidis, Bozidar Radunovic, and Milan Vojnovic.  Fennel: Streaming graph partitioning for massive scale graphs.  In WSDM, 2014.
  * [39] Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio.  Graph attention networks.  arXiv preprint arXiv:1710.10903, 2017.
  * [40] Roger Waleffe, Jason Mohoney, Theodoros Rekatsinas, and Shivaram Venkataraman.  Mariusgnn: Resource-efficient out-of-core training of graph neural networks.  In Eurosys, 2023.
  * [41] Minjie Yu Wang.  Deep graph library: Towards efficient and scalable deep learning on graphs.  In ICLR, 2019.
  * [42] Yuke Wang, Boyuan Feng, Gushu Li, Shuangchen Li, Lei Deng, Yuan Xie, and Yufei Ding.  Gnnadvisor: An adaptive and efficient runtime system for gnnacceleration on gpus.  In OSDI, 2021.
  * [43] Zeke Wang, Hongjing Huang, Jie Zhang, Fei Wu, and Gustavo Alonso.  FpgaNIC: An FPGA-based versatile 100gb SmartNIC for GPUs.  In ATC, 2022.
  * [44] Wikipedia.  MaxCliqueDyn.  https://en.wikipedia.org/wiki/MaxCliqueDyn_maximum_clique_algorithm, 2022\.
  * [45] Jaewon Yang and Jure Leskovec.  Defining and evaluating network communities based on ground-truth.  In SIGKDD, 2012.
  * [46] Jianbang Yang, Dahai Tang, Xiaoniu Song, Lei Wang, Qiang Yin, Rong Chen, Wenyuan Yu, and Jingren Zhou.  Gnnlab: A factored system for sample-based gnn training over gpus.  In Eurosys, 2022.
  * [47] Rex Ying, Ruining He, Kaifeng Chen, Pong Eksombatchai, William L Hamilton, and Jure Leskovec.  Graph convolutional neural networks for web-scale recommender systems.  In SIGKDD, 2018.
  * [48] Zhongbao Yu, Jiaqi Zhang, Xin Qi, and Chao Chen.  Application research of graph neural networks in the financial risk control.
  * [49] Hanqing Zeng, Hongkuan Zhou, Ajitesh Srivastava, Rajgopal Kannan, and Viktor Prasanna.  Graphsaint: Graph sampling based inductive learning method.  arXiv preprint arXiv:1907.04931, 2019.
  * [50] Dalong Zhang, Xin Huang, Ziqi Liu, Zhiyang Hu, Xianzheng Song, Zhibang Ge, Zhiqiang Zhang, Lin Wang, Jun Zhou, Yang Shuang, and Yuan Qi.  Agl: a scalable system for industrial-purpose graph machine learning.  arXiv preprint arXiv:2003.02454, 2020.
  * [51] Da Zheng, Chao Ma, Minjie Wang, Jinjing Zhou, Qidong Su, Xiang Song, Quan Gan, Zheng Zhang, and George Karypis.  Distdgl: distributed graph neural network training for billion-scale graphs.  In 2020 IEEE/ACM 10th Workshop on Irregular Applications: Architectures and Algorithms (IA3), 2020.
  * [52] Da Zheng, Xiang Song, Chengru Yang, Dominique LaSalle, and George Karypis.  Distributed hybrid cpu and gpu training for graph neural networks on billion-scale heterogeneous graphs.  In SIGKDD, 2022.
  * [53] Da Zheng, Xiang Song, Chengru Yang, Dominique LaSalle, Qidong Su, Minjie Wang, Chao Ma, and George Karypis.  Distributed hybrid cpu and gpu training for graph neural networks on billion-scale graphs.  arXiv preprint arXiv:2112.15345, 2021.
  * [54] Rong Zhu, Kun Zhao, Hongxia Yang, Wei Lin, Chang Zhou, Baole Ai, Yong Li, and Jingren Zhou.  Aligraph: a comprehensive graph neural network platform.  VLDB, 2019.


11institutetext: 1Memorial Sloan Kettering Cancer Center, New York, NY, USA
2Kyungpook National University, South Korea
11email: {keey, kangm2, jeongs1<EMAIL_ADDRESS><EMAIL_ADDRESS>
# Joint Multi-Echo/Respiratory Motion-Resolved Compressed Sensing
Reconstruction of Free-Breathing Non-Cartesian Abdominal MRI

Youngwook Kee1 MungSoo Kang1 Seongho Jeong1,2 and Gerald Behr1

###### Abstract

We propose a novel respiratory motion-resolved MR image reconstruction method
that jointly treats multi-echo k-space raw data. Continuously acquired non-
Cartesian multi-echo/multi-coil k-space data with free breathing are
sorted/binned into the motion states from end-expiratory to end-inspiratory
phases based on a respiratory motion signal. Temporal total variation applied
to the motion state dimension of each echo is then coupled in the $\ell_{2}$
sense for joint reconstruction of the multiple echoes. Reconstructed source
images of the proposed method are compared with conventional echo-by-echo
motion-resolved reconstruction, and R2* of the proposed and echo-by-echo
methods are compared with respect to a clinical reference. We demonstrate that
inconsistency between echoes is successfully suppressed in the proposed joint
reconstruction method, producing high-quality source images and R2*
measurements compared to clinical reference.

###### Keywords:

Multi-Echo Non-Cartesian MRI Motion-Resolved Image Reconstruction
Collaborative/Vectorial Total Variation.

##  1 Introduction

In multi-echo gradient-echo (mGRE) MRI, k-space raw data are acquired at
multiple time points (echo times or TEs). Reconstructed images at different
TEs are utilized in quantitative tissue parameter mapping, such as proton
density fat fraction (PDFF) [3, 18], R2* relaxation rate [8], or quantitative
susceptibility mapping (QSM) [20, 13]. Clinical applications include fatty
liver disease [15] and hepatic iron overload assessment [9]. Recent advances
in non-Cartesian k-space sampling trajectories have facilitated free-breathing
mGRE MRI [14, 1], obviating the need for patients to hold their breath which
is challenging for children and some adults. In addition, respiratory motion-
resolved compressed sensing (CS) reconstruction, such as XD-GRASP [5], has
demonstrated to enable motion-resolved PDFF/R2*/QSM [19, 10, 11] mitigating
the confounding factor of respiratory motion.

Despite these advances, applying respiratory motion-resolved CS reconstruction
to highly-undersampled mGRE data is not straightforward. The question arises
as to what mathematical machinery needs to be devised/incorporated for joint
reconstruction of multiple echoes to prevent potential “echo inconsistency”
produced by echo-by-echo CS reconstruction that necessarily requires iterative
optimization. Little attention has been paid to the optimal treatment of
highly-undersampled mGRE imaging data, which is crucial in motion-tolerant
quantitative body MRI. In this paper, inspired by vectorial/collaborative
total variation (TV) in color image denoising/deblurring/inpainting [4, 6], we
propose a novel treatment for mGRE image reconstruction where all echoes are
jointly reconstructed.

The remainder of this paper is organized as follows: Section 2 describes the
details of the proposed joint multi-echo/multi-coil/respiratory motion-
resolved image reconstruction method with numerical implementation, Section 3
describes the methods and materials used to demonstrate the effectiveness of
the proposed method, Section 4 reports the results and compares them with
existing methods, and finally, Section 5 povides a discussion of our findings
and a conclusion.

##  2 Theory

###  2.1 Problem Formulation

Let $N$ be the image size, $C$ be the number of coils, $E$ be the number of
echoes, $T$ be the number of motion states, and $M$ be the number of
measurements in k-space. We are interested in reconstructing
$u_{k}\in\mathbb{C}^{N\times T}$ in image space from
$y_{j,k}\in\mathbb{C}^{\lfloor M/T\rfloor\times T}$ in k-space for
$k=1,\dots,E$ and $j=1,\dots,C$. Notice that the total number of k-space
measurements $M$ is assumed to be sorted/binned into $\lfloor M/T\rfloor\times
T$ from end-expiratory to end-inspiratory motion state according to a motion
signal [5, 10, 11]. Therefore, each motion state is highly undersampled by a
factor of $T$, and the following model-based MRI reconstruction is considered:

$\displaystyle\underset{u_{1},\dots,u_{E}}{\text{minimize}}\,\frac{1}{2}\sum_{j=1}^{C}\sum_{k=1}^{E}||\sqrt{D}(FS_{j}u_{k}-y_{j,k})||_{F}^{2}+\lambda\cdot{\cal
R}(u_{1},\dots,u_{E}),$ (1)

where $S_{j}:\mathbb{C}^{N\times T}\to\mathbb{C}^{N\times T}$ is the $j$-th
coil sensitivity map, $F:\mathbb{C}^{N\times T}\to\mathbb{C}^{\lfloor
M/T\rfloor\times T}$ is the nonuniform Fourier transform, and
$D:\mathbb{C}^{\lfloor M/T\rfloor\times T}\to\mathbb{C}^{\lfloor
M/T\rfloor\times T}$ is the density compensation factor [16, 17].
$||\cdot||_{F}$ is the Frobenius norm in complex space and ${\cal
R}(u_{1},\dots,u_{E})$ is the regularization term that can be defined via
$\ell_{2}$ coupling along the echo as follows.

$\displaystyle{\cal R}(u_{1},\dots,u_{E})$ $\displaystyle=\sum_{{\bf
x}}\left(|\partial_{t}u_{1}({\bf x})|^{2}+|\partial_{t}u_{2}({\bf
x})|^{2}+\cdots+|\partial_{t}u_{E}({\bf x})|^{2}\right)^{1/2}$
$\displaystyle=||\partial_{t}{\bf u}||_{2,1},$ (2)

where ${\bf u}=(u_{1},\dots,u_{E})$ and $\partial_{t}$ is the partial
derivative or finite difference along the motion dimension (time). This is
known as collaborative/vectorial TV [4, 6], initially considered in RGB color
image processing. _Here, we make a key observation that multiple echoes in
mGRE data can be combined in the $\ell_{2}$ sense for joint multi-
echo/respiratory motion-resolved CS reconstruction (henceforth, joint TE/MR
recon). This approach takes into account the echo signal evolution, which is
currently absent in the conventional echo-by-echo motion-resolved
reconstruction (henceforth, echo-by-echo MR recon)_.

###  2.2 Numerical Optimization

The collaborative/vectorial TV with $\ell_{2}$ coupling is convex in $\bf u$,
and we rewrite (1) with the dualization of the data consistency term [2, 12]
as follows.

$\displaystyle\min_{u_{1},\dots,u_{E}}\max_{\begin{subarray}{c}(\xi_{1},\dots,\xi_{E})\in
K\\\
\zeta_{1,1},\dots,\zeta_{C,E}\end{subarray}}\left(\sum_{j=1}^{C}\sum_{k=1}^{E}\langle\sqrt{D}(FS_{j}u_{k}-y_{j,k}),\zeta_{j,k}\rangle-\frac{1}{2}||\zeta_{j,k}||_{2}^{2}\right)+\sum_{k=1}^{E}\langle\partial_{t}u_{k},\xi_{k}\rangle,$
(3)

where the convex set $K:=\\{(\xi_{1},\dots,\xi_{E})\in\mathbb{C}^{N\times
T}\times\cdots\times\mathbb{C}^{N\times
T}:(|\xi_{1}|^{2}+\cdots+|\xi_{E}|^{2})^{1/2}\leq 1/\lambda\\}$. The primal-
dual hybrid gradient (PDHG) algorithm [2] is suitable to solve (9). By
applying gradient ascent for the dual variables and gradient descent for the
primal variable, we obtain the following update equations:

$\displaystyle\xi_{k}^{n+1}$
$\displaystyle\leftarrow\operatorname{proj}_{K}(\xi_{k}^{n}+\sigma\partial_{t}\bar{u}_{k}^{n}),$
(4) $\displaystyle\zeta_{j,k}^{n+1}$
$\displaystyle\leftarrow\operatorname{prox}(\zeta_{j,k}^{n}+\sigma(\sqrt{D}(FS_{j}\bar{u}_{k}^{n}-y_{j,k}))),$
(5) $\displaystyle u_{k}^{n+1}$ $\displaystyle\leftarrow
u_{k}^{n}-\tau\left(\sum_{j=1}^{C}(\sqrt{D}FS_{j})^{H}\zeta_{j,k}^{n}+\partial_{t}^{H}\xi_{k}^{n}\right),$
(6) $\displaystyle\bar{u}_{k}^{n+1}$ $\displaystyle\leftarrow
2u_{k}^{n+1}-u_{k}^{n},$ (7)

for all $k$ and $j$. The projection onto $K$ and proximal operator are
performed by

$\displaystyle\xi_{k}\leftarrow\frac{\tilde{\xi}_{k}}{\max(1,(|\xi_{1}|^{2}+\cdots+|\xi_{E}|^{2})^{1/2}/\lambda)},\quad\zeta_{j,k}\leftarrow\frac{\tilde{\zeta}_{j,k}}{1+\sigma}.$
(8)

The step sizes $\sigma$ and $\tau$ are both set to 1/8. The additional update
step for the primal variable in (7) is an extragradient step for convergence
[2].

#### Remark I.

The calculation of $(\sqrt{D}FS_{j})^{H}\zeta_{j,k}^{n}$ for $j=1,\dots,C$,
which is the most computationally intensive step, can be distributed across
multiple GPUs using the Message Passing Interface (MPI) to achieve speedup.

#### Remark II.

The $\ell_{1}$ coupling between echoes for ${\cal R}(u_{1},\dots,u_{E})$ can
be used instead, which is expressed as

$\displaystyle||\partial_{t}{\bf u}||_{1,1}=\sum_{{\bf
x}}|\partial_{t}u_{1}({\bf x})|+|\partial_{t}u_{2}({\bf
x})|+\cdots+|\partial_{t}u_{E}({\bf x})|.$ (9)

Notice that there is _no effective coupling_ between the echoes as the
gradient of the objective function with respect to each echo is independent
from each other, possibly causing “echo inconsistency” in the final image. The
only change in the PDHG algorithm is the convex set $K$ on which the dual
variable $\zeta_{j,k}$ is projected, which is performed as
$\xi_{k}\leftarrow\tilde{\xi}_{k}/\max(1,|\xi_{k}|/\lambda)$.

##  3 Methods

###  3.1 Synthetic Image (Toy Example)

Figure 1: Synthetic Image Experiment (see Section 3.1 and 4 for description.).

To illustrate the concept of collaborative/vectorial TV in the idea of joint
TE/MR recon, we designed a synthetic image reconstruction experiment as
follows: 2D radial and variable density spiral k-space sampling trajectories
were designed using the following parameters: Gmax = 80 mT/m, SR = 200 T/m/s,
FOV = 25 cm, in-plane resolution = 1$\times$1 mm2, R = 6 (acceleration). The
generated trajectories in Fig. 1(A) were used to sample the k-space (Fourier
space) data of the input RBG image shown in Fig. 1(B) using the nonuniform
Fourier transform (NUFFT). Notice that there is misalignment between the RGB
color channels along the vertical axis, which was introduced by shifting the G
and B channels. This is illustrated in Fig. 1(I-L), where transparent line
profiles are drawn along the vertical dashed line in Fig. 1(B).

The R, G, and B channels can be considered as multi-echo images, where each
channel represents the images at TE1, TE2, and TE3, respectively. The
intensity profiles of each color channel, or equivalently each echo, can be
viewed as representing different motion states. We formulate the following
reconstruction problem: Let $m_{j}:(\Omega_{{\cal
F}}\subset\mathbb{R}^{2})\to\mathbb{C}$ be the $j$-th channel of the vector-
valued function in Fourier space (k-space) $m:(\Omega_{{\cal
F}}\subset\mathbb{R}^{2})\to\mathbb{C}^{3},(k_{x},k_{y})\mapsto(m_{1}(k_{x},k_{y}),m_{2}(k_{x},k_{y}),$
$m_{3}(k_{x},k_{y}))$. The goal is to reconstruct ${\bf
u}:(\Omega\subset\mathbb{R}^{2})\to\mathbb{C}^{3},(x,y)\mapsto(u_{1}(x,y),u_{2}(x,y),u_{3}(x,y))$
from the undersampled input Fourier samples $m$ by solving

$\displaystyle\underset{u_{1},u_{2},u_{3}}{\text{minimize}}\,\frac{1}{2}\sum_{k=1}^{3}\int_{\Omega_{\cal
F}}\|\sqrt{d}({\cal
F}u_{k}-m_{k})\|_{2}^{2}+\lambda\int_{\Omega}||\partial_{y}{\bf u}||_{1},$
(10)

where ${\cal F}:\mathbb{C}\to\mathbb{C}$ is the nonuniform Fourier transform
operator, and $d$ is the density compensation factor. The channels (or echoes)
can be coupled with the $\ell_{2}$ norm as follows.

$\displaystyle\int_{\Omega}||\partial_{y}{\bf
u}||_{2}=\int_{\Omega}\left(|\partial_{y}u_{1}|^{2}+|\partial_{y}u_{2}|^{2}+|\partial_{y}u_{3}|^{2}\right)^{1/2}.$
(11)

Note that the above variational formulation of generic image reconstruction
problem (10) and (11) coincides with (1) and (2.1) for MRI reconstruction when
$C=1$ (single coil). To evaluate the effectiveness of the proposed approach,
we used PDHG to solve (10) and compared the results with gridding and channel-
by-channel (or echo-by-echo) CS reconstructions.

###  3.2 Human Subjects and Data Acquisition

With IRB approval and informed consent, 7 subjects were recruited, including 4
adult patients with known and suspected iron overload and 3 healthy adult
volunteers. The subjects underwent imaging on a 3T clinical MRI scanner (Signa
Premier, GE Healthcare, Waukesha, WI) using a 3D multi-echo cones MRI method
with free breathing (FB Cones) based on the implementation described in [7,
14]. Subsequently, a commercially available Cartesian-based mGRE sequence
(IDEAL-IQ, GE Healthcare, Waukesha, WI) was performed with a single breath-
hold as a clinical reference (BH Cartesian).

One of the healthy subjects underwent two FB Cones acquisitions, one with
regular breathing and the other with deep breathing. The imaging parameters
for FB Cones were as follows: Initial TE/$\Delta$TE/TR =
0.032/1.4-1.5/11.4-11.5 ms; $\\#$TEs = 6; FA = 3∘; resolution =
2$\times$2$\times$2 mm3; rBW = 1106-1250 Hz/Px; readout duration = $\sim$1 ms.
The imaging parameters for BH Cartesian were set to the clinical standard for
iron quantification, which can be completed within a single breath-hold. These
parameters include 4X parallel imaging acceleration along the phase and slice
encoding directions and a slice thickness of 6 mm.

Figure 2: Motion-averaged, echo-by-echo MR (A, C, E, and G) and join TE/MR (B,
D, F, and H) reconstructions of end-expiration/inspiration of a healthy
volunteer with deep breathing. The areas indicated by the color arrows are
magnified in separate images and labeled with the same color bounding box.

###  3.3 Image Quality Assessment and ROI-based R2* Measurements

To assess image quality, organ interfaces, such as the liver dome, pulmonary
vasculature/bronchi, and liver/kidney interfaces, were closely examined
visually and/or qualitatively. To facilitate this process, areas of interest
were magnified, and the sharpness of the root sum of squares of the
reconstructed multi-echo images (T2*w) was visually inspected using the
‘imgradient3’ function in MATLAB. R2* maps were computed using a magnitude-
based exponential fitting from reconstructed images from FB Cones as well as
BH Cartesian. To quantify the R2* values in the liver parenchyma, we placed a
6-voxel radius region of interest (ROI) in the 3 consecutive coronal slices of
the liver, avoiding large vessels that could affect the R2* measurements.
Then, mean and std were calculated.

##  4 Results

Figure 3: T2*w, magnitude of gradients, and R2* maps for motion-averaged,
echo-by-echo MR, and joint TE/MR reconstructions of a patient with iron
overload are shown. The second and third columns show end-expiration, and the
fourth and fifth columns show end-inspiration. ROI-based R2* measures are
shown at the bottom of the R2* maps.

Fig. 1(C-H) shows reconstructed RGB color images from the Fourier samples
“acquired” with the radial and variable density spiral trajectories in Fig.
1(A). The color noise introduced by misalignment in the R, G, and B channels
of the input image (Fig. 1(B)) was effectively suppressed in the joint
reconstructions (Fig. 1(G, H)) regardless of the trajectory type. In contrast,
gridding (Fig. 1(C, D)) and channel-wise reconstructions (Fig. 1(E, F))
preserved the color noise/misalignment. This difference is better appreciated
in the line profiles (Fig. 1(I-L)) extracted from a vertical line passing
through the center of the images, which is identical to the white dotted line
in Fig. 1(B). The misaligned edges between the R, G, and B channels in the
input image were reconstructed to have common edges in the joint
reconstruction (Fig. 1(J, L)), as opposed to the channel-wise reconstruction
(Fig. 1(I, K)).

Fig. 2 shows reconstructed images of a healthy volunteer who performed deep
breathing during the acquisition. Both echo-by-echo motion-resolved (echo-by-
echo MR) and joint multi-echo/motion-resolved (joint TE/MR) methods
demonstrated the capability to resolve end-expiratory/-inspiratory motion
phases, in contrast to motion-averaged reconstruction. To facilitate a visual
comparison between echo-by-echo MR and joint TE/MR reconstructions, areas of
interest are indicated by colored arrows and magnified for closer examination.
The joint TE/MR reconstruction provided a clearer visualization of the liver
dome (Fig. 2(A, B, E, and F)) and pulomary vasculature/bronchi (Fig. 2(G, H))
when compared to the echo-by-echo reconstruction. In addition, the spurious
noise present in the echo-by-echo MR reconstruction (Fig. 2(C)) was
successfully suppressed in the joint TE/MR reconstruction (Fig. 2(D)).

Fig. 3 displays T2*w, magnitude of gradients, and R2* maps of a patient with
iron overload, and compares motion-averaged, echo-by-echo MR, and joint TE/MR
reconstructions. The echo-by-echo MR and joint TE/MR reconstructions exhibit
sharper image quality than the motion-averaged reconstruction, which shows a
blurred liver dome. In the area of interest indicated by the red arrows, both
end-expiratory and end-inspiratory phases of the joint TE/MR reconstruction
exhibit visually sharper image quality than the echo-by-echo MR
reconstruction. ROI-based R2* value obtained from the joint TE/MR
reconstruction was closest to that obtained from BH Cartesian (230 s-1; not
shown).

The ROI-based R2* measurements of all the human subjects imaged for this study
are reported in Table 1. As shown, the joint TE/MR reconstruction (second
column from the right; show in bold) produced the lowest R2* values compared
to the other reconstruction methods and was closest to those obtained from the
clinical standard BH Cartesian method.

Table 1: ROI-based R2* measurements of healthy volunteers (HVs) and patients. | FB Cones | FB Cones | FB Cones | BH Cartesian
---|---|---|---|---
| motion-averaged | echo-by-echo MR | joint TE/MR | reference
HV$\\#$1 | 106.72 $\pm$ 19.05 | 70.18 $\pm$ 15.46 | 65.51 $\pm$ 12.41 | 41.89 $\pm$ 14.81
HV$\\#$2 | 94.58 $\pm$ 21.98 | 76.19 $\pm$ 23.45 | 71.73 $\pm$ 22.51 | 58.37 $\pm$ 17.34
HV$\\#$3 (regular) | 117.36 $\pm$ 25.89 | 87.59 $\pm$ 14.34 | 80.79 $\pm$ 13.89 | 62.59 $\pm$ 18.93
HV$\\#$3 (deep) | 246.08 $\pm$ 38.41 | 96.13 $\pm$ 25.53 | 80.75 $\pm$ 17.68 | 39.35 $\pm$ 16.79
Patient$\\#$1 | 261.37 $\pm$ 28.54 | 258.94 $\pm$ 34.15 | 252.25 $\pm$ 31.83 | 226.18 $\pm$ 32.07
Patient$\\#$2 | 311.98 $\pm$ 49.22 | 280.19 $\pm$ 49.81 | 264.52 $\pm$ 47.50 | 224.97 $\pm$ 30.48
Patient$\\#$3 | 397.15 $\pm$ 51.95 | 393.47 $\pm$ 66.39 | 378.08 $\pm$ 64.77 | 304.43 $\pm$ 58.74
Patient$\\#$4 | 489.12 $\pm$ 44.01 | 468.00 $\pm$ 52.00 | 445.23 $\pm$ 56.12 | 400.32 $\pm$ 43.79

##  5 Conclusion

In this paper, we have proposed a novel approach for handling multi-echo
k-space raw data in respiratory motion-resolved CS reconstruction of highly-
undersampled free-breathing 3D cones acquisition for liver R2* MRI. By making
a key observation that multi-echo images in MRI are somewhat similar to the R,
G, and B channels in a color image and demonstrating its effectiveness in a
synthetic image reconstruction experiment, we have utilized
vectorial/collaborative TV to suppress inconsistencies between echoes in _in
vivo_ experiments. The proposed joint TE/MR reconstruction has been shown to
successfully suppress “echo inconsistency” compared to echo-by-echo
reconstructions, enabling 5D (3D space + 1D echo + 1D motion state) image
reconstruction from continuously acquired non-Cartesian multi-echo k-space
data. Future work will include exploring other convex/non-convex combinations
of the echoes.

In conclusion, the proposed joint TE/MR reconstruction method successfully
suppressed “echo inconsistency” observed in echo-by-echo motion-resolved CS
reconstruction of free-breathing 3D multi-echo cones liver MRI.

## References

  * [1] Armstrong, T., Dregely, I., Stemmer, A., Han, F., Natsuaki, Y., Sung, K., Wu, H.H.: Free-breathing liver fat quantification using a multiecho 3 d stack-of-radial technique. Magnetic resonance in medicine 79(1), 370–382 (2018)
  * [2] Chambolle, A., Pock, T.: A first-order primal-dual algorithm for convex problems with applications to imaging. Journal of mathematical imaging and vision 40(1), 120–145 (2011)
  * [3] Dixon, W.T.: Simple proton spectroscopic imaging. Radiology 153(1), 189–194 (1984)
  * [4] Duran, J., Moeller, M., Sbert, C., Cremers, D.: Collaborative total variation: a general framework for vectorial tv models. SIAM Journal on Imaging Sciences 9(1), 116–151 (2016)
  * [5] Feng, L., Axel, L., Chandarana, H., Block, K.T., Sodickson, D.K., Otazo, R.: Xd-grasp: golden-angle radial mri with reconstruction of extra motion-state dimensions using compressed sensing. Magnetic resonance in medicine 75(2), 775–788 (2016)
  * [6] Goldluecke, B., Strekalovskiy, E., Cremers, D.: The natural vectorial total variation which arises from geometric measure theory. SIAM Journal on Imaging Sciences 5(2), 537–563 (2012)
  * [7] Gurney, P.T., Hargreaves, B.A., Nishimura, D.G.: Design and analysis of a practical 3d cones trajectory. Magnetic Resonance in Medicine 55(3), 575–582 (2006)
  * [8] Hernando, D., Kramer, J.H., Reeder, S.B.: Multipeak fat-corrected complex r2* relaxometry: theory, optimization, and clinical validation. Magnetic resonance in medicine 70(5), 1319–1331 (2013)
  * [9] Hernando, D., Levin, Y.S., Sirlin, C.B., Reeder, S.B.: Quantification of liver iron with mri: state of the art and remaining challenges. Journal of Magnetic Resonance Imaging 40(5), 1003–1021 (2014)
  * [10] Kang, M., Jafari, R., Behr, G., Otazo, R., Kee, Y.: Free-breathing liver qsm with high isotropic resolution using respiratory motion-resolved 3d multi-echo cones mri. In: Proceedings of the 2022 ISMRM & ISMRT Annual Meeting, London, England, UK (Abstract $\\#$0116)
  * [11] Kang, M., Wen, Y., Carl, M., Behr, G., Otazo, R., Kee, Y.: Free-breathing r2* mapping for hepatic iron quantification using respiratory motion-resolved 3d multi-echo ute cones mri. In: Proceedings of the 2022 ISMRM & ISMRT Annual Meeting, London, England, UK (Abstract $\\#$1338)
  * [12] Kee, Y., Deh, K., Dimov, A., Spincemaille, P., Wang, Y.: Primal-dual and forward gradient implementation for quantitative susceptibility mapping. Magnetic resonance in medicine 78(6), 2416–2427 (2017)
  * [13] Kee, Y., Liu, Z., Zhou, L., Dimov, A., Cho, J., De Rochefort, L., Seo, J.K., Wang, Y.: Quantitative susceptibility mapping (qsm) algorithms: mathematical rationale and computational implementations. IEEE Transactions on Biomedical Engineering 64(11), 2531–2545 (2017)
  * [14] Kee, Y., Sandino, C.M., Syed, A.B., Cheng, J.Y., Shimakawa, A., Colgan, T.J., Hernando, D., Vasanawala, S.S.: Free-breathing mapping of hepatic iron overload in children using 3d multi-echo ute cones mri. Magnetic Resonance in Medicine 85(5), 2608–2621 (2021)
  * [15] Li, Q., Dhyani, M., Grajo, J.R., Sirlin, C., Samir, A.E.: Current status of imaging in nonalcoholic fatty liver disease. World journal of hepatology 10(8), 530 (2018)
  * [16] Ong, F., Uecker, M., Lustig, M.: Accelerating non-cartesian mri reconstruction convergence using k-space preconditioning. IEEE transactions on medical imaging 39(5), 1646–1654 (2019)
  * [17] Pruessmann, K.P., Weiger, M., Börnert, P., Boesiger, P.: Advances in sensitivity encoding with arbitrary k-space trajectories. Magnetic Resonance in Medicine: An Official Journal of the International Society for Magnetic Resonance in Medicine 46(4), 638–651 (2001)
  * [18] Reeder, S.B., Pineda, A.R., Wen, Z., Shimakawa, A., Yu, H., Brittain, J.H., Gold, G.E., Beaulieu, C.H., Pelc, N.J.: Iterative decomposition of water and fat with echo asymmetry and least-squares estimation (ideal): application with fast spin-echo imaging. Magnetic Resonance in Medicine: An Official Journal of the International Society for Magnetic Resonance in Medicine 54(3), 636–644 (2005)
  * [19] Schneider, M., Benkert, T., Solomon, E., Nickel, D., Fenchel, M., Kiefer, B., Maier, A., Chandarana, H., Block, K.T.: Free-breathing fat and r2* quantification in the liver using a stack-of-stars multi-echo acquisition with respiratory-resolved model-based reconstruction. Magnetic resonance in medicine 84(5), 2592–2605 (2020)
  * [20] Wang, Y., Liu, T.: Quantitative susceptibility mapping (qsm): decoding mri data for a tissue magnetic biomarker. Magnetic resonance in medicine 73(1), 82–101 (2015)

### Suplementary Materials

We provide the results of a phantom study performed as supplementary
materials, which strongly supports the conclusion drawn from the _in vivo_
human subject study in the main manuscript.

Figure 4: Configuration of a phantom containing vials with gadolinium (Gd)
solutions at different concentrations is illustrated in (A). The Gd solution
(Gadavist, Bayer Pharmaceuticals, Wayne, NJ) was diluted with saline and 1$\%$
agarose gel solution to eight concentrations ranging from 25 to 200 mM. Each
of the diluted Gd solutions (25, 50, 75, 100, 125, 150, 175, and 200 mM) was
placed in a 50 mL Falcon tube and contained within a polystyrene foam
container. During acquisition, periodic table motion was applied with the
parameters prescribed in (B). Reconstructed images (at TE1) using the proposed
method with five motion states are shown in (C). Note that the first column in
(C) may be considered the end-expiratory motion state in the _in vivo_ study
reported in the main manuscript. In addition to the proposed method, echo-by-
echo MR reconstruction and gridding reconstruction without table motion were
also performed.

Figure 5: Reconstructed images (gridding without table motion, echo-by-echo MR
with table motion, and join TE/MR recons with table motion) at TE2 and TE3 are
shown in the first and second columns of the top row. For the echo-by-echo MR
and joint TE/MR reconstructions, the motion state where all the vials are
inside the prescribed FOV, or that corresponds to end-expiratory motion state,
is shown. As indicated by the red arrows, the echo-by-echo MR reconstruction
exhibts poor SNR in the vials with high concentrations of Gd solution ([Gd]),
compared to joint TE/MR reconstruction, which shows similar image quality to
the gridding without motion. As also indicated by the green arrows, the
residual _unresolved_ motion in the echo-by-echo MR reconstruction is clearly
suppressed in the joint TE/MR reconstruction. These reconstruction artifacts
manifest as inaccuracies in R2* measurements, as shown in the [Gd]-R2* scatter
plot (third column of the top row) as well as in the Bland-Altman plots
(bottom row). The ROI measurements from the joint TE/MR reconstruction are
closer to those of the gridding reconstruction without bias than the echo-by-
echo MR reconstruction. Note that the vertical axis of the Bland-Altman plot
is (joint TE/MR or echo-by-echo MR $-$ gridding)/gridding $\times$ 100 [$\%$].