[论文解读] CosmoFlow: Using Deep Learning to Learn the Universe at Scale

Amrita Mathuriya, Deborah Bard+15 more|arXiv (Cornell University)|Aug 14, 2018

Parallel Computing and Optimization Techniques参考文献 14被引用 24

一句话总结

CosmoFlow 提出了一种基于 TensorFlow 的高度可扩展深度学习框架，用于从 3D 暗物质模拟中预测宇宙学参数（$\Omega_M$，$\sigma_8$，$n_s$）。通过优化 3D 卷积神经网络，并在 Cori 超算机的 8192 个节点上实现完全同步的数据并行训练，该框架在持续性能上达到 3.5 Pflop/s，平行效率达 77%，从而实现了参数估计前所未有的精度。

ABSTRACT

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 $Ω_M$, $σ_8$ and n$_s$ with unprecedented accuracy.

研究动机与目标

• 解决从大规模 3D 暗物质模拟中进行宇宙学参数估计时面临的计算瓶颈。
• 在百亿亿次 HPC 平台上，利用深度学习实现对宇宙学参数（$\Omega_M$，$\sigma_8$，$n_s$）的高精度回归。
• 针对基于 CPU 的超算机，对完整软件栈——网络架构、I/O、通信及原语——进行优化，以实现深度学习的可扩展性。
• 在超算机上使用 TensorFlow 实现大规模、完全同步的数据并行训练，并保持高效率。
• 通过大规模深度学习实现对宇宙学数据的快速、精确科学探索。

提出的方法

• 将 Ravanbakhsh 等人（2017）提出的 3D 卷积神经网络架构适配于 $128^3$ 体素输入数据，用于预测三个宇宙学参数。
• 在 TensorFlow 框架内使用 MKL-DNN 实现优化的 3D 卷积与池化原语，以实现面向 CPU 的高性能计算。
• 利用 Cray PE 机器学习插件，在数千个节点之间实现基于 MPI 的高效数据并行。
• 采用基于高速 SSD 的文件系统，缓解在 Lustre 上观察到的 I/O 瓶颈，提升扩展效率。
• 使用同步随机梯度下降（SSGD），在节点间进行完整参数平均，以确保大规模训练下的收敛性。
• 采用 HPC 最佳实践，对整个堆栈——从数据 I/O、模型训练到节点间通信——进行优化，适配 Intel Xeon Phi（KNL）处理器。

实验结果

研究问题

• RQ1 在 3D 暗物质模拟上训练的深度学习模型能否准确预测宇宙学参数 $\Omega_M$，$\sigma_8$ 和 $n_s$？
• RQ2 在拥有 8192 个节点的超算机上训练宇宙学深度学习模型时，可达到怎样的性能水平和扩展效率？
• RQ3 I/O 系统的选择（如 Lustre 与基于 SSD 的系统）对大规模深度学习工作负载的训练可扩展性和性能有何影响？
• RQ4 在极端规模（8192 个节点）下，完全同步的数据并行训练在多大程度上能保持收敛性和准确性？
• RQ5 经过优化的基于 CPU 的深度学习框架能否在科学应用中实现百亿亿次级别的性能？

主要发现

• CosmoFlow 在 Cori 超算机的 8192 个节点上，使用 Intel Xeon Phi 处理器，实现了 3.5 Pflop/s 的持续单精度性能。
• 系统在大规模下实现了 77% 的并行效率，标志着首次在超算机上使用 TensorFlow 实现大规模、完全同步的数据并行训练深度学习模型。
• 在 8192 个节点上训练已接近收敛，尽管计算负载极高，但仍产生了具有科学意义的预测结果。
• 模型在 2048 个节点上实现完全收敛，为从暗物质分布中估计宇宙学参数提供了当前最佳的科学精度。
• 采用高速带宽的 SSD 基础文件系统显著提升了 I/O 性能，并缓解了在 Lustre 上观察到的可扩展性下降问题。
• 单节点性能在 KNL 节点上达到 535 Gflop/s，全栈优化使系统能够高效处理超过 1.4TB 的宇宙学模拟数据。