NVIDIA DGX SuperPOD助力京东探索研究院Vega-MT模型大赛夺魁！

模型训练任务缩短到两周内完成,且在算力和扩展性上都有倍数提升。

（ 图片来源于京东 ）

借助 NVIDIA DGX SuperPOD,京东探索研究院训练了参数量近五十亿的 Vega-MT 模型,在 2022 年第 17 届国际机器翻译大赛(WMT)上大放异彩。Vega-MT 在中文-英语(BLEU 33.5, chrF 0.611)、英文-中文(BLEU 49.7, chrF 0.446)、德语-英语(BLEU 33.7, chrF 0.585)、英语-德语(BLEU 37.8, chrF 0.643)、捷克语-英语(BLEU 54.9, chrF 0.744)、英语-捷克语(BLEU 41.4, chrF 0.651)和英语-俄罗斯语(BLEU 32.7, chrF 0.584)在七个翻译赛道上荣获冠军。

作为大规模 AI 基础架构,NVIDIA DGX SuperPOD 具备完整、先进的基础架构。相较于原来的 V100 集群,DGX SuperPOD 除了在单卡算力上取得了接近两倍的提升外,在算力扩展性方面也呈线性增长,即扩展性方面相较于从前也取得了两倍的提升。在多节点的情况下总计获得了约 4 倍的提升。因此,原本需要数月才能完成类似模型(相当规模、相当复杂度)的训练任务缩短到了两周,研究人员也有了更充裕的时间针对模型优化。

客户简介及应用背景

京东作为以供应链为基础的技术与服务企业。京东探索研究院(JD Explore Academy)秉承“以技术为本,致力于更高效和可持续的世界”的集团使命,是以京东集团以各事业群与业务单元的技术发展为基础,集合全集团资源和能力,成立的专注前沿科技探索的研发部门,是实现研究和协同创新的生态平台。探索研究院深耕泛人工智能 3 大领域,包括“量子机器学习”、“可信人工智能”、“超级深度学习”,从基础理论层面实现颠覆式创新,助力数智化产业发展及社会变革,以原创性科技赋能京东集团零售、物流、健康、科技等全产业链场景,打造源头性科技高地,实现从量变到质变的跨越式发展,引领行业砥砺前行。

国际机器翻译大赛(WMT)是全球学术界公认的国际顶级机器翻译比赛,由国际计算语言学协会(ACL)举办,是协会旗下的顶尖赛事。自 2006 年至今,WMT 的每次比赛都是全球各大高校、科技公司与学术机构展示自身机器翻译实力的平台,更见证了机器翻译技术的不断进步。

京东探索研究院在 WMT 大赛上这一重大成就进一步验证了自然语言处理大模型在理解、生成、跨语种建模上的优越性。

客户挑战

机器翻译面临诸多挑战:常见的几个语种使用广泛数据资源丰富,小语种在跨境电商中非常必要但数据不充分,小数据集的训练面临挑战;同时,挖掘语种与语种的联系也是难点之一,因为语言生成的复杂性、多义性,表达的多样性、文化背景,语种之间的差异等都是机器翻译比赛中绕不开的难题。

从 2018 年的 GPT-1 的 1.1 亿参数,到现今大规模语言模型拥有万亿参数,大模型在多种语言任务上准确率的显著提高,有助于我们构建对自然语言有着更加丰富理解的智能系统。

Vega-MT 采用了诸多先进技术,包括多方向预训练(multidirectional pre-training)、超大模型(Extremely Large Transformer)、循环翻译(cycle translation)和双向自训练(bidirectional self-training),来充分挖掘双语数据、单语数据的知识。此外,还使用了噪声信道重排序和泛化性微调等策略来增强 Vega-MT 系统的鲁棒性和译文的信达雅水平。

但是,在训练大模型时,我们依然会面临诸多困难。在此之前对于一般任务一张 GPU 即可胜任模型的训练,但是在大模型场景,需要多节点协同才能完成最终的训练任务,这也对现有的 GPU 计算集群提出了新的挑战。以知名的 GPT-3 为例,它使用了 45 TB 的训练数据,最高达到 1750 亿的模型参数量;在采用混合精度时,它共占用约 2.8 TB 的显存,需要超过 35 张 GPU 才能将模型全部放下。

因此,训练的挑战集中在单卡算力和多卡多节点通信上,训练也会跨越多个节点。此时数据传输、任务调度、并行优化、资源利用率等方面就显得尤为重要。

