【阅读记录-章节5】从零构建大语言模型-星语课程网

返回
顶部

修改密码

* 旧密码：

* 新密码：

* 确认新密码：

首页 > 知识摘要 > 人工智能 > 正文

【阅读记录-章节5】从零构建大语言模型

2025-02-05 15:29:25

来源：本站编辑

阅读量：

model.train() # 设置模型为训练模式- device (torch.device)- device (torch.device)) # 计算训练集损失) # 使用模型生成新的Token ID) # 计算训练集和验证集的损失# 使用 generate_text_simple 函数生成Token IDs并使用该模型进行文本生成# 返回训练损失、验证损失和已见Token数量的记录model.train() # 重新设置模型为训练模式# - torch.log 函数计算每

5\. Pretraining on unlabeled data
---------------------------------

到目前为止，我们已经实现了**数据采样和注意力机制**，并**编写了大型语言模型（LLM）的架构代码**。接下来，我们将实现一个**训练函数**并对LLM进行**预训练**。同时，我们将学习**基本的模型评估技术**，以衡量生成文本的质量，这是在训练过程中优化LLM所必需的。此外，我们还将讨论**如何加载预训练的权重**，为我们的LLM在微调阶段提供一个坚实的起点。图5.1展示了我们的整体计划，重点介绍了本章将讨论的内容。
![](/u/cms/www/202502/0514561427cq.png)
> **权重参数的补充介绍**
> 
> *   在大型语言模型（LLM）和其他深度学习模型的上下文中，**权重参数**是模型学习和调整的核心。这些权重也被称为权重参数，或简称为参数，指的是可训练的参数，学习过程会对其进行调整，以优化模型的性能。
> 
> **权重参数的作用**
> 
> *   **决定模型表现**：权重参数决定了模型在输入数据上的表现，通过训练过程不断调整这些参数，以最小化预测误差，从而提升模型的准确性和生成文本的质量。

> **在PyTorch中的存储与访问**
> 
> *   **存储位置**：在像PyTorch这样的深度学习框架中，权重通常存储在线性层（`torch.nn.Linear`）中。例如，我们在第3章中使用线性层实现了多头注意力模块，在第4章中实现了GPT模型。
> *   **访问方式**：
>     *   **单个层的权重**：初始化一个层（例如 `new_layer = torch.nn.Linear(...)`）后，可以通过 `.weight` 属性访问其权重，如 `new_layer.weight`。
>     *   **所有可训练参数**：为了方便起见，PyTorch 允许通过 `model.parameters()` 方法直接访问模型的所有可训练参数，包括权重和偏置。这在实现训练循环时尤为重要，因为我们需要迭代这些参数以更新模型。

### 5.1 Evaluating generative text models

在本章中，我们将在简要回顾第4章的文本生成内容后，设置我们的大型[言模型（LLM）进行文本生成，并讨论评估生成文本质量的基本方法。随后，我们将计算训练损失和验证损失。图5.2展示了本章涵盖的主题，前三个步骤已被突出显示
![](/u/cms/www/202502/05145653f31p.png)
#### 5.1.1 Evaluating generative text models

首先，让我们设置LLM并简要回顾我们在第4章中实现的文本生成过程。我们从初始化GPT模型开始，稍后将使用`GPTModel`类和`GPT_CONFIG_124M`字典对其进行评估和训练（参见[第4章](https://blog.csdn.net/weixin_44329069/article/details/144110715?spm=1001.2014.3001.5501)）：

import torch
    from chapter04 import GPTModel
    
    GPT_CONFIG_124M = {
        "vocab_size": 50257,
        "context_length": 256,
        "emb_dim": 768,
        "n_heads": 12,
        "n_layers": 12,
        "drop_rate": 0.1,
        "qkv_bias": False
    }
    
    torch.manual_seed(123)
    model = GPTModel(GPT_CONFIG_124M)
    model.eval()

我们将上下文长度（`context_length`）从1024个令牌缩短至256个令牌。相比上一章，唯一的调整是减少了上下文长度，这一修改降低了训练模型的计算需求，使得在标准笔记本电脑上进行训练成为可能。

原始的GPT-2模型拥有1.24亿个参数，配置为处理最多1024个令牌。在训练过程完成后，我们将更新上下文长度设置，并加载预训练权重，以使模型能够处理配置为1024个令牌上下文长度的模型。使用`GPTModel`实例，我们采用第4章中的`generate_text_simple`函数，并引入两个便捷函数：`text_to_token_ids`和`token_ids_to_text`。这些函数有助于在文本和令牌表示之间进行转换，这是我们将在本章中贯穿使用的技术。  
![](/u/cms/www/202502/05145722uc6l.png)
图5.3展示了使用GPT模型的三步文本生成过程：

1.  **令牌化**：分词器将输入文本转换为一系列令牌ID（参见第2章）。
2.  **模型生成**：模型接收这些令牌ID并生成相应的logits，这些logits是表示词汇表中每个令牌概率分布的向量（参见第4章）。
3.  **解码**：这些logits被转换回令牌ID，分词器将其解码为人类可读的文本，完成从文本输入到文本输出的循环。

我们可以按照以下代码实现文本生成过程：

**Listing 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)
        return encoded_tensor
    
    def token_ids_to_text(token_ids, tokenizer):
        flat = token_ids.squeeze(0)
        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))

使用这段代码，模型生成了以下文本：

Output text:
    Every effort moves you rentingetic wasn? refres RexMeCHicular stren

显然，模型尚未生成连贯的文本，因为它尚未经过训练。为了定义什么使得文本“连贯”或“高质量”，我们必须实现一种数值方法来评估生成的内容。这种方法将使我们能够在整个训练过程中监控和提升模型的性能。

接下来，我们将计算生成输出的损失指标。这个损失将作为训练进展的指标。此外，在后续章节中，当我们对LLM进行微调时，我们将回顾评估模型质量的其他方法。

#### 5.1.2 Calculating the text generation loss

在训练大型语言模型（LLM）时，评估生成文本的质量是至关重要的。通过计算文本生成的损失，我们可以量化模型的性能，并指导训练过程的改进。本文将通过一个实际的例子，逐步讲解如何评估模型生成文本的质量。

我们将从回顾数据的加载方式和 `generate_text_simple` 函数生成文本的过程开始。
![](/u/cms/www/202502/05145754g3c7.png)
如图5.4所示，文本生成过程可以分为五个步骤：

