Toward Training a Large 3D Cosmological CNN
with Hybrid Parallelization
1,2,?

23,244

Yosuke Oyama , Naoya Maruyama , Nikoli Dryden , Peter Harrington , Jan Balewski ,5,1
342
Satoshi Matsuoka , Marc Snir , Peter Nugent , and Brian Van Essen
1Tokyo Institute of Technology, 2Lawrence Livermore National Laboratory, 3University of Illinois at Urbana-Champaign,4
Lawrence Berkeley National Laboratory, 5Riken Center for Computational Science?
oyama.y.aa@m.titech.ac.jp

Overview

Strong scaling

We present a case study of large-scale training of a 3D convolutional neural networkmodel for cosmological analyses of dark matter distributions. This work extends
existing work [1] for predicting cosmological parameters using CNNs for better
prediction accuracy and performance by exploiting finer-grained parallelism in distributed
convolutions. We show significant improvements using the latest complex cosmologicaldataset in both strong scaling and weak scaling, achieving 1.42 PFlop/s (111x of speedup
over one computing node) on a single training task with a mini-batch size of 128 byusing 512 Tesla V100 GPUs. Our framework enables to train a huge CNN whose input
size is 4 × 5123 that was previously unfeasible due to its memory pressure.

2.28x of speedup on 4 nodes (16 GPUs) over one node with a mini-batch size of 1
The main bottleneck is I/O of excessive data size (1 GiB/sample) via PCIe and inter-node data
shuffleN = 1, W = 512

N = 8, W = 512

0.3

2.28x

[s]

[s]

0.3
0.2

0.2

0.1

0.1

0.0

0.0

Background: Hybrid-parallel training

1248{Dp , 1}
{Dp , Hp }0.4

16 328

Number of nodes (Dp /4)

1632641Number of nodes{Dp , 1}0.4

24Seq. data load
ForwardBackward
Update8161280.00

More amount of parallelism: The mini-batch size (N) × layer size (W n , where n is the number oflayer dimensions)Less influence on resulting inference accuracy by keeping the mini-batch size small

0.10

0.15

0.20

Number of nodes (8Dp Hp /4)

0.25

0.30

0.35

Time [s](a) Strong scaling with two different mini-batch sizes

Hybrid-parallelism has advantages over data/model-parallelism

0.05(b) Time breakdown with N = 1, W = 512

Figure 2. Strong scaling of the CosmoFlow network. {Dp , Hp } represents the depth and the height
dimensions are distributed among Dp and Hp process groups respectively.

Table 1. Comparison of parallel strategies for training a CNN.

Weak scaling

Data
Samples
O(N)
7
7
33+

119x and 120x of speedup on 128 nodes from 1 node with the W = 128 cubes andW = 256 cubes datasets respectively
111x of speedup over 1 node with W = 512 by exploiting hybrid-parallelism even if
layer-wise communication is introducedHybrid
Samples & Layers
O(NW n )3333+

1.19x of speedup by increasing the number of nodes from 64 to 128 with N = 64

The CosmoFlow dataset [1, 2] is composed of a set of 3D dark matter distributions
along with their cosmological parameters10{1, 1} (Synthetic){1, 1} (Conduit){1, 1} (Direct I/O)

410

Speed [samples/s]

CosmoFlowNGPU = 1, W = 2564119x103

NGPU = 1/Dp Hp , W = 512{1, 1} (Synthetic){1, 1} (Conduit){1, 1} (Direct I/O)

3Speed [samples/s]

NGPU = 8, W = 128Speed [samples/s]

What to parallelize
Available parallelism
GPU memory pressure
Influence on accuracy
Weak-scaling
Strong-scaling

ModelLayers
O(W n )
33+N/A
310120x1021248163264 128

102111x110100101

We use a CNN composed of seven 3D convolutional layers and three fully-connected layers
A single network cannot be trained with the original input size due to its memory pressure{2, 2} (Synthetic){4, 1} (Synthetic){8, 1} (Synthetic){2, 2} (Conduit){4, 1} (Conduit){8, 1} (Conduit){2, 2} (Direct I/O)
{4, 1} (Direct I/O)
{8, 1} (Direct I/O)

12

Number of nodes

48163264 128

12

Number of nodes

48163264 128

Number of nodes

