在家复制GPT-2

几个月前，我开始从事一个研究项目，试图从头开始训练自己的，更有效的语言模型。我可以从Tensorflow Reseach Cloud访问128核TPUv3 Pod，并使用它来预训练一个价格为124 $ M的参数GPT-2模型，使其困惑度接近OpenAI的结果（我的预训练模型被训练了约1/8 $ th OpenAI为其训练模型并在OpenWebText上获得$ 21 $ ppl，而OpenAI的模型为$ 17 $ ppl的迭代次数），然后使用一种语言来预先训练ALBERT风格的GPT-2（我称之为ALGPT2）语言模型分解的输入嵌入和逐层参数共享，这将使模型中的参数数量从$ 124 $ M减少到约$ 12 $ M。

不幸的是，ALGPT-2的表现不如GPT-2（ALGPT-2在OpenWebText上获得了31 $ ppl，而我的训练有素的GPT-2模型则为21 $ ppl），但是我在写这一系列博文来回顾过去几个月我学到的一切。

我想在今年春天自己进行的这类“研究项目”中要做的主要事情就是开发和训练价格更高的参数为GPT-2的$ 124 $ M版本。我想预训练参数版本$ 1.5 $ B的GPT-2，但是由于我只能使用TPU吊舱一个星期，所以我不得不选择可以及时训练的模型。 $ 100 $ k的迭代训练运行时间约为$ 20 $，这使我有足够的时间运行多个实验。相比之下，完全按照OpenAI的培训程序进行培训，并进行完整的$ 800 $ k迭代培训将占用几乎整个一周的时间，并用完了我的大部分配额。

通过预训练GPT-2的$ 124 $ M参数版本，使其几乎接近OpenAI的结果，我几乎可以复制它（我使用的预训练模型的训练量约为OpenAI训练其迭代次数的$ 1/8 $ th）模型，并在标准OpenWebText数据集上获得了21 $的困惑度（ppl），而OpenAI的模型则为$ 17 $ ppl），

我做一个效率更高的变压器的想法并没有真正实现，因为我的训练有素的变压器最终比同等的GPT-2模型差20美元/ ppl，但是我想写下两三个月的经验无论如何都在努力。

关于我自己：我是滑铁卢大学的一名即将毕业的软件工程专业的学生，​​这篇文章应该是我在2020年3月至5月左右从事的NLP研究项目的写照（中间是在2020年首次Covid-19锁定中，我目前正在2020年7月15日撰写这份报告，同时等待我的Pix2PixHD模型在colab上为我正在进行的新项目训练数百个纪元）。

我很幸运能在变压器大行其道之前就开始学习NLP，我还记得当word2vec和LSTM在许多NLP任务上成为SOTA时，看到NLP领域在仅仅一个数年之久，从只有少数几个层且单位价格在512美元左右的LSTM被认为是大型网络和训练成本高昂，到训练顶部具有关注层的LSTM到原始的变压器编码器/解码器网络，到ULMFIT和ELMO，然后是BERT，RoBERTa，GPT-2和T5，到几个月前，诸如Sparse变压器，Reformer和Synthesizers的新型，更高效的自我关注替代品爆炸了，并且现在是GPT-3，IMO有潜力真正改变NLP的整个领域。

就在几年前，我们在数据集上训练了浅层递归网络，然后在大型数据集上对大型转换器语言模型进行了预训练，并在针对特定任务的数据集上进行了微调。现在的整个想法是，在一个巨大的数据集上训练一个巨大的语言模型，然后通过将某些特定任务的示例放在输入之前，以几次学习的形式对模型进行调节，感觉好像它真的可以使NLP模型变得很多与当今相比，更易于访问且更易于生产，并使人机对话更加真实。

我已经徘徊了很长时间，让我们开始这篇文章的主题。

模型根部的嵌入层将给定令牌索引的单热点向量（所有GPT-2模型使用的词汇量为$ 50257 $）映射到$ 768 $维度向量（此博客中的所有GPT-2数字） post将针对GPT-2的$ 124 $ m参数版本）。

嵌入矩阵后面是一堆自我注意和前馈层，每层输出一个768维矢量（保持每一层的输出数量恒定），该矢量构成了变压器的主要部分。

然后，在自我注意层堆栈之后进行输出嵌入（将输入和输出嵌入的权重捆绑在一起，使训练更容易），该映射将$ 768 $维向量映射为变压器最后一层的输出到相同的$ 50257 $维向量，表示词汇表中每个标记是序列中下一个标记的概率。

