Graph neural networks (GNNs) are a powerful approach for machine learning on graph datasets. Such datasets often consist of millions of modestly-sized graphs, making them well-suited for data-parallel training. However, existing methods show poor scaling due to load imbalances and kernel overheads. We propose an optimized 2D scatter-gather based represen-tation of GNNs that is amenable to distributed, data-parallel training without changing the underlying mathematics of the GNN. By padding graph data to a fixed size on each process, we can simplify data ingestion, make use of efficient compute kernels, equally distribute computation load, and reduce overheads. We benchmark edge-conditioned GNNs with the PCQM4M-LSC and OGB-PPA datasets. Our implementation shows better runtime performance than the state-of-the-art, with a $12\times$ strong-scaling speedup on 16 GPUs and an $89.4\times\ \text{weak}$ -scaling speedup on 100 GPUs.

Transformer models have revolutionized the field of Natural Language Processing (NLP) and they achieve state-of-the-art performance in applications like machine translation, question answering, regression, and summarization. However, training Transformers is challenging because of their large memory and compute requirements. The literature contains several approaches to parallelize training, like layer parallelism and pipeline parallelism, but they are optimized to benefit out-of-core models and they don’t exploit the inherent parallelism in Transformer models. Other work uses model parallelism to achieve weak scaling by increasing the model size. In this paper, we propose sub-graph parallelism that provides a significant performance improvement over pure data parallelism with a fixed number of resources, and as an additional technique for strong- and weak-scaling without increasing model capacity. Our technique accelerates the training of Transformer models and we generalize the concept to any neural network with multiple branches. We optimize the communication for sub-graph parallelism and combine it with data parallelism to scale performance up to 1024 GPUs. To decrease communication overheads, we propose a topology-aware scheme that limits inter-node communication. Finally, we empirically compare sub-graph parallelism with pure data parallelism and demonstrate its performance benefits in end-to-end training.

Monumental advances in deep learning have led to unprecedented achievements across various domains. While the performance of deep neural networks is indubitable, the architectural design and interpretability of such models are nontrivial. Research has been introduced to automate the design of neural network architectures through neural architecture search (NAS). Recent progress has made these methods more pragmatic by exploiting distributed computation and novel optimization algorithms. However, there is little work in optimizing architectures for interpretability. To this end, we propose a multi-objective distributed NAS framework that optimizes for both task performance and "introspectability," a surrogate metric for aspects of interpretability. We leverage the non-dominated sorting genetic algorithm (NSGA-II) and explainable AI (XAI) techniques to reward architectures that can be better comprehended by domain experts. The framework is evaluated on several image classification datasets. We demonstrate that jointly optimizing for task error and introspectability leads to more disentangled and debuggable architectures that perform within tolerable error.

We improved the quality and reduced the time to produce machine learned models for use in small molecule antiviral design. Our globally asynchronous multi-level parallel training approach strong scales to all of Sierra with up to 97.7% efficiency. We trained a novel, character-based Wasserstein autoencoder that produces a higher quality model trained on 1.613 billion compounds in 23 minutes while the previous state of the art takes a day on 1 million compounds. Reducing training time from a day to minutes shifts the model creation bottleneck from computer job turnaround time to human innovation time. Our implementation achieves 318 PFLOPs for 17.1% of half-precision peak. We will incorporate this model into our molecular design loop enabling the generation of more diverse compounds; searching for novel, candidate antiviral drugs improves and reduces the time to synthesize compounds to be tested in the lab.

Accelerating large-scale CNN training is needed to keep training times reasonable as datasets grow larger and models become more complex. Existing frameworks primarily scale using dataparallelism, but this is limited by the mini-batch size, which cannot grow arbitrarily. We introduce three algorithms that partition channel or filter data to exploit parallelism beyond the sample dimension. Further, they partition the parameters of convolutional layers, replacing global allreduces with segmented allreduces-smaller, concurrent allreduces among disjoint processor sets. These algorithms enable strong scaling, reduced communication overhead, and reduced memory pressure, enabling training of very wide CNNs. We demonstrate improved strong and weak scaling, including up to 4.1x reductions in training time for residual networks and 4x reductions in allreduce overhead. We also show that wider models provide improved accuracy on ImageNet. We study the current limitations of our algorithms and provide a direction for future optimizations of large-scale deep learning frameworks.

Training deep neural networks on large scientific data is a challenging task that requires enormous compute power, especially if no pre-trained models exist to initialize the process. We present a novel tournament method to train traditional as well as generative adversarial networks built on LBANN, a scalable deep learning framework optimized for HPC systems. LBANN combines multiple levels of parallelism and exploits some of the worlds largest supercomputers. We demonstrate our framework by creating a complex predictive model based on multi-variate data from high-energydensity physics containing hundreds of millions of images and hundreds of millions of scalar values derived from tens of millions of simulations of inertial confinement fusion. Our approach combines an HPC workflow and extends LBANN with optimized data ingestion and the new tournament-style training algorithm to produce a scalable neural network architecture using a CORAL-class supercomputer. Experimental results show that 64 trainers (1024 GPUs) achieve a speedup of 70.2× over a single trainer (16 GPUs) baseline, and an effective 109% parallel efficiency.