1.  **输入文本转换为Token IDs**：将输入文本转换为对应的Token IDs序列。
2.  **模型预测下一个Token的概率分布**：将Token IDs输入模型，计算每个位置下一个Token的概率分布（Logits），并通过Softmax函数转换为概率。
3.  **选择最可能的下一个Token**：对概率分布应用Argmax函数，选取概率最高的Token ID。
4.  **更新Token序列**：将选取的Token ID添加到输入序列中，重复步骤2和3，生成完整的输出序列。
5.  **将Token IDs转换回文本**：将生成的Token IDs序列转换回可读的文本。

> **需要注意的是**，图5.4中的示例为了简化，仅使用了一个包含7个Token的小词汇表。然而，在实际中，我们的GPT模型使用了包含50,257个词的更大词汇表，因此Token IDs的范围是0到50,256。

为了更好地理解，我们使用两个输入示例：

inputs = torch.tensor([
        [16833, 3626, 6100],    # ["every effort moves"]
        [40, 1107, 588]         # ["I really like"]
    ])

对应的目标输出（Targets）为：

targets = torch.tensor([
        [3626, 6100, 345],      # ["effort moves you"]
        [1107, 588, 11311]      # ["really like chocolate"]
    ])

注意，目标输出是将输入序列右移一个位置得到的。这种方式在第2章的数据加载器实现中已介绍过，主要用于训练模型预测序列中的下一个Token。

我们将输入数据输入模型，计算Logits，并通过Softmax函数转换为概率：

with torch.no_grad():
        logits = model(inputs)
        probas = torch.softmax(logits, dim=-1)
    print(probas.shape)

输出的概率张量形状为：

torch.Size([2, 3, 50257])

*   第一个维度`2`表示批次大小，即有两个输入示例。
*   第二个维度`3`表示每个输入序列中的Token数量。
*   第三个维度`50257`表示词汇表的大小，即每个位置上可能的下一个Token的概率分布。

通过对概率分布应用Argmax函数，我们可以得到模型预测的下一个Token IDs：

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

输出的Token IDs为：

Token IDs:
    tensor([[[16657],
             [  339],
             [42826]],
    
            [[49906],
             [29669],
             [41751]]])

我们使用Tokenizer将预测的Token IDs转换回可读文本：

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

输出结果：

Targets batch 1: effort moves you
    Outputs batch 1: Armed heNetflix

可以看到，模型生成的文本与目标文本有明显差异，这是因为模型还没有经过训练，输出的结果较为随机。

##### 评估模型生成文本的质量
![](/u/cms/www/202502/05145822fxhp.png)
现在，我们希望通过计算损失（如图5.5所示）来数值化地评估模型生成文本的性能。这不仅有助于衡量生成文本的质量，也是实现训练函数的基础，我们将使用它来更新模型的权重，以改进生成的文本。

我们实现的文本评估过程的一部分，如图5.5所示，是测量生成的Tokens与正确预测（目标）之间的“距离”。我们稍后实现的训练函数将使用这些信息来调整模型的权重，使其生成的文本更接近（或理想情况下匹配）目标文本。  
![](/u/cms/www/202502/05145845kcdr.png)
模型训练的目标是提高与正确目标Token IDs对应的索引位置上的Softmax概率，如图5.6所示。我们接下来将实现的评估指标也使用了这个Softmax概率来数值化地评估模型生成的输出：在正确位置上的概率越高，效果越好。

请记住，图5.6显示的是针对一个包含7个Token的小词汇表的Softmax概率，以便将所有内容放在一张图中。这意味着初始的随机值大约在1/7左右，约等于0.14。然而，我们的GPT-2模型使用的词汇表有50,257个Token，因此大多数初始概率将徘徊在0.00002（1/50,257）左右。

为了量化模型的性能，我们需要计算生成文本与目标文本之间的差异。这通常通过计算损失函数（如交叉熵损失）来实现。

我们可以提取模型在目标Token位置上的概率值：

# 设置文本索引为0，表示第一个文本样本
    text_idx = 0
    
    # 提取第一个文本样本中每个位置上目标Token的概率
    # probas的形状为[批次大小, 序列长度, 词汇表大小]
    # targets[text_idx]包含了第一个文本样本的目标Token IDs
    # [0, 1, 2]表示序列中的三个位置
    target_probas_1 = probas[text_idx, [0, 1, 2], targets[text_idx]]
    
    # 打印第一个文本样本的目标Token概率
    print("Text 1:", target_probas_1)
    
    # 设置文本索引为1，表示第二个文本样本
    text_idx = 1
    
    # 提取第二个文本样本中每个位置上目标Token的概率
    # targets[text_idx]包含了第二个文本样本的目标Token IDs
    target_probas_2 = probas[text_idx, [0, 1, 2], targets[text_idx]]
    
    # 打印第二个文本样本的目标Token概率
    print("Text 2:", target_probas_2)

输出结果：

Text 1: tensor([7.4541e-05, 3.1061e-05, 1.1563e-05])
    Text 2: tensor([1.0337e-05, 5.6776e-05, 4.7559e-06])

这些概率值表示模型在每个位置上预测目标Token的概率。由于模型尚未训练，这些概率值非常低。

训练模型的目标是**最大化正确Token的概率**，即提高目标Token在概率分布中的概率值。这可以通过最小化损失函数来实现。

常用的损失函数是**交叉熵损失**，它可以衡量预测分布与目标分布之间的差异。具体的计算方法将在后续的训练函数中实现。

> **反向传播**
> 
> *   我们如何最大化与目标 Token 对应的 Softmax 概率值？总体而言，我们需要更新模型的权重，使模型对我们希望生成的相应 Token ID 输出更高的概率值。权重的更新是通过一种称为**反向传播**的过程完成的，这是训练深度神经网络的标准技术（有关反向传播和模型训练的更多详细信息，请参见附录 A 的 A.3 至 A.7 节）。
> *   反向传播需要一个**损失函数**，该函数计算模型预测输出（在这里是对应目标 Token ID 的概率值）与实际期望输出之间的差异。这个损失函数衡量了模型的预测与目标值之间的偏差有多大。

接下来，我们将为两个示例批次 `target_probas_1` 和 `target_probas_2` 计算损失。主要步骤如图5.7所示。
![](/u/cms/www/202502/05145915owo1.png)
**计算概率分数的对数**

由于我们已经应用了步骤1到3以获得 `target_probas_1` 和 `target_probas_2`，接下来进行步骤4，对概率分数取对数：