ALBERT（精简BERT）是一本论文，着眼于BERT，并通过以下四种主要方法确定了提高效率和减少模型中参数数量的一些方法：因式分解，分层参数共享，句子顺序预测辅助损失，并消除辍学现象。

GPT-2的嵌入具有很多参数。它实际上只是尺寸为50257美元乘以768美元的矩阵。这意味着单独的输入嵌入将消耗几乎$ 50257 \ times 768 = \ space \ sim 38,000,000 $参数，这是模型中$ 128 $ M总参数中的很大一部分。

ALBERT作者提出了一种中间分解大小为$ 128 $的分解式嵌入：一种嵌入大小为$ 50257 \乘以128 $，另一种嵌入大小为$ 128 \乘以768 $。通过将较大的嵌入矩阵分解为两个较小的矩阵，嵌入中使用的参数总数从大约$ 38 $ M到大约$ 6 $ M。

作者尝试使用不同的中间嵌入大小，并在$ 128 $上作为参数和性能之间的良好折衷。

类BERT（nn。Module）：def __init__（self，n_layers）：super（）。 __在自身 。块= nn。 ModuleList（[[_的范围（n_layers）中的_的Block（）]））// ... def forward（self，x）：// ...表示self中的block。块：x =块（x）// ...

类ALBERT（nn。Module）：def __init__（self，n_layers）：super（）。 __在自身 。 n_layers = n_layers个self。 block = Block（）// ... def forward（self，x）：// ... for _在self中。 n_layers：x =块（x）// ...

通过仅定义一个变压器块并在其上循环n_layers次，ALBERT节省了GPU内存，该内存将用于存储所有图层的参数。

由于我们通常使用$ 32 $位浮点数将参数存储在GPU上，因此将$ 1.5 $ B参数GPT-2存储在GPU上将消耗约$ 6 $ GB的GPU内存-这是$ 16 $ GB的很大一部分考虑到存储模型激活所需的内存以及优化器所需的任何动量参数，普通V100 GPU上的内存已被耗尽。相反，如果您在$ 1.5 $ B参数GPT-2中的所有转换器层之间共享参数，则生成的模型将仅具有$ 37 $ M参数，而参数共享版本将仅使用约$ 148 $ MB的GPU内存。

作者尝试将参数共享应用于BERT，并发现它会降低性能，但使训练越来越大的模型变得更加容易。

在像JAX这样的机器学习框架中，默认情况下，在使用XLA编译代码时，该框架会展开并内联循环，展开和内联循环的大小会使计算图真正变大，并且编译时间很长。因此，建议您在这种情况下使用lax.scan（）之类的功能。

ALBERT作者添加了辅助损失以帮助训练。由于语言建模通常是自动回归的，因此我没有将其用于自定义模型。

ALBERT作者从BERT中删除了所有遗漏，并发现它显着提高了性能。

这几乎就是我的想法：以GPT-2，添加分解式嵌入，在所有转换器层之间共享参数，删除辍学（我实际上错过了有关ALBERT删除辍学的部分，直到我工作很远，但我确实运行了一两次运行而不会出现辍学，以了解其工作原理），并在大型数据集上进行预训练数十万次迭代。

我无法自己对GPT-2之类的东西进行预训练，因此我将其应用于Tensorflow研究云（TFRC）。

TFRC强调要帮助来自非传统背景的研究人员，这对于那些不是“传统”机器学习研究人员的人来说都是一个了不起的资源。他们愿意给我一个17岁，没有正规教育或证书（甚至没有高中文凭：/）的我免费使用功能强大的TPU集群的机会。能够参与该计划对我真的很有帮助，尤其是因为我没有专用的GPU，并且通常依靠Colab的免费GPU来训练我的模型。

我给TFRC小组发了电子邮件，询问我是否可以从5美元的独立TPUv3（每个具有8个内核）升级到TPU吊舱以预训练大型语言模型。第二天（！）我收到一封电子邮件，说我可以使用可抢先的128核TPUv3 Pod持续7天，很不幸，它不足以让我预训练$ 1.5 $ B参数模型，但足以在$ 124 $ M的模型上进行一些训练。

因此，对于设置，我将完成设置虚拟机和TPU Pod以及预处理数据集所需的所有步骤。

当我从事这个项目时，我设置了两个虚拟机。一个具有大量RAM和CPU内核以快速处理数据，另一个具有运行TPU训练脚本的小型实例。关于TPU和TPU单元训练的好处之一是，只要您的数据已作为一组TFRecord文件进行了预处理，就不需要真正强大的VM实例，它可以为您节省很多金钱/计算额度。

您可以查看有关我用来设置VM和预处理数据集的每个命令的完整列表。

我使用TF2.1的n-1-standard-16实例来处理OpenWebText数据集。确保将实例与SSD而不是默认HDD一起使用，因为处理数据集涉及处理很多非常小的文本文件，并且主要受驱动器io速度的限制。我犯了一个使用HDD的错误，仅提取数据集的TAR存档大约需要7个小时。我将所有数据放在〜/ data / openwebtext /的文件夹中，因此，如果您想将数据下载到其他地方，请对其进行修改。

TIL：大多数常见的linux实用工具（如ls，mv和cat）实际上并未针对在OpenWebText中处理近1000万个文件进行过优化。仅计算数据集中的文本文件数可能需要几分钟。

下载OpenWebText数据集（实际上只是一堆包含许多文本文件的tar归档文件的tar归档文件）并解压缩它：

数据集压缩后约为12GB，未压缩时约为53GB，仅包含约800万个文本文件。

我将数据集中的前$ 100,000 $文件移动到单独的目录中以创建验证集：

我在训练集的$ 1 $ M文件子集中训练了一个词汇量为$ 50,257 $（与GPT-2相同）的字节级BPE令牌生成器（我不确定GPT-2是否在整个培训中都对令牌生成器进行了培训数据集或仅在一个子集上，但我知道CTRL纸会在其训练集的5％划分上训练其标记器。）我使用了Hugginface的快速基于Rust的令牌生成器库及其ByteLevelBPETokenizer令牌生成器。

来训练令牌生成器，或者只是看一下它的主要细节（它只是训练令牌生成器并将其以及配置文件保存到磁盘）：

import os import glob从tokenizers导入json import ByteLevelBPETokenizer path = glob。 glob（os。path。join（＆＃39; ./ data / openwebtext＆＃39;，＆＃39; *＆＃39;））[：1000000] tok = ByteLevelBPETokenizer（）tok。训练（文件=路径，vocab_size = args。vocab_size，special_tokens = args。control_codes）。保存（＆＃39; ./ tokenizer /＆＃39;）tokenizer_config = {＆＃34; max_len＆＃34; ：1024}，并以fp：json的格式打开（操作系统路径。 。转储（tokenizer_config，fp）

TPU Pods希望您的数据可以在GCP云存储桶中以一组TFRecord文件的形式获得，这些文件将下载到TPU板中每个内置的强大VM中，该虚拟机将负责反序列化文件并将其馈送到TPU芯片。确保您的GCP存储桶和TPU插盒位于同一计算区域中，否则，通过在计算区域之间传输数百GB的数据，您将很快承担很多费用。

在使用TPU Pod时，这件事没有得到很好的记录（这实际上不适用于TPU）：TPU Pod会创建大量（100 GB）的日志，这些日志会发送到Stackdriver，在其中收取大约50的费用每个摄取的GiB日志超出特定限制的美分（我认为大约是每月50GiB）。在短短几天的培训中，我最终被收取了约100美元的IIRC费用。幸运的是，我仍然拥有大部分的免费GCP积分，因此这对我来说并不是一个大问题，但请务必关闭提取TPU的日志。

当我可以访问TPU吊舱时，我遇到了一个问题，即我的代码可以在单个TPU上正常运行，但是会抛出超出范围的错误：在TPU吊舱上运行时出现序列结束错误。我为此苦苦思索了很长时间，直到看了一下Kaggle的讨论贴，说TPU希望每个TPU板（8个核）都获得自己的TFrecord文件（到那时，我将火车分成8组）我应该将TFRecord文件拆分为16个（每个板128个核心/ 8个核心）的TFRecord文件。

TPU非常适合扩展到庞大的模型和庞大的数据集，但是您需要了解很多TPU特定的信息（尤其是针对TPU Pod的信息），这些信息并未包含在文档中并且不容易找到。_**

python3 make_tfrecords.py --path ./data/openwebtext/ --save_path ./train/ --files_per_tfrecord 500000 \ --use_control_codes --seq_len 1024 --min_seq_len --tokenizer ./tokenizer/

python3 make_tfrecords.py --path ./data/openwebtext-valid/ --save_path ./val/ --files_per_tfrecord 50000 \ --use_control_codes --seq_len 1024 --min_seq_len --tokenizer ./tokenizer/

从火车转换原始文本文件并进行验证，将其分成两组$ 16 $ TFRecord文件。

我对数据集中文本字段的平均长度进行了快速分析，$ 67 $％的文件的令牌少于$ 1024 $，$ 35 $％的文件的令牌少于$ 512 $，只有$ 10 $％的文件的令牌少于$ 256 $代币。这意味着，如果我想使数据集尽可能整洁，并使模型的每个输入序列都包含一个连续的$ 1024 $令牌流，则数据集的大小会小很多。因此，每个人都准备了一个＆lt; | endoftext |＆gt;这样的令牌。到每个序列的开头，并将长度小于$ 1024 $的序列连接在一起。具体操作方法的细节（例如，您将数据集视为单个令牌流，然后将其分解为长度为$ 1024 $的序列，还是跟踪小于$ 1024 $的序列，然后将它们串联在一起形成一个单序列）确实不应该对模型的性能产生太大的影响，但是您可以在此处查看我的实现。

我的版本没有充分利用快速，多线程的batch_encode_plus（）方式并行对大型数据集进行令牌化，因为它只在文件的每一行中保留第一个context_len令牌，这使得处理带有或少于$ 1024 $令牌的文件成为可能更难。因此，对数据集进行标记化大约需要花费$ 8 $小时，这是我想要改进的地方。

火车设置为大约$ 26 $ GB，包含约$ 8 $ M个文本文件，这些文本文件已转换为将近$ 7 $ M tfrecord个示例，每个示例带有$ 1024 $令牌（与GPT-2相同）。验证集大约为$ 300 $ MB，包括大约$ 100 $ K个文本文件，这些文本文件已转换为大约$ 90 $ K tfrecord个示例，每个示例都带有$ 1024 $令牌（也与GPT-2相同）。

由于我使用的是TPU，因此您现在可以实际使用的唯一真正的库就是Tensorflow。我不想经历学习如何在TF2中制作自定义训练循环和内容的学习过程，所以我只是坚持使用Keras。您可以在这里查看我的训练脚本（非常简短）。这很简单，因此我不会复制整个培训脚本，但是我将讨论一些小代码段。

我通常喜欢在脚本中添加一个ptvsd断点，以便在将其推送到我的VM之前可以使用vscode在本地调试培训脚本。

通常，当您在Keras上使用TPU时，会将TPU的IP地址和端口传递给TPUClusterResolver，但是当使用TPU Pod时，会将TPU本身的名称传递给解析器。

当我复制GPT-2的$ 124 $ M参数版本时，我尝试使用了OpenAI所使用的许多原始超参数，但是我必须修改一些内容，以便可以及时进行训练。

注意：出于某些原因，GPT-2论文的作者没有确切说明他们用于训练模型的学习率，而是指出“手动调整每个模型的学习率以获得5％的最佳困惑度”。 WebText的示例”。

OpenAI以512美元的批量大小训练他们的模型，总共进行了800美元的K迭代（通过训练集得出总共约60美元的时期）。

我为GPT-2模型训练的价格为OpenAI为其训练的迭代次数的$ 1/8（总共约$ 100 $ K迭代），因为在我的128-核心TPU Pod。如果我想训练GPT-2进行与OpenAI相同的迭代次数，那么一次训练运行将耗尽我一周访问Pod的大部分时间。

由于我的TPU吊舱是可抢占的，并且每隔$ 24小时会重置一次，因此我通常必须至少恢复一次训练运行，这就是为什么所有这些图表通常都进行两次或更多次训练的原因。

因此，这是我的模型，非常接近复制GPT-2，在将近$ 90 $ K的训练迭代结束时，训练的困惑度约为$ 21.5 $。为了进行比较，GPT-2在约$ 800 $ K的训练迭代后获得了约$ 17.5 $ ppl的训练困惑，因此相差仅约$ 4 $ ppl。

我制作了一个colab笔记本，展示了如何使用预训练的GPT-2模型生成文本

我想使用节省内存的Adafactor优化器来简化大型语言模型的训练，但是我所有的Adafactor训练运行都比使用AdamW差很多（〜5ppl IIRC）（这可能是由于未使用Adafactor的动量参数或相对更新规模，因此我希望尽快对此进行调查）。

我最初使用的是Adam的默认学习率$ 1e-4 $，但是很快我发现可以使用更高的学习率（如$ 1e-3 $）来更快地训练模型。

GPT-3的第二部分列出了OpenAI团队在训练GPT-3时用于不同尺寸模型的学习率。他们对$ 124 $ M版本的模型使用$ 6e-4 $的学习率，并随模型大小降低学习率。

您可以看一下部分培训，以了解不同学习率的培训之间的区别。

自从我用 ......