从零实现ChatGPT:第四章在无标签数据上预训练

准备深入学习transformer，并参考一些资料和论文实现一个大语言模型，顺便做一个教程，今天是第四部分。
本系列禁止转载，主要是为了有不同见解的同学可以方便联系我，我的邮箱 fanzexuan135@163.com

第4章：在无标签数据上预训练

本章将介绍如何在无标签数据上预训练语言模型。我们将学习如何：

• 在训练过程中计算训练集和验证集的损失，以评估生成文本的质量
• 实现训练函数并预训练语言模型
• 保存和加载模型权重以继续训练
• 从OpenAI加载预训练权重

在前面的章节中，我们实现了数据采样、注意力机制，并编写了语言模型架构的代码。本章的核心重点是实现训练函数并从头开始预训练语言模型，如图3-1所示。

图3-1 编写语言模型的三个主要阶段示意图：在通用文本数据集上预训练，然后在标注数据集上微调。本章重点介绍语言模型的预训练，包括实现训练代码、评估性能以及保存和加载模型权重。

如图3-1所示，我们还将学习一些基本的模型评估技术，用于衡量生成文本的质量，这是优化语言模型训练过程的必要条件。此外，我们将讨论如何加载预训练权重，为后续章节中对语言模型的微调提供一个良好的起点。

权重参数

在语言模型和其他深度学习模型中，权重指的是学习过程中调整的可训练参数。这些权重也被称为权重参数或简称参数。在PyTorch等框架中，这些权重存储在线性层中，例如我们在第2章实现多头注意力模块和第2章实现GPT模型时使用的那些。初始化一个层后（new_layer = torch.nn.Linear(...)），我们可以通过weight属性访问其权重：new_layer.weight。此外，为了方便起见，PyTorch允许通过model.parameters()方法直接访问模型的所有可训练参数，包括权重和偏置，我们将在后面实现模型训练时使用它。

评估生成文本模型

我们首先基于前一章的代码设置语言模型进行文本生成，并讨论评估生成文本质量的基本方法。本节的内容概述如图3-2所示。

图3-2 本章涵盖的主题概览。我们首先回顾上一章的文本生成，并实现基本的模型评估技术，可在预训练阶段使用。

如图3-2所示，下一小节将回顾我们在上一章末设置的文本生成，然后在后续小节中深入探讨文本评估和计算训练集与验证集损失。

使用GPT生成文本

这一部分我们基于第2章的代码设置语言模型，并简要回顾了文本生成过程。首先，使用第2章的GPTModel类和GPT_CONFIG_M字典初始化将在本章评估和训练的GPT模型：

import torch
from chapter2 import GPTModel

GPT_CONFIG_M = {
    'vocab_size': 50257,
    'context_length': 512,  # 修改为512 tokens  
    'emb_dim': 768, 
    'n_heads': 12,
    'n_layers': 12,
    'dropout_rate': 0.1,
    'tkv_bias': False,
}

torch.manual_seed(1234)  
model = GPTModel(GPT_CONFIG_M)
model.eval()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

相比上一章的GPT_CONFIG_M字典，我们唯一的调整是将上下文长度context_length减少到512个token。这一修改减少了训练模型的计算需求，使得在标准笔记本电脑上进行训练成为可能。

最初，具有1.1亿参数的GPT模型配置为可以处理多达2048个token。在本章末尾的训练过程之后，我们将更新上下文大小设置，并加载预训练权重以使用配置为1024个token上下文长度的模型。

使用GPT模型实例，我们采用上一章介绍的generate_text_simple函数，并引入两个方便的函数text_to_token_ids和token_ids_to_text。这些函数方便在文本和token表示之间进行转换，这是我们将在本章中使用的技巧。为了更清楚地理解，图3-3说明了这个过程，然后我们深入研究代码。

图3-3 生成文本涉及将文本编码为语言模型处理为Logit向量的token ID。然后将Logit向量转换回token ID并解码为文本表示。

图3-3说明了使用GPT模型的三步文本生成过程。首先，分词器将输入文本转换为一系列token ID，如第1章所述。其次，模型接收这些token ID并生成相应的logits，它们是表示词汇表中每个token概率分布的向量，如第2章所述。第三，这些logits被转换回token ID，分词器将其解码为人类可读的文本，从而完成从文本输入到文本输出的循环。

在代码中，我们实现文本生成过程如下：

import tiktoken
from chapter2 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)  # add batch dimension
    return encoded_tensor

def token_ids_to_text(token_ids, tokenizer):
    flat = token_ids.squeeze(0)  # remove batch dimension    
    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 = 20,
    context_size = GPT_CONFIG_M['context_length'],
)

print(f"Output text:\n{token_ids_to_text(token_ids, tokenizer)}")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

使用以上代码，模型生成以下文本：

Output text:  
Every effort moves yourentingetic was n't refres5ex MeCHicularstren
1
2

根据输出，很明显模型还不能生成连贯的文本，因为它还没有经过训练。为了定义什么是连贯或高质量的文本，我们必须实现一种数值方法来评估生成的内容。这种方法将使我们能够在整个训练过程中监控和提高模型的性能。

下一节介绍了如何计算生成输出的损失度量。训练过程中监控这个损失，以指示训练进度。此外，在后续章节关于微调语言模型时，我们将回顾评估模型质量的其他方法。

计算文本生成损失

本节探讨在训练过程中通过计算所谓的文本生成损失来数值评估文本质量的技术。我们将通过一个实际的例子一步一步地讨论这个主题，使概念清晰易懂。首先简要回顾第1章如何加载数据，以及第2章如何通过generate_text_simple函数生成文本。

图3-4说明了本节涵盖的整体流程，从输入文本到语言模型生成的文本。

图3-4 本节涵盖的主题概览。我们从回顾上一章的文本生成开始，并在后续小节中实现基本的模型评估技术。

如图3-4所示，下一小节回顾了我们在上一章末设置的文本生成，然后在后续小节中深入探讨文本评估和计算训练集与验证集损失。

使用GPT生成文本

这一部分我们设置语言模型，并简要回顾了第2章实现的文本生成过程。首先使用第2章的GPTModel类和GPT_CONFIG_M字典初始化将在本章评估和训练的GPT模型：

import torch
from chapter2 import GPTModel

GPT_CONFIG_M = {
    'vocab_size': 50257,  
    'context_length': 512,  # 修改为512 tokens
    'emb_dim': 768,
    'n_heads': 12, 
    'n_layers': 12,
    'dropout_rate': 0.1,
    'tkv_bias': False,
}

torch.manual_seed(1234)
model = GPTModel(GPT_CONFIG_M)
model.eval()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

考虑到GPT_CONFIG_M字典，与上一章相比，我们唯一的调整是将上下文长度context_length减少到512个token。这一修改减少了训练模型的计算需求，使得在标准笔记本电脑上进行训练成为可能。

最初，具有1.1亿参数的GPT模型配置为可以处理多达2048个token。在本章末尾的训练过程之后，我们将更新上下文大小设置，并加载预训练权重以使用配置为1024个token上下文长度的模型。

