# SOLAR: A Highly Optimized Data Loading Framework for Distributed Training of
CNN-based Scientific Surrogates

Baixi Sun  Indiana University BloomingtonINUSA<EMAIL_ADDRESS>,  Xiaodong
Yu111Authors to whom correspondence should be addressed: Xiaodong Yu
<EMAIL_ADDRESS>and Dingwen Tao (ditao@iu.edu). Argonne National Lab LemontILUSA
<EMAIL_ADDRESS>,  Chengming Zhang  Indiana University BloomingtonINUSA
<EMAIL_ADDRESS>,  Jiannan Tian  Indiana University BloomingtonINUSA<EMAIL_ADDRESS>,
Sian Jin  Indiana University BloomingtonINUSA<EMAIL_ADDRESS>,  Kamil Iskra
Argonne National Lab LemontILUSA<EMAIL_ADDRESS>,  Tao Zhou  Argonne
National Lab LemontILUSA<EMAIL_ADDRESS>,  Tekin Bicer  Argonne National Lab
LemontILUSA<EMAIL_ADDRESS>,  Pete Beckman  Argonne National Lab LemontILUSA
<EMAIL_ADDRESS>and  Dingwen Tao111Authors to whom correspondence should be
addressed: Xiaodong Yu<EMAIL_ADDRESS>and Dingwen Tao (ditao@iu.edu). Indiana
University BloomingtonINUSA<EMAIL_ADDRESS>
###### Abstract.

CNN-based surrogates have become prevalent in scientific applications to
replace conventional time-consuming physical approaches. Although these
surrogates can yield satisfactory results with significantly lower computation
costs over small training datasets, our benchmarking results show that data-
loading overhead becomes the major performance bottleneck when training
surrogates with large datasets. In practice, surrogates are usually trained
with high-resolution scientific data, which can easily reach the terabyte
scale. Several state-of-the-art data loaders are proposed to improve the
loading throughput in general CNN training; however, they are sub-optimal when
applied to the surrogate training. In this work, we propose SOLAR, a surrogate
data loader, that can ultimately increase loading throughput during the
training. It leverages our three key observations during the benchmarking and
contains three novel designs. Specifically, SOLAR first generates a pre-
determined shuffled index list and accordingly optimizes the global access
order and the buffer eviction scheme to maximize the data reuse and the buffer
hit rate. It then proposes a tradeoff between lightweight computational
imbalance and heavyweight loading workload imbalance to speed up the overall
training. It finally optimizes its data access pattern with HDF5 to achieve a
better parallel I/O throughput. Our evaluation with three scientific
surrogates and 32 GPUs illustrates that SOLAR can achieve up to 24.4$\times$
speedup over PyTorch Data Loader and 3.52$\times$ speedup over state-of-the-
art data loaders.

PVLDB Reference Format:
PVLDB, 16(1): XXX-XXX, 2022.
doi:XX.XX/XXX.XX ††This work is licensed under the Creative Commons BY-NC-ND
4.0 International License. Visit https://creativecommons.org/licenses/by-nc-
nd/4.0/ to view a copy of this license. For any use beyond those covered by
this license, obtain permission by emailing<EMAIL_ADDRESS>Copyright is held
by the owner/author(s). Publication rights licensed to the VLDB Endowment.
Proceedings of the VLDB Endowment, Vol. 16, No. 1 ISSN 2150-8097.
doi:XX.XX/XXX.XX

PVLDB Artifact Availability:
The source code, data, and/or other artifacts have been made available at
%leave␣empty␣if␣no␣availability␣url␣should␣be␣sethttps://github.com/ANL-FAIR-
SBI/SOLAR (per request by email).

##  1\. Introduction

CNN Surrogates Deep neural networks (DNNs) such as convolutional neural
networks (CNNs) are critical to traditional artificial intelligence (AI) tasks
such as image and vision recognition (He et al., 2016; Voulodimos et al.,
2018), recommender systems (He et al., 2020; Zhang et al., 2019), and natural
language processing (Collins and Duffy, 2001; Alshemali and Kalita, 2020).
Very recently, domain scientists such as physicists and chemists also leverage
neural network techniques to build surrogates for their scientific
applications to substitute conventional physical solutions (Obiols-Sales et
al., 2020; Dong et al., 2020, 2019; Yin et al., 2022; Shi et al., 2022).
Conventional solutions built upon physical principles usually consist of
iterative processes with complex operations. They can provide accurate
results, however, are quite time-consuming, even with the accelerations on
high-performance computing (HPC) platforms. For instance, the traditional
iterative phase retrieval algorithm for ptychographic image reconstruction
takes hours or even days to converge on multiple GPUs with only a medium-scale
dataset (Yu et al., 2022). To accelerate the reconstruction, physicists
develop a CNN-based autoencoder (named PtychoNN (Pty, [n. d.])) to approximate
the reconstructed images from raw X-ray data and achieve 300$\times$ speedup
compared to the traditional physical approach (Cherukara et al., 2020).
Another example is CosmoFlow (Mathuriya et al., 2018), which is a CNN-based
model used to predict the cosmological parameters in the universe and is more
than 6,000$\times$ faster than the traditional physical approach when
processing large 3D dark matter parameters on a single compute node. All these
models, called CNN-based scientific surrogate models,111We interchangeably
use (CNN) surrogate models and surrogates in this paper. allow scientists to
quickly get acceptable approximate results without performing traditional
time-consuming physical algorithms.

Performance Bottleneck Although CNN surrogates have been approved to be
efficient with small datasets, their training performances with large-scale
data on the HPC platforms remain unclear. To obtain insights, we first
benchmark the training of some representative surrogates. Based on the
benchmarking results (see §3), we observe that, in contrast to the training in
conventional AI tasks, data loading is the major bottleneck of surrogate
training (occupies 98% of training time in some cases). There are two possible
reasons to justify this observation:  (1) The networks in surrogate models are
relatively simple and have a small number of parameters. For example, PtychoNN
(Cherukara et al., 2020) consists of only 1.2 million parameters.  (2) The
training sets of surrogates usually are high-resolution scientific data (Kurth
et al., 2018; Mathuriya et al., 2018), and thus the total dataset sizes are
immense. For instance, the training dataset of CosmoFlow (Mathuriya et al.,
2018) is over 10 TB consisting of $512^{3}$ high-resolution 3D samples (Oyama
et al., 2020).  As a comparison, ResNet-50, a classical network for image
recognition, has 23 million parameters and is typically trained only on a
150GB dataset (e.g., ImageNet (Deng et al., 2009)). Accordingly, while the
computations for updating networks dominate the training performance of
classical AI tasks, the top challenge in surrogate training is to efficiently
manage and load the huge scientific training dataset (Zhang et al., 2021a).

Research Challenges A common approach for reducing the data loading cost is to
buffer the dataset in the near-processor memories, so that loading data from
remote file systems (e.g., PFS) becomes a one-time effort. However, in many
cases, especially with large-scale training datasets, the memory space is not
enough to buffer the entire dataset, and buffer misses will happen. Because of
the random data shuffling, the buffer misses are actually frequent when
loading each local mini-batch. Therefore, we still need to load samples from
either PFS or neighbor nodes (in distributed training) in each training step.
Although the buffering scheme can alleviate the data loading overhead, it
causes three issues:  (i) The buffer hit rate cannot be guaranteed due to data
shuffling in each step.  (ii) The loading workload for each device can be
severely imbalanced.  (iii) It leads to frequent random accesses to small data
fragments, making the parallel I/O libraries (e.g., HDF5 (The HDF Group,
2010)) inefficient.

Existing Works There are a few state-of-the-art data loaders have been
proposed to tackle some of the above issues caused by the buffering scheme.
(1) NoPFS (Dryden et al., 2021) utilizes a heuristic performance model to
predict the data to be used by the next epoch, and accordingly determine the
data eviction scheme to increase the buffer hit rate.  (2) DeepIO (Zhu et al.,
2018) limits the data shuffling within the buffer of each compute node to
eliminate the buffer misses and achieve maximum data reuse.  (3) Lobster (Liu
et al., 2022) improves NoPFS and balances the loading workload by giving
different numbers of loading processes to different devices.  Although these
data loaders can achieve improved performance in general training tasks, they
are sub-optimal when applied to surrogate training with large scientific
datasets.  (1) NoPFS can only approximately optimize the prefetching for only
the next epoch, and thus can not maximize the data reuse in most cases.  (2)
DeepIO significantly reduces the randomness and accordingly will sharply
degrade the training accuracy of surrogates, which is not acceptable as the
outputs of scientific surrogates usually are high-resolution.  (3) Lobster
trades the loading process imbalance for a balanced loading workload. This
tradeoff is less efficient in surrogate training, as data loading takes the
majority time of the entire training, hence the process imbalance exists
during most training time, and its overhead is nontrivial.  Moreover, there is
no existing work addressing the third issue.

