使用 PyTorch 学习生成式人工智能——从零构建一个生成式预训练Transformer

228 阅读38分钟

本章内容包括

  • 从零构建生成式预训练Transformer
  • 因果自注意力机制
  • 从预训练模型中提取和加载权重
  • 使用GPT-2生成连贯文本(GPT-2是ChatGPT和GPT-4的前身)

生成式预训练Transformer 2(GPT-2)是OpenAI于2019年2月发布的先进大型语言模型(LLM),标志着自然语言处理(NLP)领域的一个重要里程碑,并为后续更复杂模型的开发铺平了道路,包括ChatGPT和GPT-4。

GPT-2是在其前身GPT-1的基础上改进而来,旨在基于给定的提示生成连贯且上下文相关的文本,展现出在多种风格和主题上模拟人类文本生成的惊人能力。发布时,OpenAI最初决定不公开发布参数最多(15亿参数)的最强版本(本章将从零构建的版本即为此),主要担忧模型可能被滥用,例如生成误导性新闻、在线冒充他人,或自动生成辱骂和虚假内容。这一决定引发了AI及科技社区关于AI伦理、创新与安全之间平衡的广泛讨论。

随后,OpenAI采取分阶段发布策略,逐步开放较小版本,并监控影响,探索安全部署方法。最终在2019年11月,OpenAI发布了完整模型、相关数据集及检测模型生成文本的工具,促进了负责任AI使用的讨论。基于此,你将学习如何从GPT-2提取预训练权重并加载到自己构建的模型中。

GPT-2基于第9和第10章讨论的Transformer架构,但与之前构建的英译法翻译器不同,GPT-2是仅含解码器(decoder-only)的Transformer,即模型中没有编码器堆栈。英译法中,编码器负责捕捉英文短语含义,传递给解码器生成法语翻译;而文本生成任务中,模型无需编码器理解另一种语言,而是仅基于句中已有词元依次生成文本。和其他Transformer模型一样,GPT-2利用自注意力机制并行处理输入,大幅提升训练大型语言模型的效率和效果。

GPT-2预训练于大规模文本语料,任务是根据上下文词预测下一个词,从而学习广泛的语言模式、语法和知识。

本章你将从零构建GPT-2XL(GPT-2的最大版本)。之后,学习如何从Hugging Face(一个AI社区,托管并协作开发机器学习模型、数据集及应用)提取预训练权重并加载至你的GPT-2模型。你将使用GPT-2通过输入提示生成文本,模型根据可能的下一个词元概率进行采样,能基于输入提示生成连贯且上下文相关的段落文本。此外,和第8章一样,你还可以通过调整温度和Top-K采样控制生成文本的创造性。

尽管GPT-2是NLP的重要进步,但你应适度预期,认识其内在局限。不可将GPT-2直接与ChatGPT或GPT-4相比,因为GPT-2XL仅有15亿参数,而ChatGPT约有1750亿参数,GPT-4估计达到1.76万亿参数。GPT-2的主要限制之一是缺乏对生成内容的真正理解。模型基于训练数据中词的概率分布预测下一个词,生成的文本语法正确、表面合理,但不具备词义的真实理解,可能导致内容不准确、荒谬或表层化。

另一个关键限制是GPT-2的上下文感知能力有限。它能维持短文本的连贯性,但处理较长文本时容易出现连贯性丢失、自相矛盾或内容无关。我们应谨慎对待其生成需要持续上下文关注和细节掌控的长篇内容的能力。因此,尽管GPT-2在NLP领域迈出了重要一步,阅读其生成内容时仍需保持理性怀疑,设定现实预期。

11.1 GPT-2架构与因果自注意力机制

GPT-2是一个纯解码器(decoder-only)的Transformer模型(它基于句中之前的词元生成文本,无需编码器去理解另一种语言),与第9章和第10章中讨论的英译法翻译器的解码器部分相似。与双语翻译模型不同,GPT-2没有编码器,因此在输出生成过程中不包含来自编码器的输入。该模型完全依赖序列中之前的词元来生成输出。

本节将介绍GPT-2的架构,并深入探讨GPT-2模型的核心机制——因果自注意力(causal self-attention)。

11.1.1 GPT-2的架构

GPT-2有四个不同规模的版本:小型(Small,S)、中型(Medium,M)、大型(Large,L)和超大型(Extra-Large,XL),它们的能力各不相同。本章主要关注最强大的版本GPT-2XL。最小的GPT-2模型约有1.24亿参数,而超大型版本约有15亿参数,是GPT-2系列中参数最多、能力最强的版本。GPT-2XL能够理解复杂的上下文,生成连贯且细腻的文本。

GPT-2由多个相同的解码器模块组成。超大型版本拥有48个解码器模块,其余三个版本分别拥有12、24和36个解码器模块。每个解码器模块包含两个不同的子层:第一子层是因果自注意力层,稍后将详细讲解;第二子层是基础的逐位置全连接前馈网络,类似于英译法翻译器中的编码器和解码器模块。每个子层都包含层归一化(layer normalization)和残差连接(residual connection),以稳定训练过程。

图11.1展示了GPT-2的架构示意图。

image.png

图11.1 GPT-2模型架构示意图。GPT-2是一个仅包含解码器的Transformer,由N个相同的解码器层组成。每个解码器模块包含两个子层:第一个子层是因果自注意力层,第二个子层是前馈网络。每个子层都采用层归一化(layer normalization)和残差连接(residual connection)。输入首先经过词嵌入和位置编码,然后两者相加的结果被传入解码器。解码器输出再经过层归一化和线性层处理。