# 对目标概率取对数并合并两个批次
    
    # 1. 使用 torch.cat 函数将两个目标概率张量 (target_probas_1 和 target_probas_2) 在第一个维度（即批次维度）上进行拼接
    #    - target_probas_1 和 target_probas_2 分别对应两个不同的输入批次
    #    - 拼接后的张量形状为 [6]，因为每个批次有3个概率值，总共2个批次
    combined_probas = torch.cat((target_probas_1, target_probas_2))  
    
    # 2. 对拼接后的概率值取自然对数
    #    - torch.log 函数计算每个元素的自然对数
    #    - 对数转换有助于将乘法操作转化为加法，简化后续的数学优化过程
    #    - 取对数后，较小的概率值会变得更负，便于计算损失
    log_probas = torch.log(combined_probas)  
    
    # 3. 打印对数概率值以供检查
    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中的所有六个步骤。

> **交叉熵损失**
> 
> *   交叉熵损失是机器学习和深度学习中常用的一种度量方法，用于衡量两个概率分布之间的差异——通常是真实标签的分布（这里是数据集中的Token）与模型预测的分布（例如，LLM生成的Token概率）。在机器学习的框架中，特别是像PyTorch这样的框架，`cross_entropy` 函数计算离散结果的交叉熵，这类似于模型生成的Token概率下目标Token的负平均对数概率，使得“交叉熵”和“负平均对数概率”这两个术语在实际中相关且常常可以互换使用。

在应用 `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 张量有三个维度：批次大小、Token 数量和词汇表大小。Targets 张量有两个维度：批次大小和 Token 数量。

对于 PyTorch 中的 `cross_entropy` 损失函数，我们需要通过合并批次维度来将这些张量展平：

# 将Logits和Targets展平以适应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 生成的 Token IDs，Logits 包含了在进入 Softmax 函数之前模型未缩放的输出。

之前，我们应用了 Softmax 函数，选择了对应目标 ID 的概率分数，并计算了负的平均对数概率。现在，PyTorch 的 `cross_entropy` 函数将为我们完成所有这些步骤：

# 计算交叉熵损失
    loss = torch.nn.functional.cross_entropy(logits_flat, targets_flat)
    print(loss)

输出结果为：

tensor(10.7940)

这个损失值与我们手动应用图5.7中的各个步骤时得到的结果相同。

> **困惑度（Perplexity）**
> 
> *   困惑度是常与交叉熵损失一起使用的度量，用于评估模型在语言建模等任务中的性能。它提供了一种更易于理解的方式来理解模型在预测序列中下一个 Token 时的不确定性。
> *   困惑度衡量模型预测的概率分布与数据集中实际单词分布的匹配程度。与损失类似，较低的困惑度表明模型的预测更接近实际分布。
> *   困惑度通常被认为比原始的损失值更具可解释性，因为它表示模型在每一步预测下一个 Token 时不确定的有效词汇量。在给定的示例中，这意味着模型在词汇表中的48,725个Token中对生成下一个 Token 的选择感到不确定。

我们已经为两个小型文本输入计算了损失，用于说明目的。接下来，我们将把损失计算应用到整个训练集和验证集上，以进一步评估和优化模型的性能。

#### 5.1.3 Calculating the training and validation set losses

在训练大型语言模型（LLM）之前，我们必须首先准备训练和验证数据集。接下来，如图5.8所示，我们将计算训练集和验证集的交叉熵损失，这是模型训练过程中一个重要的组成部分。  
![](/u/cms/www/202502/05145942xeyw.png)
为了计算训练和验证数据集的损失，我们使用一个非常小的文本数据集——埃迪丝·沃顿（Edith Wharton）的短篇小说《裁决》（The Verdict），这是我们在第2章中已经使用过的。选择公共领域的文本可以避免任何使用权相关的问题。此外，使用这样一个小的数据集可以让代码示例在标准笔记本电脑上几分钟内执行完成，即使没有高端GPU，这对于教育目的尤其有利。

> **注意**：感兴趣的读者还可以使用本书的补充代码准备一个由古腾堡计划（Project Gutenberg）中超过60,000本公共领域书籍组成的大规模数据集，并在其上训练LLM（详细信息见附录D）。

以下代码加载《裁决》短篇小说：

# 加载《裁决》短篇小说
    import os
    import urllib.request
    
    if not os.path.exists("the-verdict.txt"):
        url = ("https://raw.githubusercontent.com/rasbt/"
               "LLMs-from-scratch/main/ch02/01_main-chapter-code/"
               "the-verdict.txt")
        file_path = "the-verdict.txt"
        urllib.request.urlretrieve(url, file_path)
        
    with open(file_path, "r", encoding="utf-8") as file:
        text_data = file.read()

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

# 计算数据集中的字符数和Token数
    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个Token，文本可能看起来太小，不足以训练一个LLM，但如前所述，这是为了教育目的，以便我们可以在几分钟内运行代码而不是数周。此外，稍后我们将从OpenAI加载预训练权重到我们的GPTModel代码中。

接下来，我们将数据集划分为训练集和验证集，并使用第2章的数据加载器准备LLM训练的批次。这个过程在图5.9中进行了可视化。  
![](/u/cms/www/202502/0515000717v7.png)
由于空间限制，我们使用了`max_length=6`。然而，对于实际的数据加载器，我们将`max_length`设置为LLM支持的256 Token上下文长度，以便在训练过程中让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`代码：

# 从第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])

基于上述代码输出，我们有九个训练集批次，每个批次包含两个样本和256个Token。由于我们仅分配了10%的数据用于验证，因此只有一个验证批次，包含两个输入示例。如预期，输入数据（x）和目标数据（y）具有相同的形状（批次大小乘以每批次的Token数），因为目标是将输入序列右移一个位置得到的，如第2章所述。

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

def calc_loss_batch(input_batch, target_batch, model, device):
        """
        计算单个批次的交叉熵损失
    
        参数：
        - input_batch (Tensor): 输入批次的Token IDs
        - target_batch (Tensor): 目标批次的Token IDs
        - model (nn.Module): 训练中的LLM模型
        - device (torch.device): 计算设备（CPU或GPU）
    
        返回：
        - loss (Tensor): 该批次的交叉熵损失
        """
        # 将输入和目标批次移动到指定设备
        input_batch = input_batch.to(device)
        target_batch = target_batch.to(device)
        
        # 获取模型的Logits输出
        logits = model(input_batch)
        
        # 计算交叉熵损失
        loss = torch.nn.functional.cross_entropy(
            logits.flatten(0, 1),  # 展平Logits以适应cross_entropy函数
            target_batch.flatten()  # 展平目标批次
        )
        
        return loss