使用GPT模型实例，我们采用上一章介绍的generate_text_simple函数，并引入两个方便的函数text_to_token_ids和token_ids_to_text。这些函数方便在文本和token表示之间进行转换，这是我们将在本章中使用的技巧。为了更清楚地理解，图3-5说明了这个过程，然后我们深入研究代码。

图3-5 生成文本涉及将文本编码为语言模型处理为logit向量的token ID。然后将logit向量转换回token ID并解码为文本表示。

图3-5说明了使用GPT模型的三步文本生成过程。首先，分词器将输入文本转换为一系列token ID，如第1章所述。其次，模型接收这些token ID并生成相应的logit，它们是表示词汇表中每个token概率分布的向量，如第2章所述。第三，这些logit被转换回token ID，分词器将其解码为人类可读的文本，从而完成从文本输入到文本输出的循环。

在代码中，我们实现文本生成过程如下：

import tiktoken
from chapter2 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)  # add batch dimension
    return encoded_tensor

def token_ids_to_text(token_ids, tokenizer):  
    flat = token_ids.squeeze(0)  # remove batch dimension
    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 = 20,
    context_size = GPT_CONFIG_M['context_length'],
)

print(f"Output text:\n{token_ids_to_text(token_ids, tokenizer)}")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

使用上述代码，模型生成以下文本：

Output text:
Every effort moves yourentingetic was n't refres5ex MeCHicularstren  
1
2

基于输出，很明显模型还不能生成连贯的文本，因为它还没有经过训练。为了定义什么是连贯或高质量的文本，我们必须实现一个数值方法来评估生成的内容。这种方法将使我们能够在整个训练过程中监控和提高模型的性能。

下一节介绍了如何计算生成输出的损失度量。这个损失作为模型训练过程中进度和成功的指标。此外，在后续章节关于微调语言模型中，我们将回顾评估模型质量的其他方法。

计算文本生成损失

本节探讨在训练过程中通过计算所谓的文本生成损失来数值评估文本质量的技术。我们将通过一个实际的例子一步一步地讨论这个主题，使概念清晰易懂。首先简要回顾第1章如何加载数据，以及第2章如何通过generate_text_simple函数生成文本。

图3-6说明了从输入文本到语言模型生成文本的整体流程。

图3-6 对于左侧显示的每个输入token，我们计算一个包含与词汇表中每个token对应的概率分数的向量。每个向量中概率分数最高的索引位置代表最可能的下一个token ID。与概率分数最高相关联的这些token ID被选择并映射回模型生成的文本。

图3-6概述了generate_text_simple函数在内部所做的文本生成过程。我们需要在计算文本质量的损失之前执行这些相同的初始步骤，这将在本节后面介绍。

图3-6概述了具有小型token词汇表的文本生成过程，以便将此图像放在单个页面上。然而，我们的GPT模型使用的词汇表要大得多，由50257个单词组成，因此以下代码中的token ID将从0到50256，而不是0到8。此外，图3-6仅示意了一个文本示例（“every effort moves”）以简化说明。在下面的动手代码示例中，实现了图3-6中的步骤，我们将使用两个输入示例（“every effort moves"和"I really like”）作为GPT模型的输入。

考虑到这两个输入示例，它们已经映射到了相应的token ID：

inputs = torch.tensor([[8982, 4387, 10554, 21862],   # every effort moves
                       [314,  1309,  2744,     0]])  # I really like
1
2

与这些输入相匹配，targets包含我们希望模型生成的token ID：

targets = torch.tensor([[4387, 10554, 21862, 18529],   # effort moves you 
                        [1309,  2744,     0, 11914]])  # really like chocolate
1
2

请注意，目标是输入向前移动一个位置的结果，这是我们在第1章实现数据加载器时讨论的概念。这种移位策略对于教模型预测序列中的下一个token至关重要。

当我们将输入输入到模型中以计算两个输入示例的logit向量时，每个包含三个token，并应用softmax函数将这些logit值转换为概率分数，这对应于图3-6中的步骤2：

with torch.no_grad():
    logits = model(inputs)
    probas = torch.softmax(logits, dim=-1)  # Probability of each token in vocab
    print(probas.shape)  
1
2
3
4

生成的概率分数（probas）张量的维度如下：

torch.Size([2, 4, 50257])
1

第一个数字对应于inputs中的两个示例（行），也称为批次大小。第二个数字对应于每行中的token数。最后一个数字对应于嵌入维度，它由词汇表的大小决定，如前几章所述。

在将logit转换为概率后，第2章的generate_text_simple函数然后将这些概率分数转换回文本，如图3-6中的步骤3-5所示。

我们可以通过应用argmax函数将概率分数转换回token ID来实现步骤3和4：

token_ids = torch.argmax(probas, dim=-1, keepdim=True)
print("Token IDs:\n", token_ids)
1
2

给定我们有2个输入批次，每个包含3个token，应用argmax函数到概率分数（图3-6中的步骤3）产生2组输出，每个有3个预测的token ID：

Token IDs:
tensor([[[14991],  #First batch
         [11],
         [49970]],

        [[22],     #Second batch
         [311], 
         [47499]]])
1
2
3
4
5
6
7
8

最后，步骤5将token ID转换回文本：