GPT-2首先将一串词元的索引输入词嵌入和位置编码层,得到输入嵌入(稍后将详细讲解此过程)。输入嵌入依次通过N个解码器模块,随后输出经过层归一化和线性层。GPT-2输出的维度等于词汇表中独特词元的数量(所有GPT-2版本的词汇表大小均为50,257个词元)。该模型设计的目标是基于序列中所有先前的词元预测下一个词元。

为了训练GPT-2,OpenAI使用了一个名为WebText的数据集,该数据集通过自动抓取互联网内容构建,包含各种文本,如Reddit中高票数的链接,旨在覆盖广泛的人类语言和主题。估计该数据集包含约40GB文本。

训练数据被切分为固定长度的序列(所有GPT-2版本为1024词元)作为输入。序列向右移动一个词元后作为训练时的输出。由于模型采用因果自注意力机制,训练过程中序列中未来词元被屏蔽(即隐藏),因此模型实际被训练成基于序列中所有先前词元预测下一个词元。

11.1.2 GPT-2中的词嵌入和位置编码

GPT-2使用一种名为Byte Pair Encoding(BPE)的子词分词方法,将文本拆分成单独的词元(多数情况下是完整单词或标点符号,但对罕见词拆分成音节)。这些词元映射为介于0到50,256之间的索引,因为词汇表大小为50,257。GPT-2通过词嵌入将训练文本转化为向量表示,捕捉文本的语义,类似于前两章中你所做的工作。

举例说明,短语“this is a prompt”经过BPE分词后变为四个词元 ['this', ' is', ' a', ' prompt']。每个词元用大小为50,257的独热编码表示。GPT-2将它们传入词嵌入层,将高维稀疏的独热向量压缩为低维连续浮点向量,GPT-2XL中向量长度为1600(其他三个版本分别为768、1024和1280)。因此,短语“this is a prompt”通过词嵌入后被表示为4×1600的矩阵,而非原始的4×50,257矩阵。词嵌入显著减少了模型参数数量,提高了训练效率。图11.2左侧展示了词嵌入的工作原理。

image.png

图11.2 GPT-2首先将序列中每个词元表示为一个长度为50,276的独热向量。序列的词元表示经过词嵌入层压缩成维度为1600的嵌入向量。GPT-2同样将序列中每个位置表示为长度为1024的独热向量。序列的位置表示经过位置编码层压缩成维度同样为1600的嵌入。词嵌入和位置编码相加形成输入嵌入。

GPT-2和其他Transformer类似,输入数据并行处理,因此无法天然识别输入的序列顺序。为解决这一问题,需要将位置编码加到输入嵌入中。GPT-2采用了不同于2017年“Attention Is All You Need”论文中的位置编码方法。GPT-2的位置编码方式与词嵌入类似。考虑模型能处理最多1024个词元的输入序列,序列中每个位置初始用同样大小的独热向量表示。例如,序列“this is a prompt”中,第一个词元用第一个元素为1其余为0的独热向量表示,第二个词元用第二个元素为1其余为0的向量表示。因此,该短语的位置表示形成一个4×1024的矩阵,如图11.2右上角所示。

为了生成位置编码,序列的位置表示经过一个1024×1600维度的线性神经网络,该网络权重随机初始化,训练中不断优化。结果,序列中每个词元的位置编码是一个1600维向量,与词嵌入向量维度一致。序列的输入嵌入是词嵌入和位置编码的加和,如图11.2底部所示。对于“this is a prompt”这一短语,词嵌入和位置编码均为4×1600矩阵,输入嵌入是这两个矩阵的和,仍为4×1600维度。

11.1.3 GPT-2中的因果自注意力

因果自注意力是GPT-2(以及整个GPT系列模型)的关键机制,使模型能够基于之前生成的词元序列条件生成文本。它类似于我们在第9、10章讨论的英译法翻译器解码器每层第一个子层中的掩码自注意力,但实现细节略有差异。

注:“因果”一词指模型确保对某个词元的预测只能受到该词元之前词元的影响,尊重文本生成的因果(时间前进)顺序。这对于生成连贯且上下文相关的文本输出至关重要。

自注意力是一种机制,使输入序列中的每个词元都能关注序列中的所有其他词元。在GPT-2等Transformer模型中,自注意力允许模型在处理某个词元时权衡其他词元的重要性,从而捕获句子中词与词之间的上下文关系。

为确保因果性,GPT-2的自注意力机制经过修改,使得任何给定词元只能关注自己及之前的词元。这通过在计算注意力时屏蔽序列中未来的词元(即当前词元后面的词元)实现,保证模型在预测下一个词元时无法“看到”未来词元。例如,在“this is a prompt”这句话中,当模型用“this”预测“is”时,掩码会屏蔽后三个词。实现时,计算注意力分数时将未来词元对应位置设为负无穷,softmax后未来词元权重为零,有效将其从注意力计算中移除。

下面用具体代码示例说明因果自注意力的实现过程。短语“this is a prompt”经过词嵌入和位置编码后的输入嵌入是一个4×1600的矩阵。我们将其传入GPT-2中的N个解码器层。每个解码器层先通过因果自注意力子层,具体步骤如下。输入嵌入通过三个神经网络分别生成查询Q、键K和值V,代码如下:

import torch
import torch.nn as nn

torch.manual_seed(42)
x = torch.randn((1, 4, 1600))                      # ①
c_attn = nn.Linear(1600, 1600 * 3)                 # ②
B, T, C = x.size()
q, k, v = c_attn(x).split(1600, dim=2)            # ③
print(f"the shape of Q vector is {q.size()}")
print(f"the shape of K vector is {k.size()}")
print(f"the shape of V vector is {v.size()}")      # ④

① 创建输入嵌入x
② 创建一个线性层,将输入维度1600映射到4800(1600×3)
③ 输入x经过线性层后,分割成Q、K、V三部分
④ 输出Q、K、V的形状