我们可以使用`calc_loss_batch`实用函数来实现以下`calc_loss_loader`函数，该函数计算通过给定数据加载器采样的所有批次的损失：

def calc_loss_loader(data_loader, model, device, num_batches=None):
        """
        计算通过数据加载器（data_loader）返回的多个批次的平均交叉熵损失。
    
        参数：
        - data_loader (DataLoader): 数据加载器，提供输入和目标批次。
        - model (nn.Module): 要评估的模型。
        - device (torch.device): 计算设备（CPU或GPU）。
        - num_batches (int, 可选): 要处理的批次数。如果为None，则处理所有批次。
    
        返回：
        - avg_loss (float): 所有处理批次的平均交叉熵损失。
        """
    
        # 初始化总损失为0
        total_loss = 0.0
    
        # 如果数据加载器为空，返回NaN
        if len(data_loader) == 0:
            return float("nan")
        
        # 如果未指定处理的批次数，则处理所有批次
        elif num_batches is None:
            num_batches = len(data_loader)
        
        # 如果指定了批次数，则取较小的值以避免超出数据加载器的范围
        else:
            num_batches = min(num_batches, len(data_loader))
        
        # 迭代数据加载器中的批次
        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()
            else:
                # 如果已经处理了指定的批次数，退出循环
                break
        
        # 计算并返回平均损失
        avg_loss = total_loss / num_batches if num_batches > 0 else float("nan")
        return avg_loss

现在，我们已经准备好了计算训练和验证集损失的工具函数，接下来我们将整个流程整合起来，计算训练和验证集的交叉熵损失。

# 指定计算设备（GPU 或 CPU）
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # 将模型移动到指定设备
    model.to(device)
    
    # 禁用梯度计算，因为我们只需要前向传播来计算损失
    with torch.no_grad():
        # 计算训练集的平均交叉熵损失
        train_loss = calc_loss_loader(train_loader, model, device)
        
        # 计算验证集的平均交叉熵损失
        val_loss = calc_loss_loader(val_loader, model, device)
        
    # 打印训练集和验证集的损失
    print("Training loss:", train_loss)
    print("Validation loss:", val_loss)

**输出示例**：

Training Loss: 10.7940
    Validation Loss: 10.7940

这些损失值反映了模型在训练集和验证集上的表现。较低的损失值表示模型的预测与目标更接近。

现在我们有了一种衡量生成文本质量的方法，我们将训练LLM以减少这一损失，从而使其在生成文本方面表现得更好，如图5.10所示。  
![](/u/cms/www/202502/05150037h44w.png)
接下来，我们将专注于对LLM进行预训练。模型训练完成后，我们还将实现替代的文本生成策略，并学习如何保存和加载预训练模型的权重。

### 5.2 Training an LLM

终于到了实现预训练LLM（我们的GPTModel）代码的时候。为此，我们将专注于一个简单明了的训练循环，以保持代码简洁易读。

> **注意**：感兴趣的读者可以在附录D中了解更多高级技术，包括学习率预热（learning rate warmup）、余弦退火（cosine annealing）和梯度裁剪（gradient clipping）。
![](/u/cms/www/202502/0515005772fv.png)
图5.11中的流程图展示了一个典型的PyTorch神经网络训练工作流程，我们将使用它来训练LLM。该流程图概述了八个步骤，从每个epoch的迭代、处理批次、重置梯度、计算损失和新梯度、更新权重，到最后的监控步骤，如打印损失和生成文本样本。

> **注意**：如果您对使用PyTorch训练深度神经网络相对陌生，并且对这些步骤中的任何一个不熟悉，建议阅读附录A的A.5至A.8节。

我们可以通过以下代码中的 `train_model_simple` 函数来实现这个训练流程。

def train_model_simple(model, train_loader, val_loader,
                           optimizer, device, num_epochs,
                           eval_freq, eval_iter, start_context, tokenizer):
        """
        计算通过数据加载器（data_loader）返回的多个批次的平均交叉熵损失。
        
        参数：
        - model (nn.Module): 要评估的模型。
        - train_loader (DataLoader): 训练数据加载器。
        - val_loader (DataLoader): 验证数据加载器。
        - optimizer (Optimizer): 优化器，用于更新模型权重。
        - device (torch.device): 计算设备（CPU或GPU）。
        - num_epochs (int): 训练的总轮数。
        - eval_freq (int): 评估的频率（每多少步进行一次评估）。
        - eval_iter (int): 在每次评估中处理的批次数。
        - start_context (str): 用于生成文本样本的起始文本片段。
        - tokenizer (Tokenizer): 分词器，用于编码和解码文本。
        
        返回：
        - train_losses (list): 训练集的损失记录。
        - val_losses (list): 验证集的损失记录。
        - track_tokens_seen (list): 已见Token数量的记录。
        """
        # 初始化用于跟踪损失和已见Token数量的列表
        train_losses, val_losses, track_tokens_seen = [], [], []
        tokens_seen, global_step = 0, -1
        
        # 遍历每个epoch
        for epoch in range(num_epochs):
            model.train()  # 设置模型为训练模式
            # 开始主训练循环
            for input_batch, target_batch in train_loader:
                optimizer.zero_grad()  # 清零之前批次的梯度
                loss = calc_loss_batch(
                    input_batch, target_batch, model, device
                )  # 计算当前批次的损失
                loss.backward()  # 反向传播计算梯度
                optimizer.step()  # 更新模型权重
                tokens_seen += input_batch.numel()  # 累计已见Token数量
                global_step += 1  # 更新全局步数
                
                # 如果当前步数是评估频率的倍数，则进行评估
                if global_step % eval_freq == 0:
                    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)  # 记录已见Token数量
                    
                    # 打印当前epoch、步数及损失
                    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(
                        model, tokenizer, device, start_context
                    )
        
        # 返回训练损失、验证损失和已见Token数量的记录
        return train_losses, val_losses, track_tokens_seen

**evaluate\_model** 函数对应于图5.11中的第7步。它在每次模型更新后打印训练集和验证集的损失，以便我们评估训练是否改善了模型。具体来说，`evaluate_model` 函数在计算训练集和验证集的损失时，确保模型处于评估模式，并禁用梯度跟踪和Dropout，以获得稳定和可复现的结果。

