不使用库，不走捷径：用 PyTorch 从零开始构建 LLM

这是一份不废话的指南，教你从零开始构建、训练和微调 Transformer 架构。

OpenAI 近期发布了备受期待的开源 GPT-OSS 模型，这不禁我们反思我们一路走来的历程。几年前，甚至在 ChatGPT 出现之前，我记得读过一篇关于 GPT 模型（可能是 GPT-2）的文章，它能够自己撰写文章和诗歌，而当时这些都只是实验。如今，它已经成为我日常生活中不可或缺的一部分。这一切都始于 2017 年谷歌研究院发表的里程碑式论文《注意力机制就是一切》（Attention is All You Need）。该论文提出了 Transformer 架构，并很快在 2018 年催生了第一个 GPT 模型——GPT-1（生成式预训练 Transformer）。

在过去的八年中，大型语言模型的发展可谓突飞猛进。如今的大型语言模型具备多模态能力、高级推理技巧和创新的架构改进。然而，它们的核心仍然依赖于Transformer框架。由于如今大型语言模型可以通过用户友好的框架和API轻松上手，许多开发者往往忽略了这种底层设计的精妙之处。

本文将深入剖析 Transformer 架构，逐一解读其各个组成部分。读完本文，你将从零开始构建自己的 LLM（语言学习模型），并创作全新的 Coldplay 歌曲（说不定你还会开始为 Coldplay 谱写歌词呢！）。

分词器

你提供的任何文本都会被 LLM 的分词器分解成称为token的更小单元，token的范围可以从一个字符到整个单词。

考虑以下文本：“Hold my math!” 根据模型的设计，token可以是单词、子词，甚至是字符。

词级分词：["Hold", "my", "math", "!"]

子词级分词：["Hold", "my", "ma", "th", "!"]

字符级分词：["H", "o", "l", "d", " ", "m", "y", " ", "m", "a", "t", "h", "!"]

下一个Token预测器

大型语言模型从基本定义上来说就是下一个token预测器。给定输入的token，模型学习分析并预测下一个token的概率。

它们一次只能处理固定数量的token，每一步生成一个token。你看到的长回复，正是通过高效地重复这个过程实现的。这是通过在每次预测后反复向前滑动输入窗口，并在遇到一个eostoken或达到某个长度限制时停止来实现的。

事实上，你可能会认为对于 LLM API 调用，你应该这样编写代码。

messages = [ 
    { 
"role": "system" , 
"content": "你是一位富有创意的故事讲述者。"
     }, 
    { 
"role": "user" , 
"content": "写一个创意故事"
     }, 
]

但是，经过库的处理后，输入到大型语言模型的输入将会不同。

""" 
<|im_start|>system
你是一位富有创意的故事讲述者。
<|im_end|> 
<|im_start|>user
写一个创意故事
<|im_end|> 
<|im_start|>assistant
"""

经过预训练（你将在后续章节中看到），大型语言模型会根据指令进行微调，使其能够用于人机交互。这称为指令调优。否则，大型语言模型将仅仅是一个简单的文本生成器。上述格式展示了输入传递给经过指令调优的模型的方式，尽管<|im_start|>user某些模型的标签可能有所不同。

Attention is All You Need

那么，LLM是如何决定生成什么内容的呢？它可以自由生成任何随机内容。但如何才能生成有意义的输出呢？为此，神经网络需要学习并利用上下文信息，不仅包括最后一个token，还包括输入句子的其他部分。

考虑以下输入示例：The cat chased the

如果模型只看最后一个token（"the"），它可能会预测几乎任何内容："banana"，，，等等"man"。"moon"

但是，如果它使用完整的上下文（"The", "cat", "chased", "the"），它就知道"mouse"比"moon"或"banana"更有可能。

这不仅适用于完成任务，也适用于翻译。

英语：

I eat a red apple.

法语：

Je mange une pomme rouge.

匹配的词语是：

English:  [I]     [eat]     [a]      [red]     [apple]
           ↓        ↓        ↓         ↓          ↓
French:   [Je]   [mange]   [une]    [rouge]    [pomme]    

正如你所料，在法语中，red应该放在apple之后。因此，逐字翻译在这里行不通。相反，模型应该学习并理解在这种情况下词序会发生这样的变化。

在序列建模的早期，循环神经网络（RNN）是处理文本的首选方法。它们按顺序逐个处理单词，并将“隐藏状态”向前传递，因此每一步都带有对先前信息的记忆。这种方法适用于较短的依赖关系，但随着序列长度的增加，来自先前单词的重要信息往往会逐渐消失。

