通过nanoGPT深入挖掘生成式预训练Transformer
欢迎来到我的项目的第二部分，我将通过使用TinyStories数据集和在一台老旧的游戏笔记本电脑上训练的nanoGPT深入研究Transformer和基于GPT的模型的复杂性。在第一部分中，我为输入到字符级生成模型的数据集做了准备。您可以在下面找到第一部分的链接。
学习Transformer的代码 第一部分
一系列的第一部分，我将努力通过使用nanoGPT首先学习Transformer的代码
towardsdatascience.com
在本文中，我旨在对GPT模型及其组成部分以及其在nanoGPT中的实现进行详细分析。我选择了nanoGPT，因为它具有简单直接的Python实现GPT模型的代码，大约只有300行，以及类似易于理解的训练脚本。具备必要的背景知识后，通过阅读源代码，人们可以快速理解GPT模型。坦率地说，当我第一次审查代码时，我缺乏这种理解。其中一些材料仍然使我困惑。然而，我希望通过我所学到的一切，这篇解释能够为那些希望在内部获得GPT样式模型工作原理直观理解的人提供一个起点。
为了准备这篇文章，我阅读了各种论文。最初，我以为只阅读开创性论文“Attention is All You Need”就足以让我对LLM有所了解。这是一个天真的假设。虽然这篇论文确实介绍了Transformer模型，但是后续的论文将其适应到了更高级的任务，如文本生成。 “AIAYN”只是一个更广泛主题的简介。尽管如此，我回想起HackerNews上提供了一个阅读清单，以全面了解LLMs。经过快速搜索，我在这里找到了这篇文章。我没有按顺序阅读所有内容，但我打算在完成本系列后重新阅读这个阅读清单继续我的学习之旅。
话虽如此，让我们深入了解。为了详细理解GPT模型，我们必须从Transformer开始。Transformer采用了称为缩放点积注意力的自注意机制。下面的解释源自这篇关于缩放点积注意力的有见地的文章，我建议您阅读以获得更深入的理解。基本上，对于输入序列的每个元素（第i个元素），我们希望将输入序列与序列中的所有元素的加权平均值乘以第i个元素。这些权重是通过将第i个元素的向量与整个输入向量进行点积，然后应用softmax函数计算得出的，使得权重的取值介于0到1之间。在原始的“Attention is All You Need”论文中，这些输入被命名为查询（整个序列），键（第i个元素的向量）和值（也是整个序列）。传递给注意机制的权重被初始化为随机值，并在神经网络中的进一步传递中进行学习。
nanoGPT实现了缩放点积注意力，并将其扩展为多头注意力，即同时进行多个注意力操作。它还将其实现为torch.nn.Module，这使得它可以与其他网络层组合使用。
import torchimport torch.nn as nnfrom torch.nn import functional as Fclass CausalSelfAttention(nn.Module): def __init__(self, config): super().__init__() assert config.n_embd % config.n_head == 0 # 所有头的键、查询和值投影，但在一个批次中 self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias) # 输出投影 self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias) # 正则化 self.attn_dropout = nn.Dropout(config.dropout) self.resid_dropout = nn.Dropout(config.dropout) self.n_head = config.n_head self.n_embd = config.n_embd self.dropout = config.dropout # 闪电注意力让GPU发疯，但仅在PyTorch >= 2.0中支持 self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') if not self.flash: print("警告：使用缓慢的注意力。闪电注意力需要PyTorch >= 2.0") # 因果掩码以确保注意力仅应用于输入序列中的左侧 self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size)) .view(1, 1, config.block_size, config.block_size)) def forward(self, x): B, T, C = x.size() # 批量大小，序列长度，嵌入维度（n_embd） # 为批量中的所有头计算查询、键和值，并将头部前移为批次维度 q, k, v = self.c_attn(x).split(self.n_embd, dim=2) k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) # 因果自注意力；自注意：(B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T) if self.flash: # 使用Flash Attention CUDA内核进行高效注意力计算 y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True) else: # 注意力的手动实现 att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf')) att = F.softmax(att, dim=-1) att = self.attn_dropout(att) y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs) y = y.transpose(1, 2).contiguous().view(B, T, C) # 重新组装所有头部输出 # 输出投影 y = self.resid_dropout(self.c_proj(y)) return y让我们进一步解析这段代码，从构造函数开始。首先，我们验证注意力头的数量（n_heads）是否能够均匀地整除嵌入的维度（n_embed）。这一点很重要，因为当将嵌入划分为每个头的部分时，我们希望覆盖整个嵌入空间而没有任何间隙。接下来，我们初始化两个线性层，c_att和c_proj：其中c_att是保存由缩放点积注意力计算组成的矩阵工作空间的层，而c_proj则存储计算结果的最终值。在c_att中，嵌入维度被扩大了三倍，因为我们需要为注意力的三个主要组件（查询，键和值）留出空间。
我们还有两个Dropout层，attn_dropout和resid_dropout。Dropout层根据给定的概率随机将输入矩阵的元素置零。根据PyTorch文档的说法，这有助于减少模型的过拟合。在config.dropout中的值是在Dropout层中丢弃给定样本的概率。
我们通过验证用户是否可以访问PyTorch 2.0来完成构造函数。如果可以访问，该类将使用它的优化版本的缩放点积注意力；否则，我们将设置一个偏置掩码。这个掩码是注意力机制的可选掩码功能的组成部分。torch.tril方法将矩阵的上三角部分转换为零。当与torch.ones方法结合使用时，它有效地生成了一个由1和0组成的掩码，注意力机制使用该掩码为给定的输入样本生成预期的输出。
接下来，我们进入该类的forward方法，其中应用了注意力算法。首先，我们确定输入矩阵的大小，并将其分为三个维度：Batch大小，Time（或样本数量），Corpus（或嵌入大小）。nanoGPT采用了分批学习过程，当我们研究使用这个注意力层的变压器模型时，我们将更详细地探讨它。现在，只需要理解我们正在处理批处理的数据即可。然后，我们将输入x传递给线性变换层c_att，将维度从n_embed扩展到三倍的n_embed。该转换的输出被分割成我们的q，k，v变量，它们是我们输入到注意力算法中的变量。随后，使用view方法将这些变量中的数据重新组织成PyTorch scaled_dot_product_attention函数所期望的格式。
当无法使用优化函数时，代码将默认使用手动实现的缩放点积注意力。它首先对q和k矩阵进行点积运算，其中k被转置以适应点积函数，然后将结果按k的大小的平方根进行缩放。然后，使用先前创建的偏置缓冲区对缩放的输出进行掩码处理，将0替换为负无穷。接下来，对att矩阵应用softmax函数，将负无穷转换回0，并确保所有其他值在0和1之间进行缩放。然后，在得到att矩阵和v的点积之前，应用一个Dropout层以避免过拟合。
无论使用的是哪种缩放点积实现，多头输出都会被并排重新组织，然后通过最后一个Dropout层传递，并返回结果。这是不到50行Python/PyTorch代码的完整实现。如果您对上述代码不完全理解，我建议花些时间进行复习，然后再继续阅读本文的其余部分。
在我们深入讨论整合了一切的GPT模块之前，我们还需要两个构建模块。第一个是一个简单的多层感知机（MLP）——在“Attention is All You Need”论文中称为前馈网络（feed-forward network）——和注意力块，它将注意力层与MLP结合起来，完成论文中表示的基本变压器架构。这两个构建模块在下面的nanoGPT代码片段中实现。
class MLP(nn.Module): """多层感知机""" def __init__(self, config): super().__init__() self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias) self.gelu = nn.GELU() self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias) self.dropout = nn.Dropout(config.dropout) def forward(self, x): x = self.c_fc(x) x = self.gelu(x) x = self.c_proj(x) x = self.dropout(x) return xclass Block(nn.Module): def __init__(self, config): super().__init__() self.ln_1 = LayerNorm(config.n_embd, bias=config.bias) self.attn = CausalSelfAttention(config) self.ln_2 = LayerNorm(config.n_embd, bias=config.bias) self.mlp = MLP(config) def forward(self, x): x = x + self.attn(self.ln_1(x)) x = x + self.mlp(self.ln_2(x)) return x尽管看起来代码行数很少，但是 MLP 层给模型增加了额外的复杂性。本质上，线性层将每个输入层与输出层的每个元素进行连接，使用线性变换在它们之间传递值。在上述代码中，我们从嵌入大小 n_embed 开始，然后在输出中将其增加四倍。这个四倍增加是任意的；MLP 模块的目的是通过添加更多的节点来增强网络的计算能力。只要 MLP 开头的维度增加并且 MLP 结尾的维度减少是等价的，得到相同的初始输入/最终输出维度，那么缩放数字只是另一个超参数。另一个关键的要素是激活函数。这个 MLP 实现由两个线性层和 GELU 激活函数连接而成。原始论文使用的是 ReLU 函数，但是 nanoGPT 使用 GELU 以确保与 GPT2 模型检查点的兼容性。
接下来，我们来看 Block 模块。这个模块完成了我们在“Attention”论文中概述的 transformer block。本质上，它将输入通过一个归一化层传递给注意力层，然后将结果添加回输入。此添加的输出再次进行归一化，然后传递给 MLP，并再次添加到自身中。这个过程实现了 transformer 中描述的解码器部分。对于文本生成，通常只使用解码器，因为它不需要将解码器的输出与除了输入序列之外的任何东西进行条件化。transformer 最初是为机器翻译而设计的，它需要考虑输入令牌编码和输出令牌编码。然而，对于文本生成，只使用单个令牌编码，消除了通过编码器进行跨注意力的需要。nanoGPT 的作者 Andrej Karpathy 在本系列的第一篇文章中提供了详细的解释。
最后，我们来到主要的组件：GPT 模型。大约 300 行的文件中，大部分都是用于 GPT 模块的。它管理诸如模型微调和为模型训练设计的实用程序等有益功能（本系列文章的下一篇将讨论此主题）。因此，我在下面提供了 nanoGPT 存储库中提供的简化版本。
class GPT(nn.Module): def __init__(self, config): super().__init__() assert config.vocab_size is not None assert config.block_size is not None self.config = config self.transformer = nn.ModuleDict(dict( wte = nn.Embedding(config.vocab_size, config.n_embd), wpe = nn.Embedding(config.block_size, config.n_embd), drop = nn.Dropout(config.dropout), h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), ln_f = LayerNorm(config.n_embd, bias=config.bias), )) self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) # 使用 torch.compile() 时，使用权重绑定会生成一些警告： # "UserWarning: functional_call was passed multiple values for tied weights. # This behavior is deprecated and will be an error in future versions" # 不确定这是什么意思，到目前为止似乎是无害的。TODO 调查一下 self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying # 初始化所有权重 self.apply(self._init_weights) # 根据 GPT-2 论文，对残差投影应用特殊的缩放初始化 for pn, p in self.named_parameters(): if pn.endswith('c_proj.weight'): torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer)) def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) def forward(self, idx, targets=None): device = idx.device b, t = idx.size() assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}" pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t) # 正向传播 GPT 模型本身 tok_emb = self.transformer.wte(idx) # shape (b, t, n_embd) 的令牌嵌入 pos_emb = self.transformer.wpe(pos) # shape (t, n_embd) 的位置嵌入 x = self.transformer.drop(tok_emb + pos_emb) for block in self.transformer.h: x = block(x) x = self.transformer.ln_f(x) if targets is not None: # 如果我们有一些预期的目标，还要计算损失 logits = self.lm_head(x) loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1) else: # 推断时的迷你优化：只在最后一个位置上前向 lm_head logits = self.lm_head(x[:, [-1], :]) # 注意：使用列表 [-1] 以保留时间维度 loss = None return logits, loss @torch.no_grad() def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None): """ 使用索引条件序列 idx（shape 为 (b,t) 的 LongTensor）并完成序列 max_new_tokens 次，每次将预测结果反馈到模型中。 大多数情况下，您可能希望确保处于 model.eval() 模式下进行操作。 """ for _ in range(max_new_tokens): # 如果序列上下文变得过长，我们必须在 block_size 处进行裁剪 idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:] # 将模型前向传播以获取序列中索引的 logits logits, _ = self(idx_cond) # 提取最终步骤的 logits 并按所需的温度进行缩放 logits = logits[:, -1, :] / temperature # 可选地裁剪 logits 以仅保留前 k 个选项 if top_k is not None: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits[logits < v[:, [-1]]] = -float('Inf') # 应用 softmax 将 logits 转换为（归一化的）概率 probs = F.softmax(logits, dim=-1) # 从分布中采样 idx_next = torch.multinomial(probs, num_samples=1) # 将采样的索引附加到运行序列中并继续 idx = torch.cat((idx, idx_next), dim=1) return idx让我们从类的构造函数开始。不同的层被组装成一个PyTorch的ModuleDict，提供了一些结构。我们首先有两个Embedding层：一个用于标记嵌入，一个用于位置嵌入。nn.Embedding模块被设计成稀疏地填充值，优化了存储能力。在此之后，我们有一个dropout层，然后是n_layer个Block模块，形成我们的attention层，然后是另一个单独的dropout层。lm_head线性层将attention Blocks的输出降低到词汇大小，作为GPT的主要输出之一，除了损失值。
一旦层被定义，就需要进行额外的设置，然后我们才能开始训练模块。在这里，Andrej将位置编码的权重与输出层的权重关联起来。根据代码注释中链接的论文，这样做可以减少模型的最终参数，同时提高其性能。构造函数还初始化了模型的权重。由于这些权重将在训练过程中学习，它们被初始化为随机数的高斯分布，模块的偏置被设置为0。最后，使用了GPT-2论文中的一个修改，即将任何残差层的权重缩放为层数的平方根。
在网络的前向传播过程中，批处理大小和样本数（这里用t表示）从输入大小中获取。然后，在训练设备上为将要成为位置嵌入的存储器创建内存。接下来，我们将输入标记嵌入到标记嵌入层wte中。在此之后，计算在wpe层上的位置嵌入。这些嵌入相加后通过一个dropout层。然后将结果传递到每个n_layer块中并进行归一化。最终结果传递给线性层lm_head，将嵌入权重降低为词汇表中每个标记的概率得分。
在计算损失时（例如在训练过程中），我们使用交叉熵计算预测标记与实际标记之间的差异。如果没有，则损失为None。损失和标记概率都作为前馈函数的一部分返回。
与以前的模块不同，GPT模块还有其他方法。对我们来说最重要的是generate函数，这对于使用生成模型的人来说是熟悉的。给定一组输入标记idx、一定数量的max_new_tokens和一个温度temperature，它生成max_new_tokens个标记。让我们深入了解它是如何实现的。首先，将输入标记修剪到适合block_size的范围内（其他人称之为上下文长度），如果需要，从输入的末尾进行采样。接下来，将标记馈送到网络，并根据输入的温度对输出进行缩放。温度越高，模型越有创造力和幻觉。较高的温度也会导致输出不那么可预测。然后，应用softmax函数将模型输出权重转换为0到1之间的概率。采样函数用于从概率中选择下一个标记，然后将该标记添加到输入向量中，该向量被反馈到GPT模型中用于下一个字符。
感谢您耐心阅读这篇全面的文章。尽管查看带注释的源代码是了解代码段功能的有价值的方法，但没有什么能取代个人操作代码的各个部分和参数。为此，我提供了一个链接到nanoGPT存储库中完整的model.py源代码
nanoGPT/model.py at master · karpathy/nanoGPT
The simplest, fastest repository for training/finetuning VoAGI-sized GPTs. – nanoGPT/model.py at master ·…
github.com
在接下来的文章中，我们将探索nanoGPT的train.py脚本，并在TinyStories数据集上训练一个字符级模型。请关注我的VoAGI，以确保您不会错过！
我利用了大量的资源来创建这篇文章，其中许多资源已经在本文和上一篇文章中进行了链接。然而，如果我不与您分享这些资源，以便进一步探索任何主题或对概念进行替代解释，那么我就会忽略我的职责。
- 让我们从头开始构建GPT：用代码详细解释 – YouTube
- LLM阅读清单 – 博客
- “Attention is All You Need” – 论文
- “语言模型是无监督多任务学习者” – GPT-2论文
- 解释和说明多层感知机 – VoAGI
- 权重绑定 – Papers With Code
- 图解变换器神经网络指南：逐步解释 – YouTube
使用GPT-4和自定义LangChain脚本进行编辑。