print(f"Targets batch: {token_ids_to_text(targets[0], tokenizer)}")
print(f"Outputs batch: {token_ids_to_text(token_ids[0].flatten(), tokenize
1
2

当我们解码这些token时，我们发现这些输出token与我们希望模型生成的目标token完全不同：

Targets batch: effort moves you
Outputs batch: Armed he Netflix
1
2

模型生成与目标文本不同的随机文本，因为它还没有经过训练。我们现在介绍如何通过所谓的损失数值评估模型生成文本的性能，如图3-7所示。这不仅有助于衡量生成文本的质量，而且还是我们稍后实现的训练函数的构建块，我们用它来更新模型权重以改进生成的文本。

图3-7 在本节的其余部分中，我们现在实现文本评估函数。在下一节中，我们将此评估函数应用于整个用于模型训练的数据集。

如图3-7所示，我们在本节剩余部分实现的文本评估过程的一部分是衡量生成的token与正确预测（目标）的距离。训练函数将使用此信息来调整模型权重，以生成更类似于（或理想情况下与目标文本匹配）目标文本的文本。

回传

如何最大化与目标token对应的softmax概率值？答案是通过一个称为反向传播的过程来更新模型权重，以便模型输出更高的值用于我们想要生成的各自token ID。权重更新是通过深度神经网络训练的标准技术反向传播完成的（有关反向传播和模型训练的更多详细信息，请参见附录A中的A.1至A.4节）。

反向传播需要一个损失函数，该函数计算模型预测输出（这里是与目标token ID对应的概率）与实际期望输出之间的差异。此损失函数衡量模型预测与目标值的偏差程度。

在本节的其余部分，我们计算两个示例批次的概率分数target_probas_1和target_probas_2的损失。训练函数将使用此信息来调整模型权重，以生成更类似于（或理想情况下与目标文本匹配）目标文本的文本。

主要步骤如图3-8所示。

图3-8 计算损失涉及几个步骤。步骤1-2计算与目标张量对应的token概率。然后在步骤3-4中通过对数和平均对这些概率进行转换。

由于我们已经应用了图3-8中列出的步骤1-2来获得target_probas_1和target_probas_2，我们继续步骤3，应用对数到概率分数：

log_probas = torch.log(torch.cat((target_probas_1, target_probas_2)))
print(log_probas)
1
2

这导致以下值：

tensor([-13.1186, -13.1157, -13.1212])
1

使用对数概率分数在数学优化中比直接处理分数更容易管理。我在参考文献部分的附录B中的一个讲座中进一步详细介绍了这个主题。

接下来，我们通过计算平均值将这些对数概率组合成单个分数（图3-8中的步骤4）：

avg_log_probas = torch.mean(log_probas)
print(avg_log_probas)
1
2

生成的平均对数概率分数如下：

tensor(-13.1185)
1

目标是通过更新模型权重作为训练过程的一部分，将平均对数概率尽可能接近于0。

然而，在深度学习中，常见的做法不是将平均对数概率推高到0，而是将负平均对数概率降低到0。负平均对数概率只是平均对数概率乘以-1，这对应于图3-8中的步骤5：

neg_avg_log_probas = -avg_log_probas  
print(neg_avg_log_probas)
1
2

这打印出tensor(13.1185)。

这个负值转变为0的术语在深度学习中称为交叉熵损失。

PyTorch在这里派上用场，因为它已经有一个内置的cross_entropy函数，可以为我们处理图3-8中的所有这些步骤：

交叉熵损失

在其核心，交叉熵损失是机器学习和深度学习中流行的度量，用于衡量两个概率分布之间的差异，通常是数据集中标签（这里是token）的真实分布和模型预测的分布（例如，语言模型生成的token概率）。

在机器学习的上下文中，特别是在PyTorch等框架中，cross_entropy函数为离散结果计算这个度量，这类似于给定模型生成的token概率下目标token的负平均对数概率，使交叉熵和负平均对数概率相关并在实践中经常互换使用。

在我们应用交叉熵函数之前，让我们简要回顾一下logits和targets张量的形状：

print("Logits shape:", logits.shape)
print("Targets shape:", targets.shape)
1
2

生成的形状如下：

Logits shape: torch.Size([2, 4, 50257])
Targets shape: torch.Size([2, 4])  
1
2

如我们所见，logits张量有三个维度：批次大小、token数和词汇表大小。targets张量有两个维度：批次大小和token数。

对于PyTorch中的cross_entropy_loss函数，我们希望通过组合批次维度来展平这些张量：

logits_flat = logits.flatten(0, 1)
targets_flat = targets.flatten()

print("Flattened logits:", logits_flat.shape)  
print("Flattened targets:", targets_flat.shape)
1
2
3
4
5

生成的张量维度如下：

Flattened logits: torch.Size([8, 50257])
Flattened targets: torch.Size([8])
1
2

请记住，目标是我们希望语言模型生成的token ID，logit包含softmax函数处理以获得概率分数之前的未缩放模型输出。

以前，我们应用了softmax函数，选择了与目标ID对应的概率分数，并计算了负平均对数概率。PyTorch的cross_entropy函数将为我们处理所有这些步骤：

loss = torch.nn.functional.cross_entropy(logits_flat, targets_flat)
print(loss)
1
2

生成的损失与我们之前手动应用图3-8中显示的各个步骤时获得的相同：

tensor(13.1185)
1

困惑度

困惑度是一个与交叉熵损失一起使用的度量，用于评估语言建模等任务中模型的性能。它可以提供一种更易于解释的方式来理解模型在预测序列中下一个token时的不确定性。

困惑度衡量模型预测的概率分布与数据集中实际词分布的匹配程度。与损失类似，较低的困惑度表明模型预测更接近实际分布。

困惑度可以计算为perplexity = torch.exp(loss)，当应用于先前计算的损失时，返回tensor(496.3506)。

困惑度通常被认为比原始损失值更易于解释，因为它表示模型在每一步中不确定的有效词汇量。在给定的示例中，这将转化为模型不确定在词汇表中的496个单词或token中的哪一个作为下一个token生成。

在本节中，我们计算了两个小文本输入的损失以进行说明。在下一节中，我们将损失计算应用于整个训练和验证集。

计算训练集和验证集损失

在本节中，我们首先准备将用于训练语言模型的训练和验证数据集。然后我们计算训练和验证集的交叉熵，如图3-9所示，这是模型训练过程的一个重要组成部分。

图3-9 在上一节中计算交叉熵损失之后，我们现在将此损失计算应用于用于模型训练的整个文本数据集。

如图3-9所示，为了计算训练和验证数据集上的损失，我们使用一个非常小的文本数据集，Edith Wharton的短篇小说《The Verdict》，我们在第1章中已经使用过。通过选择公共领域的文本，我们规避了任何与使用权相关的问题。此外，我们使用如此小的数据集的原因是，它允许在标准笔记本电脑上在几分钟内执行代码示例，即使没有高端GPU，这对于教育目的特别有利。

感兴趣的读者还可以使用本书的补充代码从Project Gutenberg准备一个包含50,000多本公共领域书籍的更大规模数据集，并在这些书籍上训练语言模型（有关详细信息，请参见附录D）。

预训练语言模型的成本

为了将我们项目的规模放在视角中，考虑相对流行的、公开可用的具有65亿参数的Llama模型的训练。该模型需要在昂贵的A100 GPU上进行180,000 GPU小时的训练，处理1.4万亿个token。在撰写本文时，在AWS上运行一个具有8个A100的云服务器每小时的成本约为30美元。粗略估计，这样一个语言模型的总训练成本约为540万美元（计算为180,000小时除以8，然后乘以30美元）。

以下代码加载我们在第1章中使用的《The Verdict》短篇小说：周。此外，我们将在本章末尾将预训练的权重从OpenAI加载到我们的GPT模型代码中。

接下来，我们将数据集分为训练集和验证集，并使用第1章的数据加载器为语言模型训练准备批次。这个过程在图3-10中可视化。

图3-10 在准备数据加载器时，我们将输入文本分成训练集和验证集部分。然后我们对文本进行标记化（为简单起见仅显示训练集部分），并将标记化的文本划分为用户指定长度的块（此处为20）。最后，我们打乱行并将分块的文本组织成批次（此处批次大小为3），我们可以使用这些批次进行模型训练。

出于可视化目的，图3-10使用最大长度为20（由于空间限制）。但是，对于我们正在实现的实际数据加载器，我们将max_length设置为等于语言模型支持的token上下文长度，以便语言模型在训练期间看到更长的文本。

使用可变长度进行训练

为了简单和效率，我们使用大小相似的块对训练数据进行训练。然而，在实践中，用可变长度的输入训练语言模型也是有益的，以帮助语言模型在使用时更好地泛化到不同类型的输入。

为了实现图3-10中可视化的数据拆分和加载，我们首先定义一个train_ratio，将数据的80%用于训练，剩余的20%用作验证数据，用于训练期间的模型评估：

train_ratio = 0.8
split_idx = int(train_ratio * len(text_data))
train_data = text_data[:split_idx]
val_data = text_data[split_idx:]
1
2
3
4

使用train_data和val_data子集，我们现在可以创建相应的数据加载器，重用第1章的create_dataloader_v2代码：

from chapter1 import create_dataloader_v2
torch.manual_seed(42)

train_loader = create_dataloader_v2(
    train_data,
    batch_size = 32, 
    max_length = GPT_CONFIG_M['context_length'],
    stride = GPT_CONFIG_M['context_length'],
    drop_last = True,
    shuffle = True,
)

val_loader = create_dataloader_v2(  
    val_data,
    batch_size = 32,
    max_length = GPT_CONFIG_M['context_length'], 
    stride = GPT_CONFIG_M['context_length'],
    drop_last = False,  
    shuffle = False,
)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

我们在上述代码中使用了相对较小的批次大小，以减少计算资源需求，因为我们使用的是非常小的数据集。在实践中，使用批次大小为128或更大的批次训练语言模型并不罕见。

作为可选的检查，我们可以遍历数据加载器以确保它们被正确创建：

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)
1
2
3
4
5
6
7