此处我们创建一个4×1600的矩阵x,大小与“this is a prompt”的输入嵌入相同。然后将x传入三个线性层(尺寸均为1600×1600)得到查询Q、键K和值V。运行以上代码,输出:

the shape of Q vector is torch.Size([1, 4, 1600])
the shape of K vector is torch.Size([1, 4, 1600])
the shape of V vector is torch.Size([1, 4, 1600])

Q、K、V均为4×1600的形状。接下来,我们将其拆分为25个并行的注意力头。每个头关注输入的不同部分或方面,帮助模型捕获更广泛的信息,形成更细致且有上下文感知的理解。因此我们得到25组Q、K、V:

hs = C // 25
k = k.view(B, T, 25, hs).transpose(1, 2)
q = q.view(B, T, 25, hs).transpose(1, 2)
v = v.view(B, T, 25, hs).transpose(1, 2)            # ①
print(f"the shape of Q vector is {q.size()}")
print(f"the shape of K vector is {k.size()}")
print(f"the shape of V vector is {v.size()}")        # ②

① 将Q、K、V拆分为25个头
② 输出多头Q、K、V的形状

运行以上代码,输出:

the shape of Q vector is torch.Size([1, 25, 4, 64])
the shape of K vector is torch.Size([1, 25, 4, 64])
the shape of V vector is torch.Size([1, 25, 4, 64])

此时Q、K、V的形状变为25×4×64,表示有25个头,每个头有4×64的查询、键和值。

接下来计算每个头的缩放注意力分数:

import math
scaled_att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
print(scaled_att[0, 0])

缩放注意力分数是Q与K的点积,再除以键向量维度的平方根(此处为1600/25=64)。每个头的注意力分数是4×4矩阵。打印第一个头的分数:

tensor([[ 0.2334,  0.1385, -0.1305,  0.2664],
        [ 0.2916,  0.1044,  0.0095,  0.0993],
        [ 0.8250,  0.2454,  0.0214,  0.8667],
        [-0.1557,  0.2034,  0.2172, -0.2740]], grad_fn=<SelectBackward0>)

第一个头的缩放注意力分数如图11.3左下表所示。

练习11.1
张量scaled_att包含25个头的缩放注意力分数。之前打印了第一个头的分数。请问如何打印第二个头的缩放注意力分数?

接着,我们对缩放注意力分数应用掩码,屏蔽序列中的未来词元:

mask = torch.tril(torch.ones(4, 4))               # ①
print(mask)
masked_scaled_att = scaled_att.masked_fill(mask == 0, float('-inf'))   # ②
print(masked_scaled_att[0, 0])

① 创建一个下三角掩码矩阵
② 将掩码应用于缩放注意力分数,将未来词元对应位置设为负无穷

以上代码通过掩码确保模型在计算注意力时不会看到未来词元,从而实现因果性。

image.png

图11.3 因果自注意力中如何计算掩码注意力权重。掩码被应用于缩放后的注意力分数,使得对应未来词元(矩阵主对角线上方的值)的位置变为负无穷(–∞)。随后,对掩码后的缩放注意力分数应用softmax函数,得到掩码注意力权重。该掩码确保模型预测某个词元时,只能受到序列中该词元之前词元的影响,而不会被未来词元干扰。这对于生成连贯且上下文相关的文本输出至关重要。

运行上述代码时,会看到以下输出:

tensor([[1., 0., 0., 0.],
        [1., 1., 0., 0.],
        [1., 1., 1., 0.],
        [1., 1., 1., 1.]])
tensor([[ 0.2334,    -inf,    -inf,    -inf],
        [ 0.2916,  0.1044,    -inf,    -inf],
        [ 0.8250,  0.2454,  0.0214,    -inf],
        [-0.1557,  0.2034,  0.2172, -0.2740]], grad_fn=<SelectBackward0>)

掩码是一个4×4矩阵,如图11.3顶部所示。掩码的下三角(主对角线及以下部分)为1,上三角(主对角线以上部分)为0。掩码应用到缩放后的注意力分数时,矩阵上半部分的值变为–∞(图11.3中部底部)。因此,当我们对缩放注意力分数应用softmax后,注意力权重矩阵的上半部分被赋值为0(图11.3右下角):

import torch.nn.functional as F
att = F.softmax(masked_scaled_att, dim=-1)
print(att[0,0])

打印第一个注意力头的权重值为:

tensor([[1.0000, 0.0000, 0.0000, 0.0000],
        [0.5467, 0.4533, 0.0000, 0.0000],
        [0.4980, 0.2790, 0.2230, 0.0000],
        [0.2095, 0.3001, 0.3042, 0.1862]], grad_fn=<SelectBackward0>)

第一行表示在第一个时间步,词元“this”只关注自身,不关注任何未来词元。同理,第二行表示“this is”两个词元互相关注,但不关注未来的“a prompt”。

注:此数值示例中的权重未经训练,不应将这些注意力权重值字面理解。我们仅用作说明因果自注意力的工作原理。

练习11.2
我们已打印第一个头的注意力权重。如何打印最后一个(即第25个)头的注意力权重?

最后,我们计算每个头的注意力向量,方法是将注意力权重与对应的值向量做点积。25个头的注意力向量再合并成一个整体注意力向量:

y = att @ v
y = y.transpose(1, 2).contiguous().view(B, T, C)
print(y.shape)

输出结果为:

torch.Size([1, 4, 1600])

因果自注意力的最终输出是一个4×1600的矩阵,与因果自注意力子层的输入维度相同。解码器层设计保证输入和输出维度一致,这使得我们可以堆叠多个解码器层来增强模型的表示能力,并在训练过程中实现层次化特征提取。

11.2 从零构建GPT-2XL

