[논문 리뷰] Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization

Paras Jain, Ajay N. Jain외 6명|arXiv (Cornell University)|2019. 10. 07.

Parallel Computing and Optimization Techniques인용 수 51

한 줄 요약

본 논문은 텐서 재재계산을 제한된 최적화 문제로 공식화하고, MILP 또는 근사적 LP 반올림을 통해 최적의 재재계산 스케줄을 계산하여 훈련 메모리를 최소한의 오버헤드로 감소시키고 더 큰 입력 크기를 가능하게 하는 Checkmate를 제시한다.

ABSTRACT

We formalize the problem of trading-off DNN training time and memory requirements as the tensor rematerialization optimization problem, a generalization of prior checkpointing strategies. We introduce Checkmate, a system that solves for optimal rematerialization schedules in reasonable times (under an hour) using off-the-shelf MILP solvers or near-optimal schedules with an approximation algorithm, then uses these schedules to accelerate millions of training iterations. Our method scales to complex, realistic architectures and is hardware-aware through the use of accelerator-specific, profile-based cost models. In addition to reducing training cost, Checkmate enables real-world networks to be trained with up to 5.1x larger input sizes. Checkmate is an open-source project, available at https://github.com/parasj/checkmate.

연구 동기 및 목표

• 임의의 DAG에 대한 텐서 재재계산을 메모리 제약 최적화 문제로 형식화한다.
• 메모리 예산 아래에서 학습 시간을 최소화하는 MILP 수식을 개발한다.
• 현실적인 네트워크로 확장하기 위한 LP 반올림을 통한 빠른 근사화를 제공한다.
• 타깃 가속기에 대해 프로파일링하여 Checkmate를 하드웨어 및 비용 의식적으로 만든다.
• 실질적인 DNN 학습 효율성을 향상시키는 오픈 소스 시스템을 제공한다.

제안 방법

• 학습 DAG를 메모리 비용이 있는 계산 그래프로 모델링하고 노드별 메모리 M_v와 비용 C_v를 정의한다.
• 실행을 T 단계로 전개하고 재계산을 나타내는 이진 변수 R_t,i와 체크포인트 유지 여부를 나타내는 S_t,i를 정의한다.
• 메모리 제약 MILP(및 선형화된 변형)을 구성하여 M_budget 아래에서 총 계산 비용을 최소화한다.
• 동적 메모리 사용을 포착하기 위해 메모리 계정 U_t,k와 해제 지시자 Free_t,i,k를 도입하고, 다항 제약을 불 대수식으로 선형화한다.
• 해결 공간을 줄이고 MILP 해법의 성능을 향상시키기 위해 프런티어-진행(frontier-advancing) 단계 분할을 사용한다.
• 정확한 MILP가 다루기 어려운 경우 근사적으로 근사적 해를 얻기 위한 두 단계 LP 기반 반올림 알고리즘을 제공한다.

실험 결과

연구 질문

• RQ1 임의의 신경망 그래프에 대해 메모리 제약 하에서 텐서 재재계산을 최적화 문제로 어떻게 형식화할 수 있는가?
• RQ2 최적 재재계산을 적용할 때 메모리 절감과 계산 오버헤드 사이의 트레이드오프는 무엇인가?
• RQ3 LP 기반 근사 방법이 실용적인 시간 한계 내에서 복잡한 아키텍처에 대해 거의 최적의 스케줄을 제공할 수 있는가?
• RQ4 Checkmate가 현실적인 네트워크와 하드웨어 프로파일에 얼마나 잘 확장되는가?
• RQ5 최적 재재계산이 기존 방법에 비해 더 큰 입력이나 배치 크기로 학습을 가능하게 하는가?

주요 결과

• 최적 재재계산은 메모리 사용을 줄이고 이미지 분류 및 분할 아키텍처 전반에서 modest한 오버헤드로 더 큰 입력 크기를 가능하게 한다.
• 이 방법은 임의의 DAG와 비균일한 레이어 비용을 지원하며 메모리 효율성 및 계산 시간 측면에서 기준선보다 우수하다.
• Checkmate는 동일한 가속기 메모리에서 최대 5.1x 더 큰 입력 크기를 가능하게 한다.
• 완전한 MILP는 표준 해법기로 중형 네트워크를 한 시간 이내에 해결할 수 있으며 프런티어-진행 파티션이 해결 속도를 크게 향상시킨다.
• 두 단계 반올림은 많은 벤치마크에서 다항 시간 내에 거의 최적의 정수 해를 제공한다.
• 이 시스템은 TensorFlow 2.0으로 구현되었으며 오픈 소스이다(GitHub: parasj/checkmate).