[논문 리뷰] Using Machine Learning at Scale in HPC Simulations with SmartSim: An Application to Ocean Climate Modeling

Sam Partee, Matthew J. Ellis외 5명|arXiv (Cornell University)|2021. 04. 13.

Oceanographic and Atmospheric Processes참고 문헌 31인용 수 13

한 줄 요약

이 논문은 고성능 계산(HPC) 시뮬레이션 내에서 확장 가능하고 온라인 기계학습 추론을 가능하게 하는 소프트웨어 프레임워크인 SmartSim을 제안한다. PyTorch, TensorFlow 및 ONNX 모델을 Fortran 기반의 해양 기후 모델(MOM6)에 통합하여 12년간 9700억 번의 추론을 수행한 12개의 글로벌 해양 시뮬레이션을 구현하였으며, 런타임 오버헤드를 최소화하면서도 안정적인 성능을 달성하였다.

ABSTRACT

We demonstrate the first climate-scale, numerical ocean simulations improved through distributed, online inference of Deep Neural Networks (DNN) using SmartSim. SmartSim is a library dedicated to enabling online analysis and Machine Learning (ML) for traditional HPC simulations. In this paper, we detail the SmartSim architecture and provide benchmarks including online inference with a shared ML model on heterogeneous HPC systems. We demonstrate the capability of SmartSim by using it to run a 12-member ensemble of global-scale, high-resolution ocean simulations, each spanning 19 compute nodes, all communicating with the same ML architecture at each simulation timestep. In total, 970 billion inferences are collectively served by running the ensemble for a total of 120 simulated years. Finally, we show our solution is stable over the full duration of the model integrations, and that the inclusion of machine learning has minimal impact on the simulation runtimes.

연구 동기 및 목표

• 파일 I/O 병목 현상 없이 Fortran/C/C++로 작성된 전통적 HPC 시뮬레이션과 기계학습을 통합하는 도전 과제 해결.
• 장기 실행 기간 동안 실시간으로 분산된 추론을 수행할 수 있도록 딥 뉴럴 네트워크의 실시간 분산 추론 지원.
• HPC 워크로드와 현대 기계학습 라이브러리 간의 언어 및 데이터 교환 장벽 극복.
• 기후 규모 해상도에서 기계학습이 보강된 HPC 시뮬레이션의 확장성과 안정성 입증.
• 다양한 HPC 시스템과 여러 기계학습 프레임워크를 지원하는 유연하고 라이브러리에 종속되지 않는 인터페이스 제공.

제안 방법

• SmartSim의 인fra구조 및 클라이언트 라이브러리를 활용하여 기계학습 추론을 시뮬레이션 로직에서 분리.
• SmartRedis API를 사용하여 Fortran 시뮬레이션에서 기계학습 모델을 저지연, 원격으로 실행.
• JIT 트레이싱된 PyTorch 모델을 MOM6에 통합하여 기존의 난류 운동에너지 파arameterization을 대체.
• 12개의 앙상블 멤버 각각이 19개 노드에서 실행되도록 공유된 기계학습 모델을 배포하고, 각 타임스텝에서 동기화된 추론을 수행.
• Redis 기반의 데이터 교환을 활용하여 파일 I/O 오버헤드를 제거하고 시뮬레이션 및 기계학습 구성 요소 간 실시간 데이터 흐름을 가능하게 함.
• CPU 및 GPU 워크로드를 모두 지원하여 다양한 하드웨어와 소프트웨어 스택 간의 상호운용성을 확보하는 이질적 HPC 시스템 지원.

실험 결과

연구 질문

• RQ1 파일 I/O 병목 현상 없이 대규모, 생산용 HPC 기후 시뮬레이션에 기계학습을 통합할 수 있는가? 성능이나 시뮬레이션 안정성에 영향을 주지 않는가?
• RQ2 대규모로 분산된 HPC 시뮬레이션 인스턴스 간에 단일 기계학습 모델이 얼마나 효율적으로 실시간 추론을 제공할 수 있는가?
• RQ3 공유된 모델 아키텍처를 사용하여 글로벌 해양 모델에 온라인 기계학습 추론을 통합할 경우 성능 오버헤드는 얼마나 되는가?
• RQ4 SmartSim 프레임워크는 Fortran 기반 시뮬레이션과 PyTorch 및 TensorFlow와 같은 현대 기계학습 라이브러리 간에 원활하고 언어에 종속되지 않는 통신을 가능하게 하는가?
• RQ5 SmartSim을 통한 기계학습 추론 통합이 장기간 기후 시뮬레이션 동안 수치적 안정성을 유지하는가?

주요 결과

• SmartSim-EKE 앙상블은 10년씩 12개의 글로벌 고해상도 해양 시뮬레이션을 성공적으로 완료하였으며, 수치적 불안정성이 관측되지 않았다.
• 120년의 시뮬레이션 기간 동안 총 9700억 번의 추론 연산이 수행되어 엑사스케일에 준하는 규모에서의 확장성을 입증하였다.
• SmartSim을 통한 기계학습 추론 통합으로 인한 성능 오버헤드는 극히 미미하였으며, 시뮬레이션 런타임은 대부분 영향을 받지 않았다.
• 다양한 하드웨어와 소프트웨어 스택을 지원하는 이질적 HPC 시스템에서 수천 개의 프로세서에서 선형 확장성을 달성하였다.
• 모든 앙상블 멤버에 공통된 기계학습 모델을 사용함으로써 시뮬레이션 별 모델 중복 없이 일관되고 동기화된 추론을 가능하게 하였다.
• 해당 솔루션은 생산용으로 사용 가능하며 완전히 재현 가능하며, 모든 소스 코드, 데이터셋 및 훈련된 모델이 GitHub에 공개되어 있다.