[논문 리뷰] FlowPM: Distributed TensorFlow Implementation of the FastPM Cosmological N-body Solver

Chirag Modi, François Lanusse외 1명|arXiv (Cornell University)|2020. 10. 22.

Galaxies: Formation, Evolution, Phenomena참고 문헌 53인용 수 46

한 줄 요약

FlowPM은 Mesh-TensorFlow에서 새로운 다중 그리드, 다해상도 피라미드 체계를 사용하는 TensorFlow 기반의 GPU 가속, 미분 가능, 분산 Particle-Mesh N-body 해결책을 구현하여 Python FastPM 대비 약 10배의 벽시계 속도 향상을 달성하고 우주론 추론을 위한 미분 가능 순방향 모델링을 가능하게 한다.

ABSTRACT

We present FlowPM, a Particle-Mesh (PM) cosmological N-body code implemented in Mesh-TensorFlow for GPU-accelerated, distributed, and differentiable simulations. We implement and validate the accuracy of a novel multi-grid scheme based on multiresolution pyramids to compute large scale forces efficiently on distributed platforms. We explore the scaling of the simulation on large-scale supercomputers and compare it with corresponding python based PM code, finding on an average 10x speed-up in terms of wallclock time. We also demonstrate how this novel tool can be used for efficiently solving large scale cosmological inference problems, in particular reconstruction of cosmological fields in a forward model Bayesian framework with hybrid PM and neural network forward model. We provide skeleton code for these examples and the entire code is publicly available at https://github.com/modichirag/flowpm.

연구 동기 및 목표

• Distributed GPU를 활용한 미분 가능 PM 프레임워크를 통해 빠르고 확장 가능한 우주론 N-body 시뮬레이션을 가능하게 한다.
• 상호 프로세스 간 통신을 줄이기 위한 새로운 다해상도 피라미드 기반 다중 격자 힘 추정법을 개발하고 검증한다.
• 분산 모델-병렬 N-body 계산을 위한 Mesh-TensorFlow 기반 아키텍처를 제공한다.
• 하이브리드 PM/ML 순방향 모델과의 결합을 통한 순방향 모델링 및 우주론 필드 재구성에의 적용 가능성을 보여준다.

제안 방법

• GPU 가속과 미분 가능성을 위한 TensorFlow에서의 FastPM 스타일 PM 진화를 구현한다.
• 저해상도 글로벌 격자에서의 장거리 힘과 고해상도 로컬 격자에서의 단거리 힘을 나누기 위해 다해상도 피라미드를 사용하는 2단계 다중 격자 힘 추정을 도입한다.
• Mesh-TensorFlow를 이용해 분산 격자 청크의 슬라이스별 연산이 가능하도록 패딩이 있는 홀로 교환 프로토콜을 사용한다.
• 격자 차원의 분산 분배를 제어하는 프로세서 망에 텐서 연산을 매핑하기 위해 Mesh-TensorFlow를 채택한다.
• 차원 스케일링을 위한 피라미드 구성에 대해 3D 컨볼루션과 bspline 커널을 이용한 평활한 다운샘플/업샘플(REDUCE/EXPAND) 연산을 구성하고 평가한다.
• 정확도를 표준 FastPM 코드와 비교하고 서로 다른 격자/메시 구성에서 Cori GPU의 스케일링을 평가한다.

실험 결과

연구 질문

• RQ1 PM 우주론 N-body 해결책을 미분 가능하고 GPU 가속 프레임워크에서 구현할 수 있는가?
• RQ2 다중 격자, 다해상도 피라미드 체계가 정확도를 유지하면서 분산 FFT 통신을 줄일 수 있는가?
• RQ3 FlowPM의 대형 GPU 메시에서의 성능 및 확장성은 Python FastPM 구현과 비교해 어떤 특징을 가지는가?
• RQ4 FlowPM을 순방향 모델링 우주론 추론 워크플로우(예: 혼합 PM/ML 구성 요소를 이용한 재구성)에 효과적으로 통합할 수 있는가?

주요 결과

• FlowPM은 테스트 구성에서 FastPM에 비해 잔차가 낮은 정확도를 보이며(전달 함수 및 교차상관은 0.01% 이내).
• 두 단계 다중 격자 피라미드 체계가 글로벌 격자에서 장거리 힘을, 로컬 격자에서 단거리 힘을 처리함으로써 분산 FFT 통신을 줄인다.
• FlowPM은 Cori GPUs의 8-V100 노드에서 규모 확장 가능하며, 비교 구성에서 Python FastPM 대비 평균 약 10배의 벽시계 속도 향상을 보인다.
• Mesh-TensorFlow 기반 구현은 분산 모델-병렬 시뮬레이션과 경계 영역을 프로세스 슬라이스 간에 관리하기 위한 홀로 교환을 가능하게 한다.
• 이 프레임워크는 순방향 모델링 및 재구성 작업에 적합한 미분 가능 시뮬레이션을 지원하며, 신경망 순방향 모델과의 toy 예제를 통해 시연되었다.