Our Contributions To ultimately mitigate the aforementioned data loading
issues in surrogate training, in this work, we design a data loading
framework, called SOLAR, that can drastically improve the loading efficiency
during the distributed training of CNN surrogates. Our design is fundamentally
different from the state-of-the-art data loaders (detailed differences will be
elaborated in §4) and is based on our key observations during surrogate
training:  (i) The shuffled sample indices for all epochs can be determined
prior to the training, and reordering the sample indices within the global
batch makes no change to the randomness and keeps the same training accuracy.
(ii) The computation imbalance is much less harmful than the data loading
workload imbalance during the surrogate training.  (iii) A chunked large data
load can be more efficient than multiple random small data loads with HDF5,
even if the overall loaded size is bigger.  Accordingly, we propose three
novel designs in our framework:  (i) We avoid the runtime shuffling after each
epoch and generate a complete shuffled index list for all epochs ahead of
training. This pre-determined list enables two orthogonal offline
optimizations: (1) re-arranging the epoch order to optimize the buffer
eviction scheme, and (2) changing the mapping from data samples to devices to
increase the data locality. Compared to NoPFS, this design leverages fully
determined sample indices and can maximize data reuse globally.  (ii) We trade
the computation imbalance for a balanced data-loading workload. Different from
Lobster, our tradeoff cost is negligible since computation is only a tiny
portion of the surrogate training.  (iii) We aggregate multiple single sample
accesses within a certain locality into a chunk load to improve the parallel
HDF5 loading throughput.  The contributions of our entire work are summarized
as follows.

  * •

We conduct performance benchmarking on three surrogate models (PtychoNN,
AutoPhaseNN, and CosmoFlow) and identify their major performance issues during
the training.

  * •

We propose SOLAR, a data loading framework using fundamentally different
designs compared to the existing works and can ultimately mitigate the
surrogate training issues.

  * •

We generate the pre-determined shuffled index list in SOLAR and accordingly
introduce two optimizations. This design allows SOLAR to change the epoch
order to optimize the buffer eviction scheme globally and maximize the buffer
hit rate and data reuse.

  * •

We propose a tradeoff between lightweight computational imbalance and
heavyweight data-loading imbalance in SOLAR to improve the overall training
performance.

  * •

We optimize the data access pattern with HDF5 to maximize SOLAR’s parallel I/O
throughput.

  * •

We evaluate SOLAR using three surrogates over five datasets with various
system configurations. The results demonstrate that SOLAR achieves up to
24.43$\times$ reduction in data-loading time and 3.03$\times$ speedup in
overall training performance.

The remaining paper is organized as follows. First, we discuss the research
background in §2 and motivation in §3. Second, we present our design of SOLAR
in §4 and our evaluation results in §5. Third, we discuss related work in §6.
Finally, we conclude our work and discuss future plans in §7.

##  2\. Background

###  2.1. DNN-Based Surrogates

Figure 1. Workflow of ptychographic image reconstruction. The CNN-based
surrogate model achieves hundreds of times speedup over iterative phase
retrieval algorithms (Cherukara et al., 2020).

DNN-based surrogate models are designed to substitute heavy-cost traditional
iterative algorithms in scientific simulations. For example, Figure 1 shows
the workflow of reconstructing images from scanned diffraction patterns.
Specifically, AI-based surrogate models achieve 300$\times$ speedup in the
image reconstruction task compared to the conventional physical approaches,
with minor loss of the image resolution.

However, the problem size is becoming huge due to the growing size of
scientific datasets. For example, the coherent imaging data can be generated
in terabytes per hour by the Advanced Photon Source Upgrade (APS-U) (Noirot et
al., 2018). When the dataset size cannot fit into the memory of a single
device (i.e., CPU, GPU), a straightforward solution is to use mini-batch
training, which loads a small batch of data at a time for each step. Since the
training time would be long with a single device, it often requires
parallelizing the workload to multiple devices using data parallelism,
especially for training DNN-based surrogates with large-scale input data.

###  2.2. Distributed Training with Data Parallelism

Data parallelism techniques are widely utilized for training on large datasets
(Mathuriya et al., 2018; Kurth et al., 2018; Deng et al., 2009; Farrell et
al., 2021). When applying data parallelism on HPC clusters, the training
typically includes three stages: (1) I/O: loading the data from a remote
parallel file system (i.e., GPFS, Lustre) to host memory; (2) Computation:
performing forward and backward phases to calculate the local gradient on each
device; (3) Communication: synchronizing averaged gradients across multiple
devices to update model weights. Among these stages, I/O is a bottleneck for
distributed training with large-scale datasets (Dryden et al., 2021; Pumma et
al., 2019; Zhu et al., 2018). Thus, SOLAR is designed primarily to optimize
the data loading stage in training DNN-based surrogates.

The access order is crucial for data reuse in data loading (Pumma et al.,
2017). In most DNN training, each data sample is accessed only once during an
epoch. Thus, to reuse the data, the buffer needs to keep the samples that will
be used soon and evict those that will be used further in the future (Dryden
et al., 2021). However, random shuffling makes it hard to achieve a good reuse
rate and reduce data loading time. Existing approaches such as DeepIO (Zhu et
al., 2018) change the access order to achieve a better data reuse rate.
However, these methods reduce the randomization of the dataset and degrade the
training accuracy. On the other hand, an existing work (Yang and Cong, 2019)
proved that by adjusting the access order in a particular way, the
synchronized gradients will be the same as the original random shuffle.
Specifically, it proves that reordering the data samples within the global
batch produces the same gradients after synchronization. This inspires us to
create a better data reuse method without affecting the training accuracy.

###  2.3. I/O Frameworks for Distributed Training

DataLoader in PyTorch (Paszke et al., 2019) is the default data loading tool
and for data parallelism. It has two optimizations: (1) using a subprocess to
read data in order to fully utilize the bandwidth; (2) prefetching the
training data to overlap data loading and computation. However, when the data
loading dominates the epoch, GPUs are idle and starve due to a lack of data.
To this end, some works store the data that are accessed in the current epoch
to reuse it in the next epoch.

Dryden et al. (Dryden et al., 2021) proposed NoPFS, a prefetching and
buffering strategy to reduce the data loading time by reusing the accessed
data samples. The optimization is twofold: (1) combining the clairvoyant
method with the distribution of samples on multiple nodes to optimize data
eviction policy and (2) designing a performance model to leverage the multi-
layer architecture of the HPC cluster. Liu et al. proposed Lobster (Liu et
al., 2022), a dynamic thread arrangement and data eviction strategy that
further improves the data-loading performance by balancing the data-loading
process across multiple GPUs. Their optimizations are twofold: (1) rearranging
the threads to read and preprocess the data samples to balance the data
loading, and (2) tuning the data sample eviction to achieve better data reuse.
However, both methods focus only on the online buffer eviction strategy and do
not obtain the same observations as SOLAR; thus, their performance
improvements in data loading are limited for training surrogate models.

##  3\. Research Motivation

In this section, we present the benchmarking results of the distributed
surrogate training. We first illustrate the scalability of training a typical
CNN-based surrogate with three DL frameworks. We then break down the training
time and exhibit that data loading is the major bottleneck taking up to 98.6%
of the training time. These results motivate us to explore more efficient
data-loading solutions for surrogate training.

We conduct the benchmarking on ThetaGPU (the, [n. d.]) (each node has eight
A100 GPUs) at the Argonne Leadership Computing Facility, following the
guidelines from MLPerf HPC (Farrell et al., 2021).

Scalability Study We first perform a study to show the scalability of
distributed training over surrogate models with widely used deep learning
frameworks. Specifically, we implement PtychoNN with TensorFlow mirrored
strategy (Abadi et al., 2016), Horovod (Sergeev and Del Balso, 2018), and
PyTorch DDP (Li et al., 2020). We train PtychoNN on a 17 GB dataset (will be
detailed in §5) with different numbers of GPUs. Specifically, we take the
average epoch time among 99 epochs, excluding the warmup time in the first
epoch. Figure 2 shows that these three frameworks can provide similar
scalability from a single to 8 GPUs. Therefore, we choose Pytorch DDP as the
distributed computation tool for the experiments in the rest of this paper.

Figure 2. Training time in seconds using different numbers of GPUs with
different frameworks.  Figure 3. Time breakdown on tested datasets when using
prefetch to overlap data loading with computation.  Table 1. Breakdown of
computation time and data loading time (in seconds) when training PtychoNN on
1.2 TB dataset.

#GPU | 32 | 64 | 128
---|---|---|---
data loading | 307.7 |  1.00$\times$ | 159.7 |  1.93$\times$ | 80.2 |  3.83$\times$
| 98.5% |  | 98.6% |  | 98.6% |
computation | 4.7 |  1.00$\times$ | 2.3 |  2$\times$ | 1.1 |  4.13$\times$
| 1.5% |  | 1.4% |  | 1.4% |
total | 312.3 |  1.00$\times$ | 162.0 |  1.93$\times$ | 81.4 |  3.84$\times$

Performance Breakdown We then break down the overall performance to
investigate the bottlenecks in distributed training of surrogate models. We
train three surrogate models on the corresponding datasets, i.e., pythcoNN
(Cherukara et al., 2020) on the Coherent Diffraction (CD) 321GB dataset and
1.2TB dataset, AutoPhaseNN (Yao et al., 2022) on the 151GB Bragg Coherent
Diffraction Imaging (BCDI) dataset, and CosmoFlow (Mathuriya et al., 2018) on
a 1TB subset of the 3D universe simulation (CosmoFlow) dataset (datasets will
be detailed in §5).

The breakdown times are shown in Figure 3 and Table 1. Specifically, Figure 3
shows the time breakdown to train three surrogates on three datasets using
different numbers of GPUs. We also synthesize a 1.2TB dataset for ptychoNN to
indicate how different sizes of training datasets impact the training
performance of the same surrogate model, as shown in Table 1.

