Press "Enter" to skip to content

32分钟训练神经机器翻译,速度提升45倍

深度模型的训练时间通常对研究者而言都是很大的挑战,我们需要花数小时甚至数天才能知道某个小改进到底好不好。然而如果我们从一开始就考虑降低模型训练时间,那么很多概念都能迅速验证。在 Facebook 开发者的试验中,他们采用了低精度和大批量等一系列加速训练方法,并成功地将需要 24 小时训练的 NMT 降低到 32 分钟。该项目相关的分布式训练代码已开源。

项目地址:https://github.com/pytorch/fairseq

我们想让用户用自己喜欢的语言体验我们的产品,同时与世界各地的人们建立联系。为此,我们使用神经机器学习(NMT)自动翻译帖子和内容中的文本。我们之前关于这一课题的研究 fairseq 已经开源,这是一个序列到序列的学习库,可供任何人训练 NMT 模型,完成自动摘要或其他文本生成任务。
随着 NMT 模型在从大规模单语数据(只有一种语言的数据)中的学习越来越成功,训练速度变得越来越重要。为了适应这样的发展,我们必须想办法大大减少训练时间。直到最近,NMT 模型的这种训练在单台机器上还需要数周,对于快速实验和部署来说,这一速度实在太慢。
由于在精度和训练设置上做了几项改变,我们在短短 32 分钟内就能训练一个强大的 NMT 模型(参见《Scaling Neural Machine Translation》),速度是原来的 45 倍(原来需要 24 小时)。在之后的一项研究《Understanding Back-Translation at Scale》中,我们展示了这一速度大幅提升的训练设置如何让我们利用单语文本训练准确得多的模型。

  • 《Scaling Neural Machine Translation》:https://arxiv.org/abs/1806.00187

  • 《Understanding Back-Translation at Scale》(EMNLP 2018):https://arxiv.org/abs/1808.09381

加速训练

我们首先在一台配备 8 块 Volta GPU 的英伟达 DGX-1 机器上把模型的训练时间从近 24 小时缩减至不到 5 小时。神经网络包含数百万个参数,这些参数在训练过程中会被调整并通常以 32 位浮点精度存储。然后,我们把训练的精度从 32 位转换为 16 位,这样不仅能减少 GPU 内存占用,还让我们能够使用英伟达最新 Volta GPU 架构提供的经过高度优化的 Tensor Core。降低了精度(16 位)的训练有时会由于浮点「溢出」而导致模型质量降低。在本研究中,我们实施了一个已知的自动监控和防止溢出程序,并将训练时间从近 24 小时减至 8.25 小时,速度提升 2.9 倍而没有降低模型质量。
接下来,我们通过所谓的累积更新,延迟了模型更新来减少 GPU 之间的通信。通过同步训练模型,从而每个 GPU 可以维持一个完整、相同的模型副本,但是会处理训练数据的不同部分。处理时需要反向传播——在网络上被分为 forward pass 和 backward pass(前向/反向),以计算训练模型所需的统计量。在处理每个小批量(mini-batch),GPU 会同步地相互交流结果(梯度)。这样可能导致效率低下:首先,在 GPU 之间发送数据会消耗时间;其次,处理较快的工作站必须等待速度较慢的工作站(即所谓散兵游勇),然后训练才能继续。
后者对于文本相关模型来说是一大挑战,不同的句子长度会使这个问题更加严重。我们对这两个问题的解决方案是在工作线程互相通信之前,让每个工作线程处理并累积多个小批量的结果,从而延迟同步点。这将训练时间缩减至 7.5 小时,且无损模型质量。
累积更新还能增加有效的批大小,即每一个训练步所使用的数据量。在我们的研究设置中,批大小增大为 16 倍。这使得学习率可以翻倍,训练时间减少到 5.2 个小时,比原始系统实现了 4.6 倍的加速。

