在Graphcore IPU上训练稀疏的大型语言模型

2021/07/08

在Graphcore IPU上训练稀疏的大型语言模型

作者:

稀疏神经网络研究推动创新的良性循环。算法改进使软件库能够有效地表示稀疏模式。然后，程序员可以使用这些库最大化地利用可用硬件，并因此减少计算要求。同时，硬件创新同样支持更高效、更灵活的软件库，使研究人员能够探索新的算法创新。

一段时间以来，这一创新周期在各个层面都取得了许多令人振奋的进展。例如Lottery Ticket Hypothesis，也就是彩票假设（Frankle和Carbin，2018年；Frankle等人，2019年；Ramanujan等人，2019年；Pensia等人，2020年）强调了算法洞察力，即一个已训练模型的一小部分权重子集足以高度准确地表示学习到的函数（Han等人，2015年）。在这个想法的推动下，人们提出了新的训练方法来“操纵彩票”，即在不产生密集训练的计算成本的情况下，找到高准确性的稀疏子网络（Evci等人，2020年；Jayakumar等人，2020年）。

除了这项工作之外，研究人员还付出了大量的努力来扩展主要机器学习框架对稀疏计算的软件支持。OpenAI推动将其块稀疏工具箱集成到PyTorch中就是这种改进支持的一个例子。

硬件层面也在进行重大创新。专门为支持稀疏计算而设计的新加速器，例如Graphcore的智能处理器（IPU)，正被行业和学术界的创新者广泛使用。

考虑到最近人工智能模型尺寸上的爆炸式增长，所有这些进展都是非常及时的。像GPT3这样当前广泛使用的大型、基于Transformer的语言模型是最先进架构中模型尺寸不断增加的一个相关例子。虽然这些越来越大的模型带来了持续的准确性改进，但训练和推理对计算的要求也变高了很多。在这种情况下，改进后的稀疏方法在限制成本和能源消耗方面具有巨大潜力，有助于模型的普及和减轻高功耗引起的环境问题。

在我们的案例研究中，正如在SNN Workshop 2021上展示的那样，我们在IPU处理器上使用可扩展语言模型预训练的具体示例来探索稀疏计算的潜力。具体而言，我们实施了BERT语言模型的强大的稀疏预训练基线，并在浮点运算（FLOP）基础上进行比较时证明了其和密集模型相比的优势。除了对本案例研究中涉及的硬件、软件和算法的各个方面进行的探索之外，我们希望我们的发现可以作为该领域进一步研究的切实起点。

现在让我们重点介绍我们研讨会贡献的一些主要发现。我们将从IPU的概述以及使其非常适合于稀疏计算的特性开始。然后，我们将介绍用于实施稀疏算法的开源软件库。最后，我们将讨论一些实验结果。这些结果揭示了可以助力有效执行的一些算法特性。

硬件：IPU架构和稀疏计算

Graphcore IPU是一种新颖的大规模并行处理器，专为机器学习工作负载而设计。第二代芯片拥有1472个完全独立的处理器内核（Tile），每个内核有6个线程，可以同时运行8832条独立的并发指令（多指令多数据，MIMD）。值得注意的是，每个tile都有自己的处理器内存储块，每张卡可以以每秒180 TB的速度访问，这是片外DRAM存储访问带宽的100多倍。IPU中计算和存储的集中不仅使其更节能，而且还意味着它非常适合高度不规则的计算，这些计算必须使用细粒度、灵活的方法进行并行化。

由于IPU是专为大量并行而设计的，因此非常适合稀疏工作负载。网络拓扑的算法稀疏化通常会导致具有细粒度的、高度不规则的计算图，因此IPU的一个优势就是，即使是最小的计算结构，IPU也能够有效并行化。

即使稀疏结构的粒度很高，IPU架构也可以保持稀疏训练和推理的计算效率。尺寸为16的乘法累加单元使得IPU可以实现良好的线程利用率和性能，即使对于1×1、4×4、8×8或16×16这样非常小的块尺寸也是如此。IPU还支持许多旨在加速常见稀疏原语计算的指令，这有助于进一步提高吞吐量。