In Figure 3 we can observe that the data loading is the major performance
bottleneck. Specifically, when using four GPUs, data loading takes 83.1%,
77.3%, and 43.2% of the total training time with PtychoNN, AutoPhaseNN, and
CosmoFlow, respectively. Along with the increase of the GPU number (i.e., a
weak scaling), the computation time reduces quickly whereas the data loading
cost drops very slowly (even increasing in some cases). Consequently, the
data-loading issue becomes more severe when using more computation resources.
For instance, When we increase the GPU number from four to 16 for CosmoFlow
training with the 3D universe simulation dataset, the percentage of data
loading time increases from 43.2% to 73.4% of the total training time. In
addition, in Table 1, we can observe that the data-loading problem also can
become worse when the dataset gets larger, as it takes up to 98.6% of the
total training time when training ptychoNN on the 1.2TB dataset. We note that
we use more GPUs for training with the 1.2TB data in order to keep a
reasonable computation workload on each GPU.

Moreover, we observe that our achieved data-loading throughput with parallel
HDF5 is lower than that of reported by the HDF5 Group (Byna et al., 2020).
This is because random accesses with small data degrade the data-loading
efficiency, which will be discussed in detail in §4.4.

Motivation The above results and observations motivate us to design an
efficient framework to optimize data loading time for CNN-based surrogates by
developing offline scheduling and runtime buffering strategies.

##  4\. Proposed Design of SOLAR

###  4.1. Overview

Figure 4. Our proposed offline scheduling: (a) Obtain access order based on
random shuffle before training; (b) Adjust epoch order by solving a graph
problem to maximize data reuse; (c) Re-arrange node-to-sample mapping to
optimize data locality.

To reduce the data-loading overhead, a common solution is to buffer the loaded
samples in a faster memory space (e.g., DRAM) for reuses in the following
epochs. The buffering approach theoretically can reduce the number of samples
to be re-loaded from PFS. However, it usually cannot reach the expected
throughput improvement due to the following issues:  (1) The buffer hit rate
could be low due to the sample shuffling after each epoch.  (2) The loading
workload of each device is imbalanced because of the buffer hit rate
variability on different devices.  (3) Reloading evicted samples from PFS
becomes random access to small fragmented data, which aligns poorly with the
desired access pattern of parallel HDF5.

Our SOLAR framework addresses the above issues based on our key observations
during the surrogate training:  (1) We can predetermine the shuffled sample
indices for all epochs and reorder the indices within a global batch, without
affecting the training accuracy. (2) The impact of computation imbalance is
trivial compared to that of data-loading workload imbalance.  (3) Loading a
chunked large data load can be faster than loading several small data
fragments with HDF5, even though the chunked one loads redundant data.
Accordingly, SOLAR proposes three novel designs (one for each issue):  (1) It
generates a shuffled index list prior to the training, and then re-arranges
the epoch order and changes the sample-device mapping according to the list to
maximize the data reuse.  (2) It trades the computation imbalance for a
balanced data-loading workload for better overall training performance.  (3)
It uses a large chunked loading to replace small fragmented loadings within a
certain locality to improve the loading throughput.

We implement these designs in SOLAR as an offline scheduler and a runtime
buffering strategy, which will be demonstrated in Figure 4 and Figure 5. The
offline scheduler includes optimizing the access order and computing the
temporal (when to load) and spatial (where the data is located) information
about each data sample. The runtime buffering strategy includes data eviction
optimization and an optimized chunk-loading mechanism. Also, the workflow and
architecture for the runtime buffering strategy is illustrated in Figure 5.
Table 2 lists the notation used throughout this paper. To our best knowledge,
SOLAR is the first work obtaining the above observations and insights in
surrogate training, and accordingly design optimizations for surrogate data
loading.

Table 2. Notation used in this paper. Variable | Unit | Definition
---|---|---
$|\text{Buffer}|$ | image | Buffer Size
$|\text{Batch}|_{l}$ | image | Local Batch Size
$|\text{Batch}|_{g}$ | image | Global Batch Size
$|\text{Dset}|$ | image | Dataset Size
$|\text{chunk}|$ | image | Optimal chunk size in aggregated chunk loading
$G$ |  | A graph in which each vertex represents an epoch
$p$ |  |  A path within graph $G$
$N_{u,v}$ | image |  The number of samples to load from epoch $u$ to epoch $v$
$n$ |  | A compute node on HPC cluster
$E$ |  | Number of epochs
$e_{i}$ |  |  The $i$-th epoch

###  4.2. Maximizing Data Reuse

####  4.2.1. Optimizing Epoch Order

Limitations of Existing Work NoPFS (Dryden et al., 2021) tries to increase the
data reuse by optimizing the buffer eviction scheme. However, since it
generates the shuffle list at runtime after each epoch, it has to utilize the
prediction of data sample distribution with a performance model to determine
the runtime prefetching and caching strategy. Meanwhile, it considers only the
data reuse problem between each pair of epochs. By doing so, the solution
provided by NoPFS to optimize the data-loading performance is suboptimal. In
comparison, SOLAR generates the shuffle list for all epochs offline.
Specifically, our benefits are two-fold: (1) instead of considering only the
data reuse in the next epoch, SOLAR generates the shuffle list for all epochs
so that we can have a global view for optimization; (2) instead of determining
when and where to fetch data online and in a complicated fashion, we perform
offline access order optimizations that maximize the data reuse rate.222Note
that the access order optimization is a one-time overhead.

Observation To address the first issue, we leverage a key observation that
instead of generating during runtime, the shuffled index list can be
determined ahead of training. Since the data reuse rate varies when we change
the order of epochs, we can optimize the epoch order according to the list to
maximize the data reuse. Here, the epoch order is defined as the order of
using the shuffled index lists, where each shuffled index list indicates the
data samples’ access order for an epoch. In distributed training, each data
sample is accessed only once in each epoch. That is to say, we can only reuse
the loaded data in the next epoch. In practice, the result of the random
shuffle (shuffled index list) is determined with a given random seed.
Furthermore, parameters like the number of epochs, number of devices, and
random seed are pre-defined. Thus, instead of performing the random shuffle at
the beginning of each epoch, we can produce the shuffled index list for all
epochs before training and then design offline optimizations based on these
shuffled index lists. In terms of data reuse, the data samples must be loaded
if they have not already been stored in the current process during the last
epoch. Similar to the “cache hit,” a high data reuse rate between two epochs
is preferred to reduce the data loading operations. Through experiments, we
observe that the data reuse rate can vary with the epoch order. Thus, we
propose an algorithm to arrange the epoch order to maximize the data reuse.

Proposed Optimization With the above important observation, we design an
optimization that maximizes data reuse by adjusting the order of epochs.
First, we generate the shuffled index list for all epochs using the defined
random seed, as shown in Figure 4 (a). Instead of performing the shuffle
operation for each epoch, the ahead-of-time random shuffle procedure is
conducted offline to obtain the access order for all epochs before training.
By doing so, we are able to perform offline optimizations according to the
access order. Next, we define the problem as finding an epoch order such that
the data reuse is maximized. Then, we abstract this problem into a graph
problem. Consider the shuffle list as a graph $G$ whose nodes represent all
epochs, and the edge weight denotes the loading cost from one epoch to the
other; we are solving the problem by finding the shortest path that visits
each node exactly once. Let $N_{u,v}$ denote the number of data samples to
load from epoch $u$ to epoch $v$; we compute the edge weight by

(1) $N_{u,v}=\mathbf{card}(\text{Buffer}_{v}-\text{Buffer}_{u}),$

where $\textit{Buffer}_{v}$ means the set of data in epoch $v$’s last buffer,
$\textit{Buffer}_{u}$ means the set of data in epoch $u$’s first buffer, and
$\mathbf{card}$ means the cardinality of a set. Note that $N_{u,v}$ cannot be
the same as $N_{v,u}$ by definition. Given a buffer size
$|\textit{Buffer\hskip 1.0pt}|$, the first buffer means the set of first
$|\textit{Buffer\hskip 1.0pt}|$ elements in epoch $u$, and so forth. To find
an optimized solution to schedule the order of epochs, we have

(2)
$\min\left(\text{cost}_{\text{I/O}}\right)=\min\Big{(}\sum^{p}N_{u,v}\Big{)},$

where $p$ is a path and $\text{cost}_{\text{I/O}}$ denotes the modeled total
cost of path $p$. Thus, the problem now is to find the shortest path in graph
$G$ such that each vertex in the graph is visited exactly once.

This problem is similar to the well-known Traveling Salesman Problem (TSP)
(Flood, 1956). The differences are: (1) It does not need to go back to the
origin vertex, which is defined as path-TSP and proved to be NP-complete
(Papadimitriou, 1977); (2) The edges are bi-directional in terms of weights.
Leveraging heuristic algorithms can enable a near-optimal solution to this
problem.

Implementation Details We utilize Particle Swarm Optimization (PSO) to solve
the TSP problem. PSO (Kennedy and Eberhart, 1995) is suitable for providing
the near-optimal solution for this problem through observation, and it takes a
graph as the input and outputs the shortest path that visits each vertex
exactly once.

Specifically, PSO is an approximation algorithm that originated from the
social behavior of swarms. First, each agent starts from a random vertex
within the graph, and its fitness function is defined. The output of the
fitness function refers to the distance between the current solution of the
agent and the best solution among agents (Shi et al., 2007). Second, the swarm
of agents moves synchronously so that the local-best solution is aggregated
and updated to replace the global best solution at each iteration. Finally,
the algorithm will terminate if all the vertices in the graph have been
visited and will output the shortest path as well as the length of the
shortest path.

####  4.2.2. Boosting Data Locality

