[ASPLOS 2020] SwapAdvisor: Push Deep Learning Beyond the GPU Memory Limit via Smart Swapping

2020. 3. 1. 20:11Research

뉴럴넷의 깊이(layer의 수)와 넓이(레이어 당 파라미터 수)가 커질 수록 더 좋은 accuracy를 얻을 수 있다는 건 여러 연구를 통해서 밝혀진 사실이다. 하지만 GPU 메모리의 한계로 인해서 마냥 큰 모델을 만들 수만은 없다. 이 논문에서는 GPU 메모리 한계를 극복하기 위한 한 가지 방법으로 dataflow graph가 주어졌을 때 operator scheduling, memory allocation, swap decision을 자동으로 생성한다.

 

1. Introduction

이전의 연구들은 GPU 메모리 사용량을 줄이기 위해서 더 적은 floating point precision을 사용하거나, quantization을 통해서 더 sparse한 모델 파라미터를 생성하는 compression 기법을 제시하였다. 하지만 이런 방법은 model accuracy가 낮아지고, hyperparameter 튜닝에 시간을 써야한다는 점에서 한계가 있다.

 

이 논문에서는 GPU와 CPU 메모리 간에 tensor data를 swapping하며 사용하는 tensor swapping 기법으로 accuracy drop 없이 효율적인 메모리 사용을 가능케 하는 방법을 제시한다. Tensor swapping이 가능한 데에는 다음과 같은 이유가 있다.

 

1. CPU 메모리는 GPU 메모리 보다 훨씬 싸고 더 큰 용량을 가지고 있다.

2. GPU는 computation(텐서 계산)과  communication을 동시에 효율적으로 처리할 수 있다.

3. CPU-GPU 간의 communication bandwidth는 상당히 괜찮은 편이다.

 

Tensorflow와 같은 여러 DL framework에서는 symbolic graph execution을 통해 모델을 표현하기 때문에 실제 프로그램 실행 이전에 dataflow graph를 통해서 어떠한 연산이 이루어질지 미리 파악할 수 있다. 따라서 이를 바탕으로하여 언제, 무엇을 CPU-GPU 메모리 간에 swap 해야할지 결정할 수도 있을 것이다.

 

물론 dataflow graph만으로 완벽한 swapping scheduling이 가능한 것은 아니고, 어떻게 operator들의 실행이 scheduling 되었는지, 어떻게 메모리 할당이 이루어졌는지에 대한 정보 역시 필요하다.

 

논문에 따르면 현재 방식에서는 control flow가 없는 static dataflow graph가 주어졌을 때 single GPU에 한하여만 지원이 된다고 한다.

 

 

2. Background

DNN에서 메모리가 주로 사용되는 곳은 다음 세 가지로 나눌 수 있다.

 