def evaluate_model(model, train_loader, val_loader, device, eval_iter):
        """
        计算训练集和验证集的平均交叉熵损失
        
        参数：
        - model (nn.Module): 要评估的模型
        - train_loader (DataLoader): 训练数据加载器
        - val_loader (DataLoader): 验证数据加载器
        - device (torch.device): 计算设备（CPU或GPU）
        - eval_iter (int): 在每个数据加载器中要处理的批次数
        
        返回：
        - train_loss (float): 训练集的平均损失
        - val_loss (float): 验证集的平均损失
        """
        model.eval()  # 设置模型为评估模式
        with torch.no_grad():  # 禁用梯度计算
            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

类似于`evaluate_model`，`generate_and_print_sample`函数是一个方便的函数，用于跟踪模型在训练过程中的改进情况。具体来说，`generate_and_print_sample`函数接受一个文本片段（`start_context`）作为输入，将其转换为Token ID，然后将其输入到LLM中，使用我们之前使用的`generate_text_simple`函数生成一个文本样本：

def generate_and_print_sample(model, tokenizer, device, start_context):
        """
        生成并打印文本样本
        
        参数：
        - model (nn.Module): 已训练的模型
        - tokenizer (Tokenizer): 分词器，用于编码和解码文本
        - device (torch.device): 计算设备（CPU或GPU）
        - start_context (str): 生成文本的起始文本片段
        
        功能：
        - 生成并打印一个文本样本，用于评估模型在训练过程中的生成能力
        """
        model.eval()  # 设置模型为评估模式
        context_size = model.pos_emb.weight.shape[0]  # 获取模型的上下文大小
        encoded = text_to_token_ids(start_context, tokenizer).to(device)  # 将起始文本编码为Token ID，并移动到设备上
        with torch.no_grad():  # 禁用梯度计算
            token_ids = generate_text_simple(
                model=model, idx=encoded,
                max_new_tokens=50, context_size=context_size
            )  # 使用模型生成新的Token ID
        decoded_text = token_ids_to_text(token_ids, tokenizer)  # 将Token ID解码回文本
        print(decoded_text.replace("\n", " "))  # 打印生成的文本样本
        model.train()  # 重新设置模型为训练模式

`evaluate_model`函数提供了模型训练进度的数值估计，而`generate_and_print_sample`函数则提供了一个具体的文本示例，以便在训练过程中评估模型的生成能力。

> **AdamW 优化器**
> 
> *   Adam优化器是训练深度神经网络的一个常用选择。然而，在我们的训练循环中，我们选择了AdamW优化器。AdamW是Adam的一个变种，它改进了权重衰减方法，旨在通过惩罚较大的权重来最小化模型复杂度并防止过拟合。这种调整使AdamW能够实现更有效的正则化和更好的泛化能力，因此在LLM的训练中经常使用。

让我们通过一个实际的例子来看看这一切是如何工作的：使用AdamW优化器和我们之前定义的`train_model_simple`函数，训练一个GPTModel实例10个epoch。

import torch
    
    # 设置随机种子以确保结果可复现
    torch.manual_seed(123)
    
    # 初始化GPTModel模型（假设GPT_CONFIG_124M已定义）
    model = GPTModel(GPT_CONFIG_124M)
    model.to(device)  # 将模型移动到指定设备
    
    # 定义AdamW优化器
    optimizer = torch.optim.AdamW(
        model.parameters(),
        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
    )