Limitations of Existing Work DeepIO (Zhu et al., 2018) tries to maximize the
data reuse by increasing the data locality. It limits the sample shuffling
within each device so that the buffered data can ultimately be reused.
However, this approach seriously hurts the randomness and sharply reduces the
training accuracy. This limits DeepIO being applicable in some cases, because
(1) For some tasks like image classification on Imagenet, the accuracy is over
90%, thus trading off training accuracy is not a good choice. (2) The
relationship between randomness and training accuracy is unknown for many
neural networks and datasets, thus, we cannot know if it is worth utilizing
the framework until trying it for each network and dataset.

Observation To boost the data locality without sacrificing training accuracy,
we design our optimization based on a key observation: with the determined
shuffled index list, changing the access order and node-to-sample mapping
within a global batch will still result in the same gradient after
synchronization (as proved in (Yang and Cong, 2019)). Here, the node-to-sample
mapping is defined as the relationship of which data sample to access in
training on which node. Since the dataset is fully randomized in each epoch,
the node-to-sample mapping will vary across epochs. It incurs extra data
loading from the file system or point-to-point communication to exchange the
needed data samples. Meanwhile, generating the mini-batch from locally
buffered data samples significantly affects the training accuracy (Zhu et al.,
2018). Luckily, data parallelism is a synchronous training mechanism that
averages the local gradients computed from mini-batches of data. These mini-
batches from multiple devices are together named the global mini-batch or
global batch. Specifically, changing the training order of data samples within
a global batch does not affect the gradient results after synchronization
(Yang and Cong, 2019). Consider the following update function in data
parallelism:

(3) $w_{t+1}=w_{t}-\textstyle\eta{N_{B}}^{-1}\sum^{N}_{k=1}\sum_{x\in
B_{k,t}}\nabla f(x,w_{t})$

where $w_{t}$ denotes the weight at iteration $t$, $\eta$ denotes the learning
rate, $N$ is the number of devices and $B_{k,t}$ is the mini-batch on device
$k$ of iteration $t$, and $\nabla f(x,w_{t})$ is the local gradient. From the
equation, we can see that so long as the data samples in a global batch
($\sum^{N}_{k=1}\sum_{x\in B_{k,t}}\nabla f(x,w_{t})$) are trained together,
changing the node to sample mapping within the global batch produces the same
gradient result after gradient averaging. Thus, we can leverage this
characteristic of data parallelism to achieve a higher in-memory data reuse
rate without affecting the training accuracy.

Proposed Optimization With the key observations mentioned above, we are able
to re-arrange the node-to-sample mapping to optimize data locality and thus
improve the data reuse rate. Specifically, we propose to improve the data
reuse rate by updating the node-to-sample mapping based on locality, instead
of requesting intra-node data movement via point-to-point communication, as
shown in Figure 4 (c). Through offline computation, we can acquire the
knowledge of node-to-sample mapping in the current epoch. With this
information, we can update the mapping relationship within each global mini-
batch for the next epoch to benefit data reuse. For example, for the original
shuffle list, sample $s$ is trained on epoch $e_{i}$ on node $n_{x}$. For the
next epoch $e_{i+1}$, it should be trained in step $j$ on node $n_{y}$. We
arrange the access order to make the sample $s$ trained on node $n_{x}$ still
on epoch $e_{i+1}$ and step $j$.

The node-to-sample mapping optimization in SOLAR is inspired by an interesting
thought: How do you identify that the color you see is the same as others? For
example, everyone names colors like “Red,” “Green,” and “Blue,” but how do
people know that they see the same color in their vision? One answer can be:
Learning name-to-color mapping is enough. It is similar to the data locality
problem in distributed training in that we only need to remap the samples and
nodes within a global batch.

Implementation Details The data locality optimization is implemented by re-
arranging the access order of each epoch according to its previous epoch. By
doing so, we can avoid node-to-node communication. Specifically, we only need
to arrange the access order within the global mini-batch by computing the
samples accessed and stored in the last epoch, and remapping the buffered
samples to the corresponding node for the coming epoch. The rest of the data
samples will be evicted during runtime after the buffer is populated because
their interval between the last accessed time and the next access time is
large. By doing so, we efficiently utilize the limited buffer storage space
and achieve maximum data reuse.

Figure 5. Data loading from PFS to GPUs using our proposed runtime buffering
strategy with two optimizations.

###  4.3. Loading Workload Balancing

Limitations of Existing Work There are a few existing works that try to solve
the second (imbalance) issue. Locality aware (Yang and Cong, 2019) utilizes
intra-node communication to balance the data loading during runtime. By doing
so, it has to schedule the sending source and receive target for each data
sample which is an NP-Complete problem and can only be approximately solved.
This limitation will lead to performance loss since we still have extra I/O
overhead when the approximation algorithm is not accurate. Lobster (Liu et
al., 2022) tackles the imbalance problem by re-assigning preprocessing threads
for data loading. It essentially trades the thread imbalance for a balanced
loading workload. In surrogate training, since data loading takes the majority
of the training time, the thread imbalance will exist during most of the
training process, and thus its system overhead will be amplified. Besides,
surrogates, like PtychoNN, do not require complex data preprocessing and data
augmentation, hence cannot directly utilize Lobster.

Figure 6. An example illustrating performance improvements with data load
balancing.

Observation In order to overcome the imbalance issue, we leverage an
observation that is important for our load-balancing optimization:
computational imbalance is less harmful to overall training time compared to
the load imbalance. In data parallelism, due to the random shuffle, the number
of data samples that can be reused on each node is not the same. In other
words, the number of samples that need to be loaded from the PFS is different
across nodes, resulting in imbalanced I/O. Furthermore, this imbalance causes
a large amount of idle time for GPUs as they must wait for the data to be
loaded from the file system to the host memory, formatted into a batch, and
transferred to GPUs. The importance of data load balancing to training
performance is illustrated in Figure 6. We can see that when the data loading
time is significantly larger than the computation time, prefetch cannot
overlap with the computation. After the first optimization in the middle
example, the number of samples loaded from the parallel file system is
reduced, thus the I/O time is reduced. Meanwhile, the data imbalance is
introduced, which limits the data loading efficiency. In the bottom example,
after data load balancing, the overall training time is reduced.

Figure 7. Comparison of computation time on each GPU between original training
and our solution with the load-balancing trade-off, demonstrating that the
imbalanced batch size does not affect performance.

The voyage to explore possible data load balancing optimization designs is
stormy. Fortunately, when we performed micro-benchmarking experiments to
observe the computation time for balanced and imbalanced training batch sizes,
we discovered that the computation time does not differ much between the two
cases. Specifically, 16 A100 GPUs are used in the experiments, and we first
fixed the batch size for each GPU and trained on that setting; Then we trained
with modified batch sizes where each batch size is the defined batch size
minus the rank index of each GPU. We observe that varying the batch size when
training results in a similar computation time compared to the fixed batch
size training, as shown in Figure 7. The black line shows the average
computation time of 10 epochs when the batch size is set to 64. The gray line
shows the average computation time of 10 epochs when batch size is set to
$|\textit{batch}|=64-\textit{rank}$. Note that the computation time of both
balanced and imbalanced solutions are relatively close to each other,
considering some system variance between GPUs. This means we can trade off the
compute balance for the load balance to improve the overall performance.

Proposed Optimization With the key observation mentioned above, we design the
data load balancing optimization that trades computational balance for load
balance. Specifically, to balance the data loading across nodes, our goal is
to let each node load a similar amount of data during every step. After
designing a buffer for data reuse, the data samples within a mini-batch have
two states: (1) to fetch from the PFS, and (2) to fetch from the in-memory
buffer. The number of data samples in the state (1) on multiple GPUs is not
the same, thus the data loading is not balanced for the data parallelism. From
our observation, we design an optimization to distribute the number of data
samples evenly in the state (1) across multiple GPUs. By doing so, the
training batch sizes on the GPUs are different. As shown in equation 3, since
we did not change the global batch, this optimization will not affect training
accuracy. Furthermore, we observe that changing the training batch size does
not affect computation time much, and we demonstrate the importance of data
load balancing in Figure 6. This is the trade-off: we diminish the computation
balance to achieve load balance.

Implementation Details During offline scheduling, we compute the indices of
the data that need to be fetched from the PFS. During runtime, we use those
indices to perform the computation balance trade-off, as shown in Figure 5.
Once the mini-batch of data is ready on a GPU, the computation part of the
training begins.

###  4.4. Aggregated Chunk Loading

HDF5 (The HDF Group, 2010) is an I/O library designed for HPC applications and
achieves high read/write throughput. However, we observe that the throughput
of data loading with parallel HDF5 is low in our training. Specifically, while
loading the same amount of data in total, random access to many small files
incurs higher loading time compared to chunked loading. This raises the
question of how different access patterns can affect the data-loading
performance. Specifically, there are four flavors of access patterns, (1)
random access, in which a process reads one sample randomly until all samples
in the dataset have been accessed once, (2) sequential-stride access, in which
a process iteratively reads samples with a fixed stride, (3) chunk-cycle
loading, in which a process loads samples one by one in its assigned chunk,
and (4) full-chunk loading, in which a process loads multiple consecutive
samples in its assigned chunk in one go.

Figure 8. An example of offsets visited in an HDF5 file with different access patterns. Table 3. I/O time (in seconds) of different access patterns. Pattern | Time | Norm’ed | Speedup
---|---|---|---
Random Access | 645.864 |  203.42$\times$ |  1.00$\times$
Sequential Stride Access | 84.421 |  26.59$\times$ |  7.65$\times$
Chunk Cycle Loading | 30.537 |  9.62$\times$ |  21.15$\times$
Full Chunk Loading | 3.175 |  1.00$\times$ |  203.42$\times$

