[논문 리뷰] The Case for Strong Scaling in Deep Learning: Training Large 3D CNNs with Hybrid Parallelism

Yosuke Oyama, Naoya Maruyama외 7명|arXiv (Cornell University)|2020. 07. 25.

Advanced Neural Network Applications참고 문헌 54인용 수 9

한 줄 요약

이 논문은 데이터 병렬성과 공간 병렬성을 조합하여 단일 샘플을 여러 GPU에 분산하는 방식으로, 대규모 3D 합성곱 신경망(3D CNNs)에 대한 강력한 스케일링을 가능하게 하는 엔드 투 엔드 하이브리드-병렬 학습 프레임워크를 제안한다. 미니배치 크기 이외의 병렬성 확장을 통해 I/O를 최적화하고 공간 분할을 적용함으로써 높은 성능을 달성하며, 이는 전체 해상도의 512³ 천체물리 시뮬레이션 데이터에 대한 학습을 가능하게 하여, CosmoFlow의 예측 정확도를 이전의 낮은 해상도 학습 대비 10배 향상시킨다.

ABSTRACT

We present scalable hybrid-parallel algorithms for training large-scale 3D convolutional neural networks. Deep learning-based emerging scientific workflows often require model training with large, high-dimensional samples, which can make training much more costly and even infeasible due to excessive memory usage. We solve these challenges by extensively applying hybrid parallelism throughout the end-to-end training pipeline, including both computations and I/O. Our hybrid-parallel algorithm extends the standard data parallelism with spatial parallelism, which partitions a single sample in the spatial domain, realizing strong scaling beyond the mini-batch dimension with a larger aggregated memory capacity. We evaluate our proposed training algorithms with two challenging 3D CNNs, CosmoFlow and 3D U-Net. Our comprehensive performance studies show that good weak and strong scaling can be achieved for both networks using up 2K GPUs. More importantly, we enable training of CosmoFlow with much larger samples than previously possible, realizing an order-of-magnitude improvement in prediction accuracy.

연구 동기 및 목표

• 우주론적 시뮬레이션과 3D 의료 영상과 같은 고차원 과학적 데이터에서 대규모 3D CNN을 학습할 때 발생하는 메모리 및 계산 병목 현상을 해결하기 위해.
• 큰 샘플로 인해 GPU당 메모리 압박이 심해지는 바탕이 되는 전통적 데이터 병렬성의 한계를 극복하기 위해.
• 미니배치 크기를 늘리지 않고도 학습 시간을 단축할 수 있는 강력한 스케일링을 달성하기 위해, 개별 샘플을 GPU 간에 공간적으로 분할하는 공간 병렬성을 도입하기 위해.
• 데이터 로딩에 동일한 하이브리드 병렬 기법을 적용하여 I/O 파이프라인을 종단 간 스케일링함으로써, 큰 샘플에 대한 I/O 병목 현상을 최소화하기 위해.
• 전체 해상도 학습이 장거리 특징 학습이 필요한 과학적 과제에서 모델 정확도를 크게 향상시킬 수 있음을 입증하기 위해.

제안 방법

• 단일 3D 샘플을 여러 GPU에 분산시켜 GPU당 메모리 사용량을 줄이는 방식으로, 데이터 병렬성과 공간 병렬성을 융합한 하이브리드 병렬 알고리즘을 도입한다.
• 특징 맵의 공간 차원(예: 높이, 너비, 깊이)을 처리 요소 간에 나누어 3D 합성곱에 대한 공간 분할을 적용함으로써, 미니배치 한계를 초월한 강력한 스케일링을 가능하게 한다.
• 큰 데이터 샘플에 대한 I/O 파이프라인을 공간 분할을 통해 최적화하여, GPU 수에 비례하여 확장 가능한 집합적이고 비동기적 I/O를 실현한다.
• 하이브리드 병렬 학습과 I/O를 지원하기 위해 커스텀 커널과 통신 패턴을 추가한 LBANN 딥러닝 프레임워크를 확장하여, 낮은 통신 오버헤드를 확보한다.
• 채널/특징 병렬성에서 사용하는 전체 텐서 통신 대비 데이터 이동을 최소화하기 위해, 공간 분할 간의 할로( halo) 교환을 활용하여 기울기 통신을 수행한다.
• 모델 기반 성능 분 析을 통해 3D CNN의 계산 및 I/O 병목 현상을 식별하고 완화한다.

실험 결과

연구 질문

• RQ1 개별 샘플의 공간 분할을 도입함으로써 데이터 병렬성 외의 병렬성을 확장함으로써 3D CNN 학습에서 강력한 스케일링을 달성할 수 있는가?
• RQ2 공간 병렬성과 데이터 병렬성을 융합한 하이브리드 병렬 학습이 최대 2,048개 GPU에서 고성능 및 낮은 통신 오버헤드를 유지하면서도 효율적인 스케일링을 가능하게 하는가?
• RQ3 계산에 사용된 하이브리드 병렬 기법을 동일하게 적용함으로써, 큰 3D 데이터 샘플에 대한 I/O 파이프라인이 효과적으로 스케일링될 수 있는가?
• RQ4 메모리 및 계산 요구량이 더 높음에도 불구하고, 전체 해상도의 3D 데이터(예: 512³)로 학습할 경우 낮은 해상도 데이터로 학습하는 것보다 모델 정확도가 크게 향상되는가?
• RQ5 CosmoFlow 및 3D U-Net과 같은 3D CNN의 성능 특성과 스케일링 효율성은 하이브리드 병렬 학습 및 I/O를 적용했을 때 어떻게 되는가?

주요 결과

• 512³ 전체 해상도 샘플을 사용해 CosmoFlow를 학습한 결과, 평균 제곱오차(MSE)가 0.0018로, 128³ 샘플로 학습한 경우(MSE = 0.0289)에 비해 예측 정확도가 10배 향상되었다.
• 64의 미니배치 크기로 학습한 CosmoFlow 모델은 512에서 2048개 GPU로 확장할 때 1.77배의 성능 향상을 기록하여, 미니배치 크기를 늘리지 않고도 강력한 스케일링을 입증했다.
• 3D U-Net 모델은 256에서 512개 GPU로 확장할 때 1.42배의 성능 향상을 기록하여, 다양한 아키텍처와 워크로드에서 강력한 스케일링 행동을 보였다.
• CosmoFlow와 3D U-Net 모두 2,048개 GPU까지 양호한 약한 및 강한 스케일링을 보였으며, 성능 분석을 통해 계산 및 I/O 구성 요소 전반에서 높은 스케일링 효율성이 확인되었다.
• 하이브리드 병렬 I/O 파이프라인은 GPU 수에 비례하여 성공적으로 스케일링되었으며, 과학적 3D 데이터에서 흔한 기가바이트 크기의 샘플이라도 I/O 병목 현상을 최소화하여 효과적으로 처리할 수 있었다.