执行`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
    Intermediate results removed to save space
    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,
    ...
    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

正如我们所见，训练损失显著改善，从9.781降至0.391。模型的语言能力有了相当大的提升。开始时，模型只能在起始上下文后附加逗号（“Every effort moves you,”）或重复词语“and”。在训练结束时，它能够生成语法正确的文本。

类似于训练集损失，我们可以看到验证集损失从高值（9.933）开始，在训练过程中逐渐降低。然而，它从未像训练集损失那样降低，在第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):
        """
        绘制训练集和验证集的损失曲线
        
        参数：
        - epochs_seen (Tensor): 已见的epoch数量
        - tokens_seen (int): 已见的Token数量
        - train_losses (list): 训练集的损失记录
        - val_losses (list): 验证集的损失记录
        """
        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")
        
        # 设置x轴为整数
        ax1.xaxis.set_major_locator(MaxNLocator(integer=True))
        
        # 创建第二个x轴，显示已见Token数量
        ax2 = ax1.twiny()
        ax2.plot(tokens_seen, train_losses, alpha=0)
        ax2.set_xlabel("Tokens seen")
        
        fig.tight_layout()
        plt.show()
    
    # 创建一个包含epochs数量的张量
    epochs_tensor = torch.linspace(0, num_epochs, len(train_losses))
    
    # 绘制损失曲线
    plot_losses(epochs_tensor, tokens_seen, train_losses, val_losses)

生成的训练集和验证集损失图如图5.12所示。  
![](/u/cms/www/202502/05150133md2b.png)
我们可以看到，训练损失和验证损失在第一个epoch开始时都在改善。然而，损失在第二个epoch之后开始分化。这种分化以及验证损失远大于训练损失的事实表明模型正在过拟合训练数据。我们可以通过在“裁决”文本文件中搜索生成的文本片段，如“quite insensible to the irony”，来确认模型是否逐字记忆了训练数据。

由于我们使用的是一个非常小的训练数据集，并且进行了多轮训练，这种记忆化是可以预期的。通常情况下，训练模型时会使用更大规模的数据集，并且只进行一个epoch。

> **注意**：如前所述，感兴趣的读者可以尝试在古腾堡计划（Project Gutenberg）的60,000本公共领域书籍上训练模型，此时不会发生这种过拟合现象；详细信息见附录B。
![](/u/cms/www/202502/05150153fzvf.png)
如图5.13所示，我们已经完成了本章的四个目标。接下来，我们将介绍LLM的文本生成策略，以减少训练数据的记忆化并提高LLM生成文本的原创性，然后我们将讨论权重的加载和保存，以及从OpenAI的GPT模型加载预训练权重。

### 5.3 Decoding strategies to control randomness

让我们来看一下文本生成策略（也称为解码策略），以生成更具原创性的文本。首先，我们将简要回顾之前在 `generate_and_print_sample` 函数中使用的 `generate_text_simple` 函数。然后，我们将介绍两种技术：温度缩放（temperature scaling）和Top-K采样（top-k sampling），以改进该函数。

由于使用相对较小的模型进行推理不需要GPU，我们首先将模型从GPU转移回CPU。此外，在训练完成后，我们将模型设置为评估模式，以关闭诸如Dropout之类的随机组件：

# 将模型移动到CPU
    model.to("cpu")
    
    # 设置模型为评估模式，关闭Dropout等随机组件
    model.eval()

接下来，我们将GPTModel实例（model）插入到 `generate_text_simple` 函数中，该函数使用LLM一次生成一个Token：

import tiktoken
    
    # 获取GPT-2的编码器
    tokenizer = tiktoken.get_encoding("gpt2")
    
    # 使用 generate_text_simple 函数生成Token IDs
    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"]
    )
    
    # 将Token IDs解码回文本
    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

如前所述，生成的Token在每一步生成时都是根据词汇表中所有Token的最大概率分数选择的。这意味着即使我们多次在相同的起始上下文（“Every effort moves you”）上运行 `generate_text_simple` 函数，LLM也会始终生成相同的输出。

#### 5.3.1 Temperature Scaling

现在让我们来看一下温度缩放，这是一种在下一个Token生成任务中加入概率选择过程的技术。之前，在 `generate_text_simple` 函数中，我们总是使用 `torch.argmax`（也称为贪婪解码）选择概率最高的Token作为下一个Token。为了生成更具多样性的文本，我们可以将 `argmax` 替换为从概率分布中采样的函数（这里是LLM在每个Token生成步骤中为每个词汇表项生成的概率分数）。

为了通过一个具体的例子说明概率采样，我们将简要讨论使用一个非常小的词汇表进行下一个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()}

接下来，假设LLM接收到起始上下文 “every effort moves you”，并生成以下下一个Token的Logits：

import torch
    
    # 定义下一个Token的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函数选择生成的Token ID，然后通过反向词汇表将其映射回文本：

# 计算概率分布
    probas = torch.softmax(next_token_logits, dim=0)
    
    # 使用argmax选择概率最高的Token ID
    next_token_id = torch.argmax(probas).item()
    
    # 打印生成的词语
    print(inverse_vocab[next_token_id])

输出结果为：

forward

因为第四个位置（索引位置3，Python使用0索引）的Logit值最大，对应的Softmax概率分数也是最大的，所以生成的词语是 “forward”。

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

# 设置随机种子以确保结果可复现
    torch.manual_seed(123)
    
    # 使用multinomial函数根据概率分布采样Token ID
    next_token_id = torch.multinomial(probas, num_samples=1).item()
    
    # 打印采样生成的词语
    print(inverse_vocab[next_token_id])

输出结果仍然是 “forward”，就像之前一样。发生了什么？`multinomial` 函数根据概率分布按比例采样下一个Token。换句话说，"forward"仍然是最可能的Token，并且大部分时间会被 `multinomial` 选择，但并非总是如此。为了说明这一点，让我们实现一个函数，多次重复这个采样过程1000次：

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]}")
    
    # 执行采样并打印结果
    print_sampled_tokens(probas)

采样输出为：

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” 被大部分时间采样（582次），但其他词语如 “closer”、“inches” 和 “toward” 也会有一定的采样次数。这意味着，如果我们将 `argmax` 函数替换为 `multinomial` 函数，LLM有时会生成如 “every effort moves you toward”、“every effort moves you inches” 和 “every effort moves you closer” 的文本，而不是每次都生成 “every effort moves you forward”。

我们可以通过一个称为温度缩放（temperature scaling）的概念进一步控制分布和选择过程。温度缩放只是将Logits除以一个大于0的数的高大上描述：

def softmax_with_temperature(logits, temperature):
        """
        根据温度缩放后的Logits计算Softmax概率分布
        
        参数：
        - logits (Tensor): 下一个Token的Logits
        - temperature (float): 温度值，控制分布的平滑程度
        
        返回：
        - Tensor: 温度缩放后的概率分布
        """
        scaled_logits = logits / temperature  # 将Logits除以温度值
        return torch.softmax(scaled_logits, dim=0)  # 计算Softmax概率

温度大于1会导致Token概率分布更加均匀，而温度小于1会导致分布更加自信（更加尖锐或峰值更高）。让我们通过绘制原始概率和不同温度值下的概率来说明这一点：

import matplotlib.pyplot as plt
    from matplotlib.ticker import MaxNLocator
    
    # 定义温度值
    temperatures = [1, 0.1, 5]
    
    # 根据不同温度计算缩放后的概率分布
    scaled_probas = [softmax_with_temperature(next_token_logits, T) for T in temperatures]
    
    # 定义x轴
    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()
    
    # 设置x轴为整数
    ax.xaxis.set_major_locator(MaxNLocator(integer=True))
    
    # 调整布局并显示图表
    fig.tight_layout()
    plt.show()
    
![](/u/cms/www/202502/05150225qjg5.png)
如图5.14所示。温度为1时，模型会在将logits传递给softmax函数计算概率分数之前将其除以1。换句话说，使用温度为1等同于不使用任何温度缩放。在这种情况下，Token是通过PyTorch中的多项式采样函数以与原始softmax概率分数相等的概率被选中的。例如，在温度设置为1的情况下，如图5.14所示，“forward”对应的Token大约60%的时间会被选中。

此外，如图5.14所示，应用非常小的温度，例如0.1，会导致更尖锐的分布，使得多项式采样函数几乎100%的时间选择最可能的Token（这里是“forward”），接近argmax函数的行为。同样，温度为5会导致更均匀的分布，其他Token更频繁地被选中。这可以增加生成文本的多样性，但也更常导致无意义的文本。例如，使用温度为5会在大约4%的时间生成“every effort moves you pizza”这样的文本。

#### 5.3.2 Top-k sampling

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

结合概率采样和温度缩放，Top-K采样可以进一步改善文本生成结果。在Top-K采样中，我们可以将采样的Token限制在最可能的前K个Token，并通过屏蔽其他Token的概率分数将它们从选择过程中排除，如图5.15所示。  
![](/u/cms/www/202502/05150249zf6j.png)
Top-K方法通过将所有未选中的logits替换为负无穷（-inf）来实现，这样在计算softmax值时，非Top-K的Token的概率分数为0，其余概率总和为1。（细心的读者可能会记得我们在第3章3.5.1节中实现的因果注意力模块中的这个技巧。）

在代码中，我们可以按照图5.15的Top-K过程如下实现，首先选择logits值最大的K个Token：

# 定义Top-K的K值
    top_k = 3
    
    # 选择logits中最大的K个值及其对应的位置
    top_logits, top_pos = torch.topk(next_token_logits, top_k)
    print("Top logits:", top_logits)
    print("Top positions:", top_pos)

输出结果为：

Top logits: tensor([6.7500, 6.2800, 4.5100])
    Top positions: tensor([3, 7, 0])

接下来，我们使用PyTorch的`where`函数将低于Top-3中最低logit值的Token的logits值设置为负无穷（-inf）：

# 使用where函数屏蔽非Top-K的logits
    new_logits = torch.where(
        condition=next_token_logits < top_logits[-1],
        input=torch.tensor(float('-inf')),
        other=next_token_logits
    )
    
    print(new_logits)

输出结果为：

tensor([ 4.5100, -inf, -inf,  6.7500, -inf, -inf, -inf,  6.2800, -inf])

最后，我们应用softmax函数将这些新的logits转换为下一个Token的概率分布：

# 计算Top-K的概率分布
    topk_probas = torch.softmax(new_logits, dim=0)
    print(topk_probas)

输出结果为：

tensor([0.0615, 0.0000, 0.0000, 0.5775, 0.0000, 0.0000, 0.0000, 0.3610, 0.0000])

如我们所见，Top-3方法的结果是三个非零概率分数：

*   "forward"的概率为57.75%
*   "toward"的概率为36.10%
*   "closer"的概率为6.15%

我们现在可以应用温度缩放和多项式采样函数，根据这三个非零概率分数在它们之间选择下一个Token，以生成下一个Token。接下来，我们将通过修改文本生成函数来实现这一点。

#### 5.3.3 Modifying the text generation function

现在，让我们将温度采样（temperature sampling）和 top-k 采样（top-k sampling）结合起来，修改我们之前用于通过大型语言模型（LLM）生成文本的 `generate_text_simple` 函数，创建一个新的 `generate` 函数。

**一个具有更多多样性的修改后的文本生成函数**

def generate(model, idx, max_new_tokens, context_size,
                temperature=0.0, top_k=None, eos_id=None):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -context_size:]
            with torch.no_grad():
                logits = model(idx_cond)
                logits = logits[:, -1, :]
            # for 循环与之前相同：获取 logits 并仅关注最后一个时间步
            if top_k is not None:
                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:
                logits = logits / temperature  # 应用温度缩放
                probs = torch.softmax(logits, dim=-1)
                idx_next = torch.multinomial(probs, num_samples=1)
            else:
                idx_next = torch.argmax(logits, dim=-1, keepdim=True)  # 进行贪婪的下一个 token 选择
            if idx_next == eos_id:
                break  # 如果遇到序列结束 token，提前停止生成
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

让我们现在看看这个新的 `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("输出文本:\n", token_ids_to_text(token_ids, tokenizer))