We set up a microbenchmark to profile the I/O performance under different
access patterns. Specifically, we utilize the record profiling tool (UIUC, [n.
d.]) to visualize the I/O activity with the accessed offsets in the HDF5 file.

Observation Figure 8 shows the trace of different patterns. The random access
pattern results in uneven distribution of offsets on a single process. The
offset reveals the same sequential feature for sequential access and chunked
cycle loading. Table 3 shows the measured I/O time; we can see that sequential
access is faster than random access since each process spends a lot of time
seeking the data for random access. On the other hand, chunk cycle loading
takes less time than sequential access because consecutive access achieves
better I/O performance than stride access in seeking the following data.
Overall, full chunk loading achieves the best performance with a 203$\times$
speedup compared with random access.

Proposed Optimization Our optimization strategy is twofold: (1) we rearrange
the data loading to be as consecutive as possible, and (2) changing the access
order within the mini-batch maintains the randomness, equivalent to the
original random shuffle. We first find the optimal chunk size from the micro-
benchmark experiment. Then we determine if the current and the next sample to
load can be combined in the same chunked loading, based on the known indices
from our proposed offline scheduling.

Implementation Details PyTorch provides the `__getitem__` function in the
custom dataset option that accesses one data sample at a time (Paszke et al.,
2019), then the collate function is called to collate the data samples into a
batch of data. We utilize the interface by (1) sorting the indices of the
samples to be loaded in a mini-batch, (2) determining which samples can be
loaded in a chunk, and (3) loading multiple samples in chunks, as shown in
Figure 5.

##  5\. Performance Evaluation

In this section, we answer three main questions: (1) What is the performance
of SOLAR compared to state-of-the-art approaches on different datasets with
different system setups? (§5.2) (2) What is the contribution of each
optimization to SOLAR? (§5.3) (3) What is the end-to-end training performance
with SOLAR? (§5.4)

###  5.1. Experimental Setup

Our framework is implemented using PyTorch. In our experiments, we use
DGX-A100 GPU nodes with Lustre parallel file system.

Figure 9. Comparison of data-loading performance on three systems using five
datasets.

Baselines We use two baselines for comparison, namely PyTorch DataLoader and
NoPFS. PyTorch Dataloader (Li et al., 2020) is the data loading framework that
is widely used by the industry and academia. It is designed to load the data
for training and validating the neural networks implemented by PyTorch. NoPFS
(Dryden et al., 2021) is an existing data-loading framework that claims to
outperform state-of-the-art approaches, including PyTorch DataLoader, DeepIO
(Zhu et al., 2018), LBANN (Jacobs et al., 2019), and Locality-aware (Yang and
Cong, 2019).

Datasets We evaluate SOLAR on five datasets. (1) Coherent Diffraction Dataset
(CD) has 262,896 images (17 GB) in total, and the size of each image is 65KB.
We also synthesize the dataset of two sizes (321 GB, 1.2 TB) to study the
performance of using larger-scale datasets, which have 1,752,660 and
18,928,620 samples, respectively. (2) 3D X-ray Bragg coherent diffraction
imaging (BCDI) Dataset has 54,030 samples in total, each sample is 3.1MB. (3)
3D Universe Simulation (CosmoFlow) Dataset has 63,808 samples (1 TB), each one
being the 3D universe simulation data of size 17MB.

Buffer sizes The performance of SOLAR is affected by a few factors, among
which the buffer size $|\textit{Buffer\hskip 1.0pt}|$ is the major one. Thus,
we need to study the improvement of SOLAR compared to the baselines using
different buffer sizes. Specifically, there are three scenarios: (1) dataset
size $\leq$ local buffer size, (2) local buffer size $<$ dataset size $\leq$
total buffer size, (3) dataset size $>$ total buffer size.

Table 4. Number of GPUs used for different datasets. Buffer size is the number of GPUs multiplied by 8/16/40 GBs for low-/medium-/high-end systems, respectively. | CD |  |
---|---|---|---
System | 17 GB | 321 GB | 1.2 TB | BCDI | CosmoFlow
Low-end | 2 | 16 | 32 | 8 | 16
Midium-end | 2 | 16 | 32 | 8 | 16
High-end | 2 | 8 | 16 | 8 | 16

###  5.2. Evaluation of I/O Improvement

We first show the I/O improvement of SOLAR compared to PyTorch DataLoader and
NoPFS on five datasets under three different system setups. Specifically, as
shown in Table 4, the low-/medium-/high-end systems have 8/16/40 GB buffer per
GPU, respectively. To simplify the problem, we assume one GPU per node to
demonstrate the performance of SOLAR on distributed buffers.

The performance improvement of SOLAR compared to PyTorch DataLoader and NoPFS
is shown in Figure 9. Specifically, the bar shows the speedup compared to
PyTorch DataLoader. To compare SOLAR with PyTorch DataLoader, we measure data
loading time across multiple runs and take the average, then we compute the
speedup. For NoPFS, we run their simulator that is implemented based on
NoPFS’s performance model and outputs the simulated data loading time of
PyTorch DataLoader and NoPFS.333NoPFS only provides a simulator in its open-
source repository. The general trend shows that as the aggregated buffer size
increases, SOLAR achieves better speedup. For example, on the five datasets,
we observe that the speedup of SOLAR compared to PyTorch DataLoader is higher
on high-end systems than on medium-end and low-end systems. This is because
high-end systems have a larger buffer size, which is consistent with our
design.

Next, we present the evaluation results in three scenarios.

(1) dataset size $\leq$ local buffer size: In this scenario, the entire
dataset is loaded on each GPU. For example, on the CD 17 GB dataset with the
high-end system, NoPFS and SOLAR achieve no speedup because during the first
epoch, all data samples are loaded from the PFS and for the following epochs,
all data samples are loaded from the local buffer.

(2) local buffer size $<$ dataset size $\leq$ total buffer size: In this
scenario, SOLAR loads almost all the data samples from the local buffer, while
NoPFS loads half of the samples from the local buffer and half from the remote
buffers (on other nodes). Therefore, SOLAR benefits from data locality
optimization and aggregated chunk loading, so that SOLAR outperforms NoPFS.
For example, on the CD 17 GB dataset with the medium-end system, SOLAR
achieves 14.1$\times$ average (up to 24.4$\times$) speedup compared to PyTorch
DataLoader and 1.9$\times$ average (up to 3.3$\times$) speedup over NoPFS.
Also, on the BCDI dataset with high-end systems, SOLAR achieves 9.6$\times$
average (up to 10.9$\times$) speedup compared to PyTorch DataLoader and
1.9$\times$ average (up to 2.2$\times$) speedup over NoPFS.

(3) dataset size $>$ total buffer size: In this scenario, SOLAR loads the data
samples from the local buffer and the PFS, while NoPFS loads from the local
buffer, the remote buffer, and the PFS. Note that the PyTorch DataLoader does
not use the buffer. Here, SOLAR outperforms NoPFS because it achieves a better
data reuse rate and avoids inter-node data movement compared to NoPFS. For
instance, on the CD 321 GB dataset, CD 1.2 TB dataset, and CosmoFlow dataset,
SOLAR achieves up to 7.96$\times$/1.55$\times$/4.25$\times$ speedup compared
to PyTorch DataLoader, and up to 3.52$\times$/1.23$\times$/3.13$\times$
speedup over NoPFS, respectively. Meanwhile, we observe the performance
improvement limitation of two datasets (CD 321 GB and CD 1.2 TB) on the low-
end system. On the one hand, although the average speedup of SOLAR does not
achieve performance improvement as high as NoPFS, our best run outperforms
NoPFS in both cases. On the other hand, the result shows that when the buffer
size is relatively small compared to the dataset size, the performance
improvement of SOLAR can be limited compared to NoPFS. This is because when
the buffer size is small and the per-sample size is small, the access order
optimization cannot significantly improve the data-loading performance. This
shows the limitations of SOLAR and NoPFS, in that when the dataset is too
large for the buffer, the improvement is not significant compared to PyTorch
DataLoader. However, the worst case for SOLAR is that the performance is close
to NoPFS, since NoPFS does not re-arrange the access order and SOLAR does not
incur extra overhead (inter-node communication).

In addition, a recent work, called Lobster (Liu et al., 2022), develops a
solution to address I/O load imbalances among GPUs during DNN training.
However, we exclude Lobster from our comparison for three main reasons: (1)
while our work focuses on data-loading performance, Lobster also optimizes
pre-processing (such as decoding, shuffling, batching, and augmentation); (2)
Lobster’s code is not available for testing (we discuss more details of
Lobster in §6); and (3) Lobster’s performance improvement over NoPFS is up to
1.2$\times$, while SOLAR can improve the data-loading performance by
2.37$\times$ over NoPFS.

###  5.3. Evaluation of Each Optimization

Next, we break down the performance improvement to demonstrate the
contribution of each optimization. Specifically, Optim_1 is the access order
optimization, Optim_2 denotes data load balancing, and Optim_3 is aggregated
chunk loading. We also implemented a buffer with the least recently used (LRU)
strategy for PyTorch DataLoader (SOLAR–Optim_1–Optim_2–Optim_3).

Figure 10. Performance breakdown for each optimization. Optim_1 is access
order optimization, Optim_2 is data load balancing, and Optim_3 is aggregated
chunk loading.

