[논문 리뷰] Memory Optimization for Deep Networks

Aashaka Shah, Chao-Yuan Wu외 3명|arXiv (Cornell University)|2020. 10. 27.

Advanced Neural Network Applications참고 문헌 34인용 수 10

한 줄 요약

MONeT은 딥러닝 학습에서 메모리 사용량을 줄이기 위해 체크포인팅 스케줄과 메모리 효율적인 연산자 구현을 동시에 최적화하는 자동 프레임워크이다. 메모리와 계산 간의 트레이드오프를 정수계획법으로 모델링하여 파이토치 대비 3배 메모리 절감을 달성하면서도 계산 오버헤드는 9–16%에 그치며, 수동 최적화 및 자동화된 기준 대비 뛰어난 성능을 발휘한다.

ABSTRACT

Deep learning is slowly, but steadily, hitting a memory bottleneck. While the tensor computation in top-of-the-line GPUs increased by 32x over the last five years, the total available memory only grew by 2.5x. This prevents researchers from exploring larger architectures, as training large networks requires more memory for storing intermediate outputs. In this paper, we present MONeT, an automatic framework that minimizes both the memory footprint and computational overhead of deep networks. MONeT jointly optimizes the checkpointing schedule and the implementation of various operators. MONeT is able to outperform all prior hand-tuned operations as well as automated checkpointing. MONeT reduces the overall memory requirement by 3x for various PyTorch models, with a 9-16% overhead in computation. For the same computation cost, MONeT requires 1.2-1.8x less memory than current state-of-the-art automated checkpointing frameworks. Our code is available at https://github.com/utsaslab/MONeT.

연구 동기 및 목표

• 최근 5년간 GPU 메모리가 2.5배 증가한 데 비해 계산 능력은 32배 증가함에 따라 발생하는 딥러닝 학습의 증가하는 메모리 병목 현상을 해결하기 위해.
• 기존 방법들이 지역적 연산자 구현이나 전역적 체크포인팅 스케줄을 별도로 최적화하는 데 한계를 가진다는 점을 극복하기 위해.
• 고정된 메모리 예산 하에 지역적 및 전역적 메모리 절감 기법을 동시에 최적화하는 자동 프레임워크를 개발하기 위해.
• 정상 및 역전파 단계에서 피크 메모리 소비를 최소화하면서도 계산 효율성을 유지하기 위해.
• 메모리 사용량을 엄격히 제한하고 메모리와 계산 간 최적의 트레이드오프를 가능하게 하는 통합적이고 이론적으로 탄탄한 최적화 프레임워크를 제공하기 위해.

제안 방법

• 다양한 연산자 구현 방식과 체크포인팅 스케줄에 따라 정방향, 역방향 및 재계산 단계의 메모리 및 계산 비용 한계를 수식적으로 정의한다.
• 연산자 구현과 체크포인팅을 공동 최적화하는 문제를 선형화된 메모리 제약 조건을 가진 0-1 정수계획법으로 모델링한다.
• 표준 혼합정수선형계획법 솔버를 사용하여 주어진 예산 하에서 메모리 사용량을 최소화하는 최적의 구성 설정을 찾는다.
• 다양한 최적화 전략을 지원한다: 컬러레이션 알고리즘 선택, 인라인 활성화, 출력 활성화 재사용, 중간 활성화 재계산.
• 파이토치와 통합되어 최소한의 프레임워크 수정으로 종단 간 학습을 지원한다.
• 이론적 분석을 활용해 피크 메모리 소비에 대한 날카로운 상한을 유도하여 정밀한 최적화를 가능하게 한다.

실험 결과

연구 질문

• RQ1 지역적 연산자 구현과 전역적 체크포인팅 스케줄을 공동 최적화하면, 각각을 별도로 최적화하는 것보다 더 나은 메모리-계산 트레이드오프를 달성할 수 있는가?
• RQ2 역전파에서 다양한 메모리 절약 기법을 조합할 경우 피크 메모리 소비에 대한 이론적 상한은 무엇인가?
• RQ3 컬러레이션 알고리즘 선택과 체크포인팅의 공동 최적화가 전체 메모리 및 계산 효율성에 어떤 영향을 미치는가?
• RQ4 자동 선택된 메모리 효율적인 연산자 구현 방식이 수동으로 설계된 또는 히우리스틱 기반 접근 방식보다 얼마나 뛰어난가?
• RQ5 메모리 절감과 계산 오버헤드 측면에서 MONeT은 최신 자동 체크포인팅 프레임워크와 비교해 어떤가?

주요 결과

• MONeT은 ResNet, VGG, UNet, GoogleNet, MobileNet-V2 등 다양한 모델에서 파이토치 대비 3배 메모리 사용량을 절감한다.
• 유사한 계산 비용에서 MONeT은 현재 최신 자동 체크포인팅 프레임워크 대비 1.2–1.8배 더 적은 메모리가 필요하다.
• 표준 파이토치 학습 대비 계산 시간은 오직 9–16% 증가에 그친다.
• 컬러레이션 알고리즘과 체크포인팅의 공동 최적화는 뚜렷한 성능 향상을 이끌어내며, 컬러레이션 워크스페이스 메모리 특성이 속도와 메모리 트레이드오프에 크게 영향을 미친다.
• 제거 실험 결과, 컬러레이션, 출력 활성화, 중간 활성화 최적화 전략을 모두 조합할 경우 계산 오버헤드가 가장 낮아진다.
• ResNet-50에 대한 상세 사례 연구는 MONe트가 피크 메모리를 파이토치의 14.7 GB에서 8.0 GB로 줄이며 학습 안정성을 유지함을 확인한다.