PaperlayerView exact revision on arXiv

Source-linked AI summary

CosmoFlow: Using Deep Learning to Learn the Universe at Scale

Amrita Mathuriya, Deborah Bard, Peter Mendygral, Lawrence Meadows, James Arnemann, Lei Shao, Siyu He, Tuomas Karna, Daina Moise, Simon J. Pennycook, Kristyn Maschoff, Jason Sewall, Nalini Kumar, Shirley Ho, Mike Ringenburg, Prabhat, Victor Lee

arXiv:1808.04728v2astro-ph.COastro-ph.IMcs.LGphysics.comp-ph

TL;DR

Scientific data are multidimensional and voluminous, while cosmological information about dark energy is encoded in matter distributions. CosmoFlow combines an optimized TensorFlow 3D CNN with synchronous multi-node training, sustaining over 3.5 Pflop/s on 8192 Cori nodes while estimating Ω_M, σ_8, and n_s.

Abstract

from arXiv · show

Deep learning is a promising tool to determine the physical model that describes our universe. To handle the considerable computational cost of this problem, we present CosmoFlow: a highly scalable deep learning application built on top of the TensorFlow framework. CosmoFlow uses efficient implementations of 3D convolution and pooling primitives, together with improvements in threading for many element-wise operations, to improve training performance on Intel(C) Xeon Phi(TM) processors. We also utilize the Cray PE Machine Learning Plugin for efficient scaling to multiple nodes. We demonstrate fully synchronous data-parallel training on 8192 nodes of Cori with 77% parallel efficiency, achieving 3.5 Pflop/s sustained performance. To our knowledge, this is the first large-scale science application of the TensorFlow framework at supercomputer scale with fully-synchronous training. These enhancements enable us to process large 3D dark matter distribution and predict the cosmological parameters $\Omega_M$, $\sigma_8$ and n$_s$ with unprecedented accuracy.

I. OVERVIEW

Scientific deep learning must handle multidimensional, voluminous data efficiently; CosmoFlow applies an optimized, MPI-scalable 3D CNN to cosmological matter volumes and demonstrates synchronous training at supercomputer scale.

A. Deep Learning for Science

  • Scientific datasets can be multidimensional and voluminous, so deep-learning frameworks must process them efficiently at scale for rapid scientific exploration.The paper identifies this as a challenge distinct from many commercial applications.Source: A. Deep Learning for Science · p. 1

B. Deep Learning for Cosmology

  • Dark-energy effects are encoded in matter distributions, which cosmologists commonly characterize with two- or three-point correlation functions and other reduced statistics.The authors suggest methods that capture all distribution features, such as deep-learning networks, could give greater insight into dark energy.Source: B. Deep Learning for Cosmology · p. 1

C. Contributions

  • CosmoFlow adapts a prior network to 128^3-voxel inputs and predicts three cosmological parameters using a dataset generated with three parameter variations.An example input is a 128^3-voxel sub-volume from a 512 Mpc/h^3 dark-matter simulation evolved to the present day.Source: C. Contributions · p. 1, Figure 1 · p. 2
  • The project combines efficient MKL-DNN 3D-CNN primitives in optimized TensorFlow with the Cray ML Plugin for MPI-based scaling.The stated contribution includes fully synchronous data-parallel training on 8192 Cori KNL nodes at sustained 3.5 Pflop/s single-precision performance.Source: C. Contributions · pp. 1–2

A. Deep Learning in Cosmology

Prior work applied deep learning to cosmological parameter estimation and other astronomy tasks, while this paper focuses on efficient CPU-based 3D CNNs and synchronous multi-node training.

A. Deep Learning in Cosmology

  • Earlier astronomy work used deep learning for galaxy morphology and image generation, while cosmology studies estimated parameters from 3D matter distributions and 2D mass maps.The cited 2D study classified five cosmological models and reported better estimates than traditional metrics.Source: A. Deep Learning in Cosmology · p. 2
  • A prior 2D-map study took 500 wall-clock hours to train on 12,500 maps using one P100 GPU, motivating HPC for faster research turnaround.The paper describes a 20-day turnaround as insufficient for rapid development and testing of research ideas.Source: A. Deep Learning in Cosmology · p. 2

