[논문 리뷰] EcoRNN: Efficient Computing of LSTM RNN Training on GPUs

Bojian Zheng, Abhishek Tiwari외 2명|arXiv (Cornell University)|2018. 05. 22.

Advanced Neural Network Applications인용 수 5

한 줄 요약

EcoRNN는 소스 코드를 수정하지 않고도 LSTM RNN 학습 중에 특성 맵을 재계산하는 컴파일러 기반 최적화인 Echo를 제안한다. 이는 GPU 메모리 사용량을 줄이는 데 기여한다. 메모리 절감 효과와 재계산 오버헤드를 지능적으로 추정함으로써 Echo는 평균 1.89배에서 최대 3.13배의 메모리 절감을 달성하며, 이는 더 빠른 학습, 더 큰 배치 크기 또는 더 많은 레이어를 고정된 GPU 메모리 제약 조건 내에서 가능하게 한다.

ABSTRACT

The Long-Short-Term-Memory Recurrent Neural Networks (LSTM RNNs) are a popular class of machine learning models for analyzing sequential data. Their training on modern GPUs, however, is limited by the GPU memory capacity. Our profiling results of the LSTM RNN-based Neural Machine Translation (NMT) model reveal that feature maps of the attention and RNN layers form the memory bottleneck and runtime is unevenly distributed across different layers when training on GPUs. Based on these two observations, we propose to recompute the feature maps rather than stashing them persistently in the GPU memory. While the idea of feature map recomputation has been considered before, existing solutions fail to deliver satisfactory footprint reduction, as they do not address two key challenges. For each feature map recomputation to be effective and efficient, its effect on (1) the total memory footprint, and (2) the total execution time has to be carefully estimated. To this end, we propose *Echo*, a new compiler-based optimization scheme that addresses the first challenge with a practical mechanism that estimates the memory benefits of recomputation over the entire computation graph, and the second challenge by non-conservatively estimating the recomputation overhead leveraging layer specifics. *Echo* reduces the GPU memory footprint automatically and transparently without any changes required to the training source code, and is effective for models beyond LSTM RNNs. We evaluate *Echo* on numerous state-of-the-art machine learning workloads on real systems with modern GPUs and observe footprint reduction ratios of 1.89X on average and 3.13X maximum. Such reduction can be converted into faster training with a larger batch size, savings in GPU energy consumption (e.g., training with one GPU as fast as with four), and/or an increase in the maximum number of layers under the same GPU memory budget.

연구 동기 및 목표

• 신경 기계 번역과 같은 모델에서 특히 두드러지는 LSTM RNN 학습의 GPU 메모리 병목 현상을 해결하기 위해.
• 모델 소스 코드를 수정하지 않고도 메모리 프로파일을 줄이기 위해.
• 고정된 GPU 메모리 예산 내에서 더 큰 배치 크기, 더 많은 레이어, 또는 에너지 소비 감소를 위한 효율적인 학습을 가능하게 하기 위해.
• 이전의 재계산 기법들이 메모리 절감 효과와 런타임 오버헤드 사이의 균형을 이루지 못한 한계를 극복하기 위해.

제안 방법

• Echo는 계산 그래프 전반에 걸쳐 특성 맵 재계산을 자동으로 식별하고 적용하는 컴파일러 기반 기법을 사용한다.
• 실용적이고 비보수적인 메커니즘을 통해 계산 그래프 전체에서 재계산의 메모리 이점을 추정한다.
• 레이어별 특성 정보를 활용하여 재계산 오버헤드를 비보수적으로 모델링함으로써 불필요한 재계산를 방지한다.
• 이 방법은 투명하고 자동화되어 있어 모델의 소스 코드 수정이 필요하지 않다.
• Echo는 LSTM RNN을 넘어서 다른 딥러닝 워크로드로도 일반화 가능하도록 설계되어 있다.

실험 결과

연구 질문

• RQ1 어떻게 특성 맵 재계산을 효과적으로 적용하여 LSTM RNN 학습의 GPU 메모리 프로파일을 줄일 수 있는가?
• RQ2 RNN 모델의 다양한 레이어에서 메모리 절감 효과와 재계산 오버헤드 사이의 상호 관계는 어떠한가?
• RQ3 컴파일러 기반 시스템이 소스 코드 변경 없이 자동으로 메모리 사용량을 최적화할 수 있는가?
• RQ4 기존의 재계산 기법과 비교해 볼 때, 제안된 방법은 메모리 절감 효과와 성능 측면에서 어떤가?

주요 결과

• Echo는 여러 최신 기계 학습 워크로드에서 평균 1.89배의 메모리 프로파일 감소를 달성한다.
• 관찰된 최대 메모리 절감 비율은 3.13배로, 고정된 GPU 메모리 제약 조건 내에서 학습 능력을 크게 확장한다.
• 최적화 덕분에 단일 GPU로 학습할 경우, 4개의 GPU를 사용할 때와 동일한 처리량을 달성할 수 있다.
• 이 방법은 더 적은 GPU 수나 더 큰 배치 크기를 사용함으로써 에너지 소비를 줄이는 데 기여한다.
• 이 방법은 LSTM RNN 뿐만 아니라 RNN을 초월한 다른 딥러닝 모델에도 효과적이다.