生成的文本是：

输出文本:
    with random-
    Every effort moves you stand to work on surprise, a one of us had gone

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

##### 练习5.2 补充介绍

在生成文本时，**温度（temperature）** 和 **top-k** 参数的设置对生成结果的多样性和确定性有显著影响。通过调整这些参数，可以适应不同的应用场景需求。

*   **低温度（temperature < 1.0）**：
    
    *   **特点**：降低温度会减少生成过程中的随机性，使模型更倾向于选择概率更高的词汇。
    *   **应用场景**：
        *   **技术文档生成**：需要精确和一致的表述，减少语法或逻辑错误。
        *   **法律文本或合同生成**：要求高准确性和一致性，避免歧义。
        *   **问答系统**：提供准确且一致的答案，避免不相关或错误的信息。
*   **低 top-k 设置（top\_k 较小，例如 1-10）**：
    
    *   **特点**：限制模型只从概率最高的前 k 个词中选择，进一步减少生成的多样性。
    *   **应用场景**：
        *   **命令行工具或脚本生成**：需要生成精确的命令，避免错误操作。
        *   **数据填充或模板化内容生成**：确保生成内容符合特定格式或结构要求。
        *   **客户支持聊天机器人**：提供一致且可靠的回应，避免混淆用户。
*   **高温度（temperature > 1.0）**：
    
    *   **特点**：增加温度会提升生成过程中的随机性，使模型更倾向于选择概率较低的词汇，从而增加多样性和创造性。
    *   **应用场景**：
        *   **创意写作**：如小说、诗歌等，鼓励生成新颖和富有创意的内容。
        *   **广告文案生成**：需要吸引眼球和具有创意的表达方式。
        *   **对话系统中的闲聊功能**：提供多样化和有趣的回应，增强用户互动体验。
*   **高 top-k 设置（top\_k 较大，例如 50-100 或更高）**：
    
    *   **特点**：允许模型从更多可能的词汇中选择，增加生成内容的多样性和丰富性。
    *   **应用场景**：
        *   **游戏剧情生成**：需要丰富多变的剧情走向和角色对话。
        *   **艺术作品生成**：如音乐、绘画描述等，需要多样化的表达。
        *   **教育内容生成**：提供多角度的解释和示例，增强学习材料的丰富性。

通过合理调整温度和 top-k 参数，可以在生成文本的确定性与创造性之间找到平衡，满足不同应用场景的需求。

##### 练习5.3 补充介绍

要实现 `generate` 函数的确定性行为，即禁用随机采样，使其总是生成相同的输出，类似于 `generate_simple` 函数，可以通过以下设置组合来实现：

1.  **温度设置为 0**：
    
    *   **参数**：`temperature=0.0`
    *   **效果**：完全消除概率分布中的随机性，模型会选择概率最高的词汇（贪婪搜索）。
2.  **top-k 设置为 `None` 或等于词汇表大小**：
    
    *   **参数**：`top_k=None` 或 `top_k=词汇表大小`
    *   **效果**：不限制词汇选择范围，确保模型可以选择所有可能的词汇，但由于温度为 0，最终仍然选择概率最高的词汇。
3.  **使用贪婪搜索（greedy decoding）**：
    
    *   **实现**：在代码中，当 `temperature=0.0` 且 `top_k=None` 时，函数会使用 `torch.argmax` 来选择下一个词汇，确保每次生成的结果一致。
4.  **设置随机种子**（可选，但推荐）：
    
    *   **实现**：通过设置随机种子（例如 `torch.manual_seed(seed)`），确保在多次运行中模型的初始化和生成过程保持一致。
    *   **示例**：
        
            torch.manual_seed(123)

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=None,          # 不限制 top-k
        temperature=0.0,     # 禁用温度缩放，使用贪婪搜索
        eos_id=None          # 根据需要设置序列结束 token
    )
    print("输出文本:\n", token_ids_to_text(token_ids, tokenizer))

通过上述设置，`generate` 函数将始终选择概率最高的词汇，生成的文本将保持一致，确保确定性行为。这类似于 `generate_simple` 函数的行为，适用于需要高一致性和可预测性的应用场景。

