Determined AI：用免费的「分布式深度学习系统」加速NLP训练

  • 2020 年 12 月 15 日
  • AI

背景

在单个GPU上完成针对NLP任务的深度学习模型训练，一般需要很长时间才能完成。 在本文中，我们将利用Determined AI的分布式训练功能将SQuAD模型训练的BERT从数小时减少到数分钟，同时确保不会牺牲模型的准确性。

在这之前先介绍一下Determined AI，他们是一家很酷的startup，专门做深度学习的分布式学习和加速。他们的平台对个人使用者完全开源且免费，可以点击这里Star，fork和clone！

他们的开发者包括首席科学家，CMU的教授Ameet Talwalker，他开发的MLlib是Spark上的标准机器学习库。Determined AI主要有三大卖点，且支持TF，Torch和Keras。

  1. 第一是分布式训练，比如一卡多模型，多卡多模型等。
  2. 第二是参数寻优，包括网络结构搜索和寻优。
  3. 第三是对实验过程的保存和可视化，不用再担心模型跑了一半就丢失的问题。在高度优化的平台上，分布式训练可达到24倍的加速，而参数寻优只需要往常1/100的时间。

我将会写一系列文章来介绍这个平台的使用和特点，欢迎大家关注。最后的最后，求个GitHub Star啊老铁们~本文的代码示例 & 更加详细的介绍。

NLP与深度学习

在过去的几年中，像BERT和XLNet这样的神经语言模型引起了人们的广泛关注。原因是什么呢？常见NLP任务（例如实体识别和问题回答）的机器学习算法通常需要特征向量作为输入。通常，将自然语言输入表示为向量，从这一点出发我们就可以做各种NLP相关的任务。

但不能否认的是NLP训练特征向量时的开销可能会特别大，绝大部分的最先进的NLP模型训练成本都很高。 像BERT和XLNet这样的模型可以学习数亿个参数，而从头开始训练许多神经语言模型的计算成本对于大多数人来说都是过高的，对于我们这样的打工人就更是不可能完成的任务。 例如，XLNet需要在512 TPU v3上进行5.5天的培训。 在GCP上复制该代码的成本将超过50万美元。而NLP模型的复杂性和训练成本只会越来越高。比如OpenAI的GPT-3模型是在2020年5月推出的，总共有1,750亿个参数，大概的训练成本超过1200万美元 。

从一般的流程来看，我们一般都是拿预训练的模型，比如BERT作为起始点，然后用迁移学习的思路来后接我们需要的任务。一般来说前者的开销我们是没法负担的，后者虽然可能会很耗时，但一般不是特别复杂也能搞得定。

深度学习科研界的一个好处是这些复杂的预处理模型都是可以免费使用的。基本上选择TF或者Torch框架，再加上几行代码。在我们的案例中，我们将在斯坦福问题答案数据集（SQuAD）上训练BERT模型，以生成可以提取有关给定段落问题的答案的模型。 每个实验都需要不到20美元的计算成本来进行GCP训练，差不多能获得87-89％的F1，还是过得去的。

单卡训练

如果在单个GPU上为SQuAD训练BERT模型的话，我们使用了Hugging Face的Transformers库的预训练模型，代码看这里。 在单个GPU上训练大约2个epoch，在验证数据集上得出的F1分数为88.4％。

虽然这个结果不错，但是训练需要差不多7个小时， 等待结果的时间很长。而且这还是固定了超参数之后的结果。绝大部分NLP任务可能要比这个花的时间长的多。

在绝大部分场景下，往往我们可能有不止一块GPU，但如何用并行来加速训练其实并不那么容易。

用多卡进行高效分布式训练

在不改变代码的前提下，我们可以使用Determined AI系统来进行分布式训练。简单来说，我们可以吧训练工作负载分布在多个GPU（可能在多台计算机上）上，而无需更改代码。无论是在一台计算机上利用2个GPU还是在多台计算机上利用16个GPU，都只需更改配置即可。所以成本是很低的，详细可以参考这里：determined.ai/blog/stan

绝大部分时候需要调整的就像这样，以MNIST为例的话：

原始版本

def load_data():
    img_rows, img_cols = 28, 28
    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    x_train /= 255
    x_test /= 255
    return (x_train, y_train), (x_test, y_test)

Determined AI版本：

def build_training_data_loader(self) -> InputData:
        (train_images, train_labels), _ = mnist.load_data()
        train_images = train_images / 255.0

        return train_images, train_labels

    def build_validation_data_loader(self) -> InputData:
        _, (test_images, test_labels) = mnist.load_data()
        test_images = test_images / 255.0

        return test_images, test_labels

绝大部分情况下，就是稍微调整一下格式就成了~在这种前提下，我们可以在多卡上高效进行深度并行计算：

结果差不多是用2卡的时候节省时间约一半，用8卡的时候时间约为单卡的1/6左右。虽然不是完美的线性降低，但要考虑并行中communication成本，基本上scaling efficiency可以稳定在0.8左右。不仅如此，验证集的F1 score在所有试验中都是88%左右。

总结

深度学习，尤其是NLP相关的项目往往需要巨大的运算开销。如何简单的把单卡任务并行到多卡上，成为了提高训练效率的重点。在本文中，我们简单介绍了如何免费利用Determined AI平台来免费达成这个任务，从单卡的7小时训练降低到8卡的1小时训练。

感兴趣的朋友不妨看看 本文的代码示例 & 更加详细的介绍，也欢迎来试用Determined AI。下期我们会谈谈深度学习的自动调参问题~

VirMach 便宜 VPS

QNews

QNews