Figure 10 shows the performance breakdown for each optimization. First, we
observe that compared to PyTorch DataLoader, PyTorch DataLoader + LRU buffer
has a 1.2 $\times$ speedup. That shows the effect of using a buffer for data
reuse. Second, we can see that the access order optimization has a significant
effect on performance improvement because by scheduling the access order we
are able to maximize the data reuse; furthermore, compared to PyTorch
DataLoader + LRU buffer, the access order optimization achieves the most
improvement. This is because by adjusting the access order, SOLAR is able to
“intelligently” improve the data reuse rate more when using a larger buffer
size while PyTorch DataLoader + LRU buffer does not efficiently utilize the
in-memory buffer space for data reuse. Third, data load balancing optimization
further reduces the data loading time because the parallel reading idle time
is reduced when we balance the data loading workload. Fourth, the aggregated
chunk loading improves the performance further, up to a cumulative 7.5$\times$
speedup, illustrating that the aggregated chunk loading is effective. In this
experiment, the chunk size is set to 15.444The chunk size
$|\textit{chunk\hskip 0.5pt}|$ is a threshold from benchmarking. E.g.,
$|\textit{chunk\hskip 0.5pt}|=15$ means loading samples $i$ to $i+14$ in a
single chunk is faster than loading them separately.

Access Order Optimization First, we show the number of data samples loaded
from the PFS (“numPFS”) to demonstrate the contribution of access order
optimization (“Optim_1”), as shown in Figure 11. The batch size is set to 512,
and we take the maximum numPFS value from 16 GPUs. Note that for PyTorch
DataLoader, there is no reuse policy and thus it loads 512 samples per GPU on
each iteration. From the figure, we can see that numPFS is reduced up to
4.9$\times$ with the access order optimization. Figure 10 shows the
performance improvement from the reduction of numPFS.

Figure 11. Comparison of the number of samples loaded from the PFS between
PyTorch DataLoader and SOLAR.

Data Load Balancing Next, we demonstrate the contribution of load balancing.
We measure numPFS on each GPU (16 GPUs in total), as shown in Figure 12.
Specifically, the load imbalance is caused by the different number of reused
data samples on each GPU. In the imbalanced case, GPU 7 needs to load just 41
samples from the PFS, while GPU 2 needs to load 107. All other GPUs need to
wait for GPU 2 to finish loading (i.e., sync barrier imbalanced in Figure 12),
thus causing idle time as a consequence. In the balanced case, the GPU idle
time is reduced. The finish line (i.e., sync barrier balanced) is shown in
Figure 12. Thus, after load balancing the data loading is improved by
1.39$\times$ compared to the imbalanced case, and it is consistent with the
speedup of I/O time shown in Figure 10.

Figure 12. Comparison of the number of samples loaded from the PFS before and
after load balancing optimization.

Aggregated Chunk Loading Lastly, we demonstrate the contribution of aggregated
chunk loading. The number of data samples that can be loaded in chunks is
measured across multiple runs on each GPU, then divided by numPFS to compute
the percentage. As shown in Figure 13, about 7% of data samples on average (up
to 20.6% in the best case) can be loaded in chunks. Meanwhile, the worst case
is that no data samples can be loaded in chunks. Note that in the worst case,
this optimization does not harm the performance because there’s no extra
overhead.

Furthermore, we see the potential of this optimization to be applied to
distributed training tasks that do not require full randomization. For
example, MLPerfHPC performs intra-node shuffle for DeepCAM (Farrell et al.,
2021).

Figure 13. Percentage of data samples that can be loaded in chunks in
different training runs.

###  5.4. Evaluation of End-to-End Training

Figure 14. Training accuracy vs. time with PyTorch DataLoader method and our
SOLAR on PtychoNN. Figure 15. Comparison of outputs generated by PtychoNNs
trained with PyTorch DataLoader (PyTorch) and SOLAR. “True” is the ground
truth targeted in training. “I” and “Phi” denote amplitude and phase,
respectively.

In this section, we evaluate the end-to-end training performance of SOLAR on
the (CD 321GB) dataset using the high-end system. Also, we show that the epoch
order optimization can affect the training accuracy on PtychoNN. We also
performed an ablation study to demonstrate the options of SOLAR, which will be
discussed in §5.5 Specifically, we chose the 17 GB dataset to perform the
accuracy study because the 321 GB and 1.2 TB datasets are synthesized from the
17 GB dataset by repeating the same images. In order to reflect the realistic
case where the dataset cannot fit in a single device, we set the local buffer
size to 4.25 GB. Figure 14 shows the end-to-end training comparison between
the PyTorch DataLoader and SOLAR. We can see that SOLAR not only improves the
training performance but also achieves lower validation loss as well. The
time-to-solution speedup is 3.03$\times$ compared to PyTorch DataLoader.

After the model is trained, we perform the reference task using the test data.
Figure 15 shows the reference result compared to the ground truth labels (True
I and True Phi), where we can see that SOLAR does not degrade the quality of
image reconstruction and it can generate clear shapes of the amplitude (SOLAR
I) and phase (SOLAR Phi). We also surprisingly find that compared to the
PyTorch DataLoader amplitude (PyTorch I) and phase (PyTorch Phi), SOLAR even
improves the image quality in some cases.

###  5.5. Detailed Analysis

In this section, we answer two questions about SOLAR: (1) If changing the
epoch order affects the randomness, what will the performance be if the epoch
order optimization is not used? (2) How imbalanced will the training batch
size be after trading computation balance for load balance?

Effect of epoch order optimization We perform a study to demonstrate the
improvement of SOLAR with and without the Epoch Order Optimization (EOO). As
demonstrated in §4, the data locality optimization, loading workload
balancing, and aggregated chunk loading do not affect the training accuracy
because they do not change the training results after synchronization.
However, when changing the epoch order there is a concern that it changes the
randomness. On the one hand, the term “randomness” is not quantized and is
challenging to measure. On the other hand, the relationship between randomness
and training accuracy is unknown to the deep learning community. To this end,
we perform the study here to demonstrate the effect of re-arranging epoch
order on data-loading performance. First, to prove from the experiment that
rearranging the epoch order improves the data-loading performance, we
implemented a buffer that uses the Least Recently Used (LRU) buffering
strategy for PyTorch DataLoader. We observe that applying the EOO on PyTorch
DataLoader + LRU improves the performance by 25.6% compared to that without
the EOO. Second, we perform the experiment to study the effect of EOO on
SOLAR. We observe that when using the EOO in SOLAR, the performance improved
by 59.4% compared to not using the EOO.

Figure 16. Distributions of per-GPU batch sizes in different training steps
(dots are outliers).

Per-GPU batch sizes after trade-off Since in the data load balancing
optimization, we trade off computation balance with load balance, the
imbalance of training local batch size should be studied to understand how
much the training batch size is affected. Therefore, we measure the number of
data samples loaded from the PFS for each process at each step. The simulation
is performed on 16 processes and the local batch size is 512. The standard
deviation from the first step to the 10th step ranges from 7.00 to 16.42. As
shown in Figure 16, the local batch sizes are imbalanced, but the distribution
is concentrated around 512 (the original local batch size). This demonstrates
that after data load balancing, the training batch size is not very
imbalanced, meaning that trading off the computation balance for load balance
does not significantly affect the computation across GPUs.

##  6\. Related Work

Since optimizations on CPUs/GPUs and communication enable today’s DNN training
to achieve unprecedented throughput (Zhang et al., 2021b; Jin et al., 2021;
Jiang et al., 2020; Liu et al., 2021; Peng et al., 2019; Yin et al., 2021),
I/O is left behind and becomes the major bottleneck in DNN training.
Furthermore, distributed training techniques enable further acceleration of
computation. Meanwhile, I/O is unable to keep up with the improvement, and it
slows down the whole training process. Therefore, many works focus on data-
loading optimizations for distributed DNN training. We summarize the
differences between existing works and our framework in Table 5.

Table 5. Summary of techniques used in different methods.

|  PyTorch DataLoader (Paszke et al., 2019) |  Locality-Aware (Yang and Cong, 2019) |  DeepIO (Zhu et al., 2018) |  NoPFS (Dryden et al., 2021) |  Lobster (Liu et al., 2022) |  SOLAR (Ours)
---|---|---|---|---|---|---
Full Randomization | $\bullet$ | $\bullet$ |  | $\bullet$ | $\bullet$ | $\bullet$
Epoch Order Optim.  |  |  |  |  |  | $\bullet$
Node-Sample Mapping  |  | $\circ$ | $\circ$ | $\circ$ | $\circ$ | $\bullet$
Load Balancing  |  | $\circ$ | $\bullet$ |  | $\bullet$ | $\bullet$
Aggreg. Chunk Loading  |  |  |  |  |  | $\bullet$

$\bullet$ technique in use $\circ$ data loading achieved via inter-node
communication

As mentioned in Section 5.1, PyTorch Dataloader (Li et al., 2020) is a widely
used approach for efficient data loading in DNN training with PyTorch, while
NoPFS (Dryden et al., 2021) is a state-of-the-art approach to improve the I/O
performance via buffer eviction strategy based on the performance model and a
heuristic algorithm.

Yang et al. (Yang and Cong, 2019) propose a locality-aware loading strategy.
Their optimization takes three steps: (1) Global shuffle during each epoch;
(2) Determine the distribution of samples on nodes for data reuse; (3) Inter-
node communication to exchange the data samples for data load balancing.
However, the approach introduces communication overhead to reduce data loading
time for each iteration.