我们应该看到以下输出：

Train loader:
torch.Size([32, 512]) torch.Size([32, 512])
torch.Size([32, 512]) torch.Size([32, 512])
torch.Size([32, 512]) torch.Size([32, 512])
torch.Size([32, 512]) torch.Size([32, 512]) 
torch.Size([32, 512]) torch.Size([32, 512])
torch.Size([32, 512]) torch.Size([32, 512])
torch.Size([32, 512]) torch.Size([32, 512])
torch.Size([32, 512]) torch.Size([32, 512])
torch.Size([32, 512]) torch.Size([32, 512])

Validation loader:
torch.Size([32, 512]) torch.Size([32, 512])
1
2
3
4
5
6
7
8
9
10
11
12
13

基于上述代码输出，我们有9个训练集批次，每个批次有32个样本和512个token。由于我们只为验证分配了20%的数据，因此只有一个验证批次，由32个输入示例组成。

如预期的那样，输入数据（x）和目标数据（y）具有相同的形状（批次大小乘以每个批次中的token数），因为目标是输入移动一个位置，如第1章所讨论的。

接下来，我们实现一个实用函数来计算由训练和验证加载器返回的给定批次的交叉熵损失：

def calc_loss_batch(input_batch, target_batch, model, device):
    input_batch, target_batch = input_batch.to(device), target_batch.to(devi
    logits = model(input_batch)

    loss = torch.nn.functional.cross_entropy(
        logits.flatten(0, 1), target_batch.flatten())
    
    return loss
1
2
3
4
5
6
7
8

我们现在可以使用这个calc_loss_batch实用函数，它计算单个批次的损失，来实现下面的calc_loss_loader函数，它计算给定数据加载器采样的所有批次的损失：

def calc_loss_loader(data_loader, model, device, num_batches=None):
    total_loss = 0  
    if num_batches is None:  
        num_batches = len(data_loader)  # A
    else:
        num_batches = min(num_batches, len(data_loader))  # B
        
    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()  # C
        else:
            break
            
    return total_loss / num_batches  # D
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

默认情况下，calc_loss_batch函数遍历给定数据加载器中的所有批次，在total_loss变量中累积损失，然后计算损失并将其平均到批次的总数上。

或者，我们可以通过num_batches指定更小数量的批次，以加速训练期间的评估。

让我们现在看到这个calc_loss_batch函数的实际应用，将其应用于训练和验证集加载器：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # A
model.to(device)

train_loss = calc_loss_loader(train_loader, model, device)  # B
val_loss = calc_loss_loader(val_loader, model, device)

print(f"Training loss: {train_loss:.3f}")
print(f"Validation loss: {val_loss:.3f}")
1
2
3
4
5
6
7
8

生成的损失值如下：

Training loss: 10.910
Validation loss: 11.010
1
2

损失值相对较高，因为模型尚未经过训练。作为比较，如果模型学会按照它们在训练和验证集中出现的方式生成下一个token，则损失接近4。

类似于训练集损失，我们可以看到验证损失在训练开始时很高（11.010），并在训练过程中减小。但是，它永远不会变得像训练集损失那样小，并在第6个epoch后保持在10.5左右。

在更详细地讨论验证损失之前，让我们创建一个简单的图，将训练和验证集损失并排显示：

import matplotlib.pyplot as plt

def plot_losses(epochs_seen, tokens_seen, train_losses, val_losses):
    fig, ax = plt.subplots(figsize=(8, 4))
    ax.plot(epochs_seen, train_losses, label='Training loss')
    ax.plot(epochs_seen, val_losses, linestyle='--', label='Validation loss')
    ax.set_xlabel("Epochs")
    ax.set_ylabel("Loss")
    ax.legend(loc='upper right')
    
    ax2 = ax.twinx()  # A
    ax2.plot(tokens_seen, train_losses, alpha=0.3)  # B  
    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
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

生成的训练和验证损失图如图3-11所示。


图3-11 在训练开始时，我们观察到训练集和验证集损失都急剧下降，这表明模型正在学习。但是，训练集损失在第2个epoch之后继续下降，而验证损失停滞不前。这表明模型仍在学习，但在第2个epoch之后过拟合到训练集。

如图3-11所示，训练和验证损失在第1个epoch开始改善。但是，损失在第2个epoch之后开始出现分歧。这种分歧以及验证损失远高于训练损失的事实表明，模型正在过拟合训练数据。我们可以通过在《The Verdict》文本文件中搜索诸如"quite insensible to the irony"等生成的文本片段来确认模型逐字记忆训练数据。

考虑到我们使用的是非常非常小的训练数据集并在多个epoch上训练模型，这种记忆是意料之中的。通常，在更大的数据集上只训练模型一个epoch是常见的。

如前所述，感兴趣的读者可以尝试在来自Project Gutenberg的50,000本公共领域书籍上训练模型，在这种情况下不会发生过拟合（有关详细信息，请参见附录B）。

在即将到来的部分中，如图3-12所示，我们探讨语言模型用来缓解记忆效应的采样方法，从而产生更新颖的生成文本。

图3-12 在实现训练函数后，我们的模型可以生成连贯的文本。但是，它经常逐字记忆来自训练集的段落。下一节将介绍用于生成更多样化输出文本的策略。

如图3-12所示，下一节将介绍语言模型的文本生成策略，以减少训练数据的记忆并增加语言模型生成文本的原创性，然后我们将介绍权重加载和保存以及从OpenAI加载预训练权重。

解码策略以控制随机性

在本节中，我们将介绍文本生成策略（也称为解码策略）以生成更多原创文本。首先，我们简要回顾上一章的generate_text_simple函数，我们在本章早些时候的generate_and_print_sample中使用了该函数。然后我们将介绍两种技术：温度缩放和top-k采样，以改进此函数。

我们首先通过将模型从GPU转移回CPU来开始，因为相对较小的模型进行推理不需要GPU。此外，在训练之后，我们将模型设置为评估模式以关闭dropout等随机组件：

model.to('cpu')  
model.eval()
1
2

接下来，我们将GPT模型实例（model）插入到generate_text_simple函数中，该函数使用语言模型一次生成一个token：

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 = 20,
    context_size = GPT_CONFIG_M['context_length'],
)