### 5.4 Loading and saving model weights in PyTorch

到目前为止，我们已经讨论了如何通过数值评估训练进展以及从零开始预训练一个大语言模型（LLM）。尽管本例中使用的LLM和数据集都相对较小，但这一练习表明，预训练LLM的计算成本非常高。因此，能够保存LLM以便在新的会话中使用时无需重新训练显得尤为重要。  
![](/u/cms/www/202502/051503173rpe.png)
接下来，我们将讨论如何保存和加载预训练模型，这在图5.16中得到了强调。之后，我们会将一个更强大的预训练GPT模型从OpenAI加载到我们的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 的预训练，请按照以下步骤操作：

确保您已经保存了模型和优化器的状态，具体代码可以参考上述内容。如果已经保存，执行以下步骤：

import torch
    from your_model_file import GPTModel, train_model_simple  # 替换为实际文件名和函数名
    
    # 加载配置和设备
    GPT_CONFIG_124M = {
        "hidden_size": 768,
        "num_attention_heads": 12,
        "num_hidden_layers": 12,
        "vocab_size": 50257,
        "max_position_embeddings": 1024,
    }
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # 加载检查点
    checkpoint = torch.load("model_and_optimizer.pth", map_location=device)
    
    # 创建模型实例并加载权重
    model = GPTModel(GPT_CONFIG_124M)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.to(device)
    
    # 创建优化器并加载状态
    optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4, weight_decay=0.1)
    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    
    # 切换到训练模式
    model.train()

确保您有 `train_model_simple` 函数的实现，以及对应的数据集。如果数据集是分批加载的，请预处理数据集。

# 假设 train_dataset 已定义并可用
    train_dataloader = torch.utils.data.DataLoader(
        train_dataset, batch_size=16, shuffle=True
    )

调用函数继续预训练一个 epoch：

# 调用自定义的训练函数
    train_model_simple(
        model=model,
        dataloader=train_dataloader,
        optimizer=optimizer,
        device=device,
        num_epochs=1  # 继续训练一个 epoch
    )

在完成训练后，可以再次保存更新的模型和优化器状态以备后续使用：

torch.save({
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
    }, "updated_model_and_optimizer.pth")

### 5.5 Loading pretrained weights from OpenAI

之前，我们使用一个包含短篇故事书的小型数据集训练了一个小型的 GPT-2 模型。这种方法使我们能够专注于基础知识，而不需要大量的时间和计算资源。

幸运的是，OpenAI 公开分享了他们的 GPT-2 模型的权重，因此我们无需自己花费数万美元到数十万美元在大型语料库上重新训练模型。那么，让我们将这些权重加载到我们的 `GPTModel` 类中，并使用该模型进行文本生成。这里的权重指的是存储在 PyTorch 的 `Linear` 和 `Embedding` 层的 `.weight` 属性中的权重参数。例如，我们在训练模型时通过 `model.parameters()` 访问它们。在第 6 章中，我们将重用这些预训练的权重，对模型进行微调，以完成文本分类任务，并按照类似于 ChatGPT 的指令操作。

请注意，OpenAI 最初是通过 TensorFlow 保存 GPT-2 的权重的，我们需要安装 TensorFlow 才能在 Python 中加载这些权重。以下代码将使用一个名为 `tqdm` 的进度条工具来跟踪下载过程，我们也需要安装它。您可以在终端中执行以下命令来安装这些库：

pip install tensorflow>=2.15.0 tqdm>=4.66

下载代码相对较长，主要是样板代码，并不是很有趣。因此，为了节省宝贵的空间，不讨论用于从互联网获取文件的 Python 代码，我们直接从本章的在线存储库中下载 `gpt_download.py` Python 模块：

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 会话的本地目录后，您应该简要检查该文件的内容，以确保它已正确保存并包含有效的 Python 代码。

现在，我们可以从 `gpt_download.py` 文件中导入 `download_and_load_gpt2` 函数，这将把 GPT-2 的架构设置（`settings`）和权重参数（`params`）加载到我们的 Python 会话中：

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, ...]
    encoder.json: 100%|█████████████████████████| 1.04M/1.04M [00:00<00:00, 63.9kiB/s]
    hparams.json: 100%|██████████████████████████| 90.0/90.0 [00:00<00:00, ...]
    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, ...]
    model.ckpt.meta: 100%|██████████████████████| 471k/471k [00:00<00:00, 3.24MiB/s]
    vocab.bpe: 100%|████████████████████████████| 456k/456k [00:00<00:00, 1.70MiB/s]

> **注意** 如果下载代码无法正常工作，可能是由于网络连接不稳定、服务器问题或 OpenAI 更改了共享开源 GPT-2 模型权重的方式。在这种情况下，请访问本章的在线代码存储库：https://github.com/rasbt/LLMs-from-scratch，获取替代和更新的说明，如有进一步问题，请通过 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)

Token 嵌入层的权重为：

[[-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 所示，只是不同的架构元素重复的次数和嵌入维度不同。本章剩余的代码也兼容这些更大的模型。  
![](/u/cms/www/202502/05150412vvak.png)
在将 GPT-2 模型权重加载到 Python 后，我们仍然需要将它们从 `settings` 和 `params` 字典中转移到我们的 `GPTModel` 实例中。首先，我们创建一个字典，列出图 5.17 中不同 GPT 模型大小之间的差异：

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`）是否具有相同的尺寸或形状，并返回作为可训练的 PyTorch 参数的 `right` 张量：

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` 中。

**将 OpenAI 的权重加载到我们的 GPT 模型中**

import numpy as np
    
    def load_weights_into_gpt(gpt, params):
        # 设置模型的位置和 token 嵌入权重
        gpt.pos_emb.weight = assign(gpt.pos_emb.weight, params['wpe'])
        gpt.tok_emb.weight = assign(gpt.tok_emb.weight, params['wte'])
        
        # 遍历模型中的每个 Transformer 块
        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-2 模型在输出层重用了 token 嵌入权重，以减少参数总数，这个概念称为权重共享。
        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` 实现进行匹配。举一个具体的例子，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)

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

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?

我们可以确信我们正确加载了模型权重，因为模型能够产生连贯的文本。在接下来的章节中，我们将进一步使用这个预训练模型，并对其进行微调，以分类文本并按照指令操作。

上一篇：【阅读记录-章节4】...

下一篇：【阅读记录-章节6】...

已有0条评论

0/150

提交

热门评论

查看更多评论

目 录

目录