Zhu et al. (Zhu et al., 2018) propose DeepIO that provides a shuffle strategy
based on Remote Direct Memory Access (RDMA). DeepIO optimizes the data loading
on two aspects: (1) Local shuffling underlying Remote Direct Memory Access
(RDMA) that enables a node to access the data sample buffered on another node;
(2) Entropy-aware access order optimization that utilizes cross-entropy to
arrange the access order and reduce inter-node communication. However, the
optimization does not consider the NUMA effect when accessing data via RDMA
(Ren et al., 2013). Furthermore, their work trades off training accuracy to
reduce the data I/O overhead. We exclude Locality-aware and DeepIO in our
comparison, because NoPFS outperforms state-of-the-art frameworks including
Locality-aware and DeepIO, as confirmed by our experiments. Moreover, our
evaluation result shows that SOLAR achieves higher speedup than NoPFS.

Liu et al. (Liu et al., 2022) propose an approach called Lobster to optimize
data loading by (1) dynamically assigning threads to read and pre-process
training data to achieve load balance and (2) approximately arranging the data
eviction when the buffer is populated. Regarding data-loading optimization,
our approach is orthogonal to Lobster. Specifically, our approach adjusts the
access order and achieves better data reuse, while Lobster adjusts the number
of threads assigned for data reading and preprocessing to improve the data
loading throughput. Thus, SOLAR and Lobster can be combined to further improve
the end-to-end I/O performance for DNN training. Furthermore, for data-load
balance, SOLAR trades computation balance, while Lobster trades thread
balance. However, SOLAR’s trade-off strategy is better than Lobster when data
loading takes up most of the training time (e.g., 80%) because in Lobster the
number of threads is imbalanced across multiple GPUs most of the time.

Luan et al. (Luan et al., 2022) propose Exoshuffle, a shuffle system for the
MapReduce framework (Condie et al., 2010) that decouples the shuffle control
and data planes. Exoshuffle can be used to shuffle the training data stored on
each node and reduce training latency by transferring the data entirely in
memory and overlapping shuffling with computation. However, when the size of
the training data is larger than the aggregated memory size, the training data
must be kept in a remote storage system such as the PFS. Thus, Exoshuffle (in-
memory shuffling) does not work in this scenario. Moreover, it does not
consider data locality to reduce inter-node data movements.

The NVIDIA Data Loading Library (DALI) (NVIDIA, [n. d.]) is a data loading
library that improves the data preprocessing throughput by moving the task
from CPU to GPU. We do not compare with DALI because DALI is orthogonal to
SOLAR and can be applied to our work to further improve training performance.

##  7\. Conclusion and Future Work

CNN-based surrogate models are gaining more and more attention because they
can provide approximation results to replace traditional expensive numerical
algorithms. Motivated by the benchmark results showing that data loading is
the major performance bottleneck for distributed training of surrogate models,
we design SOLAR, which consists of offline scheduling and runtime buffer
strategies, to optimize the data loading performance by improving the data
reuse rate and loading the data in chunks. Experimental evaluation illustrates
that SOLAR achieves up to 24.4$\times$ and 3.3$\times$ speedup over a widely
used industry framework and a state-of-the-art approach, respectively. SOLAR
is implemented on PyTorch DataLoader and H5py (an HDF5 python interface),
which makes it easy to use with minimal changes to existing scripts.

SOLAR has two limitations: (1) It introduces a one-time offline scheduling
overhead, while this computation cost can be shared by multiple models trained
on the same dataset. Moreover, the scheduling stage can be overlapped with the
first epoch, since the dataset is initially stored on the PFS and the data
loading time on the first epoch still dominates the training time. (2) It does
not implement a multi-layer buffering strategy for HPC hierarchical storage
architectures. While our custom buffer is compatible with buffering the data
on different node-local storage layers.

Future work includes but is not limited to the following: (1) implement our
approach in combination with Lobster and DALI to further improve the
performance of data pre-processing and CPU-GPU data loading, respectively; (2)
optimize the offline and runtime scheduling overhead to further improve the
performance, then provide the framework as a python interface to improve the
performance of existing distributed training works.

## Acknowledgments

This research was supported by the U.S. Department of Energy, Office of
Science, Advanced Scientific Computing Research (ASCR), and Office of Science
User Facilities, Office of Basic Energy Sciences, under Contract DE-
AC02-06CH11357. This work was also supported by the National Science
Foundation under Grant 2303820. We gratefully acknowledge the computing
resources provided by Argonne Leadership Computing Facility.

