[论文解读] Nimble: Lightweight and Parallel GPU Task Scheduling for Deep Learning

Woosuk Kwon, Gyeong-In Yu+2 more|arXiv (Cornell University)|Dec 4, 2020

Advanced Neural Network Applications被引用 15

一句话总结

Nimble 是一种轻量级、并行的 GPU 任务调度器，专为深度学习设计，通过提前编排（AoT）调度消除调度开销，并利用自动多流调度实现独立 GPU 任务的并发执行。它在各种模型上将推理和训练速度分别提升最高达 22.34× 和 3.61×，优于 PyTorch、TensorRT 和 TVM。

ABSTRACT

Deep learning (DL) frameworks take advantage of GPUs to improve the speed of DL inference and training. Ideally, DL frameworks should be able to fully utilize the computation power of GPUs such that the running time depends on the amount of computation assigned to GPUs. Yet, we observe that in scheduling GPU tasks, existing DL frameworks suffer from inefficiencies such as large scheduling overhead and unnecessary serial execution. To this end, we propose Nimble, a DL execution engine that runs GPU tasks in parallel with minimal scheduling overhead. Nimble introduces a novel technique called ahead-of-time (AoT) scheduling. Here, the scheduling procedure finishes before executing the GPU kernel, thereby removing most of the scheduling overhead during run time. Furthermore, Nimble automatically parallelizes the execution of GPU tasks by exploiting multiple GPU streams in a single GPU. Evaluation on a variety of neural networks shows that compared to PyTorch, Nimble speeds up inference and training by up to 22.34$ imes$ and 3.61$ imes$, respectively. Moreover, Nimble outperforms state-of-the-art inference systems, TensorRT and TVM, by up to 2.81$ imes$ and 1.70$ imes$, respectively.

研究动机与目标

• 解决现有深度学习框架中调度开销过高的问题，该问题在小规模 GPU 任务中占据主导运行时间。
• 利用现代神经网络中的操作间并行性（如 NAS 和并行层设计中的结构），这些并行性在当前框架中尚未被充分利用。
• 通过预先调度 GPU 任务，在执行前消除运行时调度开销。
• 利用单个 GPU 内的多个 CUDA 流，实现高效、自动的 GPU 任务并行执行。
• 提升推理与训练的性能，特别是在小批量或低分辨率输入场景下，此时调度开销影响最为显著。

提出的方法

• 提出提前编排（AoT）调度：GPU 任务调度在执行前一次性完成，与运行时解耦。
• 在预处理阶段生成任务调度表，编码所有 GPU 任务的提交顺序、函数参数和资源依赖关系。
• 在运行时，直接根据预生成的调度表提交 GPU 任务，绕过调度逻辑，降低开销。
• 使用依赖感知的流分配算法，自动将独立的 GPU 任务分配到多个 CUDA 流。
• 通过基于追踪的分析，识别并调度可并发执行的独立操作符，实现跨流的并行执行。
• 利用静态神经网络架构和输入张量形状信息，预先计算最优的调度方案与流映射。

实验结果

研究问题

• RQ1 提前编排调度能否消除深度学习框架中的运行时调度开销？
• RQ2 在现代神经网络中，操作间并行性在多大程度上可被利用以加速 GPU 任务执行？
• RQ3 通过多个 CUDA 流实现的并行执行，在现代 GPU 上如何提升推理与训练性能？
• RQ4 调度开销对小批量和低分辨率训练工作负载的性能影响如何？
• RQ5 一个轻量级、自动化的调度器能否超越现有最先进的推理与训练系统？

主要发现

• 与 PyTorch 相比，Nimble 在深度学习推理中最高实现 22.34× 的加速，尤其在包含大量小规模 GPU 任务的模型上表现突出。
• 在训练场景中，Nimble 在小批量和低分辨率输入（如 CIFAR-10 中的数据）上最高实现 3.61× 的加速，此时调度开销最为显著。
• 在推理工作负载中，Nimble 相较于 TensorRT 最高提速 2.81×，相较于 TVM 最高提速 1.70×，展现出更优的调度效率。
• AoT 调度技术将调度开销降至接近零，即使在短时长 GPU 任务中也变得可忽略不计。
• 多流执行有效实现了独立操作符的并行化，尤其对 NASNet-A 和 DARTS 等具有高操作间并发性的模型带来显著收益。
• 在 ImageNet 和 BERT 的大批次训练中，性能提升较为有限，证实了在这些工作负载中调度开销并非主导因素。