本章内容包括:
- 计算训练集和验证集的损失,以评估训练过程中LLM生成文本的质量
- 实现训练函数并对LLM进行预训练
- 保存和加载模型权重,以便继续训练LLM
- 从OpenAI加载预训练权重
到目前为止,我们已经实现了数据采样和注意力机制,并编写了LLM架构。现在是时候实现训练函数并对LLM进行预训练了。我们将学习基本的模型评估技术,以衡量生成文本的质量,这是在训练过程中优化LLM的必要步骤。此外,我们还将讨论如何加载预训练权重,为我们的LLM提供一个稳固的起点以进行微调。图5.1展示了我们的整体计划,重点说明了本章将讨论的内容。
权重参数
在LLM和其他深度学习模型的上下文中,权重指的是学习过程中调整的可训练参数。这些权重也被称为权重参数或简单地称为参数。在PyTorch等框架中,这些权重存储在线性层中;我们在第3章中使用这些权重实现了多头注意力模块,并在第4章中实现了GPTModel。在初始化一个层后(
new_layer = torch.nn.Linear(...)),我们可以通过.weight属性访问其权重,如new_layer.weight。此外,为了方便,PyTorch允许通过model.parameters()方法直接访问模型的所有可训练参数,包括权重和偏置,我们将在后面实现模型训练时使用这一方法。
评估生成文本模型
在简要回顾第4章的文本生成之后,我们将设置LLM用于文本生成,并讨论评估生成文本质量的基本方法。随后,我们将计算训练和验证损失。图5.2展示了本章涵盖的主题,其中前面这三步被重点标注。
使用GPT生成文本
让我们设置LLM,并简要回顾我们在第4章中实现的文本生成过程。我们首先初始化将在之后评估和训练的GPT模型,使用GPTModel类和GPT_CONFIG_124M字典(见第4章):
import torch
from chapter04 import GPTModel
GPT_CONFIG_124M = {
"vocab_size": 50257,
"context_length": 256, #1
"emb_dim": 768,
"n_heads": 12,
"n_layers": 12,
"drop_rate": 0.1, #2
"qkv_bias": False
}
torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
model.eval()
#1 我们将上下文长度从1,024个标记缩短为256个标记。
#2 通常可以将dropout设置为0。
相比前一章的GPT_CONFIG_124M字典,唯一的调整是我们将上下文长度(context_length)减少到了256个标记。此修改减少了训练模型的计算需求,使得在标准笔记本电脑上进行训练成为可能。
最初,具有1.24亿参数的GPT-2模型被配置为处理最多1,024个标记。训练过程结束后,我们将更新上下文大小设置,并加载预训练权重,以使用配置为1,024标记上下文长度的模型。
使用GPTModel实例,我们采用第4章的generate_text_simple函数,并引入了两个实用函数:text_to_token_ids 和 token_ids_to_text。这些函数可以方便地在文本和标记表示之间进行转换,这项技术将在本章中频繁使用。
图5.3展示了使用GPT模型生成文本的三步过程。首先,标记器将输入文本转换为一系列标记ID(见第2章)。其次,模型接收这些标记ID并生成相应的logits,这些logits是表示词汇表中每个标记概率分布的向量(见第4章)。第三,这些logits被转换回标记ID,然后通过标记器解码为人类可读的文本,完成从文本输入到文本输出的循环。
我们可以按照下面的代码实现文本生成过程。
列表5.1 文本与标记ID转换的实用函数
import tiktoken
from chapter04 import generate_text_simple
def text_to_token_ids(text, tokenizer):
encoded = tokenizer.encode(text, allowed_special={'<|endoftext|>'})
encoded_tensor = torch.tensor(encoded).unsqueeze(0) #1
return encoded_tensor
def token_ids_to_text(token_ids, tokenizer):
flat = token_ids.squeeze(0) #2
return tokenizer.decode(flat.tolist())
start_context = "Every effort moves you"
tokenizer = tiktoken.get_encoding("gpt2")
token_ids = generate_text_simple(
model=model,
idx=text_to_token_ids(start_context, tokenizer),
max_new_tokens=10,
context_size=GPT_CONFIG_124M["context_length"]
)
print("Output text:\n", token_ids_to_text(token_ids, tokenizer))
#1 .unsqueeze(0)添加批次维度
#2 移除批次维度
使用这段代码,模型生成了以下文本:
Output text:
Every effort moves you rentingetic wasnم refres RexMeCHicular stren
显然,由于模型尚未经过训练,它还无法生成连贯的文本。要定义什么样的文本是“连贯的”或“高质量的”,我们需要实现一种数值方法来评估生成的内容。通过这种方法,我们可以在整个训练过程中监控并提升模型的性能。
接下来,我们将计算生成输出的损失度量。该损失将作为训练进度和成功的指标。此外,在后续章节中,当我们微调LLM时,还将回顾评估模型质量的其他方法。
计算文本生成损失
接下来,让我们通过计算文本生成损失来探讨在训练过程中以数值方式评估生成文本质量的技术。我们将通过一个实际示例一步步地讲解这个主题,使概念更加清晰和实用,首先简要回顾数据是如何加载的,以及如何通过generate_text_simple函数生成文本。
图5.4展示了从输入文本到LLM生成文本的整体流程,采用了五个步骤。这一文本生成过程展示了generate_text_simple函数在内部所做的工作。在计算衡量生成文本质量的损失之前,我们需要执行这些相同的初始步骤。
图5.4展示了文本生成过程,使用了一个小型的七个标记的词汇表,以便整个流程能展示在一页上。然而,我们的GPTModel使用的是一个包含50,257个词汇的更大词汇表,因此接下来代码中的标记ID范围将是0到50,256,而不是0到6。
另外,图5.4为了简化仅展示了一个文本示例(“every effort moves”)。在接下来的代码示例中,我们将使用两个输入示例供GPT模型处理(“every effort moves”和“I really like”)。
考虑这两个已经被映射到标记ID的输入示例(见图5.4,第1步):
inputs = torch.tensor([[16833, 3626, 6100], # ["every effort moves",
[40, 1107, 588]]) # "I really like"]
与这些输入相匹配,目标包含了我们希望模型生成的标记ID:
targets = torch.tensor([[3626, 6100, 345 ], # [" effort moves you",
[1107, 588, 11311]]) # " really like chocolate"]
请注意,目标是将输入向前移动了一个位置,这是我们在第2章实现数据加载器时介绍的概念。这种移位策略对于教会模型预测序列中的下一个标记至关重要。
现在,我们将输入喂入模型,以计算两个输入示例的logits向量,每个示例包含三个标记。接着我们应用softmax函数,将这些logits转换为概率分数(见图5.4,第2步):
with torch.no_grad(): #1
logits = model(inputs)
probas = torch.softmax(logits, dim=-1) #2
print(probas.shape)
#1 禁用梯度跟踪,因为我们还未进行训练
#2 每个标记在词汇表中的概率
概率分数的张量维度为:
torch.Size([2, 3, 50257])
第一个数字2对应于输入中的两个示例(行),也称为批次大小。第二个数字3对应于每个输入(行)中的标记数量。最后一个数字对应于嵌入维度,由词汇表大小决定。通过softmax函数将logits转换为概率后,generate_text_simple函数将生成的概率分数转换回文本(见图5.4,第3–5步)。
我们可以通过对概率分数应用argmax函数来完成第3步和第4步,从而获得相应的标记ID:
token_ids = torch.argmax(probas, dim=-1, keepdim=True)
print("Token IDs:\n", token_ids)
由于我们有两个输入批次,每个批次包含三个标记,应用argmax函数后(见图5.4,第3步),会生成两组输出,每组包含三个预测的标记ID:
Token IDs:
tensor([[[16657], #1
[ 339],
[42826]],
[[49906], #2
[29669],
[41751]]])
#1 第一批次
#2 第二批次
最后,第5步将标记ID转换回文本:
print(f"Targets batch 1: {token_ids_to_text(targets[0], tokenizer)}")
print(f"Outputs batch 1:"
f" {token_ids_to_text(token_ids[0].flatten(), tokenizer)}")
当我们解码这些标记时,发现输出标记与我们希望模型生成的目标标记有很大不同:
Targets batch 1: effort moves you
Outputs batch 1: Armed heNetflix
由于模型尚未经过训练,它生成了与目标文本不同的随机文本。我们现在希望通过计算损失(见图5.5)以数值方式评估模型生成的文本性能。这不仅对衡量生成文本的质量有用,也是实现训练函数的基础,用于更新模型权重,从而改进生成的文本。
文本评估过程的一部分,如图5.5所示,是测量生成的标记与正确预测(目标)之间的“距离”。我们稍后实现的训练函数将使用这些信息来调整模型权重,以生成更接近(或理想情况下,完全匹配)目标文本的内容。
模型训练的目标是增加在与正确目标标记ID对应的索引位置上的softmax概率,如图5.6所示。这个softmax概率还将用于我们接下来要实现的评估指标,以数值评估模型生成的输出:在正确位置上的概率越高,表现越好。
请记住,图5.6展示的是一个包含七个标记的紧凑词汇表的softmax概率,以便在一个图中展示所有内容。这意味着初始随机值大约会在1/7左右,也就是约0.14。然而,我们的GPT-2模型使用的词汇表有50,257个标记,因此大多数初始概率会在0.00002左右(1/50,257)。
对于每个输入文本,我们可以使用以下代码打印与目标标记对应的初始softmax概率分数:
text_idx = 0
target_probas_1 = probas[text_idx, [0, 1, 2], targets[text_idx]]
print("Text 1:", target_probas_1)
text_idx = 1
target_probas_2 = probas[text_idx, [0, 1, 2], targets[text_idx]]
print("Text 2:", target_probas_2)
每个批次的三个目标标记ID的概率为:
Text 1: tensor([7.4541e-05, 3.1061e-05, 1.1563e-05])
Text 2: tensor([1.0337e-05, 5.6776e-05, 4.7559e-06])
训练LLM的目标是最大化正确标记的概率,这涉及提高其相对于其他标记的概率。通过这种方式,我们确保LLM一致地选择目标标记——即句子中的下一个单词——作为其生成的下一个标记。
反向传播
我们如何最大化与目标标记对应的softmax概率值呢?核心思路是更新模型的权重,以便模型输出我们希望生成的标记ID对应的更高值。权重更新通过一种称为反向传播(backpropagation)的过程完成,这是训练深度神经网络的标准技术(有关反向传播和模型训练的更多细节,请参见附录A的A.3至A.7部分)。
反向传播需要一个损失函数,它用于计算模型预测输出(这里是与目标标记ID对应的概率)与实际期望输出之间的差异。这个损失函数衡量模型预测值与目标值之间的偏差。
接下来,我们将为两个示例批次的概率分数(target_probas_1 和 target_probas_2)计算损失。主要步骤如图5.7所示。由于我们已经完成了步骤1到3,获得了target_probas_1和target_probas_2,现在继续进行第4步,对概率分数应用对数运算:
log_probas = torch.log(torch.cat((target_probas_1, target_probas_2)))
print(log_probas)
这将产生以下值:
tensor([ -9.5042, -10.3796, -11.3677, -11.4798, -9.7764, -12.2561])
在数学优化中,使用概率分数的对数比直接处理分数更易于管理。虽然这本书没有详细讨论这个话题,但我在附录B中的讲座中对此进行了更深入的说明。
接下来,我们通过计算平均值将这些对数概率组合成一个分数(见图5.7中的第5步):
avg_log_probas = torch.mean(log_probas)
print(avg_log_probas)
生成的平均对数概率分数是:
tensor(-10.7940)
训练的目标是通过更新模型权重,使平均对数概率尽可能接近0。然而,在深度学习中,通常的做法并不是将平均对数概率推到0,而是将负的平均对数概率降到0。负的平均对数概率只是将平均对数概率乘以–1,这对应于图5.7中的第6步:
neg_avg_log_probas = avg_log_probas * -1
print(neg_avg_log_probas)
输出为tensor(10.7940)。在深度学习中,将这个负值(–10.7940)转换为正值(10.7940)的过程称为交叉熵损失。PyTorch非常方便,因为它已经内置了cross_entropy函数,能够为我们处理图5.7中的所有六个步骤。
交叉熵损失
交叉熵损失是机器学习和深度学习中广泛使用的度量,核心是衡量两个概率分布之间的差异——通常是标签的真实分布(例如,数据集中的标记)与模型预测的分布(如LLM生成的标记概率)之间的差异。
在机器学习的上下文中,特别是在像PyTorch这样的框架中,
cross_entropy函数用于计算离散结果的这种差异,这与模型生成的标记概率给定的目标标记的负平均对数概率类似。因此,“交叉熵”和“负平均对数概率”在实践中是相关的术语,且常常可以互换使用。
在应用cross_entropy函数之前,让我们简要回顾一下logits和targets张量的形状:
print("Logits shape:", logits.shape)
print("Targets shape:", targets.shape)
输出的形状为:
Logits shape: torch.Size([2, 3, 50257])
Targets shape: torch.Size([2, 3])
可以看到,logits张量有三个维度:批次大小、标记数量和词汇表大小。而targets张量有两个维度:批次大小和标记数量。
对于PyTorch中的cross_entropy损失函数,我们需要通过合并批次维度来将这些张量展平:
logits_flat = logits.flatten(0, 1)
targets_flat = targets.flatten()
print("Flattened logits:", logits_flat.shape)
print("Flattened targets:", targets_flat.shape)
展平后的张量维度为:
Flattened logits: torch.Size([6, 50257])
Flattened targets: torch.Size([6])
请记住,targets是我们希望LLM生成的标记ID,而logits包含在应用softmax函数之前的未缩放模型输出,用于获得概率分数。
之前,我们手动应用了softmax函数,选择了与目标ID对应的概率分数,并计算了负平均对数概率。PyTorch的cross_entropy函数将为我们处理所有这些步骤:
loss = torch.nn.functional.cross_entropy(logits_flat, targets_flat)
print(loss)
生成的损失与我们之前手动执行图5.7中的各个步骤时获得的相同:
tensor(10.7940)
困惑度
困惑度(Perplexity)是一种经常与交叉熵损失一起使用的度量,用于评估模型在语言建模等任务中的表现。它提供了一种更易解释的方式来理解模型在预测序列中下一个标记时的不确定性。
困惑度衡量模型预测的概率分布与数据集中单词的实际分布的匹配程度。类似于损失,较低的困惑度表明模型的预测更接近实际分布。
困惑度可以通过以下公式计算:
perplexity = torch.exp(loss)。在之前计算的损失上应用该公式,返回的结果是:tensor(48725.8203)困惑度通常比原始损失值更具解释性,因为它表示模型在每一步中对有效词汇表大小的不确定性。在这个例子中,这意味着模型在生成下一个标记时,不确定在词汇表中约48,725个标记中的哪个是正确的。
我们现在已经计算了两个小型文本输入的损失作为示例。接下来,我们将对整个训练集和验证集应用损失计算。
计算训练集和验证集的损失
我们首先需要准备用于训练LLM的训练集和验证集。接着,如图5.8所示,我们将为训练集和验证集计算交叉熵损失,这是模型训练过程中非常重要的一部分。
为了计算训练集和验证集的损失,我们使用了一个非常小的文本数据集,即伊迪丝·华顿(Edith Wharton)的短篇小说《裁决》(The Verdict),这是我们在第2章中已经使用过的。选择公共领域的文本可以避免与使用权相关的顾虑。此外,使用这样一个小型数据集可以在普通笔记本电脑上在几分钟内执行代码示例,即使没有高端GPU,这对于教育目的尤其有利。
注意 读者还可以使用本书的补充代码来准备一个更大规模的数据集,包括来自古登堡计划(Project Gutenberg)的超过60,000本公共领域书籍,并在这些书籍上训练LLM(详见附录D)。
预训练LLM的成本
为了理解我们项目的规模,可以参考一个相对流行的公开可用的LLM模型——拥有70亿参数的Llama 2模型的训练情况。该模型在昂贵的A100 GPU上消耗了184,320个GPU小时,处理了2万亿个标记。撰写本文时,在AWS上运行一台8 × A100的云服务器大约每小时花费30美元。粗略估计,训练这样一个LLM的总成本大约为690,000美元(计算方式为184,320小时除以8,再乘以30美元)。
以下代码加载了短篇小说《裁决》:
file_path = "the-verdict.txt"
with open(file_path, "r", encoding="utf-8") as file:
text_data = file.read()
加载数据集后,我们可以检查数据集中的字符数和标记数:
total_characters = len(text_data)
total_tokens = len(tokenizer.encode(text_data))
print("Characters:", total_characters)
print("Tokens:", total_tokens)
输出结果为:
Characters: 20479
Tokens: 5145
虽然只有5,145个标记,文本看起来可能太小而无法用于训练LLM,但如前所述,这是为了教育目的,使得我们能够在几分钟内运行代码,而不是花费数周时间。此外,稍后我们会从OpenAI加载预训练权重到我们的GPTModel代码中。
接下来,我们将数据集划分为训练集和验证集,并使用第2章中的数据加载器来准备用于LLM训练的批次。这一过程在图5.9中可视化展示。由于空间限制,我们使用max_length=6,但对于实际的数据加载器,我们将max_length设置为LLM支持的256个标记的上下文长度,以便在训练过程中LLM能够看到更长的文本。
注意 为了简化和提高效率,我们在训练过程中将数据以相似大小的块提供给模型。然而,在实际操作中,使用可变长度的输入进行训练可以帮助LLM更好地泛化不同类型的输入。
为了实现数据的划分和加载,我们首先定义一个train_ratio,将90%的数据用于训练,剩下的10%作为验证数据,在训练过程中评估模型:
train_ratio = 0.90
split_idx = int(train_ratio * len(text_data))
train_data = text_data[:split_idx]
val_data = text_data[split_idx:]
使用train_data和val_data子集,我们现在可以复用第2章中的create_dataloader_v1代码,创建相应的数据加载器:
from chapter02 import create_dataloader_v1
torch.manual_seed(123)
train_loader = create_dataloader_v1(
train_data,
batch_size=2,
max_length=GPT_CONFIG_124M["context_length"],
stride=GPT_CONFIG_124M["context_length"],
drop_last=True,
shuffle=True,
num_workers=0
)
val_loader = create_dataloader_v1(
val_data,
batch_size=2,
max_length=GPT_CONFIG_124M["context_length"],
stride=GPT_CONFIG_124M["context_length"],
drop_last=False,
shuffle=False,
num_workers=0
)
我们使用了相对较小的批次大小,以减少计算资源需求,因为我们正在处理一个非常小的数据集。实际上,训练LLM时,使用1,024或更大的批次大小并不少见。
作为可选检查,我们可以遍历数据加载器,确保它们被正确创建:
print("Train loader:")
for x, y in train_loader:
print(x.shape, y.shape)
print("\nValidation loader:")
for x, y in val_loader:
print(x.shape, y.shape)
我们应该看到以下输出:
Train loader:
torch.Size([2, 256]) torch.Size([2, 256])
torch.Size([2, 256]) torch.Size([2, 256])
torch.Size([2, 256]) torch.Size([2, 256])
torch.Size([2, 256]) torch.Size([2, 256])
torch.Size([2, 256]) torch.Size([2, 256])
torch.Size([2, 256]) torch.Size([2, 256])
torch.Size([2, 256]) torch.Size([2, 256])
torch.Size([2, 256]) torch.Size([2, 256])
torch.Size([2, 256]) torch.Size([2, 256])
Validation loader:
torch.Size([2, 256]) torch.Size([2, 256])
根据上述代码输出,我们有9个训练集批次,每个批次包含两个样本,每个样本有256个标记。由于我们只将10%的数据用于验证,因此只有一个验证批次包含两个输入示例。正如预期的那样,输入数据(x)和目标数据(y)具有相同的形状(批次大小乘以每个批次中的标记数),因为目标是将输入向前移位一个位置,如第2章所讨论的。
接下来,我们实现一个实用函数,用于计算训练和验证加载器返回的给定批次的交叉熵损失:
def calc_loss_batch(input_batch, target_batch, model, device):
input_batch = input_batch.to(device) #1
target_batch = target_batch.to(device)
logits = model(input_batch)
loss = torch.nn.functional.cross_entropy(
logits.flatten(0, 1), target_batch.flatten()
)
return loss
#1 转移到指定设备,允许将数据传输到GPU。
我们现在可以使用calc_loss_batch实用函数,它计算单个批次的损失,实现以下calc_loss_loader函数,用于计算通过给定数据加载器采样的所有批次的损失。
列表5.2 计算训练和验证损失的函数
def calc_loss_loader(data_loader, model, device, num_batches=None):
total_loss = 0.
if len(data_loader) == 0:
return float("nan")
elif num_batches is None:
num_batches = len(data_loader) #1
else:
num_batches = min(num_batches, len(data_loader)) #2
for i, (input_batch, target_batch) in enumerate(data_loader):
if i < num_batches:
loss = calc_loss_batch(
input_batch, target_batch, model, device
)
total_loss += loss.item() #3
else:
break
return total_loss / num_batches #4
#1 如果没有指定num_batches,则遍历所有批次
#2 如果num_batches超过数据加载器中的批次数量,则减少批次数量
#3 累计每个批次的损失
#4 计算所有批次的平均损失
默认情况下,calc_loss_loader函数遍历给定数据加载器中的所有批次,将损失累积到total_loss变量中,然后计算并平均所有批次的损失。作为替代方案,我们可以通过num_batches指定较小的批次数量,以加快模型训练过程中的评估速度。
现在,让我们实际应用这个calc_loss_loader函数,使用它来计算训练集和验证集加载器的损失:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device) #1
with torch.no_grad(): #2
train_loss = calc_loss_loader(train_loader, model, device) #3
val_loss = calc_loss_loader(val_loader, model, device)
print("Training loss:", train_loss)
print("Validation loss:", val_loss)
#1 如果您的机器有支持CUDA的GPU,LLM将无需修改代码直接在GPU上进行训练。
#2 禁用梯度跟踪以提高效率,因为我们还没有进行训练
#3 通过“device”设置,确保数据被加载到与LLM模型相同的设备上。
生成的损失值为:
Training loss: 10.98758347829183
Validation loss: 10.98110580444336
损失值相对较高,因为模型尚未进行训练。作为对比,如果模型能够学习生成训练集和验证集中的下一个标记,则损失会接近0。
现在我们有了一种评估生成文本质量的方法,接下来我们将训练LLM以减少这种损失,使其在生成文本方面表现得更好,如图5.10所示。
接下来,我们将重点放在对LLM的预训练上。模型训练完成后,我们将实现替代的文本生成策略,并保存和加载预训练的模型权重。
训练LLM
现在是时候实现预训练LLM(我们的GPTModel)的代码了。为此,我们将专注于一个简单的训练循环,以保持代码简洁易读。
注意 对于感兴趣的读者,可以在附录D中学习更高级的技术,包括学习率预热、余弦退火和梯度裁剪等。
图5.11中的流程图展示了典型的PyTorch神经网络训练工作流程,我们将其用于训练LLM。该流程包含八个步骤,开始于遍历每个epoch、处理批次、重置梯度、计算损失和新梯度、更新权重,并最终通过打印损失和生成文本样本来监控训练过程。
注意 如果你对使用PyTorch训练深度神经网络比较陌生,且对这些步骤感到不熟悉,可以参考附录A的A.5到A.8部分。
我们可以通过以下代码实现这个训练流程,使用train_model_simple函数:
列表5.3 预训练LLM的主要函数
def train_model_simple(model, train_loader, val_loader,
optimizer, device, num_epochs,
eval_freq, eval_iter, start_context, tokenizer):
train_losses, val_losses, track_tokens_seen = [], [], [] #1
tokens_seen, global_step = 0, -1
for epoch in range(num_epochs): #2
model.train()
for input_batch, target_batch in train_loader:
optimizer.zero_grad() #3
loss = calc_loss_batch(
input_batch, target_batch, model, device
)
loss.backward() #4
optimizer.step() #5
tokens_seen += input_batch.numel()
global_step += 1
if global_step % eval_freq == 0: #6
train_loss, val_loss = evaluate_model(
model, train_loader, val_loader, device, eval_iter)
train_losses.append(train_loss)
val_losses.append(val_loss)
track_tokens_seen.append(tokens_seen)
print(f"Ep {epoch+1} (Step {global_step:06d}): "
f"Train loss {train_loss:.3f}, "
f"Val loss {val_loss:.3f}"
)
generate_and_print_sample( #7
model, tokenizer, device, start_context
)
return train_losses, val_losses, track_tokens_seen
#1 初始化列表以跟踪损失和看到的标记
#2 开始主要训练循环
#3 重置前一个批次迭代的损失梯度
#4 计算损失梯度
#5 使用损失梯度更新模型权重
#6 可选的评估步骤
#7 在每个epoch之后打印一个文本样本
需要注意的是,train_model_simple函数使用了两个尚未定义的函数:evaluate_model 和 generate_and_print_sample。
evaluate_model函数对应图5.11中的第7步。它在每次模型更新后打印训练集和验证集的损失,以便我们评估训练是否提升了模型性能。具体来说,evaluate_model函数计算训练集和验证集上的损失,同时确保模型处于评估模式,在计算损失时禁用梯度跟踪和dropout:
def evaluate_model(model, train_loader, val_loader, device, eval_iter):
model.eval() #1
with torch.no_grad(): #2
train_loss = calc_loss_loader(
train_loader, model, device, num_batches=eval_iter
)
val_loss = calc_loss_loader(
val_loader, model, device, num_batches=eval_iter
)
model.train()
return train_loss, val_loss
#1 在评估期间禁用dropout,以确保结果稳定可重复
#2 禁用梯度跟踪,因为评估时不需要它,从而减少计算开销
与evaluate_model类似,generate_and_print_sample是一个方便的函数,用于跟踪模型在训练过程中的改进情况。特别是,generate_and_print_sample函数接受一个文本片段(start_context)作为输入,将其转换为标记ID,并使用前面使用的generate_text_simple函数将其输入LLM生成文本样本:
def generate_and_print_sample(model, tokenizer, device, start_context):
model.eval()
context_size = model.pos_emb.weight.shape[0]
encoded = text_to_token_ids(start_context, tokenizer).to(device)
with torch.no_grad():
token_ids = generate_text_simple(
model=model, idx=encoded,
max_new_tokens=50, context_size=context_size
)
decoded_text = token_ids_to_text(token_ids, tokenizer)
print(decoded_text.replace("\n", " ")) #1
model.train()
#1 使用紧凑的打印格式
evaluate_model函数为我们提供了模型训练进展的数值估计,而generate_and_print_sample文本函数则提供了一个由模型生成的具体文本示例,用以判断其在训练期间的能力。
AdamW
Adam优化器是训练深度神经网络的热门选择。然而,在我们的训练循环中,我们选择了AdamW优化器。AdamW是Adam的一个变体,它改进了权重衰减的方法,通过惩罚较大的权重来最小化模型的复杂性并防止过拟合。这一调整使得AdamW能够实现更有效的正则化和更好的泛化能力,因此,AdamW在LLM的训练中经常被使用。
让我们通过训练一个GPTModel实例来实践这些内容,训练10个epoch,使用AdamW优化器和之前定义的train_model_simple函数:
torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
model.to(device)
optimizer = torch.optim.AdamW(
model.parameters(), #1
lr=0.0004, weight_decay=0.1
)
num_epochs = 10
train_losses, val_losses, tokens_seen = train_model_simple(
model, train_loader, val_loader, optimizer, device,
num_epochs=num_epochs, eval_freq=5, eval_iter=5,
start_context="Every effort moves you", tokenizer=tokenizer
)
#1 .parameters()方法返回模型所有可训练的权重参数。
执行train_model_simple函数将启动训练过程,在一台MacBook Air或类似的笔记本电脑上大约需要5分钟完成。执行期间输出如下:
Ep 1 (Step 000000): Train loss 9.781, Val loss 9.933
Ep 1 (Step 000005): Train loss 8.111, Val loss 8.339
Every effort moves you,,,,,,,,,,,,.
Ep 2 (Step 000010): Train loss 6.661, Val loss 7.048
Ep 2 (Step 000015): Train loss 5.961, Val loss 6.616
Every effort moves you, and, and, and, and, and, and, and, and, and, and,
and, and, and, and, and, and, and, and, and, and, and, and,, and, and,
[...] #1
Ep 9 (Step 000080): Train loss 0.541, Val loss 6.393
Every effort moves you?" "Yes--quite insensible to the irony. She wanted
him vindicated--and by me!" He laughed again, and threw back the
window-curtains, I had the donkey. "There were days when I
Ep 10 (Step 000085): Train loss 0.391, Val loss 6.452
Every effort moves you know," was one of the axioms he laid down across the
Sevres and silver of an exquisitely appointed luncheon-table, when, on a
later day, I had again run over from Monte Carlo; and Mrs. Gis
#1 省略了部分中间结果以节省空间。
如我们所见,训练损失显著下降,从9.781开始收敛到0.391。模型的语言能力有了很大的提高。在开始时,模型只能将逗号附加到起始上下文(如“Every effort moves you,,,,,,,,,,,,”)或重复“and”一词。在训练结束时,它可以生成语法正确的文本。
与训练集损失相似,我们可以看到验证集损失从一开始的9.933逐渐下降,尽管它从未像训练集损失那样接近0,在第10个epoch后停留在6.452。
在进一步讨论验证集损失之前,让我们创建一个简单的图表,显示训练集和验证集损失的对比:
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
def plot_losses(epochs_seen, tokens_seen, train_losses, val_losses):
fig, ax1 = plt.subplots(figsize=(5, 3))
ax1.plot(epochs_seen, train_losses, label="Training loss")
ax1.plot(
epochs_seen, val_losses, linestyle="-.", label="Validation loss"
)
ax1.set_xlabel("Epochs")
ax1.set_ylabel("Loss")
ax1.legend(loc="upper right")
ax1.xaxis.set_major_locator(MaxNLocator(integer=True))
ax2 = ax1.twiny() #1
ax2.plot(tokens_seen, train_losses, alpha=0) #2
ax2.set_xlabel("Tokens seen")
fig.tight_layout()
plt.show()
epochs_tensor = torch.linspace(0, num_epochs, len(train_losses))
plot_losses(epochs_tensor, tokens_seen, train_losses, val_losses)
#1 创建一个共享相同y轴的第二个x轴。
#2 隐藏图形以对齐刻度。
生成的训练和验证损失图如图5.12所示。我们可以看到,训练和验证损失在第一个epoch时都开始改善。然而,损失在第二个epoch后开始分离。这种分离以及验证损失远高于训练损失表明,模型正在对训练数据过拟合。通过搜索生成的文本片段(例如,"quite insensible to the irony")在《裁决》文本文件中的出现,我们可以确认模型在记忆训练数据。
这种记忆是预期中的,因为我们使用的是非常小的训练数据集,并且训练模型进行了多个epoch。通常,模型是在更大的数据集上只训练一个epoch。
注意 如前所述,感兴趣的读者可以尝试在古登堡计划中提供的60,000本公共领域书籍上训练模型,在这种情况下不会发生过拟合;详见附录B。
如图5.13所示,我们已经完成了本章的四个目标。接下来,我们将讨论LLM的文本生成策略,以减少训练数据的记忆性并提高LLM生成文本的原创性。之后,我们将介绍权重加载以及如何保存和加载OpenAI的GPT模型预训练权重。
控制随机性的解码策略
让我们看看用于生成更具原创性的文本生成策略(也称为解码策略)。首先,我们将简要回顾之前在generate_and_print_sample函数中使用的generate_text_simple函数。然后,我们将介绍两种技术,温度缩放和top-k采样,以改进该函数。
首先,我们将模型从GPU转回CPU,因为使用相对较小的模型进行推理不需要GPU。同时,在训练后,我们将模型置于评估模式,以关闭随机组件(如dropout):
model.to("cpu")
model.eval()
接下来,我们将GPTModel实例(model)插入generate_text_simple函数中,该函数使用LLM一次生成一个标记:
tokenizer = tiktoken.get_encoding("gpt2")
token_ids = generate_text_simple(
model=model,
idx=text_to_token_ids("Every effort moves you", tokenizer),
max_new_tokens=25,
context_size=GPT_CONFIG_124M["context_length"]
)
print("Output text:\n", token_ids_to_text(token_ids, tokenizer))
生成的文本如下:
Output text:
Every effort moves you know," was one of the axioms he laid down across the
Sevres and silver of an exquisitely appointed lun
正如之前解释的那样,每次生成步骤中,生成的标记是从词汇表中选择具有最大概率分数的标记。这意味着,即使我们多次在相同的起始上下文(如“Every effort moves you”)上运行上述generate_text_simple函数,LLM也总是会生成相同的输出。
温度缩放
现在我们来看看温度缩放,这是一种为下一个标记生成任务引入概率选择过程的技术。此前,在generate_text_simple函数中,我们总是通过torch.argmax选择概率最高的标记作为下一个标记,这也被称为贪婪解码。为了生成更多样化的文本,我们可以用从概率分布中采样的方式替代argmax,即基于每次生成步骤中LLM为词汇表生成的概率分数进行采样。
为了用具体例子说明概率采样,让我们简要讨论一个使用小词汇表的下一个标记生成过程:
vocab = {
"closer": 0,
"every": 1,
"effort": 2,
"forward": 3,
"inches": 4,
"moves": 5,
"pizza": 6,
"toward": 7,
"you": 8,
}
inverse_vocab = {v: k for k, v in vocab.items()}
接下来,假设LLM给定起始上下文“every effort moves you”,并生成了以下的下一个标记的logits:
next_token_logits = torch.tensor(
[4.51, 0.89, -1.90, 6.75, 1.63, -1.62, -1.89, 6.28, 1.79]
)
如第4章所述,在generate_text_simple函数中,我们通过softmax函数将logits转换为概率,并通过argmax函数获得生成的标记ID,然后可以通过反向词汇表将其映射回文本:
probas = torch.softmax(next_token_logits, dim=0)
next_token_id = torch.argmax(probas).item()
print(inverse_vocab[next_token_id])
由于最大logit值在第四个位置(Python索引为0起,因此是索引位置3),生成的单词是“forward”。
为了实现一个概率采样过程,我们可以用PyTorch的multinomial函数替代argmax:
torch.manual_seed(123)
next_token_id = torch.multinomial(probas, num_samples=1).item()
print(inverse_vocab[next_token_id])
输出仍然是“forward”。这是怎么回事?multinomial函数根据概率分数采样下一个标记。换句话说,“forward”依然是最可能的标记,并且在大多数情况下会被multinomial选中,但不是每次都这样。为了说明这一点,让我们实现一个函数,将这种采样重复1,000次:
def print_sampled_tokens(probas):
torch.manual_seed(123)
sample = [torch.multinomial(probas, num_samples=1).item()
for i in range(1_000)]
sampled_ids = torch.bincount(torch.tensor(sample))
for i, freq in enumerate(sampled_ids):
print(f"{freq} x {inverse_vocab[i]}")
输出结果是:
73 x closer
0 x every
0 x effort
582 x forward
2 x inches
0 x moves
0 x pizza
343 x toward
正如我们所见,"forward"在大多数情况下(1,000次中582次)被采样,但其他单词如"closer"、"inches"和"toward"有时也会被采样。这意味着如果我们在generate_and_print_sample函数中用multinomial替代argmax,LLM可能会生成“every effort moves you toward”、“every effort moves you inches”或“every effort moves you closer”,而不仅仅是“every effort moves you forward”。
我们还可以通过一个称为温度缩放的概念进一步控制分布和选择过程。温度缩放是将logits除以大于0的一个数:
def softmax_with_temperature(logits, temperature):
scaled_logits = logits / temperature
return torch.softmax(scaled_logits, dim=0)
温度大于1时,会使标记概率分布更加均匀;温度小于1时,会使分布更加集中(更尖锐)。让我们通过绘制原始概率和不同温度值缩放后的概率来说明这一点:
temperatures = [1, 0.1, 5] #1
scaled_probas = [softmax_with_temperature(next_token_logits, T)
for T in temperatures]
x = torch.arange(len(vocab))
bar_width = 0.15
fig, ax = plt.subplots(figsize=(5, 3))
for i, T in enumerate(temperatures):
rects = ax.bar(x + i * bar_width, scaled_probas[i],
bar_width, label=f'Temperature = {T}')
ax.set_ylabel('Probability')
ax.set_xticks(x)
ax.set_xticklabels(vocab.keys(), rotation=90)
ax.legend()
plt.tight_layout()
plt.show()
#1 原始温度、较低温度和较高温度的对比。
生成的图如图5.14所示。
温度为1时,将logits除以1后再传递给softmax函数以计算概率分数。换句话说,使用温度1与不使用温度缩放是相同的。在这种情况下,使用PyTorch的multinomial采样函数,标记的选择概率等于原始的softmax概率分数。例如,对于温度设置为1时,"forward"对应的标记大约会在60%的情况下被选中,如图5.14所示。
此外,如图5.14所示,应用非常小的温度(如0.1)将导致更尖锐的分布,使得multinomial函数几乎100%选择最可能的标记(在这里是"forward"),接近argmax函数的行为。同样,温度为5时会产生更均匀的分布,其他标记的选择频率会更高。这可以为生成的文本增加更多的多样性,但也更容易生成无意义的文本。例如,使用温度5时,生成诸如“every effort moves you pizza”这样的文本的概率大约为4%。
练习 5.1
使用
print_sampled_tokens函数打印使用图5.14中温度缩放后的softmax概率的采样频率。在每种情况下,单词“pizza”被采样的频率是多少?你能想到一种更快且更准确的方法来确定“pizza”被采样的频率吗?
Top-k采样
我们已经实现了一种结合温度缩放的概率采样方法,以增加输出的多样性。我们看到,较高的温度值会使下一个标记的概率分布更均匀,从而减少模型反复选择最可能标记的可能性,生成更加多样化的输出。这种方法可以在生成过程中探索不太可能但可能更有趣和创造性的路径。然而,这种方法的一个缺点是有时会生成语法不正确或完全无意义的输出,如“every effort moves you pizza”。
Top-k采样结合概率采样和温度缩放,可以改进文本生成结果。在Top-k采样中,我们将采样限制在前k个最可能的标记中,并通过屏蔽其他标记的概率分数将它们排除在选择过程之外,如图5.15所示。
Top-k方法将所有未被选中的logits替换为负无穷大值(-inf),这样在计算softmax值时,非top-k标记的概率为0,剩余的概率总和为1。(细心的读者可能还记得我们在第3章的因果注意力模块中实现过这种屏蔽技巧,见3.5.1节。)
我们可以按照图5.15的方式在代码中实现Top-k过程,首先选择具有最大logit值的标记:
top_k = 3
top_logits, top_pos = torch.topk(next_token_logits, top_k)
print("Top logits:", top_logits)
print("Top positions:", top_pos)
前三个标记的logits值和标记ID,按降序排列:
Top logits: tensor([6.7500, 6.2800, 4.5100])
Top positions: tensor([3, 7, 0])
接着,我们使用PyTorch的where函数,将不在前三选择中的标记的logit值设为负无穷大(-inf):
new_logits = torch.where(
condition=next_token_logits < top_logits[-1], #1
input=torch.tensor(float('-inf')), #2
other=next_token_logits #3
)
print(new_logits)
#1 识别低于前3个logits最小值的logits
#2 将这些较小的logits设为–inf
#3 保留其他标记的原始logits
在九个标记的词汇表中,下一个标记的logits结果为:
tensor([4.5100, -inf, -inf, 6.7500, -inf, -inf, -inf, 6.2800,
-inf])
最后,我们应用softmax函数将这些logits转换为下一个标记的概率:
topk_probas = torch.softmax(new_logits, dim=0)
print(topk_probas)
正如我们所见,top-3方法的结果是三个非零概率分数:
tensor([0.0615, 0.0000, 0.0000, 0.5775, 0.0000, 0.0000, 0.0000, 0.3610,
0.0000])
现在我们可以应用温度缩放和multinomial函数进行概率采样,从这三个非零概率分数中选择下一个标记来生成文本。接下来我们将通过修改文本生成函数来实现这一点。
修改文本生成函数
现在,让我们结合温度采样和Top-k采样,对之前用于通过LLM生成文本的generate_text_simple函数进行修改,创建一个新的generate函数。
列表5.4 一个带有更多多样性的修改后的文本生成函数
def generate(model, idx, max_new_tokens, context_size,
temperature=0.0, top_k=None, eos_id=None):
for _ in range(max_new_tokens): #1
idx_cond = idx[:, -context_size:]
with torch.no_grad():
logits = model(idx_cond)
logits = logits[:, -1, :]
if top_k is not None: #2
top_logits, _ = torch.topk(logits, top_k)
min_val = top_logits[:, -1]
logits = torch.where(
logits < min_val,
torch.tensor(float('-inf')).to(logits.device),
logits
)
if temperature > 0.0: #3
logits = logits / temperature
probs = torch.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
else: #4
idx_next = torch.argmax(logits, dim=-1, keepdim=True)
if idx_next == eos_id: #5
break
idx = torch.cat((idx, idx_next), dim=1)
return idx
#1 循环与之前相同:获取logits并仅关注最后一个时间步。
#2 通过Top-k采样过滤logits。
#3 应用温度缩放。
#4 当温度缩放未启用时,执行贪婪的下一个标记选择。
#5 如果遇到序列结束标记(eos),则提前停止生成。
现在让我们看看这个新的generate函数的效果:
torch.manual_seed(123)
token_ids = generate(
model=model,
idx=text_to_token_ids("Every effort moves you", tokenizer),
max_new_tokens=15,
context_size=GPT_CONFIG_124M["context_length"],
top_k=25,
temperature=1.4
)
print("Output text:\n", token_ids_to_text(token_ids, tokenizer))
生成的文本是:
Output text:
Every effort moves you stand to work on surprise, a one of us had gone
with random-
正如我们所见,生成的文本与我们之前通过generate_simple函数生成的文本截然不同(之前生成的文本是:"Every effort moves you know," was one of the axioms he laid...),而那个文本是训练集中的一个记忆片段。
练习 5.2
尝试使用不同的温度和Top-k设置。根据你的观察,能否想到哪些应用场景需要较低的温度和Top-k设置?同样,能否想到哪些应用场景更适合较高的温度和Top-k设置?(建议在本章结束后加载OpenAI的预训练权重后,重新进行此练习。)
练习 5.3
有哪些不同的
generate函数设置组合可以强制其表现为确定性行为,即禁用随机采样,使其总是生成与generate_simple函数相似的相同输出?
在PyTorch中加载和保存模型权重
到目前为止,我们已经讨论了如何数值评估训练进展并从零开始预训练一个LLM。即使LLM和数据集相对较小,此练习仍表明预训练LLM需要大量的计算资源。因此,能够保存LLM模型是非常重要的,这样我们就不必每次在新会话中使用它时都重新训练。
现在,让我们讨论如何保存和加载预训练的模型,如图5.16所示。稍后,我们还会将OpenAI的一个更强大的预训练GPT模型加载到我们的GPTModel实例中。
幸运的是,保存PyTorch模型相对简单。推荐的方式是使用torch.save函数保存模型的state_dict,这是一个将每一层映射到其参数的字典:
torch.save(model.state_dict(), "model.pth")
"model.pth"是保存state_dict的文件名,.pth扩展名是PyTorch文件的惯例,尽管我们技术上可以使用任何文件扩展名。
之后,通过保存的state_dict,我们可以将模型权重加载到一个新的GPTModel实例中:
model = GPTModel(GPT_CONFIG_124M)
model.load_state_dict(torch.load("model.pth", map_location=device))
model.eval()
正如第4章讨论的那样,dropout有助于防止模型过拟合训练数据,通过在训练期间随机“丢弃”某些神经元。然而,在推理过程中,我们不希望随机丢弃网络已经学习的信息。使用model.eval()将模型切换到推理的评估模式,禁用模型的dropout层。如果我们计划在以后继续预训练模型——例如使用本章之前定义的train_model_simple函数——建议同时保存优化器的状态。
像AdamW这样的自适应优化器会为每个模型权重存储额外的参数。AdamW使用历史数据动态调整每个模型参数的学习率。如果没有这些参数,优化器会重置,模型可能无法达到最佳学习效果,甚至无法正确收敛,这意味着它可能失去生成连贯文本的能力。我们可以使用torch.save保存模型和优化器的state_dict内容:
torch.save({
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
},
"model_and_optimizer.pth"
)
然后,我们可以通过torch.load加载保存的数据,并使用load_state_dict方法恢复模型和优化器状态:
checkpoint = torch.load("model_and_optimizer.pth", map_location=device)
model = GPTModel(GPT_CONFIG_124M)
model.load_state_dict(checkpoint["model_state_dict"])
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4, weight_decay=0.1)
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
model.train();
练习 5.4
在保存权重之后,在新的Python会话或Jupyter notebook文件中加载模型和优化器,并使用
train_model_simple函数继续预训练一个epoch。
加载OpenAI的预训练权重
之前,我们使用一本短篇小说训练了一个小型的GPT-2模型,这使我们能专注于基础知识,而无需耗费大量时间和计算资源。幸运的是,OpenAI公开分享了他们的GPT-2模型权重,这样我们就不需要花费数万到数十万美元来训练模型。接下来,我们将这些权重加载到我们的GPTModel类中,用于文本生成。
在PyTorch中,权重是指线性层和嵌入层的.weight属性中存储的权重参数。例如,我们在训练模型时通过model.parameters()访问了它们。在第6章中,我们还将重用这些预训练权重来微调模型以进行文本分类任务,并遵循类似ChatGPT的指令。
请注意,OpenAI最初使用TensorFlow保存了GPT-2权重,因此我们需要安装TensorFlow来在Python中加载权重。接下来的代码还使用了一个名为tqdm的进度条工具来跟踪下载进度,也需要安装。
你可以在终端中执行以下命令来安装这些库:
pip install tensorflow>=2.15.0 tqdm>=4.66
下载代码相对较长,主要是模板代码,因此我们直接从本章的在线仓库下载gpt_download.py模块:
import urllib.request
url = (
"https://raw.githubusercontent.com/rasbt/"
"LLMs-from-scratch/main/ch05/"
"01_main-chapter-code/gpt_download.py"
)
filename = url.split('/')[-1]
urllib.request.urlretrieve(url, filename)
下载该文件后,可以简要检查文件内容以确保其已正确保存并包含有效的Python代码。
现在,我们可以从gpt_download.py文件中导入download_and_load_gpt2函数,加载GPT-2架构设置(settings)和权重参数(params):
from gpt_download import download_and_load_gpt2
settings, params = download_and_load_gpt2(
model_size="124M", models_dir="gpt2"
)
执行此代码将下载与124M参数GPT-2模型相关的以下七个文件:
checkpoint: 100%|███████████████████████████| 77.0/77.0 [00:00<00:00,
63.9kiB/s]
encoder.json: 100%|█████████████████████████| 1.04M/1.04M [00:00<00:00,
2.20MiB/s]
hprams.json: 100%|██████████████████████████| 90.0/90.0 [00:00<00:00,
78.3kiB/s]
model.ckpt.data-00000-of-00001: 100%|███████| 498M/498M [01:09<00:00,
7.16MiB/s]
model.ckpt.index: 100%|█████████████████████| 5.21k/5.21k [00:00<00:00,
3.24MiB/s]
model.ckpt.meta: 100%|██████████████████████| 471k/471k [00:00<00:00,
2.46MiB/s]
vocab.bpe: 100%|████████████████████████████| 456k/456k [00:00<00:00,
1.70MiB/s]
注意 如果下载代码无法正常工作,可能是由于网络连接问题、服务器问题或OpenAI更改了共享GPT-2模型权重的方式。请访问本章的在线代码仓库 github.com/rasbt/LLMs-… 获取替代和更新的说明,并通过Manning论坛提出进一步的问题。
假设前面的代码已执行完成,让我们检查settings和params的内容:
print("Settings:", settings)
print("Parameter dictionary keys:", params.keys())
内容如下:
Settings: {'n_vocab': 50257, 'n_ctx': 1024, 'n_embd': 768, 'n_head': 12,
'n_layer': 12}
Parameter dictionary keys: dict_keys(['blocks', 'b', 'g', 'wpe', 'wte'])
settings和params都是Python字典。settings字典存储了LLM的架构设置,类似于我们手动定义的GPT_CONFIG_124M设置。params字典包含实际的权重张量。我们只打印了字典的键,因为打印权重内容会占用过多屏幕空间;然而,我们可以通过print(params)或通过各自的字典键选择单个张量来检查这些权重。例如,嵌入层权重:
print(params["wte"])
print("Token embedding weight tensor dimensions:", params["wte"].shape)
输出为:
[[-0.11010301 ... -0.1363697 0.01506208 0.04531523]
[ 0.04034033 ... 0.08605453 0.00253983 0.04318958]
[-0.12746179 ... 0.08991534 -0.12972379 -0.08785918]
...
[-0.04453601 ... 0.10435229 0.09783269 -0.06952604]
[ 0.1860082 ... -0.09625227 0.07847701 -0.02245961]
[ 0.05135201 ... 0.00704835 0.15519823 0.12067825]]
Token embedding weight tensor dimensions: (50257, 768)
我们通过download_and_load_gpt2(model_size="124M", ...)设置下载并加载了最小的GPT-2模型的权重。OpenAI还分享了更大的模型权重:355M、774M和1558M。这些不同大小的GPT模型的总体架构是相同的,如图5.17所示,不同之处在于架构元素的重复次数和嵌入大小不同。本章的其余代码也兼容这些更大的模型。
将GPT-2模型权重加载到Python中后,我们仍需要将它们从settings和params字典中转移到GPTModel实例中。首先,我们创建一个字典,列出了不同大小GPT模型之间的区别,如图5.17所示:
model_configs = {
"gpt2-small (124M)": {"emb_dim": 768, "n_layers": 12, "n_heads": 12},
"gpt2-medium (355M)": {"emb_dim": 1024, "n_layers": 24, "n_heads": 16},
"gpt2-large (774M)": {"emb_dim": 1280, "n_layers": 36, "n_heads": 20},
"gpt2-xl (1558M)": {"emb_dim": 1600, "n_layers": 48, "n_heads": 25},
}
假设我们对加载最小的模型 "gpt2-small (124M)" 感兴趣。我们可以使用model_configs表中的对应设置来更新我们之前定义和使用的GPT_CONFIG_124M:
model_name = "gpt2-small (124M)"
NEW_CONFIG = GPT_CONFIG_124M.copy()
NEW_CONFIG.update(model_configs[model_name])
细心的读者可能还记得,我们之前使用的是256个token的长度,但OpenAI的原始GPT-2模型是使用1024个token长度进行训练的,因此我们需要相应地更新NEW_CONFIG:
NEW_CONFIG.update({"context_length": 1024})
此外,OpenAI在多头注意力模块的线性层中使用了偏置向量来实现查询、键和值矩阵的计算。虽然LLM中偏置向量已不常用了,但为了与预训练权重保持一致,我们仍需要启用这些偏置向量:
NEW_CONFIG.update({"qkv_bias": True})
现在,我们可以使用更新后的NEW_CONFIG字典来初始化一个新的GPTModel实例:
gpt = GPTModel(NEW_CONFIG)
gpt.eval()
默认情况下,GPTModel实例是用随机权重初始化的,用于预训练。要使用OpenAI的模型权重,我们最后需要将这些随机权重替换为加载到params字典中的权重。为此,我们首先定义一个小的assign工具函数,用于检查两个张量或数组(left和right)是否具有相同的维度或形状,并将right张量返回为可训练的PyTorch参数:
def assign(left, right):
if left.shape != right.shape:
raise ValueError(f"Shape mismatch. Left: {left.shape}, "
"Right: {right.shape}"
)
return torch.nn.Parameter(torch.tensor(right))
接着,我们定义load_weights_into_gpt函数,用于将权重从params字典加载到GPTModel实例gpt中:
import numpy as np
def load_weights_into_gpt(gpt, params):
gpt.pos_emb.weight = assign(gpt.pos_emb.weight, params['wpe'])
gpt.tok_emb.weight = assign(gpt.tok_emb.weight, params['wte'])
for b in range(len(params["blocks"])):
q_w, k_w, v_w = np.split(
(params["blocks"][b]["attn"]["c_attn"])["w"], 3, axis=-1)
gpt.trf_blocks[b].att.W_query.weight = assign(
gpt.trf_blocks[b].att.W_query.weight, q_w.T)
gpt.trf_blocks[b].att.W_key.weight = assign(
gpt.trf_blocks[b].att.W_key.weight, k_w.T)
gpt.trf_blocks[b].att.W_value.weight = assign(
gpt.trf_blocks[b].att.W_value.weight, v_w.T)
q_b, k_b, v_b = np.split(
(params["blocks"][b]["attn"]["c_attn"])["b"], 3, axis=-1)
gpt.trf_blocks[b].att.W_query.bias = assign(
gpt.trf_blocks[b].att.W_query.bias, q_b)
gpt.trf_blocks[b].att.W_key.bias = assign(
gpt.trf_blocks[b].att.W_key.bias, k_b)
gpt.trf_blocks[b].att.W_value.bias = assign(
gpt.trf_blocks[b].att.W_value.bias, v_b)
gpt.trf_blocks[b].att.out_proj.weight = assign(
gpt.trf_blocks[b].att.out_proj.weight,
params["blocks"][b]["attn"]["c_proj"]["w"].T)
gpt.trf_blocks[b].att.out_proj.bias = assign(
gpt.trf_blocks[b].att.out_proj.bias,
params["blocks"][b]["attn"]["c_proj"]["b"])
gpt.trf_blocks[b].ff.layers[0].weight = assign(
gpt.trf_blocks[b].ff.layers[0].weight,
params["blocks"][b]["mlp"]["c_fc"]["w"].T)
gpt.trf_blocks[b].ff.layers[0].bias = assign(
gpt.trf_blocks[b].ff.layers[0].bias,
params["blocks"][b]["mlp"]["c_fc"]["b"])
gpt.trf_blocks[b].ff.layers[2].weight = assign(
gpt.trf_blocks[b].ff.layers[2].weight,
params["blocks"][b]["mlp"]["c_proj"]["w"].T)
gpt.trf_blocks[b].ff.layers[2].bias = assign(
gpt.trf_blocks[b].ff.layers[2].bias,
params["blocks"][b]["mlp"]["c_proj"]["b"])
gpt.trf_blocks[b].norm1.scale = assign(
gpt.trf_blocks[b].norm1.scale,
params["blocks"][b]["ln_1"]["g"])
gpt.trf_blocks[b].norm1.shift = assign(
gpt.trf_blocks[b].norm1.shift,
params["blocks"][b]["ln_1"]["b"])
gpt.trf_blocks[b].norm2.scale = assign(
gpt.trf_blocks[b].norm2.scale,
params["blocks"][b]["ln_2"]["g"])
gpt.trf_blocks[b].norm2.shift = assign(
gpt.trf_blocks[b].norm2.shift,
params["blocks"][b]["ln_2"]["b"])
gpt.final_norm.scale = assign(gpt.final_norm.scale, params["g"])
gpt.final_norm.shift = assign(gpt.final_norm.shift, params["b"])
gpt.out_head.weight = assign(gpt.out_head.weight, params["wte"])
我们在load_weights_into_gpt函数中仔细匹配了OpenAI实现中的权重和我们的GPTModel实现。现在让我们在实践中尝试load_weights_into_gpt,并将OpenAI的模型权重加载到GPTModel实例gpt中:
load_weights_into_gpt(gpt, params)
gpt.to(device)
如果模型正确加载,我们现在可以使用它生成新文本:
torch.manual_seed(123)
token_ids = generate(
model=gpt,
idx=text_to_token_ids("Every effort moves you", tokenizer).to(device),
max_new_tokens=25,
context_size=NEW_CONFIG["context_length"],
top_k=50,
temperature=1.5
)
print("Output text:\n", token_ids_to_text(token_ids, tokenizer))
生成的文本如下:
Output text:
Every effort moves you toward finding an ideal new way to practice something!
What makes us want to be on top of that?
我们可以确信模型权重加载正确,因为模型能够生成连贯的文本。微小的错误都会导致模型无法正常工作。在接下来的章节中,我们将进一步使用这个预训练模型,并对其进行微调,以用于文本分类和执行指令。
练习 5.5
计算带有OpenAI预训练权重的GPTModel在“The Verdict”数据集上的训练集和验证集损失。
练习 5.6
尝试使用不同大小的GPT-2模型,例如,最大的1,558百万参数模型,并将其生成的文本与124百万参数模型进行比较。
总结
- 当大型语言模型(LLM)生成文本时,它一次生成一个标记。
- 默认情况下,下一个标记通过将模型输出转换为概率分数,并选择对应于最高概率的标记生成,这被称为“贪婪解码”。
- 使用概率采样和温度缩放,我们可以影响生成文本的多样性和连贯性。
- 训练集和验证集损失可以用来衡量LLM在训练过程中生成的文本质量。
- 预训练LLM的过程涉及通过调整模型的权重以最小化训练损失。
- LLM的训练循环是深度学习中的标准程序,使用常规的交叉熵损失和AdamW优化器。
- 在大规模文本语料库上预训练LLM耗时且资源密集,因此我们可以加载公开可用的权重,作为自己在大型数据集上预训练模型的替代方案。