[论文解读] Exascale Deep Learning for Scientific Inverse Problems

Nouamane Laanait, Joshua Romero+7 more|arXiv (Cornell University)|Sep 24, 2019

Machine Learning in Materials Science参考文献 30被引用 29

一句话总结

该论文提出了一种新型通信策略——去中心化的梯度规约编排与计算图感知的分组机制，用于同步分布式深度学习，在Summit超算上实现了27,600块NVIDIA V100 GPU的近线性扩展（0.93）。该方法实现了对0.5 PB科学数据集上全卷积神经网络的高效训练，达到百亿亿次性能（2.15(4) EFLOPS），并实现了从电子衍射数据中进行原子级精确的材料重构。

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}$.

研究动机与目标

• 解决超级计算机上数据并行分布式深度学习中的通信瓶颈问题，其中扩展低效性不断累积，导致训练过程碳足迹增加。
• 开发优化计算与通信重叠的通信策略，以在大规模GPU集群上实现近线性扩展。
• 实现在大规模科学数据集（0.5 PB）上训练大规模深度神经网络，以解决材料成像领域长期存在的逆问题。
• 证明使用百亿亿次深度学习近似求解科学逆问题（如电子显微镜中的相位问题）的可行性。
• 通过在超级计算基础设施上实现高性能、可扩展的训练，为将深度学习整合到科学模拟中奠定基础。

提出的方法

• 通过智能调度节点间的梯度聚合，实现去中心化的梯度规约编排，以最小化同步开销。
• 提出基于模型计算图结构的梯度张量分组机制，以根据计算图结构合并通信操作。
• 在Summit超算上，利用NCCL实现节点内通信，使用IBM Spectrum-MPI实现节点间通信，采用分层Allreduce策略。
• 利用节点本地的NVMe存储，避免全局文件系统的I/O瓶颈，实现数据加载的线性扩展，且每块GPU的数据量随数据规模线性增长。
• 设计一种分布式训练流水线，实现数据和模型副本的同步扩展，确保每块GPU处理0.5 PB数据集的一个独立子集。
• 采用混合精度训练（FP16），并充分利用V100 GPU上的Tensor Cores，以最大化计算吞吐量并减少通信量。

实验结果

研究问题

• RQ1 在拥有超过27,000块GPU的超级计算机上，新型梯度规约策略是否能实现数据并行分布式训练中的近线性扩展？
• RQ2 去中心化梯度编排与计算图感知分组的结合，如何提升计算与通信之间的重叠程度？
• RQ3 在保持高性能和模型准确率的前提下，分布式深度学习在多大程度上可扩展至0.5 PB的科学成像数据？
• RQ4 在如此大规模的分布式系统上训练的深度神经网络，是否能实现从电子显微镜逆问题中进行原子级精确的材料重构？
• RQ5 随着数据和计算规模的增加，对模型泛化能力和重构质量会产生何种影响？

主要发现

• 所提出的通信策略在Summit超算的27,600块NVIDIA V100 GPU上实现了0.93的扩展效率，证明了近线性扩展。
• 在训练过程中，系统在FP16精度下达到2.15(4) EFLOPS的峰值性能，标志着科学应用中百亿亿次深度学习的重要里程碑。
• 与128个节点相比，使用4096个节点训练显著降低了电子密度重构的平均测试误差，表明随着数据和计算规模的增加，模型泛化能力得到提升。
• 在1024个节点上训练的模型，其重构结果明显优于128个节点的模型，图像视场范围内与真实值的对齐度更高。
• 使用节点本地NVMe存储实现了数据加载的线性扩展，避免了直接从全局文件系统读取带来的10倍性能惩罚。
• 梯度规约技术已集成至Horovod（Bitvector Allreduce），并正在审核中（分组功能），表明其在本研究之外也具有广泛适用性。