print(f"Output text:\n{token_ids_to_text(token_ids, tokenizer)}")
1
2
3
4
5
6
7
8
9
10

生成的文本如下：

Output text:
Every effort moves you know was one of the axioms she laid down across the
1
2

如第3.2节所述，在每个生成步骤中，通过将模型输出转换为概率分数并选择对应于所有词汇表token中最高概率分数的token来选择生成的token。

这意味着语言模型将始终生成相同的输出，即使我们在相同的起始上下文（“Every effort moves you”）上多次运行上面的generate_text_simple函数。

以下小节介绍两个概念来控制生成文本的随机性和多样性：温度缩放和top-k采样。

温度缩放

本节介绍温度缩放，这是一种在下一个token生成任务中添加概率选择过程的技术。

以前，在generate_text_simple函数内部，我们总是使用torch.argmax（也称为贪婪解码）将具有最高概率的token采样为下一个token。为了生成具有更多变化的文本，我们可以用从概率分布（此处为语言模型在每个token生成步骤为每个词汇表条目生成的概率分数）采样的函数替换argmax。

为了用具体的例子说明概率采样，让我们简要讨论使用一个非常小的词汇表进行下一个token生成过程，以便说明：

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()}
1
2
3
4
5
6
7
8
9
10
11
12
13

接下来，假设语言模型获得起始上下文"every effort moves you"，并生成以下下一个token logit：