现在你已经了解了GPT-2的架构及其核心组件因果自注意力的工作原理,接下来我们将从零开始创建GPT-2的最大版本。

本节首先介绍GPT-2中使用的子词分词方法——字节对编码(Byte Pair Encoder,BPE)分词器,用于将文本拆分成单独的词元。然后你将学习GPT-2中前馈网络所用的GELU激活函数。接着,你将实现因果自注意力机制,并将其与前馈网络结合构成解码器模块。最后,堆叠48个解码器模块,构建GPT-2XL模型。本章的代码改编自Andrej Karpathy的优秀GitHub仓库(github.com/karpathy/mi…),如果你想深入了解GPT-2的工作原理,建议浏览该仓库。

11.2.1 BPE分词

GPT-2使用了一种叫字节对编码(BPE)的子词分词方法,源自一种数据压缩技术,已被改造用于自然语言处理中的分词任务。它广泛应用于训练大型语言模型,如GPT系列和BERT(双向编码器表示)。BPE的主要目标是在词汇表大小和分词后序列长度之间取得平衡,以高效编码文本。

BPE通过迭代合并数据集中出现频率最高的连续字符对为一个新词元,条件满足时继续合并,直到达到预期词汇表大小或无更多合并收益。BPE兼顾了字符级和词级分词的优点,有效减少词汇表规模,同时避免显著增加序列长度,这对NLP模型性能至关重要。

我们在第8章讨论过字符级、词级和子词分词方法的优缺点,并在第8章(以及第12章)亲自实现过词级分词器。本章直接借用OpenAI的分词方法,BPE的具体工作原理超出本书范围。你只需知道它将文本先转换成子词词元,再映射成对应索引。

请从Andrej Karpathy的GitHub仓库(mng.bz/861B)下载bpe.py文件,并放置于电脑中的/utils/文件夹。本章将作为本地模块使用。正如Karpathy在仓库中所述,该模块基于OpenAI的实现(mng.bz/EOlj),做了轻微修改以便理解。

以下示例展示bpe.py模块如何将文本转成词元及索引:

from utils.bpe import get_encoder

example = "This is the original text."                   # ①
bpe_encoder = get_encoder()                              # ②
response = bpe_encoder.encode_and_show_work(example)    # ③
print(response["tokens"])

① 示例文本
② 实例化bpe.py中的get_encoder类
③ 对示例文本分词并打印词元

输出为:

['This', ' is', ' the', ' original', ' text', '.']

BPE分词器将“This is the original text.”拆成6个词元,如上所示。注意,BPE分词器不会将大写字母转为小写,这有助于生成更有意义的词元,但导致词元数量显著增加。实际上,所有GPT-2版本的词汇表大小为50,276,远大于前几章的词汇表。

你也可以用bpe.py模块将词元映射为索引:

print(response['bpe_idx'])

输出:

[1212, 318, 262, 2656, 2420, 13]

上述列表对应示例文本6个词元的索引。

还可以根据索引还原文本:

from utils.bpe import BPETokenizer

tokenizer = BPETokenizer()                              # ①
out = tokenizer.decode(torch.LongTensor(response['bpe_idx']))  # ②
print(out)

① 实例化bpe.py中的BPETokenizer类
② 用tokenizer根据索引还原文本

输出:

This is the original text.

由此可见,BPE分词器成功还原了示例文本。

练习11.3
用BPE分词器将短语“this is a prompt”拆分为词元,随后将词元映射为索引,最后根据索引还原该短语。

11.2.2 高斯误差线性单元(GELU)激活函数

GELU激活函数被用于GPT-2每个解码器模块的前馈子层。GELU结合了线性和非线性激活的优点,被证实能提升深度学习任务(尤其是NLP)的模型性能。

相比常用的ReLU函数,GELU提供了更平滑的非线性曲线,训练过程中允许更细腻的调整,帮助网络更有效优化,因其反向传播时梯度更连续。为比较GELU与ReLU,先定义一个GELU类:

class GELU(nn.Module):
    def forward(self, x):
        return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * 
                         (x + 0.044715 * torch.pow(x, 3.0))))

ReLU函数因存在拐点,非处处可微分,而GELU函数处处可微,提供了更好的学习过程。下面绘制GELU和ReLU激活函数的图形以作比较。

import matplotlib.pyplot as plt
import numpy as np

genu = GELU()
def relu(x):                                           # ①
    y = torch.zeros(len(x))
    for i in range(len(x)):
        if x[i] > 0:
            y[i] = x[i]
    return y                 
xs = torch.linspace(-6, 6, 300)
ys = relu(xs)
gs = genu(xs)
fig, ax = plt.subplots(figsize=(6, 4), dpi=300)
plt.xlim(-3, 3)
plt.ylim(-0.5, 3.5)
plt.plot(xs, ys, color='blue', label="ReLU")         # ②
plt.plot(xs, gs, "--", color='red', label="GELU")    # ③
plt.legend(fontsize=15)
plt.xlabel("values of x")
plt.ylabel("values of $ReLU(x)$ and $GELU(x)$")
plt.title("The ReLU and GELU Activation Functions")
plt.show()

① 定义ReLU函数
② 用实线绘制ReLU激活函数
③ 用虚线绘制GELU激活函数

运行以上代码,将显示如图11.4所示的曲线图。

image.png

图11.4 GELU激活函数与ReLU的对比。实线表示ReLU激活函数,虚线表示GELU激活函数。ReLU函数在某些点不可微分,因为它存在拐角;而GELU函数在所有点均可微分。这种平滑性使得GELU在训练过程中为反向传播提供更连续的梯度,从而更有效地优化神经网络。

此外,GELU的数学形式使其能更有效地建模输入数据的分布。它结合了线性与高斯分布建模的特性,特别适合处理NLP任务中复杂多样的数据,有助于捕捉语言数据中的微妙模式,提升模型对文本的理解和生成能力。