B. Deep Learning on single Node

  • The work targets CPU acceleration of 3D CNNs in TensorFlow, with convolution kernels identified as the most compute-intensive DNN calculations.The authors optimize not only convolution kernels but all stages of the network topology.Source: B. Deep Learning on single Node · p. 2

C. Deep Learning on multiple nodes

  • HPC can distribute independent hyperparameter trials or train one network across nodes; this paper uses data parallelism, where nodes train on different data subsets and pool results after each step.The paper distinguishes this from model parallelism, in which nodes optimize different parts of one network using the same data.Source: C. Deep Learning on multiple nodes · p. 3
  • Synchronous scaling can slow beyond a few hundred nodes because large batches may impede convergence and a slow node can limit aggregate performance; the paper also examines I/O variability as a straggler source.Asynchronous training can require more iterations because nodes update from stale parameter estimates.Source: C. Deep Learning on multiple nodes · p. 3
  • The Cray ML Plugin uses MPI-based algorithms to address TensorFlow's default GRPC gradient-aggregation scaling limits; the authors describe their work as the first large-scale science TensorFlow application with fully synchronous training and 3D volumes at this scale.The paper also notes Horovod as an alternative using general-purpose MPI collectives.Source: C. Deep Learning on multiple nodes · p. 3

III. INNOVATIONS

CosmoFlow is a scalable TensorFlow application for predicting cosmological parameters from 3D matter distributions, adapting an earlier network to larger inputs while keeping model size manageable.

III. INNOVATIONS

  • CosmoFlow is a highly optimized, scalable TensorFlow application that predicts cosmological parameters from 3D matter distributions, with 3D CNN performance optimized for KNL processors.The work addresses multidimensional scientific data through its network, I/O, communication, and software stack.Source: §III · p. 3

A. 3D convolutional network and dataset

  • The network adapts Ravanbakhsh et al.'s CNN to a dataset an order of magnitude larger in physical volume, using 128^3-voxel inputs and an additional order of magnitude of training samples.The adapted dataset varies three cosmological parameters.Source: A. 3D convolutional network and dataset · p. 3
  • Seven convolution layers and three fully connected layers predict three parameters; an added convolution and stride-(2,2,2) average-pooling layer accommodate larger inputs while keeping parameters manageable.The final fully connected layer has three output neurons, and all convolution and fully connected layers use leaky ReLU.Source: A. 3D convolutional network and dataset · p. 3
  • For performance, convolution output channels are multiples of 16 for vectorization, and batch normalization is removed; the authors report no accuracy degradation from removal at batch size one.The passages also introduce two topology optimizations as performance considerations.Source: A. 3D convolutional network and dataset · pp. 3–4

B. SGD Optimizer

The optimizer combines Adam with layer-wise adaptive rate control and polynomial learning-rate decay. Single-node work targets 3D convolution and pooling, alongside TensorFlow element-wise operations.

B. SGD Optimizer

  • Adam is combined with LARC and polynomial (power=1) learning-rate decay; LARC clips each layer’s effective rate at Adam’s nominal rate.LARS/LARC scale updates with respect to each layer’s weight norm.Source: B. SGD Optimizer · p. 4

