OSDI'20 论文赏：ANTMAN: DYNAMIC SCALING ON GPU CLUSTERS FOR DEEP LEARNING

SOSP 与 OSDI 是系统领域的圣殿，无数研究者的梦想。OSDI 的全称是 USENIX Symposium on Operating Systems Design and Implementation，但随着时代的发展，它早已不局限在操作系统领域。在 OSDI‘20 上也出现了很多 ML System 方向的文章。今天与大家分享一下其中一篇与深度学习集群管理有关的论文 AntMan: Dynamic Scaling on GPU Clusters for Deep Learning。

这篇文章出自阿里云 PAI 团队，第一作者 Wencong Xiao，在北航和微软亚研获得博士学位后加入阿里云。之前的工作也主要集中在这一领域，比如 Gandiva: Introspective Cluster Scheduling for Deep Learning 等。

Wencong Xiao 的这一新工作通过对调度器和训练框架的联合设计，减少了同一个物理 GPU 上多个任务间的相互干扰，从某种程度上实现了对 GPU 资源的超卖，提高了 GPU 资源的利用率。这一工作目前也落地生产环境和阿里云 PAI-DLC 云原生的模型训练平台产品中。

背景知识

这一文章会比较多涉及深度学习分布式训练的内容，可以事先阅读分布式训练的方案和效率对比简单了解一下相关内容。

动机

随着深度学习在工业界的落地，在云上进行模型训练和推理的需求越来越旺盛。但是，GPU 这样的硬件设备的资源利用率，一直以来都处于相对比较低的水平。一方面因为模型训练往往涉及许多不同的环节，比如数据的预处理等。有些环节不适合在 GPU 上执行。另一方面，随着数据规模的扩大和模型复杂度的提升，分布式训练逐渐成为工业界训练场景的主流选择。同步 SGD 的分布式训练中，很多的时间花费在了等待网络 IO 上，GPU 从微观的角度来看，经常空闲。

左图是阿里巴巴内部集群（几千张卡） GPU 显存和 SM 利用率的 CDF 图，可以看出，GPU 算力的利用率非常非常低。右图是在调度分布式训练任务时，GPU 空闲的情况。在分布式训练中，通常需要任务下所有的实例都被运行时，训练才能开始。因此需要在调度上进行 gang scheduling 的支持。而 gang scheduling 会造成 GPU 的等待。任务所需要的 GPU 越多，在任务执行前 GPU 空闲的时间越长。这也很容易理解：调度器需要先把空闲的 GPU hold 住一段时间，在确定分布式训练中所有的 PS 和 Worker 都可以被运行，才会真正创建所有实例。GPU 被调度器 Hold 的这段时间没有办法被其他任务使用，所以会造成一定的浪费。

多个任务共享一个 GPU 可以提高利用率，但是多个任务会在 GPU 上（如 Memory hierarchy 等）有相互的干扰。因此，在生产环境中，通常不会采取这样的方式。目前也有一些 GPU 虚拟化的解决方案（如 Amazon Elastic Inference），但是在容器环境下，也很少有可以落地生产环境的。

除此之外，GPU 在训练过程中，SM 和显存也存在一定程度的不均衡。在左图中，DeepFM 的模型训练通常需要进行数据的预处理，这个过程只需要 CPU 参与，因此 GPU 利用率为 0%。右图也是有类似的情况。

因此，优化 GPU 的利用率的空间还是很大的。这也是这篇文章的动机。

Key Insight

文章通过实验，论述了 Key Insight：大部分模型本身占用的显存并不多，使用的显存多来自 mini-batch 过程中，在单个 mini-batch 中会被申请和释放。文章中所有的 design 基本都是围绕这一 Key Insight 展开的。

系统设计

文章联合设计了调度器和框架，让框架来在训练任务的角度支持显存和算力的动态调整，然后让调度器从集群的角度利用这一新的特性进行更有针对性的调度。

框架层的设计

框架层的修改分为两个方面，分别针对显存和算力。

在显存方面，为了实现动态的调整，一共引入了如下的修改。以 TensorFlow 为例，首先是引入了 GPUVMemAllocator（论文中提到的 UniversalAllocator）。在 GPU 的显存虚拟化 flag 被置为 False 时，使用 TensorFlow 的 BFCAllocator，如果是 True，就利用 GPUBFCAllocator 和 Host 的 BFCAllocator 创建出 VMemAllocator。在申请显存的时候，会先申请 GPU 显存，如果超出了限制，会申请 Host 内存。具体的逻辑可以参见开源代码，这部分的实现非常清晰。

    // GPUVMemAllocator will allocate host memory as backup after running out of
    // gpu device memory to avoid OOM failures
    gpu_allocator = maybe_create_gpu_vmem_allocator(gpu_allocator,
                                                        bus_id,
                                                        platform_gpu_id,
                                                        tf_gpu_id.value(),
                                                        stream_exec);