1997年提出的长短期记忆（LSTM）网络，通过使用特殊的门控机制（即输入门、遗忘门和输出门）来控制信息流，从而改进了上述方法。这些门决定保留、更新或丢弃当前输入和过去记忆中的哪些部分，使网络能够在更长的序列中保持相关信息。

注意力机制于 2017 年在《Attention Is All You Need》一书中提出，旨在解决这一局限性。注意力机制并非依赖于顺序处理，而是将输入中的每个词直接与其他所有词关联起来，计算权重以确定每个token对预测结果的影响程度。这使得模型能够高效且并行地捕捉长程依赖关系。接下来，我们将快速了解注意力机制的工作原理。

目标是衡量句子中每个token与其他token之间的关联程度或影响程度。对token对的这种关联程度称为注意力得分。收集所有token对的注意力得分，即可得到注意力矩阵。为此，我们需要每个token的三个分量。

1. 查询向量表示该token在其他token中寻找的内容。它是通过输入嵌入向量与可训练查询矩阵的叉积计算得出的Wq。
2. 关键向量表示该token提供的可搜索信息。它是通过输入嵌入向量与可训练的关键矩阵的叉积计算得出的Wk。
3. 值向量包含如果token被关注将会传递的实际信息内容。它是通过输入嵌入向量与一个可训练值矩阵的叉积计算得出的Wv。

token得分的计算方法是将查询向量与待评分词的键向量进行点积运算。为简单起见，我们将句子中的每个词都视为一个token：“The cat slept on the mat and it purred. ” 因此，如果我们对位置 3 的词“slept”进行自注意力处理，则第一个得分是 q3 与 k1 的点积。第二个得分是 q3 与 k2 的点积，依此类推。这样，每个token都会得到一个得分。每个得分都使用softmax 激活函数进行归一化，然后除以嵌入维度的平方根，d**0.5得到注意力权重。最后，将每个值向量乘以其对应的注意力权重，得到上下文向量。

这些上下文向量被组合成一个矩阵，我们称之为注意力矩阵。在实际应用中，所有这些操作都以矩阵运算的形式进行计算。正如你所看到的，“ it-cat ”和“ purred-cat ”这两个单元格的值较高，表明它们之间的相关性更高。

理论上，一个token被选为下一个token的概率应该只取决于过去的token，而与未来的token无关，否则就说不通了。为了实现这一点，我们需要在训练阶段屏蔽矩阵中所有未来的值，并调整行值，使它们的总和仍然为 1，因为这些值代表的是概率。这被称为因果注意力机制，在接下来的编码部分你会更好地理解它。

构建 Transformer 架构

在深入代码之前，我们先来了解一下整体结构。Transformer 由编码器和解码器模块构成，每个模块都包含相同的关键组件：词嵌入、位置编码、自注意力机制、多头注意力机制和前馈层。

从本质上讲，可以把它看作是一个模块化的流水线：

1. 输入的文本序列被拆分成token。
2. token被转换为数值嵌入，并带有位置编码。
3. 自我关注让每个token“关注”其他每个token，或者更简单地说，该模型会弄清楚每个token与其他token之间的关系。
4. 多头注意力机制结合了多个自注意力层，使模型能够从多个角度看待这些关系。
5. 前馈网络将组合信息转化为更强的特征。

通过堆叠许多这样的层，该模型可以学习到越来越丰富的语言表征。

以下各节将逐一剖析这些构建模块，我们将逐步在 PyTorch 中实现它们，直到从零开始构建出一个可运行的 Transformer 模型。你将重点关注 GPT 架构，它仅包含解码器部分。

分词

第一步应该是像我们之前讨论的那样，将输入文本分割成token（token）。你可以从头开始实现字符级分词或词级分词，但这里你可以使用tiktoken现成的分词工具。无论你使用哪种方法，关键步骤都是：对输入文本进行分词，构建一个将每个token映射到索引的词汇表，并确保你可以将输入文本中的token转换为数值 ID，以及将输出文本中的 ID 转换回token。

import tiktoken
def text_to_token_ids(text, tokenizer, device):
    encoded = tokenizer.encode(text, allowed_special={'<|endoftext|>'})
    encoded_tensor = torch.tensor(encoded).unsqueeze(0).to(device)  # 添加批次维度并移动到设备
return encoded_tensor 
def token_ids_to_text(token_ids, tokenizer):
    flat = token_ids.squeeze(0) # 移除批次维度
return tokenizer.decode(flat.tolist())

对于自定义词汇表大小，你可以自行实现字节对编码器。

位置编码与嵌入

位置编码对序列中每个token的位置进行编码。Transformer 模型依赖位置编码来保持序列顺序，因为与 RNN 不同，Transformer 并行处理所有token。简单的索引不足以表达长序列的信息，因此位置编码使用数学模式（通常是正弦和余弦函数）来创建一个矩阵，该矩阵嵌入了更丰富的位置信息。这使得模型既能感知序列顺序，又能利用并行计算的优势。