C. Single-node optimizations

  • Profiling identifies 3D convolution and pooling as primary hotspots; OpenMP loop-level parallelism improves TensorFlow threading compatibility with MKL-DNN.The hotspots were identified with Intel VTune Amplifier.Source: C. Single-node optimizations · p. 4
  • Convolution kernels use channel blocking, AVX512 vectorization, cache and register blocking, and threading; output width is blocked by 28 voxels.For layers with input and output channels that are multiples of 16, the 28, 16, and 16 innermost loops are unrolled and vectorized.Source: C. Single-node optimizations · p. 4
  • Backward data and weight operators use channel blocking and SIMD vectorization.For some layers, weight-gradient work is divided between channel-block teams and voxels within each team.Source: C. Single-node optimizations · p. 4
  • Average pooling is bandwidth-bound and uses 16-channel blocking, AVX512 vectorization, just-in-time assembly, and threading over output voxels.Pooling averages each channel separately, with constant weights of 1/(KS)^3 for a kernel of size KS.Source: C. Single-node optimizations · p. 5

D. Multi-node scaling

CosmoFlow uses the Cray CPE ML Plugin for fully synchronous data-parallel training. The plugin aggregates gradients across MPI ranks and integrates with TensorFlow without source modifications.

D. Multi-node scaling

  • Each synchronous training step computes local gradients, aggregates them through the CPE ML Plugin, and applies the aggregated gradients.The plugin provides MPI-based communication primitives designed for deep-learning workloads.Source: D. Multi-node scaling · p. 5
  • The plugin’s non-blocking MPI communication hides timing imbalances during gradient reduction, reducing the straggler effect.Every MPI rank computes gradients; the design uses no parameter servers or backup workers.Source: D. Multi-node scaling · p. 5
  • Communication uses user-tuned helper-thread teams; this work uses four helper threads on Cori and two on Piz Daint.Multiple communication threads can increase network utilization, particularly on Intel Xeon Phi architectures.Source: D. Multi-node scaling · p. 5
  • The plugin is called from the Python training script and added to the TensorFlow graph as a custom operation, with direct access to TensorFlow memory.This minimizes unnecessary copies of gradient data.Source: D. Multi-node scaling · p. 5

IV. SYSTEMS AND DATASETS

This section introduces the dataset configurations and the Cori and Piz Daint systems used for scaling measurements.

IV. SYSTEMS AND DATASETS

  • The section describes dataset configurations and the Cori and Piz Daint systems used for scaling measurements.Source: §IV · p. 5

A. Description of the Cori System

The scaling work uses Cori and Piz Daint and a simulation dataset of dark matter volumes with varied cosmological parameters.

A. Description of the Cori System

  • Cori measurements use KNL nodes, each with 68 cores, 16 GB of MCDRAM, and 96 GB of DDR4 memory; data is striped across DataWarp and Lustre storage.The Cori storage configuration uses 125 DataWarp nodes and 64 Lustre OSTs.Source: A. Description of the Cori System · p. 6

B. Description of the Piz Daint system

  • Piz Daint tests use hybrid nodes with an Intel Xeon processor and NVIDIA P100 GPU; training data is striped over 16 Lustre OSTs.The system also has DataWarp, which this work does not use because of availability issues.Source: B. Description of the Piz Daint system · p. 6

C. Simulation Data

  • Training simulations use MUSIC initial conditions and a modified pycola implementation of COLA, enabling a large suite of fast but accurate simulations.COLA preserves N-body accuracy at large scales while running faster than traditional N-body simulation code.Source: C. Simulation Data · p. 6
  • The dataset comprises 12,632 simulation boxes with varied ΩM, σ8, and ns values, assuming a flat universe and sampling specified parameter ranges.The ranges are 0.25 < ΩM < 0.35, 0.78 < σ8 < 0.95, and 0.9 < ns < 1.0.Source: C. Simulation Data · p. 6
  • Each simulation yields eight 128^3-voxel sub-volumes, totaling 101,056 sub-volumes that cover multiple clusters and dark matter filaments.Each sub-volume spans 256 h^-1 Mpc^3 at a resolution of 2 h^-1 Mpc^3.Source: C. Simulation Data · p. 6
  • After validation and test holdouts, 99,456 sub-volumes remain for training and are duplicated once; TFRecords contain 64 samples each.The total TFRecord data is 1.4 TB.Source: C. Simulation Data · p. 6

D. Libraries and Environment

