[논문 리뷰] Distributed Machine Learning for Computational Engineering using MPI

Kailai Xu, Weiqiang Zhu외 1명|arXiv (Cornell University)|2020. 11. 02.

Reservoir Engineering and Simulation Methods참고 문헌 42인용 수 5

한 줄 요약

이 논문은 MPI를 사용하여 딥 네ural 워크(DNNs)와 편미분방정식(PDE) 해법기를 통합하는 분산 기계학습 프레임워크를 제안한다. 데이터 통신을 계산 그래프 내의 일급 노드로 간주함으로써, 분산 메모리 시스템에서 병렬적인 전방 및 수반 계산을 가능하게 하여 강한 약한 스케일링 성능을 달성하고, 100개의 프로세서를 사용할 때 20배 이상의 성능 향상을 보이며, 대규모 PDE 제약 조건을 갖는 DNN에 대한 확장 가능한 역모델링을 입증한다.

ABSTRACT

We propose a framework for training neural networks that are coupled with partial differential equations (PDEs) in a parallel computing environment. Unlike most distributed computing frameworks for deep neural networks, our focus is to parallelize both numerical solvers and deep neural networks in forward and adjoint computations. Our parallel computing model views data communication as a node in the computational graph for numerical simulations. The advantage of our model is that data communication and computing are cleanly separated and thus provide better flexibility, modularity, and testability. We demonstrate using various large-scale problems that we can achieve substantial acceleration by using parallel solvers for PDEs in training deep neural networks that are coupled with PDEs.

연구 동기 및 목표

• 역모드 자동 미분을 사용할 때 DNN과 PDE를 결합하여 학습하는 데 발생하는 메모리 및 계산 병목 현상을 해결하기 위해.
• 계산 공 ingeneering 분야의 대규모 역문제를 위한 확장 가능한 분산 최적화를 가능하게 하기 위해.
• 일급 통신 노드를 갖는 단일 계산 그래프 추상화 아래에서 병렬 PDE 해법기와 분산 DNN 학습을 통합하기 위해.
• MPI 기반 데이터 통신 및 하이브리드 병렬 처리(MPI + 스레딩)를 통한 기울기 역전파 과정에서 발생하는 과제를 극복하기 위해.

제안 방법

• MPI 브로드캐스트, 할로 교환 등 데이터 통신 작업을 계산 그래프 내의 명시적 노드로 간주하여 종단 간 자동 미분을 가능하게 한다.
• 딥 네ural 워크(DNNs)와 유한 요소/유한 차분 기반 PDE 해법기(예: 유한 요소/유한 차분)를 통합된 계산 그래프에 통합하여 역모드 AD를 통한 기울기 계산을 수행한다.
• 기본 산술 연산이 아닌 고수준 추상화(예: 행렬 해법기, PDE 연산자)를 대상으로 하는 거시적 자동 미분 접근 방식을 사용한다.
• ADCME 라이브러리 내에서 통신 패턴을 통한 기울기 역전파를 지원하기 위해 커스터마이징된 MPI 프리미티브를 구현한다.
• 계산 및 통신 로직을 분리하여 하이브리드 MPI+스레딩 환경에서의 모듈성, 테스트 가능성 및 유연성을 향상시킨다.
• 도메인 분할을 통해 MPI 프로세스 간 계산 영역을 분할하고, 스텐실 기반 PDE 해법기의 경우 할로 교환을 수행한다.

실험 결과

연구 질문

• RQ1 분산 PDE 해법기에서의 데이터 통신을 자동 미분을 지원하는 일급 계산 그래프 노드로 모델링할 수 있는가?
• RQ2 역모드 자동 미분은 하이브리드 MPI+스레딩 환경에서 MPI 통신 패턴을 통한 기울기를 어떻게 처리할 수 있는가?
• RQ3 이 프레임워크를 사용하여 DNN과 PDE를 결합하여 학습할 때 강한 스케일링과 약한 스케일링 성능은 어느 정도 달성할 수 있는가?
• RQ4 수천 개의 코어를 통해 대규모 병렬 처리를 수행하면서도 수치 정확도와 수렴성을 유지할 수 있는가?
• RQ5 물리학 기반 학습에서 기존의 데이터 병렬 또는 모델 병렬 DNN 학습과 비교해 이 방법의 성능은 어떠한가?

주요 결과

• 음향 파동 방정식에 대해 100개의 프로세서(3200개의 코어)를 사용할 때 20배 이상의 성능 향상을 기록하였으며, 약한 스케일링에서의 성능 저하가 최소한이었다.
• 음향 파동 방정식에 대해 1에서 100개의 프로세서로 확장할 때 런타임이 오직 2배로 증가하여 뛰어난 확장성을 보였다.
• 탄성 파동 방정식에 대해서는 더 높은 메모리 및 상태 변수 복잡도에도 불구하고 100개의 프로세서에서 20배 이상의 강한 스케일링 성능을 유지하였다.
• 계산 그래프 추상화 덕분에 통신과 계산이 명확히 분리되어 모듈성과 테스트 가능성 향상이 이루어졌다.
• 이 프레임워크는 DNN 파라미터뿐 아니라 MPI 통신 작업을 통한 기울기 역전파도 성공적으로 지원하여 종단 간 기울기 계산이 가능했다.
• 이 구현은 파동 전파에서의 매개변수 식별과 같은 대규모 PDE와 DNN을 포함하는 역문제에 대해 실용적인 가능성을 입증하였다.