11.2.3 因果自注意力

如前所述,因果自注意力是GPT-2模型的核心元素。接下来,我们将在PyTorch中从零实现这一机制。

首先,我们定义本章将构建的GPT-2XL模型的超参数,为此定义一个Config()类,内容如下:

class Config():                                       # ①
    def __init__(self):
        self.n_layer = 48
        self.n_head = 25
        self.n_embd = 1600
        self.vocab_size = 50257
        self.block_size = 1024 
        self.embd_pdrop = 0.1 
        self.resid_pdrop = 0.1 
        self.attn_pdrop = 0.1                         # ②
        
config = Config()                                    # ③

① 定义Config()类
② 在类中设置模型超参数为属性
③ 实例化Config类

Config()类包含多个属性,作为GPT-2XL模型的超参数。n_layer表示解码器层数,这里为48层(“解码器模块”和“解码器层”可互换使用)。n_head表示因果自注意力时将Q、K、V拆分为25个并行头。n_embd表示词元嵌入维度为1600。vocab_size表示词汇表大小为50,257。block_size表示输入序列最大长度为1024词元。所有dropout比率均设置为0.1。

上一节详细解释了因果自注意力的工作原理,接下来定义CausalSelfAttention()类实现它:

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
        self.attn_dropout = nn.Dropout(config.attn_pdrop)
        self.resid_dropout = nn.Dropout(config.resid_pdrop)
        self.register_buffer("bias", torch.tril(torch.ones(\
                   config.block_size, config.block_size))
             .view(1, 1, config.block_size, config.block_size))  # ①
        self.n_head = config.n_head
        self.n_embd = config.n_embd

    def forward(self, x):
        B, T, C = x.size() 
        q, k ,v  = self.c_attn(x).split(self.n_embd, dim=2)     # ②
        hs = C // self.n_head
        k = k.view(B, T, self.n_head, hs).transpose(1, 2) 
        q = q.view(B, T, self.n_head, hs).transpose(1, 2) 
        v = v.view(B, T, self.n_head, hs).transpose(1, 2)       # ③
        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 
        y = y.transpose(1, 2).contiguous().view(B, T, C)        # ⑤
        y = self.resid_dropout(self.c_proj(y))
        return y

① 创建一个掩码并将其注册为buffer(不参与参数更新)
② 输入嵌入通过一个线性层生成Q、K、V
③ 将Q、K、V拆分成多个头
④ 计算每个头的掩码注意力权重
⑤ 合并所有头的注意力向量为一个整体向量

在PyTorch中,register_buffer用于注册张量为buffer。buffer中的变量不被视为模型的可训练参数,因此不会在反向传播中更新。上述代码中,我们创建了一个掩码并注册为buffer,这对后续提取和加载模型权重时有影响——我们从GPT-2XL中提取权重时将忽略掩码。

如第一节所述,输入嵌入经过三个神经网络生成查询Q、键K和值V,随后拆分为25个头并计算各头的掩码自注意力。最后,将25个头的注意力向量合并成单一向量,即CausalSelfAttention()类的输出。

11.2.4 构建GPT-2XL模型

接下来,我们将向因果自注意力子层添加一个前馈网络,组成一个解码器模块,代码如下。

代码清单11.5 构建解码器模块

class Block(nn.Module):
    def __init__(self, config):                                 # ①
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp = nn.ModuleDict(dict(
            c_fc   = nn.Linear(config.n_embd, 4 * config.n_embd),
            c_proj = nn.Linear(4 * config.n_embd, config.n_embd),
            act    = GELU(),
            dropout = nn.Dropout(config.resid_pdrop),
        ))
        m = self.mlp
        self.mlpf = lambda x: m.dropout(m.c_proj(m.act(m.c_fc(x)))) 
    def forward(self, x):
        x = x + self.attn(self.ln_1(x))                         # ②
        x = x + self.mlpf(self.ln_2(x))                         # ③
        return x

① 初始化Block类
② 模块的第一个子层是因果自注意力子层,含层归一化和残差连接
③ 模块的第二个子层是前馈网络,包含GELU激活、层归一化和残差连接

每个解码器模块由两个子层组成:第一个是因果自注意力机制,包含层归一化和残差连接;第二个是前馈网络,包含GELU激活函数、层归一化和残差连接。

我们堆叠48个解码器层,构成GPT-2XL模型的主体,代码如下。

代码清单11.6 构建GPT-2XL模型

class GPT2XL(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.block_size = config.block_size
        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.embd_pdrop),
            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = nn.LayerNorm(config.n_embd),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

    def forward(self, idx, targets=None):
        b, t = idx.size()
        pos = torch.arange(0, t, dtype=torch.long).unsqueeze(0)
        tok_emb = self.transformer.wte(idx)    
        pos_emb = self.transformer.wpe(pos)    
        x = self.transformer.drop(tok_emb + pos_emb)            # ①
        for block in self.transformer.h:
            x = block(x)                                        # ②
        x = self.transformer.ln_f(x)                            # ③
        logits = self.lm_head(x)                                # ④
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)),
                                   targets.view(-1), ignore_index=-1)
        return logits, loss

① 计算输入嵌入(词嵌入和位置编码之和)
② 输入嵌入依次通过48个解码器模块
③ 再次应用层归一化
④ 添加线性层,使输出数量等于词汇表中唯一词元的数量

我们在GPT2XL类中构建模型。模型输入是一序列词元索引。先通过词嵌入和位置编码,二者相加形成输入嵌入,接着通过48个解码器模块,之后对输出做层归一化,最后连接线性层,输出维度为50,257(词汇表大小),输出为对应词元的logits。生成文本时,我们对logits应用softmax函数,得到词汇表中每个词元的概率分布。

