[논문 리뷰] High Performance Monte Carlo Simulation of Ising Model on TPU Clusters

Kun Yang, Yifan Chen외 3명|arXiv (Cornell University)|2019. 03. 27.

Theoretical and Computational Physics참고 문헌 17인용 수 8

한 줄 요약

이 논문은 클라우드 TPU 클러스터에서 텐서플로우를 사용하여 2D 이징 모델의 고성능 몬테카를로 시뮬레이션을 제시한다. TPU의 행렬 연산과 고속 인터커넥트를 활용하여 이전 벤치마크 대비 다중 코어 성능을 250% 향상시켰다. 단일 Jupyter 노트북에 구현된 이 시스템은 bfloat16 정밀도를 유지하면서도 2048개 코어까지 강한 선형 스케일링을 보여준다.

ABSTRACT

Large-scale deep learning benefits from an emerging class of AI accelerators. Some of these accelerators' designs are general enough for compute-intensive applications beyond AI and Cloud TPU is one such example. In this paper, we demonstrate a novel approach using TensorFlow on Cloud TPU to simulate the two-dimensional Ising Model. TensorFlow and Cloud TPU framework enable the simple and readable code to express the complicated distributed algorithm without compromising the performance. Our code implementation fits into a small Jupyter Notebook and fully utilizes Cloud TPU's efficient matrix operation and dedicated high speed inter-chip connection. The performance is highly competitive: it outperforms the best published benchmarks to our knowledge by 60% in single-core and 250% in multi-core with good linear scaling. When compared to Tesla V100 GPU, the single-core performance maintains a ~10% gain. We also demonstrate that using low precision arithmetic---bfloat16---does not compromise the correctness of the simulation results.

연구 동기 및 목표

• 딥 러닝을 초월해 AI 가속기(예: 클라우드 TPU)를 과학 계산에 활용할 수 있는지 탐색하기 위해.
• TPU에서 텐서플로우를 사용하여 성능이 뛰어나고 확장 가능하며 가독성이 좋은 2D 이징 모델 몬테카를로 시뮬레이션을 개발하기 위해.
• 통계역학에서 몬테카를로 시뮬레이션의 정확성에 저하되지 않는지, 저정밀도 산술(bfloat16)의 영향을 평가하기 위해.
• 수천 개의 TPU 코어를 통해 시뮬레이션이 강한 선형 스케일링을 보이는지 확인하기 위해.

제안 방법

• 클라우드 TPU에서 텐서플로우를 사용하여 2D 이징 모델에 대해 SIMD 분산 마르코프 체인 몬테카를로(MCMC) 알고리즘을 구현하기 위해.
• TPU의 고대역폭 메모리와 2D 토로이드 메시 인터커넥트를 활용하여 효율적인 데이터 이동과 저지연 통신을 구현하기 위해.
• MXU의 활용도를 높이기 위해 배치 행렬 곱셈 대신 tf.nn.conv2d를 사용하여 근접 이웃 에너지 기여도를 효율적으로 계산하기 위해.
• 자원 활용도를 극대화하기 위해 느슨하게 팩킹된, 조밀하게 팩킹된, 슈퍼밀도 팩킹된 구성 방식을 활용하여 TPU 코어 간 작업 분배를 최적화하기 위해.
• 최소한의 코드 복잡도로 높은 성능을 달성하기 위해 텐서플로우 r1.15와 TPU v3 하드웨어를 사용하기 위해.
• 최소한의 설정으로 Google Colab을 통해 Jupyter 노트북에서 전체 시뮬레이션을 실행하여 TPU에서 상호작용 가능한 과학 계산을 가능하게 하기 위해.

실험 결과

연구 질문

• RQ1 클라우드 TPU 클러스터가 딥 러닝 외의 고성능 과학 시뮬레이션에 효과적으로 활용될 수 있는가?
• RQ2 TPU에서 텐서플로우 기반의 이징 모델 시뮬레이션 성능이 공개된 GPU 및 CPU 벤치마크와 비교해 어떻게 되는가?
• RQ3 bfloat16 정밀도를 사용할 경우 통계역학에서 몬테카를로 시뮬레이션의 정확성이 손상되는가?
• RQ4 시뮬레이션이 수천 개의 TPU 코어를 통해 얼마나 선형적으로 스케일링되는가?
• RQ5 복잡한 분산 과학 알고리즘을 성능을 희생시키지 않고도 텐서플로우에서 명확하고 효율적으로 표현할 수 있는가?

주요 결과

• TPU 기반 구현은 단일 코어 기준으로 이전에 발표된 최고의 벤치마크보다 60% 빠르고, 다중 코어 기준으로는 250% 빠르며, 뛰어난 선형 스케일링을 보였다.
• 32,768개 코어를 가진 전체 TPU v3 풀에서 시뮬레이션은 1나노초당 약 20,460.92회의 플립을 기록하여 1,000개 코어까지 강한 스케일링을 보였다.
• 1,000개 코어를 초과하면 통신 오버헤드가 심각해져 강한 스케일링 테스트에서 추가적인 스케일링 효율성이 저하되었다.
• bfloat16 정밀도를 사용하더라도 시뮬레이션 결과의 정확성이 손상되지 않았으며, 물리적 정확성이 유지되었다.
• 전체 시뮬레이션이 단일 가독성 있는 Jupyter 노트북에 포함되어 있으며, 최소한의 설정으로 Google Colab에서 실행 가능했다.
• 초기 구현 대비 배치 곱셈 대신 tf.nn.conv2d를 사용한 최적화된 버전은 약 80%의 성능 향상을 달성했다.