[论文解读] Large Scale Language Modeling: Converging on 40GB of Text in Four Hours

Raul Puri, Robert M. Kirby+2 more|arXiv (Cornell University)|Aug 3, 2018

Topic Modeling参考文献 39被引用 15

一句话总结

本论文展示了在 128 张 NVIDIA V100 GPU 上，使用混合精度训练和分布式数据并行性，仅用四小时即可在 40GB 的 Amazon Reviews 文本上训练一个 4096 维的乘法 LSTM。通过自定义学习率调度，实现了 32k 批大小的收敛，使大规模无监督语言建模适用于商业和研究部署。

ABSTRACT

Recent work has shown how to train Convolutional Neural Networks (CNNs) rapidly on large image datasets, then transfer the knowledge gained from these models to a variety of tasks. Following [Radford 2017], in this work, we demonstrate similar scalability and transfer for Recurrent Neural Networks (RNNs) for Natural Language tasks. By utilizing mixed precision arithmetic and a 32k batch size distributed across 128 NVIDIA Tesla V100 GPUs, we are able to train a character-level 4096-dimension multiplicative LSTM (mLSTM) for unsupervised text reconstruction over 3 epochs of the 40 GB Amazon Reviews dataset in four hours. This runtime compares favorably with previous work taking one month to train the same size and configuration for one epoch over the same dataset. Converging large batch RNN models can be challenging. Recent work has suggested scaling the learning rate as a function of batch size, but we find that simply scaling the learning rate as a function of batch size leads either to significantly worse convergence or immediate divergence for this problem. We provide a learning rate schedule that allows our model to converge with a 32k batch size. Since our model converges over the Amazon Reviews dataset in hours, and our compute requirement of 128 Tesla V100 GPUs, while substantial, is commercially available, this work opens up large scale unsupervised NLP training to most commercial applications and deep learning researchers. A model can be trained over most public or private text datasets overnight.

研究动机与目标

• 通过分布式训练和混合精度算术，实现 RNN 的大规模无监督语言模型预训练。
• 研究使用极大批次大小（最高达 32k）训练大型 RNN 的可行性，并实现收敛。
• 设计一种学习率调度策略，实现大批次下的稳定训练，避免发散或泛化性能差的问题。
• 证明商业规模硬件（128 张 V100）可在数小时内而非数月内完成大规模语言模型在公开 NLP 数据集上的训练。
• 通过在大规模文本语料上预训练一个稳健且可扩展的语言模型，实现下游 NLP 任务的迁移学习。

提出的方法

• 采用混合精度训练（FP16/FP32），相比 FP32，每张 GPU 的训练速度提升 4.2 倍。
• 利用 NCCL2、NVLink 和 InfiniBand，在 128 张 NVIDIA Tesla V100 GPU 上实现分布式数据并行通信，实现高带宽通信。
• 在所有 GPU 上应用 32k 批大小，相比单 GPU 训练，数据吞吐量提升 109 倍。
• 设计一种非线性的自定义学习率调度策略，避免发散，提升大批次下的收敛性能。
• 在 40GB 的 Amazon Reviews 数据集上，对字符级乘法 LSTM（mLSTM）进行 3 个周期的训练。
• 使用 Stanford Sentiment Treebank 和 IMDB 数据集，评估在情感分类任务上的迁移性能。

实验结果

研究问题

• RQ1 仅通过学习率缩放，能否使基于 RNN 的语言模型在 32k 批大小下实现稳定且收敛的训练？
• RQ2 混合精度训练在大规模 RNN 语言建模中如何影响收敛速度和模型质量？
• RQ3 何种学习率调度策略可实现大批次下的有效训练，且不降低泛化性能？
• RQ4 分布式数据并行性在 RNN 上的可扩展性如何？其在减少运行时间方面的极限是什么？
• RQ5 使用商用硬件能否在四小时内完成对 40GB 文本的大规模无监督预训练？

主要发现

• 使用 128 张 V100 GPU 和 32k 批大小，模型仅用四小时就在 40GB 的 Amazon Reviews 数据集上实现收敛。
• 混合精度训练使每张 GPU 的训练速度相比 FP32 提升 4.2 倍，显著提高吞吐量。
• 通过 128 张 GPU 的分布式数据并行性，数据吞吐量相比单 GPU 训练提升 109 倍。
• 模型在 Amazon Reviews 数据集上达到 1.038 的每字符比特数（BPC），在 SST 情感分类任务上准确率达 93.8%。
• 仅通过按批次大小缩放学习率会导致发散或收敛性能差；自定义学习率调度对稳定性至关重要。
• 即使在最大公开语料上，32k 批大小下仍不满足 $B \ll N$ 条件，表明大批次泛化存在局限性。