嵌入层是将原始token转换为神经网络可处理数据的第一步。词汇表中的每个token ID 都被映射到一个固定大小的稠密向量，称为嵌入向量。与使用稀疏独热编码表示单词不同，嵌入向量能够捕捉语义信息，因此像“king”和“queen”这样相关的词在向量空间中会比像 “king” 和 “banana” 这样不相关的词更接近。

该嵌入矩阵的维度为（词汇表大小 × 嵌入维度），其中嵌入维度是一个由你选择的超参数。位置编码矩阵的维度应与嵌入矩阵的维度相匹配，因为正如你在架构图中看到的，嵌入矩阵是通过与位置编码矩阵逐元素求和来丰富信息的。

为简化起见，假设输出的嵌入维度和注意力维度相同。对于输入，嵌入层表示整个词汇表。如果模型使用 10000 个不同的词（字面意义上的token）进行训练，则输入形状为 10000。对于 GPT-2，标准输入为 50257 个词。类似地，位置编码层的输入等于 LLM 输入层的token数（LLM 一次处理的token数）。我们称之为上下文长度。

self.embedding = torch.nn.Embedding(vocab_size, attention_dim) 
self.positional_embedding = torch.nn.Embedding(context_length, attention_dim)

在前向法中，将流量定义为：

embeddings = self.embedding(context) 
context_len = context.shape[1] 
position = torch.arange(context_len, device=context.device).unsqueeze(0) 
position_embeddings = self.positional_embedding(position) 
e = embeddings + position_embeddings

你将在我们即将构建的完整 GPT 解码器中看到它的实际应用。

自我关注：代币们如何互相八卦

现在，我们利用前面章节讨论的内容来编写自注意力模块。定义可训练的查询矩阵、键矩阵和值矩阵。

self.w_key = torch.nn.Linear(embed_dim, attention_dim, bias=bias) 
self.w_query = torch.nn.Linear(embed_dim, attention_dim, bias=bias) 
self.w_value = torch.nn.Linear(embed_dim, attention_dim, bias=bias)

计算查询向量、键向量和值向量。

k = self.w_key(x)   # (B, T, A)
q = self.w_query(x) # (B, T, A)
v = self.w_value(x) # (B, T, A)
 ```
其中，
B：批次大小，
T：上下文长度，
A：注意力维度
```

现在，使用查询向量和键向量的乘积计算注意力得分。将键向量转置以匹配维度。此外，在应用 softmax 函数之前，将结果除以嵌入维度的平方根进行归一化。

scores = (q @ k.transpose(-2, -1)) / (k.size(-1) ** 0.5)  # (B, T, T)

掩盖未来位置并应用softmax激活函数。最后，返回注意力权重与值向量的乘积。

mask = torch.triu(torch.ones(T, T, device=x.device), diagonal=1).bool()
scores = scores.masked_fill(mask, float('-1e10'))
attn = scores.softmax(dim=-1)  # (B, T, T)
final = attn @ v  # (B, T, A)

别忘了根据需要添加 dropout 来稳定训练过程。dropout 会根据你指定的百分比随机屏蔽一部分注意力权重。这里设置为 0.1，意味着 10% 的权重会被随机屏蔽。dropout 会作用于整个矩阵，而不仅仅是剩余的值。

将所有内容整合为一个模块，

classSelfAttention(torch.nn.Module):
def__init__(self, embed_dim, attention_dim, bias=False, dropout=0.1):
        super().__init__()
        self.w_key = torch.nn.Linear(embed_dim, attention_dim, bias=bias)
        self.w_query = torch.nn.Linear(embed_dim, attention_dim, bias=bias)
        self.w_value = torch.nn.Linear(embed_dim, attention_dim, bias=bias)
        self.dropout = torch.nn.Dropout(dropout)
defforward(self, x):
        B, T, _ = x.size()
        k = self.w_key(x)   # (B, T, A)
        q = self.w_query(x) # (B, T, A)
        v = self.w_value(x) # (B, T, A)
# 缩放点积注意力
        分数 = (q @ k.transpose( -2 , -1 )) / (k.size( -1 ) ** 0.5 )   # (B, T, T) 
# 因果掩码（未来位置被掩码）
        mask = torch.triu(torch.ones(T, T, device=x.device), diagonal=1).bool()
        scores = scores.masked_fill(mask, float('-1e10'))
        attn = scores.softmax(dim=-1)  # (B, T, T)
        attn = self.dropout(attn)
return attn @ v  # (B, T, A)

