Nimble: Parallel GPU Task Scheduling for DL - NIPS 논문 리뷰

 최근 AI 연구의 대부분을 차지하고 있는 Deep Learning (DL)은 기본적으로 정확도와 연산량이 Trade-off를 갖기 때문에, GPU 등 병렬처리가 가능한 Hardware를 사용할 때, 진가를 발휘한다. 요즘은 CUDA 같은 GPU 병렬처리 라이브러리와 이를 이용한 Tensorflow-GPU 버전이 사용하기 쉽게 되어 있어서, GPU를 직접 Contorl하지 않아도, GPU를 이용해서 Deep Learning 알고리즘을 돌릴 수 있다. 하지만, DL이 이용되는 시스템을 설계함에 있어서, 알고리즘이 돌아가는 GPU에서 어떤 스케쥴링과 동작이 이루어지는지를 알아보는 것도 중요하다는 생각을 했다. GPU 스케쥴링에 대한 논문은 많지만, 그 중 첫 번째로, 작년 NIPS에서 발표된 DL을 위한 Parallel GPU Scheduling 논문 [1]을 리뷰해보고자 한다.

 이 논문에서 지적하는 현재 Deep Learning Frameworks의 문제점은 GPU를 써서 Inference 및 Training을 할 때, GPU Scheduling overhead가 크고, 불필요한 Serial Processing이 진행되는 상황이 많다는 것이다. 즉, 이 논문에서 제안하는 Nimble은 GPU의 강점을 잘 살리기 위해, 적은 GPU Sheduling overhead를 유지하면서, Parallel Processing을 가능하게 해주는 방식이다. 이를 위해, GPU Kernel (명령) 실행 전에, 미리 Scheduling을 완료하는 Ahead-of-Time (AoT) 스케쥴링 방식을 사용함으로써, Scheduling overhead로 인한 속도 저하를 방지하고, GPU에서의 Multi GPU-Stream을 지원하여, 병렬 처리를 원활하게 한다.

 기존의 DL Framework (e.g., Tensorflow)에서는 아래 그림처럼, Run Time 중에 Scheduler에서 GPU 메모리 할당, Task Scheduling 등을 진행한다. 각 연산들의 종속성을 확인하고, Queue에 넣으면, Tensor의 Input, output shape을 check하고, 메모리를 할당받아 놓고, GPU 연산을 Submission하면, 정해진 순서대로 GPU가 연산을 수행하는 방식이다.

 하지만, Run Time 중 Scheduling을 진행할 경우, 아래 그림과 같은 속도 저하가 일어남을 볼 수 있다. Tensor shape check와 메모리 할당같은 내용들을 하드코딩으로 넣어서, Scheduling overhead가 거의 없는 경우에는 약 2배 빠른 Execution Time을 보이는 것을 볼 수 있다. 또한, GPU에서 Multi-Stream을 통해 Parallel Processing를 잘하면 약 3배 정도의 속도 향상을 볼 수 있는데, Fig. 3에서 보는 것처럼, Task A가 끝날 때까지, Task B에 대한 Scheduling이 끝나지 않을 수 있기 때문에, 효율적인 Parallel Processing이 힘들다. 

 이 논문에서는 Nimble이라는 Scheme을 통해 이를 해결하고자 한다. Nimble은 Train한 이후, Inference의 동작이 바뀌지 않는 Static Neural Network를 가정하고 있기 때문에, 이 시스템에서는 연산 순서와 방법 등이 고정되어 있다. 따라서 Nimble에서는 해당 Neural Network를 한 번 수행하고, Execution Trace로 기록하여, 다음 Task에 필요한 GPU Memory, 순서 등을 미리 알 수 있다. 이렇게 저장된 Excecution Trace로 미리 Scheduling을 하는 AoT를 사용하여, Scheduling으로 인한 속도 저하를 방지할 수 있다. 추가적으로, GPU에서 Multi-Stream간에 동기화를 통해, GPU에 스트림을 자동 할당하는 알고리즘도 제시되어 있다. 이렇게 제시된 두 알고리즘에 의해, Multi-Stream을 이용하되, Scheduling의 overhead가 적은 구현을 할 수 있다.

 

Conclusion

 이 논문의 Proposed Scheme인 Nimble은 Static Neural Network의 고정된 Scheudling 요소들을 저장하고, 이를 기반으로 Pre-Scheduling (AoT)를 통해, scheduling overhead를 줄인다. 이런 Scheduling과 함께 Parallel Processing이 자동적으로 Multi-stream scheduling될 수 있는 효율적인 알고리즘을 제안하고 있다. 이 논문의 Performance Evaluation에서 제시된 결과를 보면, Neural Network의 Parallel Processing 여지가 많은 구조 (동시성이 높은 구조)일 때, Nimble의 속도 (Training, Inference 포함) 향상 효과가 보인다.

 

Reference

[1] "Nimble: Lightweight and Parallel GPU Task Scheduling for Deep Learning", Woosuk Kwon, Gyeong-In Yu , Eunji Jeong, Byung-Gon Chun, Advances in Neural Information Processing systems 33 (NeurIPS 2020).