1. Model parameters: 매 iteration의 끝에 업데이트 됨. depth(# of layers) x width(size of a layer)에 따라 크기가 결정 됨.

2. Intermediate results: Activation / Loss / Gradient (Loss와 Gradient는 Training 시에만 필요)

3. Scratch space: temp 값들을 저장하는데 필요한 일시적인 저장 공간

 

이전 연구들에서는 일종의 휴리스틱을 활용해서 activation만 swap하고 parameter는 swap하지 않았다. 이 연구에서는 매뉴얼한 휴리스틱을 사용하지 않고 dataflow graph를 통해서 자동으로 swap plan을 생성한다.

 

 

3.  Challenges and Approach

좋은 swapping plan이란 GPU에서의 communication과 computation을 최대한 중첩(overlap) 시켜서 그 오버헤드를 최소화하는 plan일 것이다. 따라서 dataflow 그래프를 통해서 overlapping을 최대화 할 수 있는 swapping plan을 만드는 것이 이 연구의 지향점이다.

 

swapping plan에 영향을 주는 두 가지 중요한 요소들이 있다.

 

  1. Memory Allocation: DNN에서는 다양한 크기의 tensor들이 사용되는데 internal fragmentation을 줄이고 속도를 향상 시키기 위해서는 memory pool을 어떻게 디자인하느냐가 매우 중요하다.
  2. Operator Scheduling: 최근 DNN 모델들은 복잡한 dataflow를 갖는 형태로 발전되고 있기 때문에 execution order를 어떻게 설정 하느냐에 따라 성능에 많은 차이가 날 수 있다.

본 논문에서 제시하는 swapping plan에서는 향후에 가장 오랫동안 사용되지 않을 tensor를 CPU로 swap-out하고, 이전에 swap-out 되었던 tensor를 최대한 빨리 prefetching(swap-in)하는 방식으로 overlapping을 최대화 한다. 또한 Genetic Algorithm(유전 알고리즘)을 통해서 최적의 memory allocation과 operator scheduling의 조합을 찾아낸다. (유전 알고리즘을 사용하는 이유는 다른 NP-hard serach 휴리스틱 중에 가장 빠른 편이고 병렬화에 용이하기 때문이다.) 유전 알고리즘을 통해서 search space를 빠르게 탐색(exploration) 하기 위해서는 수행시간 등을 빠르게 평가할 수 있어야하는데 기존 프레임워크 위에서 실험을 했을 때는 그 속도가 너무 느려서 별도의 dataflow engine simulator를 사용해서 swap plan의 performance를 평가했다고 한다. (이런 한계점이 있으면 실제 효용성이 상당히 의심이 간다...)

 

유전 알고리즘: 유전 알고리즘은 어떤 미지의 함수 Y = f(x)를 최적화하는 해 x를 찾기 위해 사용되는, 진화를 모방한(Simulated evolution) 탐색 알고리즘이다.

{1, 5, 6, 8, 3, 7, 3, 5, 9, 0} 중 3개를 골라서 20으로 만드는 문제가 있다고 하자. 여기서 유전체는 각 숫자이며, 각 해의 유전자는 (1,5,3)와 같이 유전체 3개의 집합으로 이루어진다. 적합도 함수를 20과 얼마나 가까운지를 나타내는 값으로 둔다면, (1,5,3)에 대한 적합도는 f( (1,5,3) ) = 11이 된다.
먼저 첫 세대를 아무렇게나 생성한다. 첫 세대가 만약 { (1,5,3) (8,0,9) (9,9,8) (3,7,5) } 으로 형성되었다고 하자. 각각의 적합도를 구하면, { 11, 3, 6, 5 }이 되며, 이 값이 높을수록 20에서 멀기 때문에 해로서 부적당하다는 것을 의미하며, 따라서 세대를 거침에 따라 살아남을 확률이 낮게 된다.
다음 세대를 형성하기 위해, 이 세대의 개체중 2개의 유전자를 선택한다. 이때 선택은 적합도를 기준으로 확률적으로 선택한다. (룰렛 알고리즘이 자주 쓰인다) 위의 예에서 적합도가 3인 (8,0,9)는 적합도가 6인 (9,9,8)에 비해 훨씬 높은 선택 기회를 가진다. 선택된 2개의 유전자의 유전체는 랜덤한 위치에서 교환되어 새로운 세대가 형성된다. 예로 (8,0,9), (9,9,8) 이 선택되었고 교배위치가 첫번째 자리로 무작위로 결정되었다면 다음 세대의 개체는 (8,9,8), (9,0,9)가 되며, 각각의 적합도는 5, 2가 된다.

출처: 위키피디아

 

 

4. SwapAdvisor Design

아래의 그림(Figure 2)은 MXNet에서 만들어진 dataflow graph에 대해서 SwapAdvisor가 적용된 모습을 보여준다.

 

  1. Dataflow graph가 주어지면, SwapAdvisor는 initial value와 graph를 가지고 타당한(legitimate) operator schedule과 memory allocation plan을 뽑아낸다.
  2. 생성된 plan을 swap planner로 넘겨서 언제 어떤 tensor를 swap-in, swap-out 할 지 결정한다.
  3. Swap planner에서는 extra swap-in/out operator들과 control flow edge 정보까지 포함하는 augmented graph를 생성한다.
  4. Augmented graph는 dataflow simulator로 넘겨져서 전체 수행시간이 측정된다. 이 때 유전 알고리즘 기반의 search를 통해서 각 memory allocation/schedule 조합에 대한 성능 평가가 이루어지고, 보다 나은 allocation/schedule 후보를 생성해낸다.
  5. 만약 swap plan이 충분히 최적화 되었다면 최종적인 augmented graph는 프레임 워크로 보내져서 실제 수행에 사용된다.