软件：Poplar和popsparse软件库

所有硬件和架构特性都可以通过优化后的内核来定向，这些内核可以通过popsparse获取。popsparse是一个用于稀疏计算的灵活的低级别C++库。popsparse是Graphcore的开源软件库（PopLibs）之一，它建立在Poplar之上。Poplar是Graphcore的计算图构建和编译API，支持在IPU上运行的所有内容。

稀疏内核可以通过PopART API直接定向，也可以通过本研究中使用的TensorFlow和Keras层轻松集成到现有的TensorFlow模型中。支持FP16和FP32中块尺寸为1×1、4×4、8×8、16×16的稀疏性。Graphcore在GitHub上的公共示例中提供开源实施和教程。

Poplar软件栈支持静态和动态稀疏，具体取决于算法的稀疏模式是固定的还是随时间演变的。使用静态稀疏可以在编译时增加优化，这有助于提高执行效率，因此静态稀疏非常适合涉及已经训练的模型的推理工作负载。然而，许多训练算法在训练过程中缓慢地演化算法的稀疏模式；在这种情况下，动态稀疏可以提供更大的灵活性。

为了简化动态稀疏技术的实施，有一系列实用程序支持在将稀疏模式流式传输到IPU之前，在主机上轻松对其进行操作。用户可以选择最适合其实施的稀疏模式，然后将其转换为IPU内部基于CSR的稀疏格式。这使得将稀疏性集成到现有代码中并在Python中表达模式操作变得容易。

从RNN架构到GPT2的稀疏模型的完整模型实施，以及诸如Rig-L和稀疏注意力的静态和动态稀疏优化技术都是开源的。

算法：使用动态稀疏操纵彩票（The Lottery）

IPU的Poplar软件库提供了实施最先进算法稀疏训练所需要的一切。然而，对于像BERT这样的现有密集模型，哪种算法途径最好并非显而易见的。

传统上，剪枝技术是一种简单且令人惊奇的有效方法，可以加快已训练模型的推理速度。在训练期间实现类似的计算优势通常更具挑战性。一种常见的技术是从稀疏拓扑开始训练，并随着时间的推移不断演进连接模式（Evci等人，2020年；Jayakumar等人，2020年）。使用这种方法，通过剪枝和重新增长当前子网络，稀疏连接被定期修改。这意味着可以有效地探索原始过度参数化模型的可用自由度，同时保留降低计算成本的好处。

自然，这种方法的有效性高度依赖于特定的剪枝和重新分配标准。选项范围从简单的基于幅度的剪枝（Magnitude-Based Pruning）和随机重新分配到更复杂的、基于梯度的方法。

图1 | 左图：动态稀疏训练期间参数空间探索的图示。已训练网络参数的非零子集是使用参数的剪枝和重新分配来实现进化的。右图：动态稀疏训练期间剪枝率的余弦衰减以及已探索和已训练的参数比率的演变。

在这里，我们选择使用基于幅度的剪枝和随机重新增长来实施一个简单但经过良好调整的动态稀疏基线，用于1×1、4×4、8×8和16×16的不同块尺寸。具体而言，在BERT语言建模任务（Devlin等人，2018年）的基础上，我们以0.9的稀疏率对全连接编码器权重（不包括嵌入权重）进行稀疏化，这样一来这些层中仅保留10%的权重。通过对一定比例的权重（剪枝率为0.5，余弦随时间衰减，参见图1）进行基于幅度的剪枝和随机重新增长，我们定期更新稀疏模式（在预训练的第一阶段160次）。

然后我们调整学习率以产生最佳的任务性能。通过使用这种方法，我们将不同的模型尺寸稀疏化到BERT Base。图2显示了在阶段I之后实现的掩码语言模型（Masked Language Model）损失，对应于预测掩码标记的随机子集的验证性能。当与FLOP预算进行比较时，动态稀疏已训练模型在所有模型尺寸上都优于密集模型。