GPT-2 from scratch with torch用 torch 从零实现 GPT-2原文Posit AI Blog, GPT-2 from scratch with torch作者Sigrid KeydanaPosit发布日期2023-06-20授权说明原文页 “Reuse” 部分说明Text and figures are licensed under Creative Commons Attribution CC BY 4.0。原文同时说明来自其他来源的复用图片不在该许可范围内并会在图注中以 “Figure from …” 标识。本文为中文翻译保留原文结构、代码、引用关系和图注语义并按 CC BY 4.0 署名。从零实现一个语言模型可以说是形成准确直觉、理解其内部引擎如何工作的最佳方式。本文使用torch编写 GPT-2也就是原始 GPT 的直接继任者。最终你会得到一个 R 原生模型它可以直接使用 Hugging Face 上已经预训练好的 GPT-2 模型权重。目录来源与资源一个最小化 GPT-2端到端使用加载预训练权重无论你如何看待大语言模型LLM——它们有益吗危险吗会不会像加密货币一样只是短暂风潮——它们都已经存在而且就在当下。这意味着了解它们如何工作是一件好事至于需要了解到什么层次则取决于每个人自己的目标。同一天我还发布了 What are Large Language Models? What are they not?那篇文章面向更广泛的读者。本文则想面向深度学习实践者逐步走读一个用torch实现的 GPT-2Radford et al. 2019。GPT-2 是 OpenAI 一系列模型中的第二个模型越来越大训练语料也越来越庞大。你会看到一个完整的模型实现不到 250 行 R 代码。来源与资源我要展示的代码来自minhub仓库。这个仓库本身也值得单独介绍。正如 README 所强调的minhub是一组深度学习模型的最小实现灵感来自 minGPT。所有模型都设计成自包含、单文件、没有外部依赖因此很容易复制并集成到你自己的项目中。显然这让它们成为很好的学习材料但它的价值还不止于此。模型还提供了从 Hugging Face model hub 加载预训练权重的选项。更方便的是你也不必担心如何正确做分词直接从 Hugging Face 下载匹配的 tokenizer 即可。我会在本文最后一节展示它如何工作。正如minhubREADME 所说这些功能由hfhub和tok两个包提供。在minhub中gpt2.R 基本上是 Karpathy 的 MinGPT 的移植版本。实现时也参考了 Hugging Face 更复杂的 GPT-2 实现。如果你想看 Python 代码走读可以参考 https://amaarora.github.io/posts/2020-02-18-annotatedGPT2.html。那篇文章还整理了一批关于深度学习语言建模的博客和学习材料这些资料在短短几年内已经逐渐成为“经典”。一个最小化 GPT-2总体架构原始 TransformerVaswani et al. 2017由 encoder stack 和 decoder stack 两部分组成一个典型用例是机器翻译。后续模型会根据主要用途舍弃其中一侧结构。第一代 GPT 与 GPT-2 的差异主要是一些相对细微之处在架构上它只保留了 decoder stack。由于每个 decoder block 内部都接入了“自注意力”并且最开始还有 embedding 步骤因此这并不是问题外部输入与后续的内部表示在技术上没有本质区别。下面是原始 GPT 论文Radford and Narasimhan 2018中的总体架构截图这个结构对 GPT-2 仍然适用。token embedding 和 position embedding 之后是 12 次重复的 transformer block结构相同但不共享权重最终由一个依赖任务的线性层构成模型输出。在 gpt2.R 中这个全局结构及其行为由nn_gpt2_model()定义。代码被进一步模块化了所以不要因为代码和截图不完全一一对应而困惑。首先在initialize()中我们定义模块self$transformer-nn_module_dict(list(wtenn_embedding(vocab_size,n_embd),wpenn_embedding(max_pos,n_embd),dropnn_dropout(pdrop),hnn_sequential(!!!map(1:n_layer,\(x)nn_gpt2_transformer_block(n_embd,n_head,n_layer,max_pos,pdrop))),ln_fnn_layer_norm(n_embd,eps1e-5)))self$lm_head-nn_linear(n_embd,vocab_size,biasFALSE)这个模型最顶层的两个组件是transformer和lm_head后者是输出层。这种代码层面的区分有重要语义含义尤其体现在两个方面。第一transformer的定义用很简洁的方式传达了“什么构成了一个 Transformer”。之后接上的内容——在这里是lm_head——可以变化。第二更重要的是这个区分反映了深度学习自然语言处理中的基本思路或者说一种基本操作方式。学习分成两个步骤第一步也是不可或缺的一步是学习关于“语言”本身的知识这正是 LLM 做的事第二步资源消耗要小得多即适配到具体任务例如问答或文本摘要。为了看清事情发生的顺序以及每个模块发生多少次我们看看forward()tok_emb-self$transformer$wte(x)pos-torch_arange(1,x$size(2))$to(dtypelong)$unsqueeze(1)pos_emb-self$transformer$wpe(pos)x-self$transformer$drop(tok_embpos_emb)x-self$transformer$h(x)x-self$transformer$ln_f(x)x-self$lm_head(x)xtransformer中的所有模块都会被调用并执行一次这包括h。但h本身是一个顺序模块由多个 transformer block 构成。这些 block 是模型核心接下来我们就看它们。Transformer block在nn_gpt2_transformer_block()中12 个 block 中的每一个都是这样定义的self$ln_1-nn_layer_norm(n_embd,eps1e-5)self$attn-nn_gpt2_attention(n_embd,n_head,n_layer,max_pos,pdrop)self$ln_2-nn_layer_norm(n_embd,eps1e-5)self$mlp-nn_gpt2_mlp(n_embd,pdrop)在这个分辨率层次上我们看到自注意力会在每一层重新计算另一个构成性成分是前馈神经网络。此外这里还有两个执行layer normalization的模块。Layer normalization 是 transformer block 中使用的归一化类型。不同归一化算法的区别往往在于它们对哪些维度取平均。Layer normalizationBa, Kiros, and Hinton 2016可能会让一些读者意外它是对每个 batch item 做统计。也就是说对于一个模块中的每个单元都会有一个均值和一个标准差。其他维度例如图像中的空间维度和通道维度都会作为该 item 级统计计算的输入。继续放大细节我们很快会分别看注意力模块和前馈网络。不过在那之前需要先看这些层如何被调用。forward()中发生的事情只有这些x-xself$attn(self$ln_1(x))xself$mlp(self$ln_2(x))这两行值得仔细读。这里并不是简单地把前一层输出依次送入下一层而是在两个主要阶段外分别插入了 skip connection也称 residual connection。其效果是每个子模块并不替换传入的表示而是用自己的“视角”去更新它。近看 Transformer block自注意力在 GPT-2 的所有模块中自注意力看起来最吓人。但这里使用的基本算法仍然和 2014 年经典“点积注意力论文”Bahdanau, Cho, and Bengio 2014提出的思想一致注意力被概念化为相似度而相似度通过点积来度量。一个容易让人困惑的点是 self-attention 中的 “self”。这个术语首次出现在 Transformer 论文Vaswani et al. 2017中。那篇论文同时有 encoder stack 和 decoder stack。在那里“attention” 指 decoder block 如何决定关注 encoder 阶段传来的信息而 “self-attention” 是把这种技术用于 stack 内部时创造的术语也就是在一个 stack 内部的 block 之间使用注意力。到了 GPT-2只剩下现在显得有些重复命名的 self-attention。从上文继续看自注意力之所以显得复杂有两个原因。第一是 Transformer 通过 query-key-value 框架引入了 token 的“三重化”1。第二是多头注意力带来的额外批处理每一层不只运行一个注意力计算过程而是同时运行多个并行、独立的注意力头。走读代码时我会指出这两点分别在什么地方出现。我们还是从模块初始化开始。nn_gpt2_attention()这样列出它的组件# key, query, value projections for all heads, but in a batchself$c_attn-nn_linear(n_embd,3*n_embd)# output projectionself$c_proj-nn_linear(n_embd,n_embd)# regularizationself$attn_dropout-nn_dropout(pdrop)self$resid_dropout-nn_dropout(pdrop)# causal mask to ensure that attention is only applied to the left in the input sequenceself$bias-torch_ones(max_pos,max_pos)$bool()$tril()$view(c(1,1,max_pos,max_pos))|nn_buffer()除了两个 dropout 层我们还看到一个线性模块它执行上面提到的“三重化”。注意这不同于简单地复制三份相同的 token 表示。即使所有表示在初始时基本相同例如随机初始化后如此一旦模型开始训练它们也不会继续保持相同。一个名为c_proj的模块用于执行最终仿射变换。要理解它具体做什么需要看它如何被使用。一个buffer。buffer 是模块状态的一部分但不参与训练。这里的 buffer 确保注意力不会应用到“未来”的 previous-block output。基本做法是借助下三角矩阵把未来 token mask 掉。对于forward()我们把它拆成容易消化的小块来看。进入这个方法时参数x的形状正如语言模型中所预期的那样batch dimension × sequence length × embedding dimension。x$shape [1] 1 24 768接下来会发生两个批处理操作第一把表示三重化为 query、key、value第二腾出空间让指定数量的 attention head 可以一次性并行计算。先列出完整代码再解释。# batch size, sequence length, embedding dimensionality (n_embd)c(b,t,c)%-%x$shape# calculate query, key, values for all heads in batch and move head forward to be the batch dimc(q,k,v)%-%((self$c_attn(x)$split(self$n_embd,dim-1))|map(\(x)x$view(c(b,t,self$n_head,c/self$n_head)))|map(\(x)x$transpose(2,3)))首先调用self$c_attn()会为每个嵌入后的输入 token 生成 query、key 和 value 向量。split()把得到的矩阵拆成一个列表。然后map()负责第二个批处理操作三个矩阵都被 reshape并添加第四个维度这个第四维用来表示 attention head。注意这与前面把 embedding 三倍扩展的乘法过程不同。这里是把已有内容分配给不同的 head让每个 head 处理一个子集每个子集大小与使用的 head 数量成反比。最后map(\(x) x$transpose(2, 3))会交换 head 维度和 sequence-position 维度。接下来是注意力本身的计算。# causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) - (B, nh, T, T)att-q$matmul(k$transpose(-2,-1))*(1/sqrt(k$size(-1)))att-att$masked_fill(self$bias[,,1:t,1:t]0,-Inf)att-att$softmax(dim-1)att-self$attn_dropout(att)首先计算 query 和 key 之间的相似度矩阵乘法在这里实际上是批量点积。如果你对第一行最后的除法项感到好奇这个缩放操作是 GPT-2 与前代 GPT 有所不同的少数方面之一。感兴趣的话可以查看论文中相关考虑。接着应用前面提到的 mask对得到的分数做归一化并使用 dropout 正则化来鼓励稀疏性。最后计算出来的attention2需要传给后续层。这时 value 向量终于登场了——它是 query-key-value 三元组中我们还没看到发挥作用的成员。y-att$matmul(v)# (B, nh, T, T) x (B, nh, T, hs) - (B, nh, T, hs)y-y$transpose(2,3)$contiguous()$view(c(b,t,c))# re-assemble all head outputs side by side# output projectiony-self$resid_dropout(self$c_proj(y))y具体来说这里的矩阵乘法会用attention对 value 向量加权并求和。这一步会同时作用于所有 attention head也真正代表整个注意力算法的输出。剩余步骤会恢复原始输入大小。这包括把所有 head 的结果依次对齐然后应用线性层c_proj确保这些结果不会被等同、独立地对待而是以有用的方式组合起来。因此这里的 projection 实际由一个机械步骤view()和一个“智能”步骤c_proj()变换共同构成。近看 Transformer block前馈网络MLP与注意力模块相比transformer block 的第二个核心组件nn_gpt2_mlp()并没有太多神秘之处。它确实“只是”一个 MLP没有什么特殊技巧。不过仍有两点值得指出。第一你可能听说过 transformer block 中的 MLP 是 “position-wise” 的并想知道这是什么意思。考虑一个 block 中发生的事情x-xself$attn(self$ln_1(x))xself$mlp(self$ln_2(x))MLP 的输入几乎直接来自注意力模块。我们已经看到注意力模块返回的是形状为 [batch size,sequence length, embedding dimension] 的张量。在 MLP 内部也就是它的forward()中维度数量始终不变x|self$c_fc()|# nn_linear(n_embd, 4 * n_embd)self$act()|# nn_gelu(approximate tanh)self$c_proj()|# nn_linear(4 * n_embd, n_embd)self$dropout()# nn_dropout(pdrop)因此这些变换会独立地应用于序列中的每个元素。第二由于这里是激活函数唯一出现的地方可以顺便说一下。GeLU 表示 “Gaussian Error Linear Units”由 Hendrycks and Gimpel2020提出。它的想法是把类似 ReLU 的激活效果与正则化/随机性结合起来。理论上每个中间计算会根据它在高斯累积分布函数中的位置被加权也就是根据它比其他值大多少或小多少来加权。实践中正如模块实例化所显示的会使用一个近似形式。至此GPT-2 的主角——重复出现的 transformer block——就讲完了。还剩两件事之前发生了什么之后又发生了什么。从词到编码token embedding 和 position embedding严格说来如果你按要求对输入数据做分词使用 Hugging Face 中匹配的 tokenizer见下文最终得到的并不真的是“词”。但一个成熟事实仍然成立如果模型要成功提取语言知识就必须发生某种表示变换。和许多基于 Transformer 的模型一样GPT 家族用两种方式编码 token。第一种是 word embedding。回到最开始走读的顶层模块nn_gpt2_model()我们可以看到wtenn_embedding(vocab_size,n_embd)这已经很有用但由此产生的表示空间并不包含那些可能随序列位置变化的语义关系例如句法规则或短语语用。第二种编码就是用来弥补这一点的。它被称为 position embedding在nn_gpt2_model()中写作wpenn_embedding(max_pos,n_embd)又一个 embedding layer是的不过这次嵌入的不是 token而是一组预先指定的有效位置在 GPT 中范围是 1 到 1024。换句话说网络需要学习序列中的位置意味着什么。这也是不同模型差异很大的领域。原始 Transformer 使用的是一种正弦位置编码更新一些的改进可以在 GPT-NeoX 等模型中看到例如 RoPESu et al. 2021。一旦两种编码都可用它们会被直接相加见nn_gpt2_model()$forward()tok_emb-self$transformer$wte(x)pos-torch_arange(1,x$size(2))$to(dtypelong)$unsqueeze(1)pos_emb-self$transformer$wpe(pos)x-self$transformer$drop(tok_embpos_emb)得到的张量随后会传入一串 transformer block。输出应用完 transformer block 后最后一个映射由lm_head完成x-self$lm_head(x)# nn_linear(n_embd, vocab_size, bias FALSE)这是一个线性变换会把内部表示映射回离散词表索引并为每个索引分配一个分数。到这里模型本身的最终动作就结束了。接下来由样本生成过程决定如何使用这些分数。换句话说生成过程可以在几种成熟技术之间自由选择。下一节我们会看到一种相当标准的方法。至此模型走读结束。我略过了一些细节例如权重初始化如果你感兴趣可以查看 gpt2.R。端到端使用加载预训练权重很少有用户真的想从零开始训练 GPT-2。因此我们来看如何快速把它设置好用于样本生成。创建模型、加载权重、获得 tokenizerHugging Face model hub 允许你直接从 GPT-2 页面访问并下载所有必需文件包括 weights 和 tokenizer。所有文件都有版本管理这里使用最新版本。identifier-gpt2revision-e7da7f2# instantiate model and load Hugging Face weightsmodel-gpt2_from_pretrained(identifier,revision)# load matching tokenizertok-tok::tokenizer$from_pretrained(identifier)model$eval()分词decoder-only Transformer 类型的模型并不需要 prompt。但通常应用都会想把输入传给生成过程。有了tok对输入做分词就非常方便idx-torch_tensor(tok$encode(paste(No duty is imposed on the rich, rights of the poor is a hollow phrase...),Enough languishing in custody. Equality))$ids)$view(c(1,-1))idx输出torch_tensor Columns 1 to 11 2949 7077 318 10893 319 262 5527 11 2489 286 262 Columns 12 to 22 3595 318 257 20596 9546 2644 31779 2786 3929 287 10804 Columns 23 to 24 13 31428 [ CPULongType{1,24} ]生成样本样本生成是一个迭代过程模型最后一次预测会被追加到不断增长的 prompt 后面。prompt_length-idx$size(-1)for(iin1:30){# decide on maximal length of output sequence# obtain next prediction (raw score)with_no_grad({logits-model(idx1L)})last_logits-logits[,-1,]# pick highest scores (how many is up to you)c(prob,ind)%-%last_logits$topk(50)last_logits-torch_full_like(last_logits,-Inf)$scatter_(-1,ind,prob)# convert to probabilitiesprobs-nnf_softmax(last_logits,dim-1)# probabilistic samplingid_next-torch_multinomial(probs,num_samples1)-1L# stop if end of sequence predictedif(id_next$item()0){break}# append prediction to promptidx-torch_cat(list(idx,id_next),dim2)}要查看输出只需使用tok$decode()tok$decode(as.integer(idx))输出[1] No duty is imposed on the rich, rights of the poor is a hollow phrase... Enough languishing in custody. Equality is over如果想尝试文本生成只要复制这个自包含文件然后试试不同的采样参数即可。当然也可以换不同的 prompt。一如既往感谢阅读封面照片来自 Marjan Blan on Unsplash。参考文献Ba, Jimmy Lei, Jamie Ryan Kiros, and Geoffrey E. Hinton. 2016. “Layer Normalization.” https://arxiv.org/abs/1607.06450.Bahdanau, Dzmitry, Kyunghyun Cho, and Yoshua Bengio. 2014. “Neural Machine Translation by Jointly Learning to Align and Translate.” http://arxiv.org/abs/1409.0473.Hendrycks, Dan, and Kevin Gimpel. 2020. “Gaussian Error Linear Units (GELUs).” https://arxiv.org/abs/1606.08415.Radford, Alec, and Karthik Narasimhan. 2018. “Improving Language Understanding by Generative Pre-Training.”Radford, Alec, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019. “Language Models Are Unsupervised Multitask Learners.”Su, Jianlin, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. 2021. “RoFormer: Enhanced Transformer with Rotary Position Embedding.”arXiv Preprint arXiv:2104.09864.Vaswani, Ashish, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. “Attention Is All You Need.” https://arxiv.org/abs/1706.03762.复用与引用原文复用说明Text and figures are licensed under Creative Commons Attribution CC BY 4.0。来自其他来源的复用图片不在该许可范围内并会在图注中以 “Figure from …” 标识。原文建议引用Keydana (2023, June 20). Posit AI Blog: GPT-2 from scratch with torch. Retrieved from https://blogs.rstudio.com/tensorflow/posts/2023-06-20-gpt2-torch/BibTeXmisc{keydanagpt2, author {Keydana, Sigrid}, title {Posit AI Blog: GPT-2 from scratch with torch}, url {https://blogs.rstudio.com/tensorflow/posts/2023-06-20-gpt2-torch/}, year {2023} }如果你不熟悉这套术语可以参考 Jay Alammar 的经典入门文章The Illustrated Transformer。 ↩︎这里用斜体是为了提示这个词有一种特殊用法。虽然 “attention” 这个表达本身听起来有些奇怪但它经常被用来表示对原始 scores 做归一化之后得到的状态。 ↩︎