[논문 리뷰] Echo: Compiler-based GPU Memory Footprint Reduction for LSTM RNN Training
Echo는 GPU 메모리 용량을 줄이기 위해 컴파일러 기반 최적화를 사용하는 기술로, 특징 맵을 저장하는 대신 선택적으로 재계산함으로써 LSTM RNN 학습 중 메모리 사용량을 감소시킨다. 레이어별 런타임 추정과 글로벌 메모리 이점 분석을 활용하며, 최소한의 런타임 오버헤드로 평균 1.89배, 최대 3.13배의 메모리 용량 감소를 달성하여 동일 하드웨어에서 더 빠른 학습, 더 큰 배치 크기, 더 깊은 모델을 가능하게 한다.
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 메모리 소비 문제를 해결하여 배치 크기와 컴퓨팅 활용도를 제한하는 것을 목표로 한다.
- 이전의 재계산 기법이 LSTM 환경에서 메모리 이점과 런타임 오버헤드를 정확히 추정하지 못하는 한계를 극복하는 것.
- 학습 소스 코드를 수정하지 않고도 자동으로 투명하게 메모리 용량을 줄일 수 있도록 하는 것.
- LSTM RNN을 초과해 Transformer와 ResNet과 같은 다양한 모델을 동일한 GPU 메모리 예산 내에서 지원하는 것.
- 용량 감소를 통해 더 빠른 학습, 더 큰 배치 크기 또는 증가된 모델 깊이를 가능하게 하는 것.
제안 방법
- Echo는 전체 계산 그래프를 분석하고, 특징 맵을 저장하는 대신 재계산할 경우의 순수 메모리 이점을 추정하는 컴파일러 기반 접근법을 사용한다.
- 재계산로 인해 해제되는 특징 맵과 재계산을 위해 새로 할당되는 데이터를 모두 고려하여 글로벌 메모리 절감 효과를 평가하는 실용적인 메커니즘을 도입한다.
- 비보수적인 방식으로 재계산 오버헤드를 추정하며, 비완전히 연결된 레이어에서 낮은 계산 비용을 가지는 레이어별 특성들을 활용한다.
- 메모리 절감과 런타임 비용 간의 트레이드오프를 기반으로 최적의 재계산 대상 레이어를 식별하며, 계산 비용이 높은 레이어를 단순히 배제하지 않는다.
- 모델 소스 코드를 수정하거나 수동 커널 튜닝이 필요 없이 기존 딥러닝 프레임워크에 원활하게 통합된다.
- Echo는 실제 GPU에서 NMT, DeepSpeech2, Transformer, ResNet 등 여러 최신 모델을 대상으로 평가되었다.
실험 결과
연구 질문
- RQ1선택적 재계산이 LSTM RNN 학습에서 금지할 수 없는 런타임 오버헤드 없이 GPU 메모리 용량을 효과적으로 줄일 수 있는가?
- RQ2저장된 데이터 해제와 새로운 할당을 모두 고려할 때 전체 계산 그래프에서 재계산의 순수 메모리 이점을 정확히 추정할 수 있는가?
- RQ3컴파일러 기반 시스템이 복잡한 모델인 Transformer와 ResNet에서 최적의 재계산 대상을 자동으로 식별할 수 있는가?
- RQ4용량 감소를 얼마나 실질적인 이점으로 전환할 수 있는가? 예를 들어 더 큰 배치 크기나 더 깊은 모델을 가능하게 하는가?
- RQ5다양한 DNN 워크로드에서 이전 기법과 비교해 Echo의 메모리 감소 및 런타임 효율성은 어떠한가?
주요 결과
- Echo는 NMT, DeepSpeech2, Transformer, ResNet 등 다양한 모델에서 평균 1.89배, 최대 3.13배의 GPU 메모리 용량 감소를 달성한다.
- 이 방법을 통해 단일 GPU로 학습할 경우 기준 설정에서 네 개의 GPU를 사용한 것과 동일한 속도를 달성할 수 있으며, 에너지 소비를 크게 줄일 수 있다.
- 배치 크기를 동일하게 유지할 경우, RTX 2080 Ti(11GB 메모리)에서 Transformer의 최대 레이어 수를 1.83배, ResNet의 경우 4.0배로 증가시킨다.
- 이러한 감소는 최소한의 런타임 오버헤드로 달성되어 실제 학습 워크로드에 실용적임을 입증한다.
- Echo는 LSTM RNN을 초월해 Transformer와 ResNet과 같은 CNN 모델에도 효과적이며, 강력한 적용 가능성을 보인다.
- 이 기법은 동일한 메모리 예산 내에서 더 빠른 학습 수렴과 더 깊은 모델을 지원하여 딥러닝 아키텍처 설계 추세와 부합한다.
더 나은 연구,지금 바로 시작하세요
논문 읽기부터 검토까지, 연구 시간을 획기적으로 줄여보세요.
카드 등록 없음 · 무료 플랜 제공
이 리뷰는 AI가 만들고, 인간 에디터가 검토했습니다.