注:模型体积较大,本章未将模型移至GPU,导致后续文本生成速度较慢。如果你有配置了大内存(例如32GB以上)的CUDA GPU,可将模型移至GPU以加速文本生成。

接下来,实例化GPT2XL类创建模型:

model = GPT2XL(config)
num = sum(p.numel() for p in model.transformer.parameters())
print("number of parameters: %.2fM" % (num / 1e6,))

同时统计模型主体参数数量,输出为:

number of parameters: 1557.61M

说明GPT-2XL有超过15亿个参数。注意该数字不包括模型末尾线性头的参数。根据不同下游任务,可以附加不同的输出头。这里专注文本生成,附加线性头确保输出数量与词汇表唯一词元数量一致。

注:在GPT-2、ChatGPT或BERT等大型语言模型中,输出头指模型的最后一层,负责根据处理后的输入生成实际输出。输出头结构因任务不同而异。文本生成中,输出头通常为线性层,将最终隐藏状态映射为词汇表中每个词元的logits,之后通过softmax生成概率分布,用于预测下一个词元。分类任务中,输出头通常由线性层和softmax组成,将隐藏状态映射为各类别的概率。输出头结构可根据模型和任务调整,但其主要功能是将处理后的输入映射到期望的输出格式(如类别概率、词元概率等)。

最后,你可以打印GPT-2XL模型结构:

print(model)

输出示例:

GPT2XL(
  (transformer): ModuleDict(
    (wte): Embedding(50257, 1600)
    (wpe): Embedding(1024, 1600)
    (drop): Dropout(p=0.1, inplace=False)
    (h): ModuleList(
      (0-47): 48 x Block(
        (ln_1): LayerNorm((1600,), eps=1e-05, elementwise_affine=True)
        (attn): CausalSelfAttention(
          (c_attn): Linear(in_features=1600, out_features=4800, bias=True)
          (c_proj): Linear(in_features=1600, out_features=1600, bias=True)
          (attn_dropout): Dropout(p=0.1, inplace=False)
          (resid_dropout): Dropout(p=0.1, inplace=False)
        )
        (ln_2): LayerNorm((1600,), eps=1e-05, elementwise_affine=True)
        (mlp): ModuleDict(
          (c_fc): Linear(in_features=1600, out_features=6400, bias=True)
          (c_proj): Linear(in_features=6400, out_features=1600, bias=True)
          (act): GELU()
          (dropout): Dropout(p=0.1, inplace=False)
        )
      )
    )
    (ln_f): LayerNorm((1600,), eps=1e-05, elementwise_affine=True)
  )
  (lm_head): Linear(in_features=1600, out_features=50257, bias=False)
)

显示了GPT-2XL模型中详细的模块和层结构。

至此,你已成功从零构建了GPT-2XL模型!

11.3 加载预训练权重并生成文本

尽管你刚刚创建了GPT-2XL模型,但它尚未经过训练,因此不能用来生成有意义的文本。

鉴于模型参数量巨大,没有超级计算设施及大量训练数据,训练该模型几乎不可能。幸运的是,OpenAI于2019年11月5日向公众发布了包括最大版本GPT-2XL在内的GPT-2预训练权重(详见OpenAI官网声明 openai.com/research/gp… 及美国科技新闻网站The Verge报道 mng.bz/NBm7)。因此本节将加载预训练权重来生成文本。

11.3.1 加载GPT-2XL预训练参数

我们将使用Hugging Face团队开发的transformers库提取GPT-2XL的预训练权重。

首先,在Jupyter Notebook新单元格运行以下命令安装transformers库:

!pip install transformers

接着,导入transformers库中的GPT2模型,并提取GPT-2XL预训练权重:

from transformers import GPT2LMHeadModel

model_hf = GPT2LMHeadModel.from_pretrained('gpt2-xl')     # ①
sd_hf = model_hf.state_dict()                             # ②
print(model_hf)                                           # ③

① 加载预训练的GPT-2XL模型
② 提取模型权重
③ 打印原版OpenAI GPT-2XL模型结构

以上代码块输出如下:

GPT2LMHeadModel(
  (transformer): GPT2Model(
    (wte): Embedding(50257, 1600)
    (wpe): Embedding(1024, 1600)
    (drop): Dropout(p=0.1, inplace=False)
    (h): ModuleList(
      (0-47): 48 x GPT2Block(
        (ln_1): LayerNorm((1600,), eps=1e-05, elementwise_affine=True)
        (attn): GPT2Attention(
          (c_attn): Conv1D()                                  # ①
          (c_proj): Conv1D()                                  # ①
          (attn_dropout): Dropout(p=0.1, inplace=False)
          (resid_dropout): Dropout(p=0.1, inplace=False)
        )
        (ln_2): LayerNorm((1600,), eps=1e-05, elementwise_affine=True)
        (mlp): GPT2MLP(
          (c_fc): Conv1D()                                    # ①
          (c_proj): Conv1D()                                  # ①
          (act): NewGELUActivation()
          (dropout): Dropout(p=0.1, inplace=False)
        )
      )
    )
    (ln_f): LayerNorm((1600,), eps=1e-05, elementwise_affine=True)
  )
  (lm_head): Linear(in_features=1600, out_features=50257, bias=False)
)

① OpenAI使用了Conv1d层代替我们之前实现的线性层。

如果你将此模型结构与上一节中自己构建的模型结构对比,会发现两者相同,唯一区别是线性层被Conv1d层替代。正如第9、10章所述,前馈网络中我们把输入视为独立元素而非序列,因此称之为一维卷积网络。OpenAI在部分模型层用Conv1d模块替代了线性层。因此,提取Hugging Face权重并加载到我们模型时,需要对部分权重矩阵进行转置。