应用方案

在构建 AI 基础架构时,我们会面临来自方方面面的挑战,例如计算资源,网络,存储,乃至最上层用来进行任务调度的软件等,而这些方面并非独立,需要综合考量。

京东探索研究院采用的 NVIDIA DGX SuperPOD 就是一套综合完整的高性能解决方案。SuperPOD AI 集群,基于 DGX 服务器,HDR InfiniBand 200G 网卡以及 NVIDIA Quantum QM8790 交换机,计算网络和储存网络隔离,既保证了最优的算力,又可以确保节点和卡之间的高效互联,最大程度提升分布式训练的效率。

算力层面,单节点算力高达 2.4 PFLOPS,采用单节点进行训练,BERT 只需 17 分钟即可完成训练,Mask R-CNN 38 分钟,RetinaNet 也只需 83 分钟。而针对于 Transformer XL Base,181 分钟即可完成训练。同时,依托于多实例 GPU(Multi-Instance GPU, MIG)技术,可以将 GPU 分割成多个实例,每个实例都有其独立的显存,缓存以及流式多处理器,彼此之间故障隔离。这样一来可以进一步提升 GPU 的利用率,同时满足需要不同算力的任务。

网络层面,通过 Scalable Hierarchical Aggregation and Reduction Protocol(SHARP)技术,可以让聚合计算从 CPU 迁移到交换机网络,消除了节点之间多次发送数据的需求,大幅缩减了到达聚合节点的网络流量,从而显著减少执行 MPI 的时间,同时让通信效率与节点数目不再直接相关,进一步保障了算力的可扩展性。除此之外,将 CPU 从处理通信的任务中解脱出来,让宝贵的 CPU 资源聚焦于计算上,进一步提升了整体集群处理任务的能力。

储存层面,在训练模型时,往往需要从存储中多次读取训练数据,而读取操作的耗时也会在一定程度上影响训练的时效。DGX SuperPOD 使用了高性能的多层级储存架构,以平衡性能、容量和成本的需求。而借助于 GPU Direct RDMA 技术,可以绕过 CPU 直接连通 GPU,储存和网络设备,从而进行高速低时延的数据传输。

软件层面,为了构建集群以及保障集群的持久化平稳运行,上层的监控调度管理软件不可或缺。Base Command Manager 是一个集群管理系统,它可以对集群进行一系列的配置,管理用户访问,资源监控,记录日志,以及通过 slurm 进行作业任务调度。同时,NGC 上涵盖了大量 AI,HPC,数据科学相关的资源,用户可以轻松获取到功能强大的软件,容器镜像,以及各类预训练模型。

与此同时,探索研究院团队对集群进行了 7x24 小时的监控管理,确保训练任务长时间的平稳运行。监控资源利用率,也确保了每个节点上的计算资源能被充分地利用了起来。在完善的调度监控工作和 DGX SuperPOD 高可靠性的质量保证下,所有被使用到的训练节点在模型训练的 20 天(2 周 pre-training + 5天 fine-tuning)内,没有出现任何问题,训练最终圆满完成。

使用效果及影响

Vega-MT 被成功应用在了京东国庆期间对外发布的 Omni-Force AIGC 小程序中。小程序的应用是用户输入文字生成对应的图片,在 Vega-MT 的加持下,小程序可以支持多个语种的文字输入,譬如中文、英文、西班牙文等等。

京东探索研究院表示: “通过 NVIDIA DGX SuperPOD 的加持,京东探索研究院可以快速迭代模型,帮助高准确度的模型快速落地,进一步提升用户体验,降低成本,提升效果和业务收益。此次 NVIDIA DGX SuperPOD 支持我们能在 WMT 比赛中拔得头筹,不仅提高了企业的知名度,也助力了京东成为更受用户信赖的品牌”。

值得一提的是,除了 Vega-MT 外,NVIDIA DGX SuperPOD 也在探索研究院的其他大模型训练上发挥了不可或缺的作用,譬如 ViTAE,织女模型等。

关于 NVIDIA

自1993年成立以来,NVIDIA (NASDAQ: NVDA) 一直是加速计算领域的先驱。NVIDIA 1999年发明的GPU驱动了PC游戏市场的增长,并重新定义了现代计算机图形,开启了现代AI时代,推动了元宇宙的创建。NVIDIA现在是一家全栈计算公司,其数据中心规模的产品正在重塑整个行业。更多信息,请访问https://nvidianews.nvidia.com/。