图:通过 (i) 减少工作站之间的通信、(ii) 减少不同 GPU 的工作负载变化从而节约闲置时间,累加多个前向/反向步骤的梯度,实现训练加速。
最后,我们使用 16 位浮点数来进一步提高每一个工作站的批大小,从而利用释放出的额外 GPU 内存。通过将每个工作站的批大小从 3.5K 单词提升到 5K 单词,该研究将总训练时间减少到 4.9 个小时,对原始系统实现了 4.9 倍的加速。

扩展

我们对单个机器训练的优化还可应用于在多个机器上进行并行训练(即分布式训练)。我们将训练从单个 DGX-1 机器扩展到 16 个机器(即 128 块 GPU),发现可以在 37 分钟内训练完成同样的模型,与原始的单机器设置相比,速度提升了 38.6 倍。
在多个机器上训练时,另一个优化是同时发生 GPU 通信和反向传播。在反向传播之后,我们将所有信息传达给其他工作站。典型情况下,反向传播和通信按顺序执行,但是我们通过尽可能并行 GPU 通信和反向传播而节省时间。具体来说,我们在反向传播步骤结束后立刻开始在多个工作站之间同步梯度。这样 GPU 通信和反向传播发生就实现了并行,这节省了 14% 的训练时间,且使总训练时间从将近 24 个小时下降到 32 分钟,速度提升了 44.7 倍。

图:反向传播与梯度同步的并行示意图。

图:每次改进后训练总时间的变化情况。

在更多未翻译数据上训练

加速训练允许我们在更大数据集上训练模型。一般而言,训练 NMT 模型需要原文与参考译文对应的成对文本,例如双语数据等。但是,双语数据相对于网络上或其它地方常见的大量单语数据要有限地多。在第二篇论文《Understanding Back-Translation at Scale》中,我们展示了如何通过仅提供我们想要翻译的语言的文本来提升模型的翻译性能,以及利用大量的此类数据可以大幅提升模型准确度。
通过单语数据提升 NMT 模型最高效的方法之一是反向翻译(back-translation)。如果我们的目标是训练一个英语到德语的翻译模型,我们可以首先训练一个从德语翻译到英语的模型并利用该模型翻译所有的单语德语数据,然后在已有的和新的数据上简单地训练一个英语到德语的最终模型。我们的论文显示,数据如何被翻译是非常重要的,通过采样故意不总是选择最好的翻译同样也是非常重要的。

图:演示反向翻译的动画
如果我们把 2.26 亿个反向翻译的句子添加到现有的包含 500 万个句子的训练数据中,我们就能大幅提升翻译质量。下图(左)显示了这一系统在标准 WMT』14 英-德基准测试集上的准确率(fairseq & 单语数据)。该系统可以在 16 个 DGX-1 机器上经过 22.5 小时训练而成。图中还显示了 DeepL 的准确率,这是一个依赖于高质量人工翻译的专业翻译服务,此前该服务在该基准测试上表现最佳。

快速翻译

我们还提高了 fairseq 模型的翻译速度。具体来说,我们实现了很巧妙的高速缓存,或从计算和批处理中移除完成句子的几个单词而不是句子。这提高了大约 60% 的速度。下图展示了 fairseq 和其它工具集的对比。只从 32 位浮点数变到 16 位浮点数就能提升 40% 的速度。

图:在标准 WMT』14 英语-德语基准测试集上使用大型 Transformer 模型和 V100GPU 得到的翻译速度。

未来计划

扩展和提高自动化翻译技术仍将是我们工作的重点。为了让该研究项目尽快成长起来,我们共享了分布式训练的代码,并将其作为我们的 fairseq 开源项目的一部分,从而使其他研究者也可以轻松地以更快的速度训练 NMT 模型。进行分布式训练仅需稍微改动训练命令。
实现更快的 NMT 训练是一项激动人心的发现。大幅度地降低训练时间可以加快实验周转,使我们能更轻松地在更好的模型中做研究。把无标签数据应用到翻译之外的问题(如问答或自动摘要)上也是很不错的发现。

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注