[논문 리뷰] Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory

Julien Herrmann, Olivier Beaumont외 4명|arXiv (Cornell University)|2019. 11. 27.

Stochastic Gradient Optimization Techniques참고 문헌 29인용 수 13

한 줄 요약

이 논문은 제한된 메모리에서 딥 네ural 네트워크를 훈련시키기 위해 동적으로 저장할 포워드 활성화와 재계산할 활성화를 선택하는 최적의 체크포인팅 전략을 제안한다. 입력 전용 및 전체 연산력 기록 체크포인팅을 모델링하는 새로운 동적 프로그래밍 알고리즘을 사용하여, 기존 방법보다 평균 17.2% 높은 처리량을 달성하면서도, 임의의 메모리 제한과 복잡한 순차적 네트워크를 지원한다. 이는 파이토치에서 작동한다.

ABSTRACT

This paper introduces a new activation checkpointing method which allows to significantly decrease memory usage when training Deep Neural Networks with the back-propagation algorithm. Similarly to checkpoint-ing techniques coming from the literature on Automatic Differentiation, it consists in dynamically selecting the forward activations that are saved during the training phase, and then automatically recomputing missing activations from those previously recorded. We propose an original computation model that combines two types of activation savings: either only storing the layer inputs, or recording the complete history of operations that produced the outputs (this uses more memory, but requires fewer recomputations in the backward phase), and we provide an algorithm to compute the optimal computation sequence for this model. This paper also describes a PyTorch implementation that processes the entire chain, dealing with any sequential DNN whose internal layers may be arbitrarily complex and automatically executing it according to the optimal checkpointing strategy computed given a memory limit. Through extensive experiments, we show that our implementation consistently outperforms existing checkpoint-ing approaches for a large class of networks, image sizes and batch sizes.

연구 동기 및 목표

• 딥 네럴 네트워크 훈련에서 메모리 병목 현상을 해결하기 위해, 특히 IoT 시스템과 같은 저메모리 장치에서의 적용을 고려한다.
• 기존의 자동 미분 체크포인팅을 내부 레이어가 매우 복잡한 이종 DNN로 확장한다.
• 임의의 메모리 제한을 지원하고 훈련 처리량을 최적화하는 완전 자동화된 파이토치 호환 구현을 개발한다.
• 다양한 네트워크 아키텍처, 이미지 크기, 배치 크기에서 기존의 주기적 및 최적 체크포인팅 전략을 초월한다.

제안 방법

• DNN의 포워드 및 백워드 연산을 모듈의 순서로 모델링하여, 어떤 활성화를 체크포인트로 저장할지를 동적으로 선택할 수 있도록 한다.
• 두 가지 체크포인팅 유형을 도입한다: 메모리 효율적인 레이어 입력만 저장하거나, 전체 연산력 기록을 기록한다.
• 메모리 제약 조건 하에서 최적의 영속성 스케줄을 계산하기 위해 동적 프로그래밍 알고리즘을 사용하며, 최적 해의 메모리 영속성 성질을 활용한다.
• 복잡한 내부 계산을 포함한 임의의 순차적 파이토치 모듈을 지원한다.
• autograd와 원활하게 통합되어 사용자가 정의한 메모리 제한에 따라 최적 전략을 자동으로 적용한다.
• 이중 세그먼트 기반 전략과 달리 연속적인 메모리 적응을 가능하게 하며, 최소한의 오버헤드로 포워드 및 백워드 단계를 모두 지원한다.

실험 결과

연구 질문

• RQ1 내부 레이어 복잡도가 임의인 이종 딥 네럴 네트워크 체인에 대해 체크포인팅 전략을 어떻게 최적화할 수 있는가?
• RQ2 입력만 저장하는 것과 전체 연산력 기록을 저장하는 것 사이의 메모리 사용과 재계산 비용 간 최적의 트레이드오프는 무엇인가?
• RQ3 메모리 제약 조건 하에서 동적 프로그래밍 접근 방식이 기존의 주기적 및 최적 체크포인팅 전략보다 훈련 처리량 측면에서 뛰어나게 성능을 낼 수 있는가?
• RQ4 제안된 방법은 파이토치의 기본 전략과 비교해 네트워크 깊이, 이미지 크기, 배치 크기 등 다양한 조건에서 어떻게 확장되는가?
• RQ5 이 방법을 통해 저메모리 장치에서 더 큰 모델, 더 큰 이미지, 또는 더 큰 배치 크기를 훈련시킬 수 있는 정도는 어느 정도인가?

주요 결과

• 제안된 최적 전략은 모든 테스트 구성에서 가장 성능이 뛰어난 경쟁자보다 평균 17.2% 높은 훈련 처리량을 달성한다.
• 224×224 이미지와 배치 크기 8로 ResNet-1001을 훈련할 경우, 최적 전략은 9.18장/초의 처리 속도를 기록하지만, 파이토치의 기본 전략은 메모리 제한으로 인해 실패한다.
• 배치 크기 8일 경우, 파이토치의 전략은 225 GiB의 메모리가 필요하며, 이는 최적 전략의 최대 사용량보다 16배 이상 높다.
• 최적 전략은 파이토치의 기본 전략이 실패하는 경우에도 큰 모델 훈련을 가능하게 하며, ResNet-1001의 배치 크기 8에서도 적용 가능하다.
• 특히 더 높은 메모리 제한에서 기존의 revolve 알고리즘은 가용 메모리를 효율적으로 활용하지 못하는 반면, 최적 전략은 이에 비해 뛰어난 성능을 보인다.
• 다양한 네트워크 크기와 이미지 해상도에서 안정적인 성능을 유지하여 강건성과 유연성을 입증한다.