A Deep Learning Optimization Framework for Versatile GPU Workloads

Abstract

Deep Learning (DL) frameworks, such as TensorFlow and PyTorch, are widely used in various fields of applications. However, there are two limitations to these frameworks. The first limitation is that these frameworks heavily rely on the NVIDIA cuDNN for performance. Using cuDNN does not always give the best performance. One reason is that it is hard to handle every case of versatile DNN models and GPU architectures with a library that has a fixed implementation. Another reason is that cuDNN lacks kernel fusion functionality that gives a lot of chances to improve performance. The other limitation is that these frameworks do not support the training of an out-of-GPU-memory-scale model, which is essential for state-of-the-art DL techniques in the field of natural language processing.

In this thesis, we propose two performance-model-based deep learning optimization frameworks. The first is DeepCuts, a DL optimization framework for versatile workloads on a single GPU. It considers both the kernel implementation parameters and GPU architectures. It analyzes the DL workload, groups multiple DL operations into a single GPU kernel, and generates optimized GPU kernels when considering the kernel implementation parameters and GPU architecture parameters. The evaluation results under various DL workloads for inference and training indicate that DeepCuts outperforms cuDNN/cuBLAS-based implementations and the state-of-the-art DL optimization frameworks, such as TVM, TensorFlow XLA, and TensorRT.

In addition, we propose BigCuts: a DL optimization framework for out-of-GPU-memory scale model training on a GPU cluster. BigCuts takes the information of the target model to train and the target cluster system to use and finds the best way of cluster utilization using the performance model. BigCuts also considers using a storage system to train a massively scaled model whose parameters cannot be stored in CPU memory. BigCuts successfully find the best way of cluster utilization for a versatile type of models and hardware systems, unlike previous approaches such as Megatron and DeepSpeed, which lack such capability.

PyTorch, TensorFlow와 같은 딥 러닝 프레임워크는 딥 러닝 소프트웨어 에코시스템에서 가장 중요한 역할을 하는 핵심적인 소프트웨어이다. 그러나 현존하는 딥 러닝 프레임워크들은 특정 종류의 딥 러닝 워크로드에만 맞추어 최적화되어있으며, 최근 새로이 개발된 딥 러닝 모델을 활용하고자 하는 경우 성능과 활용성이 크게 떨어진다. 크게 두 가지 원인이 있다. 첫 번째로, 현존하는 딥 러닝 프레임워크들은 구현이 고정된 라이브러리(예: cuDNN)을 사용하는데, 이 때문에 다양한 종류의 딥 러닝 모델에 대해 좋은 성능을 보이기 어렵다. 또한 GPU 메모리상에 필요한 데이터가 전부 저장되는 상황만을 주로 가정하고 구현되어 있어 학습하고자 하는 모델의 크기가 GPU 메모리에 올릴 수 없을 만큼 큰 경우 학습 코드를 작성하기 어렵다.

본 논문은 이를 해결하기 위한 솔루션으로 DeepCuts와 BigCuts를 제시한다. DeepCuts는 주어진 딥 러닝 워크로드와 GPU 하드웨어 구조를 동시에 고려하여 최적화된 커널 코드를 생성하는 프레임워크이다. DeepCuts는 딥 러닝 워크로드를 표현하는 연산 그래프(computation graph)를 입력으로 받아 최적화된 GPU 커널을 생성한다. 이 과정에서 여러 개의 딥 러닝 연산을 합성하는 커널 퓨전(Kernel fusion)을 적용하며, 커널의 성능 상한을 예측하는 성능 모델을 사용하여 커널 생성을 최적화한다. DeepCuts는 cuDNN/cuBLAS를 사용하는 경우보다 높은 성능을 보이며, TVM, TensorFlow XLA, TensorRT등의 세계 최고 수준 딥 러닝 최적화 프레임워크 및 컴파일러와 비슷하거나 더 높은 성능을 보인다.

BigCuts는 사용하는 GPU 클러스터의 규모 및 종류와 상관없이 대규모의 딥 러닝 모델을 학습시키는 기법을 제시한다. 이를 위해 스토리지 시스템을 사용한 Tensor 관리 기법, 성능 모델에 기반한 구현 Design choice 선택 기법을 제시한다. 제시된 기법에 기반하여 주어진 GPU 클러스터의 구조와 학습하고자 하는 모델의 구조에 따라 최적화된 병렬화/실행 방법을 탐색할 수 있음을 보인다.