CosmoFlow is run with a specified software and hardware stack and fully synchronous training, using one sample per node to keep the global batch manageable at scale. The workflow broadcasts a shared initial model, then trains and validates through MPI-averaged results.

D. Libraries and Environment

  • The software stack uses TensorFlow r1.5, MKL-DNN and the CPE ML Plugin on Cori, and TensorFlow r1.5 with cuDNN and CUDA on Piz Daint.Each system uses one MPI rank per node; the Cori and Piz Daint runs use four and two CPE ML Plugin helper threads, respectively.Source: D. Libraries and Environment · pp. 6–7

V. EXPERIMENTS AND RESULTS

  • One sample per node keeps the global batch size manageable at full supercomputer scale.The global batch grows with node count, and convergence is directly affected by global batch size.Source: §V · p. 7

A. Workflow and Network’s FLOP rate computation

  • With batch size one, the network performs 69.33 Gflop and uses 28.15 MB of parameters.It has slightly more than seven million parameters and uses 32-bit single-precision inputs and weights.Source: A. Workflow and Network’s FLOP rate… · p. 7
  • Each rank trains and validates after ranks receive identical initial model parameters; training gradients are averaged through MPI before model updates.The workflow uses Nsamples/nranks iterations for each training or validation loop.Source: A. Workflow and Network’s FLOP rate… · p. 7
  • Table I reports time and performance for each convolution layer across forward data, backward weights, and backward data.The caption identifies the measurements’ scope but does not provide layer-specific values.Source: Table I · p. 7

B. Single-node results

Single-node profiling identifies convolution and pooling as important computational hotspots, alongside element-wise and data-reordering work. CosmoFlow reaches 535 Gflop/s on one Cori KNL node, while the reported scaling measurements include end-to-end epoch timing.

B. Single-node results

  • Larger convolutions exceed 1 TFlop/Second, while smaller convolutions contribute less to total time.Print-statement timing adds overhead on Cori, so Table I’s total time exceeds Figure 3’s.Source: B. Single-node results · p. 7
  • On one KNL node, 64 OpenMP threads, six I/O threads, and four plugin threads are distributed across master, worker, and communication roles.Only one I/O thread is active, while plugin threads mostly spin because there is no communication in this run.Source: B. Single-node results · p. 7, Figure 3 · p. 8
  • Element-wise operations and data reordering add little to the flop count, and low-load element-wise work creates OpenMP thread imbalance.Leaky ReLU, Adam+LARC with polynomial decay, and loss calculations use loop-level OpenMP parallelism.Source: B. Single-node results · pp. 7–8
  • 535 Gflop/s is achieved on one KNL node, compared with 388 Gflop/s on a Piz Daint GPU node.The KNL measurement includes I/O and CPE ML Plugin overhead.Source: B. Single-node results · p. 8

C. Multi-node scaling results

  • Multi-node speedups are measured by walltime per epoch, including training and loss averaging across MPI ranks.The average uses eight epochs from a ten-epoch run, excluding the first two, to capture end-to-end system and software performance.Source: C. Multi-node scaling results · p. 8

1) Cori:

CosmoFlow’s largest Cori run sustains over 3.5 Pflop/s with 77% full-epoch parallel efficiency on 8192 nodes. The results also show that shared-filesystem I/O constrains scaling, while DataWarp improves Cori scaling efficiency.

1) Cori:

  • At 1024 Cori nodes with training data on Lustre, scaling efficiency falls below 58%; DataWarp yields 77% efficiency at 8192 nodes.Tests with generated dummy data suggest I/O causes a significant scaling drop; Figure 4 presents epoch-timing scaling and highlights the Lustre drop.Source: 1) Cori · p. 8, Figure 4 · p. 9

2) Piz Daint:

  • On Piz Daint, Lustre scaling efficiency drops to 44% at 512 nodes, with a probable read bottleneck at 512 nodes and beyond.The authors expect Piz Daint DataWarp to improve scaling, but defer its investigation because of availability issues.Source: 2) Piz Daint · p. 8