...
Allocator* maybe_create_gpu_vmem_allocator(Allocator* gpu_allocator,
                                           int bus_id,
                                           PlatformGpuId platform_gpu_id,
                                           int tf_gpu_id,
                                           se::StreamExecutor* stream_exec) {
  bool gpu_vmem = false;
  Status status = ReadBoolFromEnvVar("TF_GPU_VMEM",
                                     true/*enabled by default*/,
                                     &gpu_vmem);
  if (!status.ok()) {
    LOG(ERROR) << "GetGPUAllocator: " << status.error_message();
  }
  if (!gpu_vmem) {
    return gpu_allocator;
  }
  SubAllocator* sub_allocator = new GpuHostAllocator(
      GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
      bus_id, {}, {});
  int64 cuda_host_mem_limit_in_mb = -1;
  status = ReadInt64FromEnvVar("TF_CUDA_HOST_MEM_LIMIT_IN_MB",
                               1LL << 16 /*64GB max by default*/,
                               &cuda_host_mem_limit_in_mb);
  if (!status.ok()) {
    LOG(ERROR) << "GetGpuHostAllocator: " << status.error_message();
  }
  int64 cuda_host_mem_limit = cuda_host_mem_limit_in_mb * (1LL << 20);
  Allocator* host_allocator =
      new BFCAllocator(sub_allocator, cuda_host_mem_limit,
                       true /*allow_growth*/,
                       strings::StrCat("GPUHost_", tf_gpu_id, "_bfc"));
  Allocator* gpu_vmem_allocator = new GPUVMemAllocator(gpu_allocator,
                                                       host_allocator,
                                                       tf_gpu_id,
                                                       stream_exec);
  return gpu_vmem_allocator;
}

为了保证训练的速度不会受到影响，Antman 做了一个自动调整显存限制的特性。当发现 Host 内存被使用的时候，会提高显存的限制阈值，这样所有的 Tensor 都可以申请在显卡上。这样只会影响一个 mini batch 的性能，后面的 mini batch 跑前向后向计算的时候，所有的 Tensor 都会被申请在显存上。这一特性的实现可以参考 gpu_adjustable_allocator。

在计算方面，如果多个任务运行在同一个 GPU 上时，主要会带来 GPU Kernel 的排队延迟，和 PCIE 总线带宽的争抢。以下图（a）和（b）来说，B 任务影响了 A 任务原本的执行，为了解决这个问题，Antman 在框架层引入了 GPU Op Manager。在原本的设计中，一旦 Kernel 的控制依赖被满足了，就会被执行。GPU Op Manager 接管了原本的流程，它会控制 GPU Kernel 被 issue 的频率。因为 GPU Op Manager 只会管控 GPU Kernel，因此 CPU 的 Op 可以照常被执行，不会被 Block。如下图（c）所示，满足了依赖的 CPU Op 可以在 GPU Op 没有被执行的时候照常执行，不会受到 GPU Op Manager 的影响。具体的实现可以参考 df28134 和 8c380d6。

调度器的设计

在调度器的设计上，Antman 并没有采取 Monolithic 的架构，而是存在两个角色：Global Scheduler 和 Local Coordinator。其中全局的调度器负责进行任务的调度，而 Local Coordinator 主要负责根据深度学习的训练任务的执行情况（任务情况，硬件指标，mini batch 的执行时间，显存和内存的使用情况等），管理训练任务的全生命周期。

Antman 根据 SLA 把任务分为 resource-guarantee 和 opportunistic 两种任务，其中前者需要保证与独占 GPU 卡类似的训练速度。Antman 的设计目标是在保证 resource-guarantee 类型任务的 SLA 的同时，提高集群的利用率。opportunistic 类型的任务主要就是用来提高集群利用率的。

全局调度器的调度算法比较简单，如下图所示。首先调度器会根据拓扑，获得一个最优的节点组合。这一步与业界主流基本一致，尽可能考虑到 NVLink 等硬件资源的拓扑结构，进行分配。如果是 resource-guarantee 的任务，有合适的节点就会直接运行。对于 opportunistic 类型的任务，Antman 会找到负载最低的节点，去执行。

Local Coordinator 最主要的职责是管理在共享的 GPU 上训练任务。在 Antman 中，一个 GPU 只会被分配给一个 resource-guarantee 的任务。所以当有 opportunistic 的任务已经在 GPU 上运行时，为了满足新来的 resource-guarantee 任务，Local Coordinator 会限制 opportunistic 任务使用的 SM 和显存。随后慢慢提高 opportunistic 的限制，确保在不影响 resource-guarantee 任务的训练速度（mini batch 的耗时）的同时，提高 opportunistic 的资源限额。

文章的实验与验证是在 Kubernetes 上进行的，配合 Kubeflow 发起训练任务，利用修改后的框架和一个自定义的调度器进行任务调度。完整的系统实现是在阿里巴巴的伏羲上完成的。

总结

阿里的这一工作通过对框架和调度器的联合设计，在某种程度上支持了任务的 SLA 级别，以及 GPU 资源在显存和算力级别的共享。可以看出，在这一方面的工作从调度器单方面的优化，逐渐走向 co-design 的方向。GPU 的虚拟化和复用也逐渐走向落地，这使得 GPU 的门槛和成本得以降低。

不过，这样的工作在落地公有云时会存在一些问题。对框架的闭源优化和修改难以回馈社区，对客户而言存在被云供应商绑定的担忧。

License

This article is licensed under CC BY-NC-SA 3.0.
Please contact me for commercial use.

OSDI'20 论文赏：AntMan: Dynamic Scaling on GPU Clusters for Deep Learning