下面以OpenAI GPT-2XL第一个解码器块中的前馈网络第一层权重为例,打印其形状:

print(model_hf.transformer.h[0].mlp.c_fc.weight.shape)

输出:

torch.Size([1600, 6400])

Conv1d层权重矩阵大小为(1600, 6400)。

再看我们刚构建模型中对应层的权重形状:

print(model.transformer.h[0].mlp.c_fc.weight.shape)

输出:

torch.Size([6400, 1600])

我们模型中线性层权重矩阵大小为(6400, 1600),正好是OpenAI权重矩阵的转置。因此,在将OpenAI GPT-2XL模型的Conv1d层权重加载到我们模型前,需要对这些矩阵转置。

接下来,将OpenAI GPT-2XL模型的参数名列为键:

keys = [k for k in sd_hf if not k.endswith('attn.masked_bias')]

注意,这里排除了以attn.masked_bias结尾的参数。OpenAI GPT-2用这些参数实现未来词元掩码。由于我们已在CausalSelfAttention()类中自建掩码并注册为PyTorch buffer,无需加载OpenAI的掩码参数。

将我们自建GPT-2XL模型的参数字典命名为sd

sd = model.state_dict()

提取OpenAI预训练权重,并加载至我们模型:

transposed = ['attn.c_attn.weight', 'attn.c_proj.weight',
              'mlp.c_fc.weight', 'mlp.c_proj.weight']          # ①
for k in keys:
    if any(k.endswith(w) for w in transposed):
        with torch.no_grad():
            sd[k].copy_(sd_hf[k].t())                          # ②
    else:
        with torch.no_grad():
            sd[k].copy_(sd_hf[k])                              # ③

① 识别OpenAI使用Conv1d替代线性层的权重名称
② 对这些权重矩阵做转置后加载到我们模型
③ 其余权重直接复制加载

这样,我们成功地将OpenAI预训练权重从Hugging Face提取并载入自建模型,转置了部分Conv1d层的权重。

现在,我们的模型已经配备了OpenAI的预训练权重,可以用来生成连贯的文本了。

11.3.2 定义generate()函数生成文本

有了OpenAI GPT-2XL模型的预训练权重,我们将使用自己从零构建的GPT2模型来生成文本。

生成文本时,我们向模型输入一个对应提示词元序列索引的序列,模型预测下一个词元的索引,并将预测结果添加到序列末尾,形成新的序列。随后模型使用新序列再次预测,如此循环,直到生成固定数量的新词元或遇到特殊词元<|endoftext|>(表示对话结束)。

GPT中的特殊词元<|endoftext|>
GPT模型训练时使用了来自多种来源的文本。训练阶段使用特殊词元<|endoftext|>来分隔不同来源的文本。生成阶段,遇到该特殊词元必须停止生成,否则模型可能开始无关的新话题,导致生成文本与当前上下文无关。

因此,我们定义了一个sample()函数,向当前序列添加指定数量的新索引。该函数以索引序列作为输入,喂入GPT-2XL模型,逐次预测一个索引并添加到序列末尾。直到达到max_new_tokens指定的新词元数或预测到<|endoftext|>为止,否则模型会随机开启无关话题。sample()函数定义如下:

model.eval()
def sample(idx, max_new_tokens, temperature=1.0, top_k=None):
    for _ in range(max_new_tokens):                            # ①
        if idx.size(1) <= config.block_size:
            idx_cond = idx  
        else:
            idx_cond = idx[:, -config.block_size:]
        logits, _ = model(idx_cond)                            # ②
        logits = logits[:, -1, :] / temperature
        if top_k is not None:
            v, _ = torch.topk(logits, top_k)
            logits[logits < v[:, [-1]]] = -float('Inf')        # ③
        probs = F.softmax(logits, dim=-1)
        idx_next = torch.multinomial(probs, num_samples=1)
        if idx_next.item() == tokenizer.encoder.encoder['<|endoftext|>']:
            break                                              # ④
        idx = torch.cat((idx, idx_next), dim=1)                # ⑤
    return idx

① 生成固定数量的新索引
② 使用GPT-2XL预测下一个索引
③ 若使用Top-K采样,将非Top-K词元的logits设为负无穷
④ 若预测到<|endoftext|>,停止生成
⑤ 将新预测的索引添加到序列末尾

sample()函数使用GPT-2XL向序列添加新索引。它包含temperature和top_k两个参数,用于调节生成文本的新颖度,和第8章中介绍的用法相同。函数返回新的索引序列。

接下来定义generate()函数,根据提示文本生成文本。它先将提示文本转为索引序列,再调用刚定义的sample()生成新的索引序列,最后将新序列转换回文本。

def generate(prompt, max_new_tokens, temperature=1.0,
             top_k=None):
    if prompt == '':
        x = torch.tensor([[tokenizer.encoder.encoder['<|endoftext|>']]],
                         dtype=torch.long)                     # ①
    else:
        x = tokenizer(prompt)                                  # ②
    y = sample(x, max_new_tokens, temperature, top_k)          # ③
    out = tokenizer.decode(y.squeeze())                        # ④
    print(out)

① 如果提示为空,则使用<|endoftext|>作为提示
② 将提示文本转换为索引序列
③ 调用sample()生成新索引序列
④ 将新索引序列解码为文本

generate()函数与第8章介绍的版本类似,不同点是这里用GPT-2XL代替了之前的LSTM模型进行预测。函数接收一个提示文本,转换为索引序列后输入模型预测下一个词元索引,生成指定数量的新索引后再将整条索引序列转回文本输出。

11.3.3 使用GPT-2XL生成文本

现在我们已经定义了generate()函数,就可以用它来生成文本了。