D. Full scale run

  • Over 3.5 Pflop/s sustained performance and 77% full-epoch parallel efficiency are achieved on 8192 Cori KNL nodes.The run completed 130 epochs, with a 3.35-second average epoch time; step-only scaling efficiency was 80%.Source: D. Full scale run · p. 8
  • The 2048-node run converges in fewer epochs than the 8192-node run, using a subset of the dataset used for the larger run.The authors tune learning-rate and decay hyperparameters and experiment with dataset variations for both runs.Source: D. Full scale run · p. 8

A. I/O

  • At 128 nodes, DataWarp takes 150 ms per step versus 179 ms on Lustre; the estimated minimum read bandwidth is 62 MB/s per compute node.The authors suspect variable Lustre bandwidth and contention may limit performance at the lowest-bandwidth OSTs.Source: A. I/O · p. 9

B. Communication

The 8192-node run produced reasonable but less accurate parameter estimates than the 2048-node target run. Communication measurements indicate high plugin bandwidth, though the loss and epoch times in Figure 5 cannot be directly compared across runs because their datasets differ in size.

B. Communication

  • Figure 5’s training and validation losses and epoch times cannot be directly compared across the 2048- and 8192-node runs because their datasets differ in size.Source: Figure 5 · p. 9
  • At 8192 nodes, estimated CPE ML Plugin bandwidth was 1.42 GB/s/node, compared with 1.7 GB/s/node at 1024 nodes and ≈10 GB/s/node Aries capacity.The 1.7 GB/s/node estimate assumes gradient aggregation alone accounts for the 33 ms step-time difference.Source: B. Communication · p. 9

A. Cosmology results

  • At 2048 nodes, relative errors were (0.0022, 0.0094, 0.0096) for (ΩM, σ8, ns); the ns error was almost 5× smaller than the best experimental uncertainty.The parameter estimates are shown in Figure 6; errors were comparable to the best experimental uncertainty for ΩM and σ8.Source: A. Cosmology results · p. 9, Figure 6 · p. 10
  • At 8192 nodes, relative uncertainties were (0.052, 0.014, 0.022) for (ΩM, σ8, ns), within a factor of 2 of today’s best experimental results,.The run was learning but had not converged to the accuracy of the 2048-node target run.Source: A. Cosmology results · p. 10

B. Implications on the application and future hardware

CosmoFlow scaled fully synchronous training to 8192 Cori nodes, sustaining 3.5 Pflop/s with 77% parallel efficiency and producing reasonable scientific predictions. The authors identify I/O, dataset size, and hyperparameter tuning as constraints, while describing further cosmological applications as within reach.

B. Implications on the application and future hardware

  • With predictions in a few minutes on 8192 nodes, the framework puts extended cosmological problems, optimized hyperparameter searches, multiple redshift snapshots, and combined probes within reach.Source: B. Implications on the application and… · p. 10
  • Scaling experiments showed a significant Lustre efficiency drop that the authors compensated for with a high-bandwidth SSD-based file system.Source: B. Implications on the application and… · p. 10
  • The 128^3-voxel, three-parameter problem would take more than 60 days on one node to reach an acceptable loss.The authors state that synchronous data-parallel training is limited by dataset size relative to MPI ranks and requires multiple runs for hyperparameter tuning.Source: B. Implications on the application and… · p. 10

VIII. CONCLUSIONS

  • CosmoFlow sustained 3.5 Pflop/s in aggregate on Cori and achieved 77% parallel efficiency on up to 8192 nodes.The conclusion reports close-to-converged 8192-node results and full convergence at 2048 nodes.Source: §VIII · p. 10, B. Implications on the application and… · p. 10
  • The authors report unprecedented cosmological-parameter accuracy, particularly for estimating ns from the dark matter distribution alone.Source: §VIII · p. 10
Loading 1808.04728v2…