Scaling CNN training is necessary to keep up with growing datasets and reduce training time. We also see an emerging need to handle datasets with very large samples, where memory requirements for training are large. Existing training frameworks use a data-parallel approach that partitions samples within a mini-batch, but limits to scaling the mini-batch size and memory consumption makes this untenable for large samples. We describe and implement new approaches to convolution, which parallelize using spatial decomposition or a combination of sample and spatial decomposition. This introduces many performance knobs for a network, so we develop a performance model for CNNs and present a method for using it to automatically determine efficient parallelization strategies. We evaluate our algorithms with microbenchmarks and image classification with ResNet-50. Our algorithms allow us to prototype a model for a mesh-tangling dataset, where sample sizes are very large. We show that our parallelization achieves excellent strong and weak scaling and enables training for previously unreachable datasets.

Training deep neural networks on large scientific data is a challenging task that requires enormous compute power, especially if no pre-trained models exist to initialize the process. We present a novel tournament method to train traditional as well as generative adversarial networks built on LBANN, a scalable deep learning framework optimized for HPC systems. LBANN combines multiple levels of parallelism and exploits some of the worlds largest supercomputers. We demonstrate our framework by creating a complex predictive model based on multi-variate data from high-energy-density physics containing hundreds of millions of images and hundreds of millions of scalar values derived from tens of millions of simulations of inertial confinement fusion. Our approach combines an HPC workflow and extends LBANN with optimized data ingestion and the new tournament-style training algorithm to produce a scalable neural network architecture using a CORAL-class supercomputer. Experimental results show that 64 trainers (1024 GPUs) achieve a speedup of 70.2 over a single trainer (16 GPUs) baseline, and an effective 109% parallel efficiency.

We identify communication as a major bottleneck for training deep neural networks on large-scale GPU clusters, taking over 10x as long as computation. To reduce this overhead, we discuss techniques to overlap communication and computation as much as possible. This leads to much of the communication being latency-bound instead of bandwidth-bound, and we find that using a combination of latency- and bandwidth-optimized allreduce algorithms significantly reduces communication costs. We also discuss a semantic mismatch between MPI and CUDA that increases overheads and limits asynchrony, and propose a solution that enables communication to be aware of CUDA streams. We implement these optimizations in the open-source Aluminum communication library, enabling optimized, asynchronous, GPU-aware communication. Aluminum demonstrates improved performance in benchmarks and end-to-end training of deep networks, for both strong and weak scaling.

The CEED co-design center is developing a number of kernels, bake-off / benchmark problems (BPs) and mini-applications (miniapps) that capture the unique requirements of high-order finite element algorithms and use them to influence standards, best practices and vendors to generate more portable, better performing code across diverse hardware designs. This is a main activity in CEED’s Hardware, Applications and Finite Element thrusts. CEED’s miniapps, a.k.a. CEEDlings, are small, standalone applications - surrogates for sub-components of interest from the selected first wave of applications (the MARBL/LLNLApp and ExaSMR projects). The BPs are even smaller standalone programs that solve a simple physical problem, including discretization, linear and/or nonlinear solve, and potentially a relatively small number of time steps. A BP may be a proxy/surrogate for part of a miniapp or another problem of general interest. In CEED, both miniapps and BPs will be used extensively in the project to evaluate and compare the performance of algorithms in specific contexts. The parts of the miniapps and BPs that play an important role in their overall performance will be identified and separated as standalone kernels. These kernels and their interaction inside BPs, miniapps and the actual applications, will be the main focus of our optimization efforts. All three standalone pieces of software (miniapps, BPs, and kernels) will serve as a basis for engaging and collaborating with hardware vendors and software technologies (ST) projects. These partners can start at the bottom of the hierarchy, e.g. with the kernels, and easily transition to BPs and miniapps as all three families will be produced in the same CEED environment. In this milestone we identified the initial kernels, bake-off problems and miniapps and delivered software and documentation for them through the new CEED website, http://ceed.exascaleproject.org. We specifically identified the Nekbone miniapp for the ExaSMR application and developed a new miniapp, Laghos, for the MARBL/LLNLApp application. We also formulated four initial benchmark problems (BP1- BP4) with clearly defined performance metrics and performed an initial bake-off comparison between the Nek5000 (spectral element) and MFEM (high-order finite element) technologies. As part of the milestone, we also engaged vendors (AMD) and ST projects (MPICH, STRUMPACK). In this document we are reporting some details and results from these R&D activities, as well as additional project-wide activities performed in Q3 of FY17, including: making CEED software available on GitHub, developing the project website, and the preparation of CEED’s first annual meeting.

We study data-parallel training of deep neural networks on high-performance computing infrastructure. The key problem with scaling data-parallel training is avoiding severe communication/computation imbalance. We explore quantizing gradient updates before communication to reduce bandwidth requirements and compare it against a baseline implementation that uses the MPI allreduce routine. We port two existing quantization approaches, one-bit and threshold, and develop our own adaptive quantization algorithm. The performance of these algorithms is evaluated and compared with MPI_Allreduce when training models for the MNIST dataset and on a synthetic benchmark. On an HPC system, MPI_Allreduce outperforms the existing quantization approaches. Our adaptive quantization is comparable or superior for large layers without sacrificing accuracy. It is 1.76 times faster than the next best approach for the largest layers in our benchmark and achieves near-linear speedup in data-parallel training.

Multi-shift triangular solves are basic linear algebra calculations with applications in eigenvector and pseudospectra computation. We propose blocked algorithms that efficiently exploit Level 3 BLAS to perform multi-shift triangular solves and safe multi-shift triangular solves. Numerical experiments indicate that computing triangular eigenvectors with a safe multi-shift triangular solve achieves speedups by a factor of 60 relative to LAPACK. This algorithm accelerates the calculation of general eigenvectors threefold. When using multi-shift triangular solves to compute pseudospectra, we report ninefold speedups relative to EigTool.