多头注意力：模型大脑中的群聊

你已经准备好了一个自注意力模块。但是对于一个在海量数据集上训练的大型模型来说，这远远不够。因此，我们并行地组合多个注意力头，创建了多头注意力机制。其核心思想是，多个注意力头可以专注于数据的不同方面，捕捉不同子空间和位置之间的关系，从而使模型能够学习更丰富、更复杂的模式。

classMultiHeadAttention(torch.nn.Module):
def__init__(self, num_heads, embed_dim, attention_dim, dropout=0.1):
        super().__init__()
        self.head_size = attention_dim//num_heads
        self.heads = torch.nn.ModuleList()
for i in range(num_heads):
            self.heads.append(SelfAttention(embed_dim=embed_dim, attention_dim=self.head_size,dropout=dropout))
defforward(self,x):
        head_outputs = []
for head in self.heads:
            head_outputs.append(head(x)) #B x T x A//num_heads
        concatenated = torch.cat(head_outputs, dim = 2)
return concatenated

总注意力维度被分配给各个注意力头，从而实现并行处理。最终得到的注意力矩阵输出被连接起来，得到完整的结果。

前馈网络

经过注意力机制处理后，每个词嵌入都会经过一个小型前馈网络。这个网络可以简单地由两个线性层组成，分别包含一个上投影层和一个下投影层，中间穿插一个非线性激活函数。当然，你也可以根据网络架构使其更加复杂。

class  FeedForward (torch.nn.Module): 
    def  __init__ ( self, attention_dim ): 
        super ().__init__() 
        self.up = torch.nn.Linear(attention_dim, attention_dim * 4 ) 
        self.gelu = torch.nn.GELU() 
        self.down = torch.nn.Linear(attention_dim * 4 , attention_dim) 
    def  forward ( self, x ): 
return self.down(self.gelu(self.up(x)))

带残差连接的解码器

现在你已经有了多头注意力模块和前馈网络，我们来构建模型的解码器部分。

仔细观察下图。你会发现，每个子层（注意力层或前馈层）并非简单地替换其输入，而是将原始输入加回子层的输出。这种简化的方法称为残差连接。它通过保留原始信号并避免梯度消失等问题，帮助模型更有效地训练。

此外，每个残差连接之后都会进行层归一化。与批量归一化（BatchNorm）不同，层归一化 是基于单个词嵌入的特征进行归一化，将神经网络层的激活值集中于均值为 0 的范围内，并将其方差归一化为 1。这确保了激活值的尺度和分布保持稳定，这对于像 Transformer 这样的深度网络至关重要，因为堆叠过多的层可能会导致训练不稳定。

按照图示将各层堆叠起来，你就会得到这个样子。

classDecoder(torch.nn.Module):
def__init__(self,num_heads,embed_dim,attention_dim, dropout=0.1):
        super().__init__()
        self.masked_multihead = MultiHeadAttention(num_heads, embed_dim, attention_dim, dropout)
        self.feed_forward = FeedForward(attention_dim)
        self.n1 = torch.nn.LayerNorm(attention_dim)
        self.n2 = torch.nn.LayerNorm(attention_dim)
defforward(self,x):
        e = self.masked_multihead(self.n1(x))
        e =  e + x
        e = self.feed_forward(self.n2(e))
return e

将所有部件组装起来：Transformer骨架

完美。现在，你只需要设计模型的输入和输出部分。输入部分是嵌入层和位置编码层，我们之前已经创建好了。它们将原始token ID 转换为包含位置信息的密集向量，为 Transformer 模块做好准备。

输出部分就更有意思了。你已经了解到，LLM 的输出必须代表整个词汇表上的概率分布。这意味着，我们期望它学习的每个token都应该有自己的输出，或者换句话说，输出的大小等于词汇表的大小。因此，我们将它们线性投影到一个词汇表大小的层上，这个层也称为LM（语言建模）头。LM 头只是一个没有激活函数的线性投影。输出仍然不是概率，但可以称为logits。在训练/推理过程中，会应用softmax 函数来获得概率分布。

import torch
from torch import nn
classGPT(nn.Module):
def__init__(self, num_heads, vocab_size, embed_dim, attention_dim, num_blocks, context_length, dropout_rate):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, attention_dim)
        self.positional_embedding = nn.Embedding(context_length, attention_dim)
        self.decoders = nn.ModuleList([
            Decoder(num_heads, attention_dim, attention_dim, dropout_rate) for _ in range(num_blocks)
        ])
        self.exit_norm = nn.LayerNorm(attention_dim)
        self.linear = nn.Linear(attention_dim, vocab_size)
