GTC 2022 | 以分时调度共享AI超算,弹性运行超大规模深度学习训练

High-Flyer    March 25, 2022

算力是驱动AI时代发展的大心脏。现今深度学习走向越来庞大的参数,应用于越来越复杂的场景,算力成为企业发展的重要需求。作为AI从业者,我们常常有过“缺卡而抢卡”的经历,无论是高校实验室里的科研还是企业里产品研发,显卡不够用于模型训练的问题时有发生。

最近,幻方在NVIDIA GTC 2022年会上分享了其对大规模深度学习训练和算力管理的思考,提出了以任务级分时调度共享AI超算的理念。GTC是面向全球科研学者、开发工程师、产业决策者的AI顶会之一,历年来发布了很多影响AI发展的重量级产品和前沿理念,受到学术界、产业界专家和精英们的广泛认可。

本文将简要为大家介绍,幻方在这次GTC大会上提出的分时调度理念和弹性深度学习训练产品

会议地址https://www.nvidia.com/gtc/session-catalog/?search=huang&tab.scheduledorondemand=1583520458947001NJiE&search=huang#/session/1639723980421001KqlP

概述

幻方根据自身量化金融业务发展的需求,逐步构建了拥有一万张A100的AI超算集群(萤火二号),在长期实践中,逐步总结研发出了高效利用算力资源的平台方法:最开始我们的研究员们各自独占几十张卡进行模型训练与测试,发现算力瓶颈很快达到,GPU完全不够用,而总体使用率不到50%,闲置资源大量浪费。之后我们逐步探索采用任务级分时调度来管理集群上的模型训练,其主要帮助我们解决如下三个问题:

  1. 当有任务需要多张GPU(1000+)训练时,如何快速方便地获取到计算资源;
  2. 如何尽可能利用起集群的全量算力,服务更多的深度学习任务;
  3. 如何用自动调度系统统一管理训练任务,把用户从繁琐的管理工作中解放出来。

任务级分时调度方法的提出是受到Linux系统CPU任务调度的启发,将显卡算力资源从传统的硬件独占方式划分转向于训练任务的集中调度,将显卡资源集中起来,在时间轴上进行分片,合理分配给集群上的每一个训练任务。用户不需要关心是否有算力资源,只需要研发好模型,调试好代码,配置好算力需求,提交集群即可。集群会自动找到空闲的算力,空闲的时间,将您的模型训练起来。

p0.png

然而想做成这样一件事情非常的困难,如何让这么多显卡资源整合起来,对外提供像一台计算机式的服务,同时发挥出优异的性能,存储、网络、计算都需要进行精心的设计和调试。幻方AI的研究员们经过长期的实践,不断地试错,逐步研发出了高性能的深度学习训练平台,幻方萤火,提供任务级的分时调度管理。用户可以按需在幻方萤火上提交任务,而集群可以持续保持着90%以上的利用率为模型训练提供服务。

幻方萤火是一系列AI基础设施的集大成者,是幻方对外开放其澎湃深度学习算力的窗口,他包括:

  1. 任务级分时调度器
  2. 高速文件读写系统 3FS
  3. 模型并行训练工具 hfreduce
  4. 高性能深度学习算子 hfai.nn

还有底层各种高性能硬件设备和软件封装。在接下来几章会为大家详细描述。

分时调度器

任务级分时调度器(task-level time-sharing scheduler)是幻方萤火超算集群的大管家,负责训练任务的接收和计算资源的调配。资源调配遵循如下策略:

  1. 在同等优先级条件下,每个任务按照“先来先服务”原则,排队依次完成训练;不同优先级条件下,高优先级可以插队率先完成训练
  2. 调度器监控任务的资源使用情况和模型规模,动态地扩展或缩减任务地GPU使用量
  3. 在计算资源不足导致的同等优先级任务排队情况,会依次轮询共享GPU算力资源,尽可能让所有排队任务都能分到算力进行训练;

p1.png