Table 2. Summary of the CosmoFlow dataset and the network architecture.
Input width (W )# of input channels (C )# of samples
Dataset size [TiB]# of conv. ops. [GFlops/sample](Forward)
Memory [GiB/sample]# of parameters [106]128
465,728
1.00
55.55
18.52
0.824
9.44

256
48,216
1.00
443.8
147.9
6.59
9.44

512
41,0271.00
3550
1183
52.7
9.44

Figure 3. Weak scaling of the CosmoFlow network. NGPU represents the number of data samples perGPU.

Detailed analysis of the training timeline
By distributing N = 64 from 64 nodes to 128 nodes,
the computation efficiency per sample is degraded as the batch size per GPU is halved
due to this computational inefficiency, part of all-reduce cannot be hidden in the main stream

Proposal: LBANN + Distconv for 3D CNNs
Distconv [3]: A hybrid-parallel CNN training implementation
Distconv distributes the computation of convolutional layers to a set of GPUs

The overhead introduced by Distconv (“DC”, blue) is nearly negligible than thecomputational kernels (“Main”, red)N = 64, W = 512, {Dp , Hp } = {4, 1}, 64 nodes

1. Perform convolution to the center part of an input tensor
2. Start a halo exchange among GPUs in the same sample group in an asynchronous stream

preload the entire dataset into CPU memory before training starts
employ an MPI-based data exchange to shuffle the data to the process that requires it

Read the dataset from the global file system or node-local SSDs (“Direct I/O”)
conv1···

NCCL
DC
Main

Two different I/O and communication-avoiding data readers
Use the Conduit [4] data exchange library as an I/O backend (“Conduit”)

Input

UpdateBackward

ForwardStream

3. Convolution is performed to the halo region in the halo stream

StreamRepeat the one-dimensional halo exchange three times to perform the three-dimensional halo exchange

ForwardN = 128, W = 512, {Dp , Hp } = {8, 1}, 128 nodes
Backward

Update

NCCL
DC
Main

0.00

0.05

0.10

0.15

0.20

0.25

0.30

0.35

Time [s]

0.00

0.05

0.10

0.15

0.20

Time [s]Figure 4. The GPU timeline of a single training iteration.

fc1,. . . ,3

conv7

Rank
Read

12

ShuffleMemory03

Conv.
Haloex.
+conv.

CPURead67

ShuffleMemory5CPUConv.

FC

Back-prop.

GPUConv.
Haloex.
+conv.

Acknowledgement

References

This research was supported by JSPS KAKENHI Grant Number JP18J22858, Japan.
This research was supported by the Exascale Computing Project (17-SC-20-SC), a joint
project of the U.S. Department of Energy’s Office of Science and National Nuclear
Security Administration, responsible for delivering a capable exascale ecosystem,including software, applications, and hardware technology, to support the nation’s
exascale computing imperative. This research used resources of the National Energy
Research Scientific Computing Center (NERSC), a U.S. Department of Energy Officeof Science User Facility operated under Contract No. DE-AC02-05CH11231. LLNLPOST-776501.

[1] A. Mathuriya, D. Bard, P. Mendygral, L. Meadows, J. Arnemann, L. Shao,
S. He, T. Kärnä, D. Moise, S. J. Pennycook, K. Maschhoff, J. Sewall,
N. Kumar, S. Ho, M. F. Ringenburg, Prabhat, and V. Lee, “Cosmoflow: Using
deep learning to learn the universe at scale,” in Proceedings of the International
Conference for High Performance Computing, Networking, Storage, andAnalysis, ser. SC ’18. Piscataway, NJ, USA: IEEE Press, 2018, pp. 65:1–65:11.
[Online]. Available: http://dl.acm.org/citation.cfm?id=3291656.3291743···

Parameter gradients aggregation(all-reduce)

Sampleexchange
4

ShuffleShuffle

Conv.

FC···

Figure 1. Overview of hybrid-parallel LBANN training.

GPUBack-prop.

[2] National Energy Research Scientific Computing Center, “Cosmoflow datasets,”2019. [Online]. Available: https://portal.nersc.gov/project/m3363[3] N. Dryden, N. Maruyama, T. Benson, T. Moon, M. Snir, and B. Van Essen,“Improving strong-scaling of cnn training by exploiting finer-grained
parallelism,” in Proceedings of the International Parallel and Distributed
Processing Symposium, ser. IPDPS ’19, 2019.
[4] Lawrence Livermore National Laboratory, “Conduit,” 2019. [Online]. Available:
https://github.com/LLNL/conduit