defforward(self, context):
        embeddings = self.embedding(context)
        context_len = context.shape[1]
        position = torch.arange(context_len, device=context.device).unsqueeze(0)
        position_embeddings = self.positional_embedding(position)
        e = embeddings + position_embeddings
for decoder in self.decoders:
            e = decoder(e)
return self.linear(self.exit_norm(e))

根据所需规模，可以添加多个解码器。对于大型 GPT 模型，解码器数量最多可达 25 个，这看似简单，但会消耗大量内存，并且需要海量数据进行训练。

目前为止一切顺利。你可以使用简单的序列生成函数来测试你的模型。

deftop_k_logits(logits, k):
    v, ix = torch.topk(logits, k)
    out = logits.clone()
    out[out < v[:, [-1]]] = float('-inf')
return out
defgenerate(model, max_new_tokens, context, context_length, temperature=1.0, top_k=None):
    res = []
for _ in range(max_new_tokens):
if context.shape[1] > context_length:
            context = context[:, -context_length:]
        logits = model(context)  # [B, T, V]
        logits = logits[:, -1,:]  # [B, V]
        logits = logits / max(temperature, 1e-3)
if top_k isnotNone:
            logits = top_k_logits(logits, top_k)
if torch.isnan(logits).any() or torch.isinf(logits).any():
raise ValueError("Logits contain NaN or Inf")
        probabilities = nn.functional.softmax(logits, dim=-1)
        probabilities = torch.clamp(probabilities, min=1e-9, max=1.0)
        next_token = torch.multinomial(probabilities, 1)  # [B, 1]
        context = torch.cat((context, next_token), dim=1)
return context
start_context = "I want something"
model = GPT(num_heads,vocab_size,embed_dim,attention_dim,num_blocks,context_length, dropout_rate).to(device)
model.eval()
token_ids = generate(
    model=model,
    context=text_to_token_ids(start_context, tokenizer, device),
    max_new_tokens=10,
    context_length=context_length
)
print("Output text:\n", token_ids_to_text(token_ids, tokenizer))

如果一切顺利，你会看到一些像这样的随机垃圾输出。

Output text:
 I want something introduceウ coaches Kard Judaism trendsCommerce rotating infiltration approach

模型预训练

简而言之，模型预训练旨在使模型能够理解并说出基本的英语。模型必须能够生成语法正确且意义大致的词序，即使上下文信息可能不多。为此，我们需要一个能够提供大量英语文本的数据集。你可以从众多公开数据集中进行选择，例如IMDb。

数据准备

将给定的数据集转换为连续文本数据，并删除产生噪声的不需要的字符。

from datasets import load_dataset 
import re 
# 加载数据
ds = load_dataset("stanfordnlp/imdb") 
# 仅保留英文（ASCII）字符
defkeep_english_only(text ):
return re.sub(r"[^\x00-\x7F]+", "", text) 
# 清理并合并文本列表
defcombine_and_clean(text_list):
# 仅保留英文
    cleaned_list = [keep_english_only(t) for t in text_list] 
# 合并成一个字符串
    combined = "" .join(cleaned_list) 
# 移除多余的空格/换行符
    combined = re.sub(r'\s+', '', combined).strip() 
return combined 
# 创建单独的合并字符串
train_text_data = combine_and_clean(ds['train']['text'])
test_text_data = combine_and_clean(ds['test']['text'])

我们需要以下格式的数据才能实现这一概念。

input_ids: [101, 102, 103, 104, 105]
untokenized input_ids: ["The", "cat", "sat", "on", "the"]
target_ids: [102, 103, 104, 105, 106]
untokenized target_ids: ["cat", "sat", "on", "the", "mat"]

因此，定义一个数据加载器，以便相应地拆分我们的数据集。

from torch.utils.data import Dataset, DataLoader
classCustomDataset(Dataset):
def__init__(self, txt, tokenizer, max_length, stride):
        self.input_ids = []
        self.target_ids = []
# 对整个文本进行分词
        token_ids = tokenizer.encode(txt, add_special_tokens= False ) 
# 使用滑动窗口将数据分块为长度为 max_length 的重叠序列
for i in range(0, len(token_ids) - max_length, stride):
            input_chunk = token_ids[i:i + max_length]
            target_chunk = token_ids[i + 1: i + max_length + 1]
            self.input_ids.append(torch.tensor(input_chunk))
            self.target_ids.append(torch.tensor(target_chunk))
def__len__(self):
return len(self.input_ids)
def__getitem__(self, idx):
return self.input_ids[idx], self.target_ids[idx]
defcreate_encoded_dataloader(txt, tokenizer, batch_size=4, max_length=128,
                         stride=128, shuffle=True, drop_last=True, num_workers=0):
