Exascale Deep Learning for Scientific Inverse Problems

TL;DR

This paper presents advanced communication strategies for distributed deep learning that enable near-linear scaling on supercomputers, demonstrated by training a neural network for materials imaging with unprecedented efficiency and accuracy.

Contribution

The paper introduces novel communication techniques that optimize overlap between computation and communication, enabling scalable training on thousands of GPUs for scientific inverse problems.

Findings

Achieved near-linear scaling up to 27,600 GPUs

Trained a neural network on 0.5 PB dataset for materials imaging

Reached a peak performance of 2.15 EFLOPS_16

Abstract

We introduce novel communication strategies in synchronous distributed Deep Learning consisting of decentralized gradient reduction orchestration and computational graph-aware grouping of gradient tensors. These new techniques produce an optimal overlap between computation and communication and result in near-linear scaling (0.93) of distributed training up to 27,600 NVIDIA V100 GPUs on the Summit Supercomputer. We demonstrate our gradient reduction techniques in the context of training a Fully Convolutional Neural Network to approximate the solution of a longstanding scientific inverse problem in materials imaging. The efficient distributed training on a dataset size of 0.5 PB, produces a model capable of an atomically-accurate reconstruction of materials, and in the process reaching a peak performance of 2.15(4) EFLOPS$_{16}$.