特别地,generate()函数支持无条件文本生成,即提示(prompt)为空时,模型会随机生成文本。这在创意写作中非常有用,生成的文本可作为灵感或创作起点。我们试试:

prompt = ""
torch.manual_seed(42)
generate(prompt, max_new_tokens=100, temperature=1.0, top_k=None)

输出示例:

<|endoftext|>Feedback from Ham Radio Recalls
  
I discovered a tune sticking in my head -- I'd heard it mentioned on several occasions, but hadn't investigated further.
  
The tune sounded familiar to a tune I'd previously heard on the 550 micro. 
During that same time period I've heard other people's receipients drone on
the idea of the DSH-94013, notably Kim Weaver's instructions in her 
Interview on Radio Ham; and both Scott Mcystem and Steve Simmons' concepts.

如你所见,上述输出连贯且语法正确,但可能不具备事实准确性。我做了快速谷歌搜索,文本似乎未从任何在线来源复制。

练习11.4
将提示设为空字符串,无条件生成文本,温度设置为0.9,最大生成新词元数为100,top_k设为40。在PyTorch中将随机种子设置为42。观察输出结果。

为了评估GPT-2XL能否根据上下文生成连贯文本,我们使用提示“I went to the kitchen and”,在该提示后生成10个词元。重复五次,以判断生成内容是否符合厨房常见活动:

prompt = "I went to the kitchen and"
for i in range(5):
    torch.manual_seed(i)
    generate(prompt, max_new_tokens=10, temperature=1.0, top_k=None)

输出示例:

I went to the kitchen and said, you're not going to believe this.
I went to the kitchen and noticed a female producer open a drawer in which was
I went to the kitchen and asked who was going to be right there and A
I went to the kitchen and took a small vial of bourbon and a little
I went to the kitchen and found the bottle of wine, and poured it into

结果表明生成文本包含了交谈、注意观察、拿取饮料等典型厨房活动,说明GPT-2XL能生成与上下文相关的内容。

接下来使用“Lexington is the second largest city in the state of Kentucky”作为提示,调用generate()函数生成最多100个新词元:

prompt = "Lexington is the second largest city in the state of Kentucky"
torch.manual_seed(42)
generate(prompt, max_new_tokens=100, temperature=1.0, top_k=None)

输出示例:

Lexington is the second largest city in the state of Kentucky. It caters to
those who want to make everything in tune with being with friends and 
enjoying a jaunt through the down to Earth lifestyle. To do so, they are 
blessed with several venues large and small to fill their every need while 
residing micro- cozy with nature within the landmarks of the city.
  
In a moment we look at ten up and coming suchache music acts from the 
Lexington area to draw upon your attention.
  
Lyrikhop
  
  
This Lexington-based group

虽然文本连贯,但事实准确性不足。GPT-2XL基本上是根据上下文预测下一个词元,上述结果显示它能记住序列前部内容并生成相关的后续词元。例如,开头讨论了肯塔基州的列克星敦市,90个词元后又提到该地区的音乐团体。

此外,如介绍中所述,GPT-2存在局限。其模型大小不到ChatGPT的1%,不到GPT-4的0.1%。GPT-3拥有1750亿参数,生成文本更连贯,但其预训练权重未公开。

接下来,我们探索温度和Top-K采样对GPT-2XL生成文本的影响。设置温度为0.9,top_k为50,其他参数保持不变:

torch.manual_seed(42)
generate(prompt, max_new_tokens=100, temperature=0.9, top_k=50)

输出示例:

Lexington is the second largest city in the state of Kentucky. It is also 
the state capital. The population of Lexington was 1,731,947 in the 2011 
Census. The city is well-known for its many parks, including Arboretum, 
Zoo, Aquarium and the Kentucky Science Center, as well as its restaurants, 
such as the famous Kentucky Derby Festival.
  
In the United States, there are at least 28 counties in this state with a
population of more than 100,000, according to the 2010 census.

生成的文本较之前更连贯,但事实依然不准确。模型编造了很多关于肯塔基州列克星敦市的人口和地标信息。

练习11.5
以“Lexington is the second largest city in the state of Kentucky”为起始提示,温度设置为1.2,top_k设为None,随机种子设为42,最大新词元数设为100,生成文本。

本章你学会了如何从零构建GPT-2(ChatGPT和GPT-4的前身),之后提取OpenAI发布的GPT-2XL预训练权重并加载至模型,并见证了模型生成的连贯文本。

由于GPT-2XL模型参数量巨大(15亿),无超级计算设备几乎无法训练。下一章你将创建一个更小的GPT模型,结构类似GPT-2,但只有约512万参数,并用海明威小说文本训练。训练后的模型将以海明威风格生成连贯文本!

总结

  • GPT-2是由OpenAI开发并于2019年2月发布的一款先进大型语言模型(LLM)。它在自然语言处理(NLP)领域具有重要里程碑意义,为后续更复杂模型的发展奠定了基础,包括其后继模型ChatGPT和GPT-4。

  • GPT-2是一种仅包含解码器的Transformer模型,模型中没有编码器部分。与其他Transformer模型类似,GPT-2利用自注意力机制并行处理输入数据,大幅提升了训练大型语言模型的效率和效果。

  • GPT-2采用了不同于2017年开创性论文《Attention Is All You Need》中所用的位置编码方法,其位置编码方式与词嵌入方法类似。

  • 在GPT-2的前馈子层中使用了GELU激活函数,GELU结合了线性与非线性激活的特性,被证明能提升深度学习任务,特别是NLP和大型语言模型训练中的性能表现。

  • 我们可以从零构建GPT-2模型,并加载OpenAI公开发布的预训练权重。你构建的GPT-2模型能够生成连贯文本,效果与原版OpenAI GPT-2模型相当。