[논문 리뷰] CosmoFlow: Using Deep Learning to Learn the Universe at Scale

Amrita Mathuriya, Deborah Bard외 15명|arXiv (Cornell University)|2018. 08. 14.

Parallel Computing and Optimization Techniques참고 문헌 14인용 수 24

한 줄 요약

CosmoFlow는 텐서플로우 기반으로 구축된 매우 확장 가능한 딥러닝 프레임워크로, 3D 어둠침재 시뮬레이션에서 천체물리학적 파arameter($\Omega_M$, $\sigma_8$, $n_s$)를 예측한다. 8192개 노드로 구성된 Cori 슈퍼컴퓨터에서 3D 컨볼루션 신경망을 최적화하고 완전한 동기화 데이터 병렬 처리를 구현함으로써 지속성 3.5 Pflop/s 성능과 77%의 병렬 효율성을 달성하여, 파arameter 추정의 사상 최고 수준의 정확도를 실현한다.

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 어둠침재 시뮬레이션에서의 천체물리학적 파arameter 추정에 있어 계산적 병목 현상을 해결하기 위해.
• 엑세스케일 HPC 플랫폼에서 CPU 기반 슈퍼컴퓨터를 활용해 천체물리학적 파arameter($\Omega_M$, $\sigma_8$, $n_s$)의 고정밀도 회귀를 가능하게 하기 위해.
• 소프트웨어 스택 전반—신경망 아키텍처, I/O, 통신, 기초 primitive—을 최적화하여 CPU 기반 슈퍼컴퓨터에서의 확장 가능한 딥러닝을 실현하기 위해.
• 텐서플로우를 사용하여 슈퍼컴퓨터에서 대규모로 완전한 동기화 데이터 병렬 학습을 구현하기 위해.
• 대규모 딥러닝을 활용해 천체물리학 데이터의 신속하고 정확한 과학적 탐색을 가능하게 하기 위해.

제안 방법

• Ravanbakhsh 등 (2017)의 3D 컨볼루션 신경망 아키텍처를 $128^3$ 바vox 입력 데이터에 적응시켜 세 가지 천체물리학적 파arameter를 예측하도록 구현하였다.
• Intel MKL-DNN를 활용해 텐서플로우 프레임워크 내부에서 CPU 최적화된 3D 컨볼루션 및 풀링 primitive를 구현하였다.
• Cray PE 머신러닝 플러그인을 활용해 수천 개 노드 간 효율적인 MPI 기반 데이터 병렬 처리를 실현하였다.
• I/O 병목 현상을 완화하기 위해 고대역폭 SSD 기반 파일 시스템을 도입하여 루스트(Lustre) 대비 성능 향상을 달성하였다.
• 모든 노드 간 전체 파aram터 평균화를 수행하는 동기화 확률적 경사 하강법(SSGD)을 사용하여 대규모에서의 수렴을 보장하였다.
• Intel Xeon Phi(KNL) 프로세서를 대상으로 HPC 최적화 기법을 적용하여 데이터 I/O, 모델 학습, 노드 간 통신에 이르기까지 전체 스택을 최적화하였다.

실험 결과

연구 질문

• RQ1 3D 어둠침재 시뮬레이션에서 학습된 딥러닝 모델이 천체물리학적 파arameter $\Omega_M$, $\sigma_8$, $n_s$를 정확하게 예측할 수 있는가?
• RQ2 8192개 노드로 구성된 슈퍼컴퓨터에서 천체물리학용 딥러닝 모델을 학습할 경우, 어떤 성능 수준과 스케일링 효율성을 달성할 수 있는가?
• RQ3 I/O 시스템 선택(예: 루스트 대비 SSD 기반)이 대규모 딥러닝 워크로드에서의 학습 스케일링 및 성능에 어떤 영향을 미치는가?
• RQ4 완전한 동기화 데이터 병렬 학습이 극한 규모(8192개 노드)에서 수렴성과 정확도를 유지하는 데 얼마나 효과적인가?
• RQ5 최적화된 CPU 기반 딥러닝 프레임워크가 과학적 응용 분야에서 엑세스케일 수준의 성능을 달성할 수 있는가?

주요 결과

• CosmoFlow는 Intel Xeon Phi 프로세서를 사용한 Cori 슈퍼컴퓨터의 8192개 노드에서 단일 정밀도 성능으로 지속성 3.5 Pflop/s를 달성하였다.
• 대규모에서 77%의 병렬 효율성을 보이며, 텐서플로우를 사용한 슈퍼컴퓨터에서의 첫 번째 대규모 완전 동기화 데이터 병렬 학습을 실현하였다.
• 8192개 노드에서의 학습은 높은 계산 부하에도 불구하고 수렴에 가까운 결과를 도출하여 과학적으로 의미 있는 예측을 제공하였다.
• 2048개 노드에서의 모델은 완전한 수렴을 달성하였으며, 어둠침재 분포에서의 천체물리학적 파arameter 추정 분야에서 최고 수준의 과학적 정확도를 확보하였다.
• 고대역폭 SSD 기반 파일 시스템의 도입으로 I/O 성능이 크게 향상되었고, 루스트에서 관찰된 스케일링 저하 현상이 완화되었다.
• 단일 노드에서의 성능은 KNL 노드 기준 535 Gflop/s에 도달하였으며, 전체 스택 최적화를 통해 1.4TB 이상의 천체물리학 시뮬레이션 데이터를 효율적으로 처리할 수 있었다.