# 创建数据集
    dataset = CustomDataset(txt, tokenizer, max_length, stride) 
# 创建数据加载
    dataloader = DataLoader(
        dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last, num_workers=num_workers, pin_memory=True)
return dataloader
total_characters = len(train_text_data)
total_tokens = len(tokenizer.encode(train_text_data))
print("Characters:", total_characters)
print("Tokens:", total_tokens)
# 检查是否足够：
if total_tokens * (0.95) < context_length:
    print("Not enough tokens for the training loader. "
"Try to lower the context_length or "
"increase the `training_ratio`")
if total_tokens * (1-0.95) <context_length:
    print("Not enough tokens for the validation loader. "
"Try to lower the context_length or "
"decrease the `training_ratio`")
train_dataloader = create_encoded_dataloader(
    train_text_data,
    tokenizer=tokenizer,
    batch_size=2,
    max_length=context_length,
    stride=context_length,
    shuffle=True,
    drop_last=True
)
test_dataloader = create_encoded_dataloader(
    test_text_data,
    tokenizer=tokenizer,
    batch_size=2,
    max_length=context_length,
    stride=context_length,
    shuffle=False,
    drop_last=True
)

训练

在训练之前，需要初始化权重，以确保模型从预定义的起点开始训练。这是 GPT 模型的标准做法。

definitialize_weights(module):
if isinstance(module, nn.Linear):
        torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
if module.bias isnotNone:
            torch.nn.init.zeros_(module.bias)
elif isinstance(module, nn.Embedding):
        torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
elif isinstance(module, nn.LayerNorm):
        torch.nn.init.ones_(module.weight)
        torch.nn.init.zeros_(module.bias)
model.apply(initialize_weights)

我们先来看损失函数。LLM（逻辑学习模型）在某种程度上执行的是多类分类，词汇表中的每个词都代表一个类别，模型输出的是所有词的概率分布。因此，我们使用交叉熵损失函数。

交叉熵奖励正确单词的高概率，惩罚正确单词概率低的情况。

请参考我们之前的例子。

Index: Token                            
0 → "The"           Inputs: ["The", "cat", "sat", "on", "the"]
1 → "cat"           Targets: ["cat", "sat", "on", "the", "mat"]  
2 → "sat"
3 → "on"
4 → "the"
5 → "mat"

计算每次预测的损失，

• 位置 0：目标 = "cat"→ P = 0.90 → L0 = −log⁡(0.9) ≈ 0.105
• 位置 1：目标 = "sat"→ P = 0.10 → L1 = −log⁡(0.1) = 2.302
• 位置 2：目标 = "on"→ P = 0.05 → L2 = −log⁡(0.05) = 2.996
• 位置 3：目标 = "the"→ P = 0.75 → L3 = −log⁡(0.75) ≈ 0.288
• 位置 4：目标 = "mat"→ P = 0.75 → L4 = −log⁡(0.75) ≈ 0.288

取平均损失，

Lavg  = (0.105+2.302+2.996+0.288+0.288)/ 5 ≈ 1.20

由于错误预测会导致巨大的损失值，少数错误预测会显著影响平均损失。基于此损失，定义评估函数。我们特意提及这一点，@torch.no_grad()是因为在计算损失期间不应更新模型权重。

criterion = nn.CrossEntropyLoss()
defcalc_loss_batch(input_batch, target_batch, model, device):
    input_batch = input_batch.to(device, non_blocking=True)
    target_batch = target_batch.to(device, non_blocking=True)
    logits = model(input_batch)  # [B, T, V]
    B, T, V = logits.shape
    loss = criterion(logits.view(B * T, V), target_batch.view(B * T))
return loss
@torch.no_grad()
defcalc_loss_loader(data_loader, model, device, num_batches=None):
if len(data_loader) == 0:
return float("nan")
    model.eval()
    total_loss = 0.0
    num_batches = len(data_loader) if num_batches isNoneelse min(num_batches, len(data_loader))
for i, (inp, tgt) in enumerate(data_loader):
if i >= num_batches:
break
        loss = calc_loss_batch(inp, tgt, model, device)
        total_loss += loss.item()
    model.train()
return total_loss / num_batches
@torch.no_grad()
defevaluate_model(model, train_loader, val_loader, device, eval_iter=1):
    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)
return train_loss, val_loss

请注意，你在计算损失时已显式设置了该值，计算完成后model.eval()请将其改回默认值。在训练过程中，最大的挑战是如何随时间调整学习率。如果学习率过高，模型将无法收敛；如果学习率过低，训练速度将非常缓慢。为了平衡这一点，我们通常使用调度器在训练过程中调整学习率。model.train()