next_token_logits = torch.tensor(
    [0.2, -0.1, 0.3, 1.2, 0

```python
file_path = "the_verdict.txt"
with open(file_path, 'r', encoding='utf-8') as file:
    text_data = file.read()
1
2
3
4
5
6
7

加载数据集后，我们可以检查数据集中的字符数和token数：

total_characters = len(text_data)
total_tokens = len(tokenizer.encode(text_data))
print(f"Characters: {total_characters}")
print(f"Tokens: {total_tokens}")    
1
2
3
4

输出如下：

Characters: 81330
Tokens: 51453
1
2

只有51,453个token，文本可能看起来太小而无法训练一个语言模型，但正如前面提到的，这是出于教育目的，以便我们可以在几分钟内运行代码，而不是几.8, 0.5, -2.3, 0.7, 0.4]
)

如前一章所述，在`generate_text_simple`内部，我们通过softmax函数将logit转换为概率，并通过`argmax`函数获得对应于生成token的token ID，然后我们可以通过反向词汇表将其映射回文本：

```python
probas = torch.softmax(next_token_logits, dim=-1)
next_token_id = torch.argmax(probas).item()
print(inverse_vocab[next_token_id])
1
2
3
4
5
6
7

由于第四个位置（索引位置3，因为Python使用0索引）的logit值最大，相应地softmax概率分数也最高，因此生成的单词是"forward"。

为了实现概率采样过程，我们现在可以用PyTorch中的multinomial函数替换argmax：

torch.manual_seed(1234)
next_token_id = torch.multinomial(probas, num_samples=1).item()  
print(inverse_vocab[next_token_id])
1
2
3

打印的输出仍然是"forward"，就像之前一样。发生了什么？multinomial函数根据其概率分数按比例采样下一个token。换句话说，"forward"仍然是最有可能的token，并且在大多数时候都会被multinomial选中，但不是所有时候。为了说明这一点，让我们实现一个函数，重复这个采样过程100次：

def print_sampled_tokens(probas):
    torch.manual_seed(1234)
    sample = [torch.multinomial(probas, num_samples=1).item() for i in range(
    sampled_ids = torch.bincount(torch.tensor(sample))
    
    for i, freq in enumerate(sampled_ids):
        print(f"{freq}x {inverse_vocab[i]}")

print_sampled_tokens(probas)        
1
2
3
4
5
6
7
8
9

采样输出如下：

1x closer
1x every
4x effort
78x forward 
10x inches
4x moves
2x pizza
1x toward
1
2
3
4
5
6
7
8

如输出所示，单词"forward"在大多数时候（100次中的78次）被采样，但其他token如"closer"、“inches"和"toward"也会在某些时候被采样。这意味着，如果我们在generate_and_print_sample函数内用multinomial函数替换argmax函数，语言模型有时会生成文本，如"every effort moves you toward”、“every effort moves you inches"和"every effort moves you closer”，而不是"every effort moves you forward"。

我们可以通过一个称为温度缩放的概念进一步控制分布和选择过程，其中温度缩放只是将logit除以大于1的数字的一种花哨说法：

def softmax_with_temperature(logits, temperature=1.0):
    scaled_logits = logits / temperature
    return torch.softmax(scaled_logits, dim=-1)  
1
2
3

大于1的温度会导致token概率更加均匀分布。小于1的温度将导致更加自信（更尖锐或更尖峰状）的分布。让我们通过绘制原始概率以及使用不同温度值缩放的概率来说明这一点：

temperatures = [0.01, 1, 2]  # Original, higher, and lower temperature 
scaled_probas = [softmax_with_temperature(next_token_logits, T) for T in tem

x = torch.arange(len(vocab))

fig, ax = plt.subplots(figsize=(10, 4))

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=45)
ax.legend()

plt.tight_layout()
plt.show()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

生成的图如图3-13所示。

图3-13 温度为1表示词汇表中每个token的未缩放概率分数。将温度降低到0.01会锐化分布，使最可能的token（此处为"forward"）将具有更高的概率分数。反之，将温度提高到2会使分布更加均匀。

温度为1除以logit 1，然后将它们传递给softmax函数以计算概率分数。换句话说，使用温度1与不使用任何温度缩放相同。在这种情况下，token的选择概率等于原始的softmax概率分数，通过PyTorch中的multinomial采样函数。

例如，对于温度设置为1，对应于"forward"的token将在大约78%的时间内被选中，如我们在图3-13中看到的。

此外，如图3-13所示，应用非常小的温度，如0.01，将导致更尖锐的分布，使multinomial函数的行为选择最可能的token（此处为"forward"）几乎100%的时间，接近argmax函数的行为。反之，温度为2会导致更均匀的分布，其中其他token被更频繁地选择。这可以为生成的文本添加更多变化，但也更频繁地导致无意义的文本。例如，使用温度2会导致"every effort moves you pizza"这样的文本大约10%的时间。

练习

使用print_sampled_tokens函数打印图3-13中显示的温度缩放的softmax概率的采样频率。在每种情况下，单词"pizza"被采样的频率是多少？你能想到一种更快、更准确的方法来确定单词"pizza"被采样的频率吗？

Top-k采样

在上一节中，我们实现了一种概率采样方法，结合温度缩放来增加输出的多样性。我们看到，更高的温度值会导致下一个token概率更加均匀分布，从而产生更多样化的输出，因为它减少了模型重复选择最可能token的可能性。这种方法允许探索不太可能但可能更有趣和更有创意的生成过程路径。然而，这种方法的一个缺点是，它有时会导致语法不正确或完全无意义的输出，例如"every effort moves you pizza"。

在本节中，我们介绍另一个概念，称为top-k采样，当与概率采样和温度缩放结合使用时，可以改善文本生成结果。

在top-k采样中，我们可以将采样的token限制为最有可能的top-k个token，并通过掩盖它们的概率分数来排除所有其他token的选择过程，如图3-14所示。

图3-14 使用top-k采样，其中k=4，我们专注于与最高logit相关联的4个token，并在应用softmax函数之前掩盖所有其他token为负无穷大（-inf）。这导致所有非top-k token被分配0的概率分布。

图3-14中概述的方法将所有未选择的logit替换为负无穷大值（-inf），以便在计算softmax值时，非top-k token的概率分数为0，剩余的概率加起来为1。仔细的读者可能还记得我们在第2章第2.1节"应用因果注意力掩码"中实现多头注意力模块时使用的这个掩码技巧。

在代码中，我们可以按如下方式实现图3-14中概述的top-k过程，从选择具有最大logit值的k个token开始：

top_k = 4
top_logits, top_pos = torch.topk(next_token_logits, top_k)

print("Top logits:", top_logits)  
print("Top positions:", top_pos)
1
2
3
4
5

按降序排列的top-k token的logit值和token ID如下：

Top logits: tensor([1.2000, 0.8000, 0.7000, 0.5000])
Top positions: tensor([3, 4, 7, 5])
1
2

随后，我们应用PyTorch的where函数将低于top选择范围内最低logit值的token的logit值设置为负无穷大（-inf）：

new_logits = torch.where(
    condition = (next_token_logits < top_logits[-1]),  # A
    input = torch.tensor(float('-inf')),  # B 
    other = next_token_logits,  # C
)

print(new_logits)  
1
2
3
4
5
6
7

下一个token的词汇表中9个token的生成logit如下：

tensor([-inf, -inf, -inf,  1.2000,  0.8000,  0.5000, -inf,  0.7000, -inf])
1

最后，让我们应用softmax函数将这些转换为下一个token概率：

top_k_probas = torch.softmax(new_logits, dim=-1)
print(top_k_probas)
1
2

如我们所见，这种top-k方法的结果是非零概率分数：

tensor([0.0000, 0.0000, 0.0000, 0.4728, 0.2167, 0.1328, 0.0000, 0.1777, 0.0000])
1

我们现在可以应用上一节中介绍的温度缩放和multinomial函数进行概率采样，以在这些非零概率分数中选择下一个token以生成下一个token。

我们在下一节中通过修改文本生成函数来做到这一点。

修改文本生成函数

前两小节介绍了两个概念来增加语言模型生成文本的多样性：温度采样和top-k采样。在本节中，我们结合并添加这些概念来修改我们之前用于通过语言模型生成文本的generate_simple函数，创建一个新的generate函数。

def generate(model, idx, max_new_tokens, context_size, temperature=1.0, top_k
    for _ in range(max_new_tokens):  # A
        idx_cond = idx[-context_size:]
        
        with torch.no_grad():  
            logits = model(idx_cond)
            logits = logits[:, -1, :]  
            
            if top_k is not None:  # B
                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 > 1:  # C
                logits = logits / temperature
                probs = torch.softmax(logits, dim=-1)
                idx_next = torch.multinomial(probs, num_samples=1)
            else:  # D
                idx_next = torch.argmax(logits, dim=-1, keepdim=True)
            
        idx = torch.cat((idx, idx_next), dim=1)
        
    return idx
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

让我们现在看到这个新的generate函数的实际应用：

torch.manual_seed(1234)

token_ids = generate(
    model = model,
    idx = text_to_token_ids("Every effort moves you", tokenizer),  
    max_new_tokens = 20,
    context_size = GPT_CONFIG_M['context_length'],
    top_k = 4,
    temperature = 1.5,
)

print(f"Output text:\n{token_ids_to_text(token_ids, tokenizer)}")
1
2
3
4
5
6
7
8
9
10
11
12

生成的文本如下：

Output text:
Every effort moves you stand to work on surprise a one of us had gone with
1
2

如我们所见，生成的文本与我们之前通过generate_simple函数在第3.2节开始时生成的文本（“Every effort moves you know was one of the axioms he laid”）非常不同，后者是从训练集中记忆的段落。

练习

尝试不同的温度和top-k设置。根据你的观察，你能想到哪些应用程序需要较低的温度和top-k设置？反之，在哪些应用中更喜欢较高的温度和top-k设置？（建议在本章末尾加载来自OpenAI的预训练权重后重新审视这个练习。）

练习

generate函数的不同组合设置是什么，可以强制确定性行为，即禁用随机采样，使其始终产生与generate_simple函数类似的相同输出？

到目前为止，我们介绍了如何预训练语言模型和使用它们生成文本。本章的最后两节将讨论我们如何保存和加载训练的语言模型，以及如何从OpenAI加载预训练权重。

在PyTorch中加载和保存模型权重

在本章中，我们讨论了如何在数值上评估训练进度并从头开始预训练语言模型。即使语言模型和数据集都相对较小，这个练习也表明预训练语言模型在计算上是昂贵的。因此，能够保存语言模型非常重要，这样我们就不必每次想在新会话中使用它时都重新运行训练。

如图3-1中的章节概述所示，我们将在本节中介绍如何保存和加载预训练模型。然后在即将到来的部分中，我们将把一个更有能力的预训练GPT模型从OpenAI加载到我们的GPT模型实例中。

图3-15 在训练和检查模型之后，通常有助于保存模型，以便我们可以在以后使用或继续训练它，这是本节的主题，然后我们在本章的最后一节加载来自OpenAI的预训练模型权重。

幸运的是，在PyTorch中保存模型相对简单。推荐的方法是使用torch.save函数保存模型的所谓state_dict，这是一个将每个层映射到其参数的字典：

torch.save(model.state_dict(), 'model.pth')
1

在上面的代码中，model.pth是保存state_dict的文件名。.pth扩展名是PyTorch文件的约定，尽管我们在技术上可以使用任何文件扩展名。

然后，在保存模型权重之后，我们可以通过state_dict将模型权重加载到新的GPT模型实例中，如下所示：

model = GPTModel(GPT_CONFIG_M)
model.load_state_dict(torch.load('model.pth'))
model.eval()
1
2
3

如第2章所述，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')
1
2
3
4

然后我们可以通过首先通过torch.load加载保存的数据，然后使用load_state_dict方法来恢复模型和优化器状态：

checkpoint = torch.load('model_and_optimizer.pth')
model = GPTModel(GPT_CONFIG_M)
model.load_state_dict(checkpoint['model_state_dict'])

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

model.train()  # Switch model to training mode 
1
2
3
4
5
6
7
8

练习

在保存权重后，在新的Python会话或Jupyter笔记本文件中加载模型和优化器，并使用train_model_simple函数继续预训练它多个epoch。

从OpenAI加载预训练权重

之前，出于教育目的，我们使用一个有限的数据集（包含一本短篇小说）训练了一个小型GPT模型。这种方法允许我们专注于基础知识，而无需大量的时间和计算资源。

幸运的是，OpenAI公开分享了他们的GPT模型的权重，从而消除了我们自己在大型语料库上重新训练模型的需要，这将花费数万到数十万美元。

在本节的其余部分，我们将这些权重加载到我们的GPTModel类中，并使用该模型进行文本生成。这里的权重指的是存储在PyTorch的Linear和Embedding层的weight属性中的权重参数。我们之前通过model.parameters()访问它们，以便在训练模型时使用。

在接下来的章节中，我们将重用这些预训练的权重来微调模型以进行文本分类任务，并像ChatGPT一样遵循指令。

请注意，OpenAI最初通过TensorFlow保存了GPT权重，我们必须安装TensorFlow才能在Python中加载权重。此外，下面的代码将使用一个名为tqdm的进度条工具来跟踪下载过程，我们也必须安装它。

你可以通过在终端中执行以下命令来安装这些库：

pip install tensorflow==2.11.0 tqdm==4.64.1 
1

下载代码相对较长，主要是样板代码，并不是很有趣。因此，与其在本章中讨论用于从互联网获取文件的Python代码，不如直接从本章的在线存储库下载gpt_download.py Python模块：

import urllib.request

url = "https://raw.githubusercontent.com/rasbt/
LLMs-from-scratch/main/ch3_pretrain/gpt_download.py"

filename = url.split("/")[-1]
urllib.request.urlretrieve(url, filename)
1
2
3
4
5
6
7

接下来，下载此文件到Python会话的本地目录后，读者可以简要检查此文件的内容，以确保它已正确保存并包含有效的Python代码。

我们现在可以从gpt_download.py文件导入download_and_load_gpt函数，如下所示，它将加载GPT架构设置（settings）和权重参数（params）到我们的Python会话中：

from gpt_download import download_and_load_gpt

settings, params = download_and_load_gpt(model_size='124M', models_dir='gpt2
1
2
3

执行上述代码将下载与1.24亿参数GPT模型相关的以下文件：

checkpoint_1001300_...............................[1KiB]
encoder.json.......................................[1MiB]
hparams.json.......................................[1KiB]
model.ckpt.data-00000-of-00001.................[440MiB]
model.ckpt.index..................................[16KiB] 
model.ckpt.meta...................................[16KiB]
vocab.bpe.........................................[1MiB]
1
2
3
4
5
6
7

更新的下载说明

如果下载代码对你不起作用，可能是由于间歇性的互联网连接、服务器问题或OpenAI共享开源GPT模型权重的方式发生变化。在这种情况下，请访问本章的在线代码存储库（https://github.com/rasbt/LLMs-from-scratch），了解替代和更新的说明，如有进一步的问题，请通过Manning论坛联系。

在前面的代码执行完成后，让我们检查settings和params的内容：

print(f"Settings: {settings}")
print(f"Parameter dictionary keys: {params.keys()}")
1
2

内容如下：

Settings: {'n_vocab': 50257, 'n_ctx': 1024, 'n_embd': 768, 'n_head': 12, 'n_
Parameter dictionary keys: dict_keys(['blocks.0', 'blocks.1', 'blocks.10', '
1
2

settings和params都是Python字典。settings字典存储语言模型架构设置，类似于我们手动定义的GPT_CONFIG_M设置。params字典包含实际的权重张量。请注意，我们只打印字典键，因为打印权重内容会占用太多屏幕空间；但是，我们可以通过打印整个字典（print(params)）或通过各自的字典键选择单个张量来检查这些权重张量，例如，嵌入层权重：

print(params['wte'])
print(f"Token embedding weight tensor dimensions: {params['wte'].shape}")
1
2

token嵌入层的权重如下：

[[-0.2433  -0.31426 -0.20737 ...  0.23305  1.1433  -0.38715]
 [ 0.10515 -0.34704 -0.12734 ... -0.14525 -0.12743 -0.12534]
 [ 0.25143  0.18864  0.41139 ... -0.21956  0.46881 -0.09653]
 ...
 [-0.27119  0.0018  -0.66206 ... -0.38255 -0.08005  0.57541]
 [-0.1208   0.13757 -0.12524 ... -0.26504 -0.49177  0.54857]
 [ 0.2821  -0.079    0.15147 ...  0.42941 -0.62312  1.1336 ]]
Token embedding weight tensor dimensions: (50257, 768)
1
2
3
4
5
6
7
8

我们通过download_and_load_gpt(model_size='124M')设置下载并加载了最小的GPT模型的权重。但请注意，OpenAI还共享了更大模型345M、774M和1558M的权重。这些不同大小的GPT模型的整体架构是相同的，如图3-16所示。

图3-16 GPT语言模型有几种不同的模型大小，参数范围从1.24亿到15.58亿。核心架构是相同的，唯一的区别是嵌入大小以及注意力头和transformer块等单个组件的重复次数。

如图3-16所示，不同大小的GPT模型的整体架构保持不变，只是不同的架构元素重复的次数不同，嵌入大小也不同。本章中剩余的代码也与这些更大的模型兼容。

加载GPT模型权重到Python后，我们仍然需要将它们从settings和params字典转移到我们的GPTModel实例中。首先，我们创建一个字典，列出不同GPT模型大小之间的差异，如图3-16所述：

model_configs = {
    'gpt-small (124M)': {'emb_dim': 768, 'n_layers': 12, 'n_heads':  12},
    'gpt-medium (345M)': {'emb_dim': 1024, 'n_layers': 24, 'n_heads': 16},
    'gpt-large (774M)': {'emb_dim': 1280, 'n_layers': 36, 'n_heads': 20},
    'gpt-xl (1558M)': {'emb_dim': 1600, 'n_layers': 48, 'n_heads':  25},
}
1
2
3
4
5
6

假设我们有兴趣加载最小的模型'gpt-small (124M)'。我们可以使用来自model_configs表的相应设置来更新我们在整个章节中定义和使用的全长GPT_CONFIG_M，如下所示：

model_name = 'gpt-small (124M)'
NEW_CONFIG = GPT_CONFIG_M.copy()
NEW_CONFIG.update(model_configs[model_name])  
1
2
3

仔细的读者可能还记得，我们之前使用512 token长度，但OpenAI的原始GPT模型是在1024 token长度上训练的，因此我们必须相应地更新NEW_CONFIG：

NEW_CONFIG.update({'context_length': 1024}) 
1

此外，OpenAI在多头注意力模块的线性层中使用偏置向量来实现查询、键和值矩阵计算。偏置向量在语言模型中不再常用，因为它们不会提高建模性能，因此是不必要的。但是，由于我们使用的是预训练权重，我们需要匹配设置以保持一致性并启用这些偏置向量：

NEW_CONFIG.update({'tkv_bias': True})
1

我们现在可以使用更新后的NEW_CONFIG字典来初始化一个新的GPTModel实例：

gpt = GPTModel(NEW_CONFIG)
gpt.eval()
1
2

默认情况下，GPTModel实例使用随机权重进行初始化以进行预训练。使用OpenAI模型权重的最后一步是用我们加载到params字典中的权重覆盖这些随机权重。为此，我们首先定义一个小的assign实用函数，用于检查两个张量或数组（left和right）是否具有相同的维度或形状，并将右侧张量作为可训练的PyTorch参数返回：

def assign(left, right):
    if left.shape != right.shape:
        raise ValueError(f"Shape mismatch: Left {left.shape} Right {right.sh
    return torch.nn.Parameter(torch.tensor(right))
1
2
3
4

接下来，我们定义一个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'])  # A
    gpt.tok_emb.weight = assign(gpt.tok_emb.weight, params['wte'])
    
    for b in range(len(params['blocks'])):  # B
        q, w, k, v = np.split( # C
            params['blocks'][b]['attn']['c_attn']['w'], axis=0) 

        gpt.trf_blocks[b].att.W_query.weight = assign(
            gpt.trf_blocks[b].att.W_query.weight, q.T)
        gpt.trf_blocks[b].att.W_key.weight = assign(
            gpt.trf_blocks[b].att.W_key.weight, k.T)  
        gpt.trf_blocks[b].att.W_value.weight = assign(
            gpt.trf_blocks[b].att.W_value.weight, v.T)
        
        q_b, k_b, v_b = np.split(
            params['blocks'][b]['attn']['c_attn']['b'], axis=0)
        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, vb)

        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].ffl.layers[0].weight = assign(
            gpt.trf_blocks[b].ffl.layers[0].weight,
            params['blocks'][b]['mlp']['c_fc']['w'].T)
        gpt.trf_blocks[b].ffl.layers[0].bias = assign(
            gpt.trf_blocks[b].ffl.layers[0].bias, 
            params['blocks'][b]['mlp']['c_fc']['b'])
        gpt.trf_blocks[b].ffl.layers[1].weight = assign(
            gpt.trf_blocks[b].ffl.layers[1].weight,
            params['blocks'][b]['mlp']['c_proj']['w'].T)    
        gpt.trf_blocks[b].ffl.layers[1].bias = assign(
            gpt.trf_blocks[b].ffl.layers[1].bias,
            params['blocks'][b]['mlp']['c_proj']['b'])

        gpt.trf_blocks[b].norm_1.scale = assign(
            gpt.trf_blocks[b].norm_1.scale,
            params['blocks'][b]['ln_1']['g'])
        gpt.trf_blocks[b].norm_1.shift = assign(
            gpt.trf_blocks[b].norm_1.shift, 
            params['blocks'][b]['ln_1']['b'])
        gpt.trf_blocks[b].norm_2.scale = assign(
            gpt.trf_blocks[b].norm_2.scale,
            params['blocks'][b]['ln_2']['g']) 
        gpt.trf_blocks[b].norm_2.shift = assign(
            gpt.trf_blocks[b].norm_2.shift,
            params['blocks'][b]['ln_2']['b'])

    gpt.final_norm.scale = assign(gpt.final_norm.scale, params['ln_f']['g'])
    gpt.final_norm.shift = assign(gpt.final_norm.shift, params['ln_f']['b'])
    gpt.out_head.weight = assign(gpt.out_head.weight, params['wte'])  # D
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

在load_weights_into_gpt函数中，我们仔细地将OpenAI实现中的权重与我们的GPTModel实现相匹配。举一个具体的例子，OpenAI将第一个transformer块的输出投影层的权重张量存储为params['blocks'][0]['attn']['c_proj']['w']。在我们的实现中，这个权重张量对应于gpt.trf_blocks[b].att.out_proj.weight，其中gpt是一个GPTModel实例。

开发load_weights_into_gpt函数需要大量猜测，因为OpenAI使用了与我们略有不同的命名约定。但是，如果我们尝试匹配维度不同的两个张量，assign函数会提醒我们。此外，如果我们在此函数中犯了错误，我们会注意到这一点，因为生成的GPT模型将无法生成连贯的文本。

让我们现在尝试实际使用load_weights_into_gpt，并将OpenAI模型权重加载到我们的GPTModel实例gpt中：

load_weights_into_gpt(gpt, params)
gpt.to(device)
1
2

如果模型加载正确，我们现在可以使用它通过我们之前的generate函数生成新文本：

torch.manual_seed(1234)

token_ids = generate(
    model = gpt,
    idx = text_to_token_ids("Every effort moves you", tokenizer),
    max_new_tokens = 50,  
    context_size = NEW_CONFIG['context_length'],  
    top_k = 4,
    temperature = 1.5,
)

print(f"Output text:\n{token_ids_to_text(token_ids, tokenizer)}")
1
2
3
4
5
6
7
8
9
10
11
12

生成的文本如下：

Output text:
Every effort moves you toward finding an ideal new way to practice somethin
What makes us want to be on top of that?
1
2
3

我们可以确信我们正确加载了模型权重，因为模型可以生成连贯的文本。在这个过程中的一个小错误就会导致模型失败。

在接下来的章节中，我将进一步使用这个预训练模型，并微调它以对文本进行分类并遵循指令。

鼓励读者尝试不同大小的GPT模型，例如最大的15.58亿参数模型，并将生成的文本与我们在本章中加载的1.24亿模型进行比较。

总结

• 当语言模型生成文本时，它们一次输出一个token。
• 默认情况下，下一个token是通过将模型输出转换为概率分数并从词汇表中选择对应于最高概率分数的token来生成的，这被称为贪婪解码。
• 使用概率采样和温度缩放，我们可以影响生成文本的多样性和连贯性。
• 训练集和验证集损失可用于衡量语言模型在训练期间生成的文本质量。
• 预训练语言模型涉及改变其权重以最小化训练损失。
• 语言模型的训练循环本身是深度学习中的标准程序，使用传统的交叉熵损失和AdamW优化器。
• 在大型文本语料库上预训练语言模型是耗时且资源密集型的，因此我们可以从OpenAI加载公开可用的权重，作为我们自己在大型数据集上预训练模型的替代方法。

以下是本章节中提到的一些参考文献和资源：

1. Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., & Sutskever, I. (2019). Language models are unsupervised multitask learners. OpenAI blog, 1(8), 9. https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf

2. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., … & Polosukhin, I. (2017). Attention is all you need. In Advances in neural information processing systems (pp. 5998-6008). https://arxiv.org/abs/1706.03762

3. Holtzman, A., Buys, J., Du, L., Forbes, M., & Choi, Y. (2019). The curious case of neural text degeneration. arXiv preprint arXiv:1904.09751. https://arxiv.org/abs/1904.09751

4. Kingma, D. P., & Ba, J. (2014). Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980. https://arxiv.org/abs/1412.6980

5. Loshchilov, I., & Hutter, F. (2017). Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101. https://arxiv.org/abs/1711.05101

6. OpenAI GPT-2 Model weights: https://github.com/openai/gpt-2

7. PyTorch Documentation: https://pytorch.org/docs/stable/index.html

8. Hugging Face Transformers library: https://huggingface.co/transformers/

9. Transformer Architecture Blog Post by Jay Alammar: http://jalammar.github.io/illustrated-transformer/

10. “The Illustrated GPT-2” Blog Post by Jay Alammar: http://jalammar.github.io/illustrated-gpt2/