## References

  * (1)
  * Pty ([n. d.]) [n. d.].  PtychoNN: Deep learning of ptychographic imaging.  https://github.com/mcherukara/PtychoNN
  * the ([n. d.]) [n. d.].  Theta – Argonne Leadership Computing Facility.  https://www.alcf.anl.gov/support-center/theta/theta-thetagpu-overview
  * Abadi et al. (2016) Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. 2016\.  TensorFlow: A system for large-scale machine learning. In _12th USENIX symposium on operating systems design and implementation (OSDI 16)_. 265–283.
  * Alshemali and Kalita (2020) Basemah Alshemali and Jugal Kalita. 2020.  Improving the reliability of deep neural networks in NLP: A review.  _Knowledge-Based Systems_ 191 (2020), 105210.
  * Byna et al. (2020) Suren Byna, M Scot Breitenfeld, Bin Dong, Quincey Koziol, Elena Pourmal, Dana Robinson, Jerome Soumagne, Houjun Tang, Venkatram Vishwanath, and Richard Warren. 2020.  ExaHDF5: Delivering efficient parallel I/O on exascale computing systems.  _Journal of Computer Science and Technology_ 35, 1 (2020), 145–160.
  * Cherukara et al. (2020) Mathew J Cherukara, Tao Zhou, Youssef Nashed, Pablo Enfedaque, Alex Hexemer, Ross J Harder, and Martin V Holt. 2020.  AI-enabled high-resolution scanning coherent diffraction imaging.  _Applied Physics Letters_ 117, 4 (2020), 044103\.
  * Collins and Duffy (2001) Michael Collins and Nigel Duffy. 2001.  Convolution kernels for natural language.  _Advances in neural information processing systems_ 14 (2001).
  * Condie et al. (2010) Tyson Condie, Neil Conway, Peter Alvaro, Joseph M Hellerstein, Khaled Elmeleegy, and Russell Sears. 2010.  MapReduce online. In _NSDI_ , Vol. 10. 20\.
  * Deng et al. (2009) Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. 2009\.  ImageNet: A large-scale hierarchical image database. In _2009 IEEE conference on computer vision and pattern recognition_. IEEE, 248–255.
  * Dong et al. (2019) Wenqian Dong, Jie Liu, Zhen Xie, and Dong Li. 2019\.  Adaptive neural network-based approximation to accelerate Eulerian fluid simulation. In _Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis_. 1–22.
  * Dong et al. (2020) Wenqian Dong, Zhen Xie, Gokcen Kestor, and Dong Li. 2020\.  Smart-PGSim: Using neural network to accelerate AC-OPF power grid simulation. In _SC20: International Conference for High Performance Computing, Networking, Storage and Analysis_. IEEE, 1–15.
  * Dryden et al. (2021) Nikoli Dryden, Roman Böhringer, Tal Ben-Nun, and Torsten Hoefler. 2021.  Clairvoyant prefetching for distributed machine learning I/O. In _Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis_. 1–15.
  * Farrell et al. (2021) Steven Farrell, Murali Emani, Jacob Balma, Lukas Drescher, Aleksandr Drozd, Andreas Fink, Geoffrey Fox, David Kanter, Thorsten Kurth, Peter Mattson, et al. 2021\.  MLPerf™ HPC: A Holistic Benchmark Suite for Scientific Machine Learning on HPC Systems. In _2021 IEEE/ACM Workshop on Machine Learning in High Performance Computing Environments (MLHPC)_. IEEE, 33–45.
  * Flood (1956) Merrill M Flood. 1956\.  The traveling-salesman problem.  _Operations research_ 4, 1 (1956), 61–75.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2016\.  Deep residual learning for image recognition. In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_. 770–778.
  * He et al. (2020) Xiangnan He, Kuan Deng, Xiang Wang, Yan Li, Yongdong Zhang, and Meng Wang. 2020\.  LightGCN: Simplifying and powering graph convolution network for recommendation. In _Proceedings of the 43rd International ACM SIGIR conference on research and development in Information Retrieval_. 639–648.
  * Jacobs et al. (2019) Sam Ade Jacobs, Brian Van Essen, David Hysom, Jae-Seung Yeom, Tim Moon, Rushil Anirudh, Jayaraman J Thiagaranjan, Shusen Liu, Peer-Timo Bremer, Jim Gaffney, et al. 2019\.  Parallelizing training of deep generative models on massive scientific datasets. In _2019 IEEE International Conference on Cluster Computing (CLUSTER)_. IEEE, 1–10.
  * Jiang et al. (2020) Yimin Jiang, Yibo Zhu, Chang Lan, Bairen Yi, Yong Cui, and Chuanxiong Guo. 2020\.  A unified architecture for accelerating distributed DNN training in heterogeneous GPU/CPU clusters. In _14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20)_. 463–479.
  * Jin et al. (2021) Sian Jin, Chengming Zhang, Xintong Jiang, Yunhe Feng, Hui Guan, Guanpeng Li, Shuaiwen Leon Song, and Dingwen Tao. 2021\.  COMET: A novel memory-efficient deep learning training framework by using error-bounded lossy compression.  _Proceedings of the VLDB Endowment_ 15, 4 (2021), 886–899.
  * Kennedy and Eberhart (1995) James Kennedy and Russell Eberhart. 1995.  Particle swarm optimization. In _Proceedings of ICNN’95-international conference on neural networks_ , Vol. 4. IEEE, 1942–1948.
  * Kurth et al. (2018) Thorsten Kurth, Sean Treichler, Joshua Romero, Mayur Mudigonda, Nathan Luehr, Everett Phillips, Ankur Mahesh, Michael Matheson, Jack Deslippe, Massimiliano Fatica, et al. 2018\.  Exascale deep learning for climate analytics. In _SC18: International Conference for High Performance Computing, Networking, Storage and Analysis_. IEEE, 649–660.
  * Li et al. (2020) Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, et al. 2020\.  PyTorch distributed: experiences on accelerating data parallel training.  _Proceedings of the VLDB Endowment_ 13, 12 (2020), 3005–3018.
  * Liu et al. (2021) Hongyuan Liu, Bogdan Nicolae, Sheng Di, Franck Cappello, and Adwait Jog. 2021.  Accelerating DNN architecture search at scale using selective weight transfer. In _2021 IEEE International Conference on Cluster Computing (CLUSTER)_. IEEE, 82–93.
  * Liu et al. (2022) Jie Liu, Bogdan Nicolae, and Dong Li. 2022.  Lobster: Load Balance-Aware I/O for Distributed DNN Training. In _ICPP’22: The 51st International Conference on Parallel Processing_.
  * Luan et al. (2022) Frank Sifei Luan, Stephanie Wang, Samyukta Yagati, Sean Kim, Kenneth Lien, SangBin Cho, Eric Liang, and Ion Stoica. 2022.  Exoshuffle: Large-Scale Shuffle at the Application Level.  _arXiv preprint arXiv:2203.05072_ (2022).
  * Mathuriya et al. (2018) Amrita Mathuriya, Deborah Bard, Peter Mendygral, Lawrence Meadows, James Arnemann, Lei Shao, Siyu He, Tuomas Kärnä, Diana Moise, Simon J Pennycook, et al. 2018\.  CosmoFlow: Using deep learning to learn the universe at scale. In _SC18: International Conference for High Performance Computing, Networking, Storage and Analysis_. IEEE, 819–829.
  * Noirot et al. (2018) Philippe Noirot, Andrzej Joachimiak, and Robert F Fischetti. 2018\.  _Workshop on Biological Science Opportunities Provided by the APS Upgrade_.  Technical Report. Argonne National Laboratory (ANL), Lemont, IL (United States).
  * NVIDIA ([n. d.]) NVIDIA. [n. d.].  NVIDIA Developer Data Loading Library (DALI).  https://docs.nvidia.com/deeplearning/dali/user-guide/docs/.  (Accessed on 10/27/2022).
  * Obiols-Sales et al. (2020) Octavi Obiols-Sales, Abhinav Vishnu, Nicholas Malaya, and Aparna Chandramowliswharan. 2020\.  CFDNet: A deep learning-based accelerator for fluid simulations. In _Proceedings of the 34th ACM international conference on supercomputing_. 1–12.
  * Oyama et al. (2020) Yosuke Oyama, Naoya Maruyama, Nikoli Dryden, Erin McCarthy, Peter Harrington, Jan Balewski, Satoshi Matsuoka, Peter Nugent, and Brian Van Essen. 2020.  The case for strong scaling in deep learning: Training large 3D CNNs with hybrid parallelism.  _IEEE Transactions on Parallel and Distributed Systems_ 32, 7 (2020), 1641–1652.
  * Papadimitriou (1977) Christos H Papadimitriou. 1977\.  The Euclidean travelling salesman problem is NP-complete.  _Theoretical computer science_ 4, 3 (1977), 237–244.
  * Paszke et al. (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. 2019\.  PyTorch: An imperative style, high-performance deep learning library.  _Advances in neural information processing systems_ 32 (2019).
  * Peng et al. (2019) Yanghua Peng, Yibo Zhu, Yangrui Chen, Yixin Bao, Bairen Yi, Chang Lan, Chuan Wu, and Chuanxiong Guo. 2019\.  A generic communication scheduler for distributed DNN training acceleration. In _Proceedings of the 27th ACM Symposium on Operating Systems Principles_. 16–29.
  * Pumma et al. (2017) Sarunya Pumma, Min Si, Wu-chun Feng, and Pavan Balaji. 2017\.  Towards scalable deep learning via I/O analysis and optimization. In _19th IEEE International Conference on High Performance Computing and Communications; 15th IEEE International Conference on Smart City; 3rd IEEE International Conference on Data Science and Systems (HPCC/SmartCity/DSS)_. IEEE, 223–230.
  * Pumma et al. (2019) Sarunya Pumma, Min Si, Wu-Chun Feng, and Pavan Balaji. 2019\.  Scalable deep learning via I/O analysis and optimization.  _ACM Transactions on Parallel Computing (TOPC)_ 6, 2 (2019), 1–34.
  * Ren et al. (2013) Yufei Ren, Tan Li, Dantong Yu, Shudong Jin, and Thomas Robertazzi. 2013.  Design and performance evaluation of NUMA-aware RDMA-based end-to-end data transfer systems. In _Proceedings of the International Conference on High Performance Computing, Networking, Storage and Analysis_. 1–10.
  * Sergeev and Del Balso (2018) Alexander Sergeev and Mike Del Balso. 2018.  Horovod: Fast and easy distributed deep learning in TensorFlow.  _arXiv preprint arXiv:1802.05799_ (2018).
  * Shi et al. (2022) Neng Shi, Jiayi Xu, Skylar W Wurster, Hanqi Guo, Jonathan Woodring, Luke P Van Roekel, and Han-Wei Shen. 2022.  GNN-Surrogate: A Hierarchical and Adaptive Graph Neural Network for Parameter Space Exploration of Unstructured-Mesh Ocean Simulations.  _IEEE Transactions on Visualization and Computer Graphics_ 28, 6 (2022), 2301–2313.
  * Shi et al. (2007) Xiaohu H Shi, Yanchun Chun Liang, Heow Pueh Lee, C Lu, and QX Wang. 2007.  Particle swarm optimization-based algorithms for TSP and generalized TSP.  _Information processing letters_ 103, 5 (2007), 169–176.
  * The HDF Group (2010) The HDF Group. 2000-2010.  Hierarchical Data Format version 5.  http://www.hdfgroup.org/HDF5
  * UIUC ([n. d.]) UIUC. [n. d.].  Multi-level tracing library.  https://github.com/uiuc-hpc/Recorder.
  * Voulodimos et al. (2018) Athanasios Voulodimos, Nikolaos Doulamis, Anastasios Doulamis, and Eftychios Protopapadakis. 2018\.  Deep learning for computer vision: A brief review.  _Computational intelligence and neuroscience_ 2018 (2018).
  * Yang and Cong (2019) Chih-Chieh Yang and Guojing Cong. 2019.  Accelerating data loading in deep neural network training. In _2019 IEEE 26th International Conference on High Performance Computing, Data, and Analytics (HiPC)_. IEEE, 235–245.
  * Yao et al. (2022) Yudong Yao, Henry Chan, Subramanian Sankaranarayanan, Prasanna Balaprakash, Ross J Harder, and Mathew J Cherukara. 2022.  AutoPhaseNN: Unsupervised physics-aware deep learning of 3D nanoscale Bragg coherent diffraction imaging.  _npj Computational Materials_ 8, 1 (2022), 1–8.
  * Yin et al. (2022) Junqi Yin, Feiyi Wang, and Mallikarjun Shankar. 2022\.  Strategies for Integrating Deep Learning Surrogate Models with HPC Simulation Applications. In _2022 IEEE International Parallel and Distributed Processing Symposium Workshops (IPDPSW)_. IEEE, 01–10.
  * Yin et al. (2021) Lujia Yin, Yiming Zhang, Zhaoning Zhang, Yuxing Peng, and Peng Zhao. 2021.  ParaX: Boosting deep learning for big data analytics on many-core CPUs.  _Proceedings of the VLDB Endowment_ 14, 6 (2021), 864–877.
  * Yu et al. (2022) Xiaodong Yu, Viktor Nikitin, Daniel J Ching, Selin Aslan, Doğa Gürsoy, and Tekin Biçer. 2022.  Scalable and accurate multi-GPU-based image reconstruction of large-scale ptychography data.  _Scientific Reports_ 12, 1 (2022), 1–16.
  * Zhang et al. (2021b) Chengming Zhang, Geng Yuan, Wei Niu, Jiannan Tian, Sian Jin, Donglin Zhuang, Zhe Jiang, Yanzhi Wang, Bin Ren, Shuaiwen Leon Song, et al. 2021b.  Clicktrain: Efficient and accurate end-to-end deep learning training via fine-grained architecture-preserving pruning. In _Proceedings of the ACM International Conference on Supercomputing_. 266–278.
  * Zhang et al. (2019) Shuai Zhang, Lina Yao, Aixin Sun, and Yi Tay. 2019\.  Deep learning based recommender system: A survey and new perspectives.  _ACM Computing Surveys (CSUR)_ 52, 1 (2019), 1–38.
  * Zhang et al. (2021a) Yuhao Zhang, Frank Mcquillan, Nandish Jayaram, Nikhil Kak, Ekta Khanna, Orhan Kislal, Domino Valdano, and Arun Kumar. 2021a.  Distributed deep learning on data systems: A comparative analysis of approaches.  _Proceedings of the VLDB Endowment_ 14, 10 (2021), 1769–1782.
  * Zhu et al. (2018) Yue Zhu, Fahim Chowdhury, Huansong Fu, Adam Moody, Kathryn Mohror, Kento Sato, and Weikuan Yu. 2018.  Entropy-aware I/O pipelining for large-scale deep learning on HPC systems. In _2018 IEEE 26th International Symposium on Modeling, Analysis, and Simulation of Computer and Telecommunication Systems (MASCOTS)_. IEEE, 145–156.