[논문 리뷰] Addressing the Memory Bottleneck in AI Model Training

David Ojika, Bhavesh Patel외 4명|arXiv (Cornell University)|2020. 03. 11.

Advanced Neural Network Applications참고 문헌 7인용 수 6

한 줄 요약

이 논문은 2세대 인텔 Xeon 스케일러블 프로세서와 대용량 시스템 메모리가 탑재된 단일 노드 서버에서 인텔 최적화된 텐서플로우를 사용하여 딥 뉴럴 네트워크를 1TB의 메모리 프로파일로 훈련시킨 최초의 사례를 보여준다. 이 방법은 분산 시스템 없이도 메모리 집약적인 AI 모델, 특히 의료 영상 분야에서의 스케일업 훈련을 가능하게 하여 기존의 메모리 병목 현상을 극복한다.

ABSTRACT

Using medical imaging as case-study, we demonstrate how Intel-optimized TensorFlow on an x86-based server equipped with 2nd Generation Intel Xeon Scalable Processors with large system memory allows for the training of memory-intensive AI/deep-learning models in a scale-up server configuration. We believe our work represents the first training of a deep neural network having large memory footprint (~ 1 TB) on a single-node server. We recommend this configuration to scientists and researchers who wish to develop large, state-of-the-art AI models but are currently limited by memory.

연구 동기 및 목표

• 대규모 AI 모델 훈련에서 발생하는 메모리 병목 현상을 해결하기 위해, 특히 의료 영상 분야에서의 응용을 목적으로 한다.
• 기존에 분산 시스템이 필요로 했던 모델을 대용량 시스템 메모리가 탑재된 단일 노드 서버에서 훈련시킬 수 있음을 입증한다.
• 제한된 인프라를 가진 연구자들에게 분산 훈련의 실용적이고 확장 가능한 대안을 제공한다.
• 충분한 메모리가 확보된 경우 최신 기술 수준의 메모리 요구가 매우 높은 모델을 단일 서버에서 훈련시킬 수 있는 가능성을 검증한다.
• 대규모 AI 모델을 개발하는 과학자들이 사용할 수 있는 하드웨어-소프트웨어 스택 권장 사항을 제공한다.

제안 방법

• 2세대 인텔 Xeon 스케일러블 프로세서를 탑재한 x86 기반 서버에서 인텔 최적화된 텐서플로우를 활용했다.
• 최대 1TB까지의 대용량 시스템 메모리를 활용하여 전체 모델과 활성화 상태를 메모리에 유지했다.
• 통신 오버헤드를 피하기 위해 분산 훈련 대신 스케일업 아키텍처를 사용했다.
• 가용 메모리의 최적 활용을 위해 텐서플로우 내부의 메모리 관리 및 데이터 플로우를 최적화했다.
• 3D 컨볼루션 네트워크의 높은 메모리 요구량으로 인해 의료 영상 분야를 사례 연구로 선정했다.
• 높은 처리량을 유지하기 위해 충분한 메모리 대역폭과 NUMA 인식 메모리 할당을 구성했다.

실험 결과

연구 질문

• RQ1 1TB의 시스템 메모리가 탑재된 단일 노드 서버에서 메모리 프로파일이 동일하게 1TB에 이르는 딥 뉴럴 네트워크를 성공적으로 훈련시킬 수 있는가?
• RQ2 분산 훈련과 비교했을 때 단일 서버에서의 스케일업 훈련은 메모리 효율성과 훈련 처리량 측면에서 어떻게 다른가?
• RQ3 모델 병렬화 없이도 안정적인 메모리 집약적 모델 훈련을 가능하게 하는 하드웨어 및 소프트웨어 스택은 무엇인가?
• RQ4 인텔 최적화된 텐서플로우는 대규모 AI 워크로드에서 메모리 활용도를 얼마나 극대화할 수 있는가?
• RQ5 클러스터 인프라 없이도 최신 기술 수준의 의료 영상 모델을 단일 서버에서 훈련시킬 수 있는가?

주요 결과

• 저자들은 1TB의 메모리 프로파일을 가진 딥 뉴럴 네트워크를 단일 노드 서버에서 성공적으로 훈련시켰으며, 이는 최초의 사례이다.
• 이 훈련은 2세대 인텔 Xeon 스케일러블 프로세서와 대용량 시스템 메모리가 탑재된 환경에서 인텔 최적화된 텐서플로우를 사용하여 달성되었다.
• 스케일업 구성으로 인해 모델 병렬화나 데이터 병렬화가 필요 없었으며, 훈련 설정이 단순화되었다.
• 분산 시스템 간의 조율 없이도 복잡한 의료 영상 모델의 안정적인 훈련이 가능했다.
• 결과적으로 충분한 메모리가 확보된 경우 대규모 최신 기술 수준의 AI 모델을 단일 서버에서 훈련시킬 수 있음을 입증했다.
• 이 연구는 인프라 제약이 있는 연구자들에게 보다 단순하고 효율적인 훈련 경로를 제공하는 실용적인 대안이 된다.