在此设置中，我们使用名为 的自定义调度器CosineWithWarmup。

classCosineWithWarmup(torch.optim.lr_scheduler._LRScheduler):
def__init__(self, optimizer, warmup_steps, total_steps, base_lr, min_lr, last_epoch=-1):
        self.warmup_steps = max(1, warmup_steps)
        self.total_steps = max(self.warmup_steps + 1, total_steps)
        self.base_lr = base_lr
        self.min_lr = min_lr
        super().__init__(optimizer, last_epoch)
defget_lr(self):
        step = self.last_epoch + 1
        lrs = []
for _ in self.base_lrs:
if step <= self.warmup_steps:
                lr = self.base_lr * step / self.warmup_steps
else:
                progress = (step - self.warmup_steps) / max(1, self.total_steps - self.warmup_steps)
                lr = self.min_lr + 0.5 * (self.base_lr - self.min_lr) * (1 + math.cos(math.pi * progress))
            lrs.append(lr)
return lrs

• Warmup阶段：在前几个步骤中，学习率从 0 线性增加到基准值。这有助于稳定训练，尤其对于像 GPT 这样的大型模型而言，否则它们可能在早期就出现偏差。
• 余弦衰减：预热后，学习率会按照余弦曲线逐渐下降，平滑地衰减至最小值（min_lr）。这可以防止学习率突然下降，并帮助模型“稳定”到一个良好的局部最小值。

从小规模开始 → 稳步上升 → 平稳下降。

settings =  { 
"learning_rate":  3e- 4 ,           
"weight_decay":  0.1 ,             # GPT 风格训练的标准值
"num_epochs":  300 ,             
"batch_size":  32 ,                # 平衡 GPU 内存和收敛速度
"warmup_steps":  1500 ,            # 预热有助于避免过早发散
"max_lr":  3 e- 4 ,                  
"min_lr":  3 e- 5 ,                  
"eval_freq":  200 ,                
"eval_iter":  20 ,                 
"gradient_clip":  1.0 ,            
"patience":  50 ,                  
"min_improvement":  1e- 4 , 
"print_interval":  1 ,             
"generate_interval":  5
}
 train_dataloader = create_encoded_dataloader (
     train_text_data, 
    tokenizer = tokenizer, 
    batch_size = settings [ "batch_size" ] , 
    max_length = context_length, 
    stride = context_length, 
    shuffle = True , 
    drop_last = True
)
 test_dataloader = create_encoded_dataloader (
     test_text_data, 
    tokenizer = tokenizer, 
    batch_size = settings [ "batch_size" ] , 
    max_length = context_length, 
    stride = context_length, 
    shuffle = False , 
    drop_last = True
)

将所有设置集中在一个模块中进行管理非常方便。接下来是训练循环，实际的学习过程就在这里进行。

deftrain_model(
    model,
    train_loader,
    val_loader,
    device,
    settings,
    save_path="checkpoints/gpt_256_256_8_8.pt",
):
    torch.manual_seed(123)
if torch.cuda.is_available():
        torch.cuda.manual_seed_all(123)
    model.to(device)
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=settings["learning_rate"],
        weight_decay=settings["weight_decay"],
        betas=(0.9, 0.95),
    )
    total_steps = settings["num_epochs"] * len(train_loader)
    scheduler = CosineWithWarmup(
        optimizer,
        warmup_steps=settings["warmup_steps"],
        total_steps=total_steps,
        base_lr=settings["max_lr"],
        min_lr=settings["min_lr"],
    )
    train_losses, val_losses, tokens_seen_track = [], [], []
    tokens_seen, global_step = 0, -1
    best_val_loss, patience_counter = float("inf"), 0
for epoch in range(settings["num_epochs"]):
        model.train()  
for step, (inp, tgt) in enumerate(train_loader):
            loss = calc_loss_batch(inp, tgt, model, device)
            loss.backward()
# gradient clipping
            torch.nn.utils.clip_grad_norm_(model.parameters(), settings["gradient_clip"])
            optimizer.step()
            optimizer.zero_grad(set_to_none=True)
            scheduler.step()
            global_step += 1
            tokens_seen += inp.numel()
# evaluation
if global_step % settings["eval_freq"] == 0:
                train_loss, val_loss = evaluate_model(
                    model, train_loader, val_loader, device,
                    eval_iter=settings["eval_iter"],
                )
                train_losses.append(train_loss)
                val_losses.append(val_loss)
                tokens_seen_track.append(tokens_seen)
                lr_now = optimizer.param_groups[0]["lr"]
                print(f"Ep {epoch+1} | step {global_step:06d} | lr {lr_now:.3e} "
f"| train {train_loss:.3f} | val {val_loss:.3f}")
# early stopping
if val_loss + settings["min_improvement"] < best_val_loss:
                    best_val_loss = val_loss
                    patience_counter = 0
                    os.makedirs(os.path.dirname(save_path) or".", exist_ok=True)
                    torch.save({
"model_state": model.state_dict(),
"optimizer_state": optimizer.state_dict(),
"epoch": epoch,
"global_step": global_step,
                    }, save_path)
                    print(f"[Checkpoint saved at step {global_step}]")