通过上述调度策略,用户完全可以按自己需求提交任务,训练模型,每项任务都可以得到公平的训练服务。而对于幻方萤火来说,集群的算力资源可以充分得到利用,使我们可以向公众开放我们的计算能力。

3FS

3FS幻方AI自研的高速读写文件系统,全称是萤火超算文件系统(Fire-Flyer File System),是幻方萤火超算实行计算存储分离架构后,存储服务中的重要一环。3FS 只服务于AI训练时计算节点中的模型批量读取样本数据这个场景,通过高速的计算存储交互减少模型训练中数据读取的开销,充分利用GPU的计算时。

p2.png

AI训练读取样本数据是一个大规模的随机读取任务,且读上来的数据不会在短时间内再次被用到,因此3FS的从设计之初就是一个比较特殊的文件系统:其不采用“读取缓存”这一最重要的工具来优化文件读取,而采用 Direct IORDMA Read 的读取方式,让模型训练在样本读取部分只用极小的 CPU 和内存开销,而充分得到超高的读取带宽。

hfreduce

hfreduce幻方AI自研的高速模型并行训练工具,是幻方萤火超算实行计算存储分离架构后,计算服务中的重要一环,根据幻方萤火的集群特性专为计算节点设计的All Reduce工具。本质上hfreduce相当于 PyTorch 中的 DistributedDataParallel(DDP),只不过使用 CPU 做加法运算以计算总梯度,而不是调用 NCCL 传递梯度到不同的显卡上,再各自计算总梯度。

p3.png

hfreduce的设计来源于幻方主要的AI场景:金融行为分析、自然语言处理、生物分子结构预测等。在这些场景中,基本是数据规模大而模型大小适中,模型的加速主要是依赖大量的数据并行。因此,幻方萤火中的机器没有配备NVLink,而是安装了 PCIe 版本的 GPU,通过软件层的设计更能给幻方AI的业务加速。

hfai.nn

为了解决AI落地的环节中,实验室代码在生产环境上常常面临着性能、准确度、资源等各种各样的问题。幻方AI针对Pytorch框架进行了深度优化,结合幻方萤火超算集群的特点,对一些常用的AI算子重新研发,提升效率,进一步提升了模型整体的训练效率,提供出了hfai.nn算子供开发者们使用,包括:

  1. hfai.nn.lstm,相比 torch.nn.lstm 提速4~5倍
  2. hfai.nn.MultiheadAttention,相比 torch.nn.MultiheadAttention 提速30%
  3. hfai.nn.LayerNorm,相比 torch.nn.LayerNorm 提速 80%

p4.png

幻方AI集结了一批NOI/ACM 金牌选手组成AI团队,持续优化深度学习核心算子,研发高性能的深度学习模型,对外提供定制的服务。

写在最后

AI 拓展了我们的能力边界,激发了我们的想象力和创造力。与大家熟悉的产品研发思维不同,我们并不是先构想出“分时调度共享AI超算”的技术,而后去实现它,而是我们作为实际的算力使用者,为了解决自己的使用困境,创造了如今的萤火超算平台。当我们受益于萤火绝佳的AI科研环境的时候,我们在想,或许它可以推而广之,让AI时代的引擎能推进更多的“天才想法”。截止到2022年第一季度,我们的萤火AI超算已经服务了50多个国内外高校实验室及科研机构,累积提供了超过270万小时的GPU算力。当前,幻方AI团队为垂直行业提供更为强大的AI定制服务,加速AI对各行各业革命性的技术颠覆。


本文作者: High-Flyer


您可以转载、不违背作品原意地摘录及引用本技术博客的内容,但必须遵守以下条款: 署名 — 您应当署名原作者,但不得以任何方式暗示幻方为您背书,亦不会对幻方的权利造成任何负面影响。 非商业性使用 — 您不得将本技术博客内容用于商业目的。 禁止演绎 — 如果基于该内容改编、转换、或者再创作,您不得公开或分发被修改内容,该内容仅可供个人使用。