else:
                    patience_counter += 1
if patience_counter >= settings["patience"]:
                        print("Early stopping triggered.")
return train_losses, val_losses, tokens_seen_track
return train_losses, val_losses, tokens_seen_track

代码中有三点值得一提。这些都是小细节，但它们对训练的稳定性和效率影响很大。

1. 梯度裁剪在反向传播过程中，模型会计算每个参数的梯度。实际上，这些梯度有时会变得非常大（梯度爆炸），尤其是在深度网络中或训练长序列数据时。如果发生这种情况，权重更新可能会导致训练不稳定，并造成损失函数发散。梯度裁剪通过限制梯度的大小来防止这种情况的发生。

2. **过早停止

   **训练LLM需要数天时间。因此，必须持续监控计算过程，如果没有取得显著进步，就应该停止训练。

3. **AdamW：

   在实际的权重更新中，我们使用AdamW AdamW，这是一种结合了 Adam 优化器优势和适当权重衰减正则化的**现代优化器。Adam 会针对每个参数单独调整学习率，这有助于模型更快收敛。“W”代表解耦权重衰减。与经典的 Adam 优化器不同，它将权重衰减与梯度更新完全分离，从而提高了模型的泛化能力。

这是我随机得到的一些输出结果。

the movie starts slow and i thought it was going to be boring, but then going to be interesting. 
the acting is okay, some are boring felt like they just gave up. 
still, it was not the worst film i’ve seen

教你的模特跟着唱（并演唱酷玩乐队的歌曲）

恭喜！经过漫长的等待，你终于拥有了自己的语言模型，它能够说一些基本的英语（尽管由于我们是在 IMDb 数据集上训练的，而且模型架构简单，所以可能不太符合逻辑）。现在，是时候教它一些 Coldplay 的风格了。为此，我们使用一个小型 Coldplay 数据集对模型进行微调。实际上，像 GPT 这样的大型语言模型就是这样构建的：

• 首先，它们会在一个庞大的通用数据集（来自书籍、网站的数十亿个token）上进行预训练。这可以教会它们语法、词汇和一般世界知识。
• 然后，研究人员会在规模较小、更专业的数据集上对模型进行微调，使其适应特定的风格或任务（例如聊天、编程、医疗问答、法律推理等）。如果没有微调，GPT 就只是一个庞大的文本预测器。而通过微调，它就能进行对话式、安全且符合我们实际需求的任务。

你可以使用这个数据集，并按照相同的预处理步骤进行准备。使用相同的训练循环进行微调，但稍微调整一下设置。

Look at the star look how the " settings_ft = { 
    "learning_rate": 1e-5,           # 降低学习率以进行微调，从而保留预训练权重
    "weight_decay": 0.01,            # 降低权重衰减
    "num_epochs": 5,                 # 由于 Coldplay 数据集较小，因此减少训练轮数
    "batch_size": 4,                 # 针对小型数据集使用较小的批次大小
    "warmup_steps": 100,             # 缩短预热时间
    "max_lr": 1e-5, 
    "min_lr": 1e-6, 
    "eval_freq": 50,                
    "eval_iter": 5,                  
    "gradient_clip": 0.5,            # 更温和的梯度裁剪
    "patience": 3,                   
    "min_improvement": 1e-4, 
    "print_interval": 1, 
    "generate_interval": 2 
} 
train_losses_ft, val_losses_ft, tokens_seen_ft = train_model( 
    model, 
    train_dataloader_ft, 
    val_dataloader_ft, 
    tokenizer, 
    device, 
    settings=settings_ft, 
    context_length=context_length, 
    save_path= "checkpoints/gpt_512_512_8_8_finetuned_coldplay.pt" , 
    sample_prompt= "Look at the star look how the"  
 )

我们再检查一次输出结果，

lights go out and the stars begin to fall i hear your voice across the night  
lights are running in circles chasing the echoes  
you are the star that keeps me alive  
Oh-ooh-oh-ooh oh, oh  
i will follow you, i will follow you

总结

好了，课程到此结束！你已经从零开始构建了自己的 Transformer 架构，并完全使用 PyTorch 进行了训练。你已经学习了从概念到代码所需的一切，其中还包含许多其他地方找不到的更详细的讲解