用 PyTorch 从零构建一个大模型
模块逐一讲解语言模型如何真正工作和构建（附代码）
AI模型价格对比 | AI工具导航 | ONNX模型库 | Vibe Coding教程 | PLC在线仿真器 | Tripo 3D | Meshy AI | ElevenLabs | KlingAI | ArtSpace | Phot.AI | InVideo
理解语言模型如何工作有两种方式。第一种是阅读它的描述：注意力机制、Transformer 模块、token 嵌入。第二种是你自己构建一个，看着它失败，修复故障，然后看着它工作。第二种方式产生完全不同质量的理解。
本文采用第二种方式。我们将使用 PyTorch 从零开始构建一个小型但完整的语言模型。不是现有模型的封装。不是微调教程。是实际的组件，按照实际的顺序组装，并在每一步解释实际的原因。
读完本文，你将拥有一个可以生成文本的工作模型。更重要的是，你将知道为什么每个部分都在那里，它在做什么，以及如果你移除它会破坏什么。
首先：确保这三个导入在文件顶部。本文中的每个代码块都假设它们已导入。
import torch
import torch.nn as nn
import torch.nn.functional as F
1、全貌：语言模型真正做什么
在编写单个类之前，你需要正确的思维模型。语言模型做一件事：给定一个词序列，它预测下一个词是什么。
就这样。这就是全部工作。
读句子 "The cat sat on the"。语言模型查看这五个词，并为词汇表中的每个词分配概率。"Mat" 获得高概率。"Quantum" 获得低概率。"Floor" 获得中等概率。模型选择最可能的下一个词，将其添加到序列中，然后重复。
本文中的所有内容——分词器、嵌入、注意力机制、Transformer 模块——都是为了让那一个预测尽可能准确。
以下是我们将构建的架构，按顺序排列：
下面每一节构建此图表的一行。当我们到达底部时，我们将它们组装成一个完整的模型。
2、分词器
神经网络只理解数字。它们无法读取 "hello" 这个词。它们可以读取数字 8730。分词器是人类语言和机器算术之间的桥梁——它将原始文本转换为数字序列，并将数字序列转换回文本。
每个语言模型都有一个分词器。分词器决定语言的最小单位是什么。在本文中，我们使用字符级分词：每个字符是一个 token。这使实现简单清晰，同时不隐藏任何重要内容。
2.1 它如何工作
分词器构建一个词汇表：一个将它见过的每个字符映射到唯一整数的字典。单词 "hello" 变成五个整数，每个字符一个。要转换回来，你反转字典。
class Tokenizer:
@staticmethod
def build_vocab(text: str) -> dict:
"""
扫描文本并为每个唯一字符分配唯一整数。
结果是词汇表：此分词器知道如何处理的完整字符集。
添加特殊 <unk> token 用于推理时遇到但训练文本中没有的字符。
"""
# sorted() 使词汇表在不同运行间确定。
# 不排序的话，相同文本可能产生不同的整数分配，
# 导致保存的模型不兼容。
vocab = {
char: idx
for idx, char in enumerate(sorted(set(text)))
}
vocab["<unk>"] = len(vocab)
return vocab
def __init__(self, vocab: dict):
self.encode_map = {str(k): int(v) for k, v in vocab.items()}
# 解码映射只是编码映射的反转。
self.decode_map = {v: k for k, v in self.encode_map.items()}
def encode(self, text: str) -> list[int]:
"""
将字符串转换为整数列表。
未知字符映射到 <unk> token 而不是崩溃。
"""
return [
self.encode_map.get(char, self.encode_map["<unk>"])
for char in text
]
def decode(self, indices: list[int]) -> str:
"""
将整数列表转换回字符串。
"""
return "".join([
self.decode_map.get(idx, "<unk>")
for idx in indices
])
2.2 为什么使用字符级
词级分词有更简单的词汇表，但它在遇到从未见过的词时会崩溃。字符级分词词汇表很小（大约100个字符覆盖大多数英文文本）并且能处理任何输入而不会崩溃。出于学习目的，字符级是正确的选择：词汇表适合内存，代码透明，没有被外部库隐藏。
试一试：
text = "Hello, PyTorch."
vocab = Tokenizer.build_vocab(text)
tokenizer = Tokenizer(vocab)
encoded = tokenizer.encode("Hello")
print(encoded) # 类似 [3, 4, 8, 8, 11] 取决于词汇表顺序
decoded = tokenizer.decode(encoded)
print(decoded) # Hello
3、嵌入层
分词后，每个词是一个整数。问题是单个整数不携带任何关于含义的信息。整数42并不比整数1"更接近"整数43。但单词 "car" 在概念上应该接近单词 "vehicle"。单个整数无法表示这种关系。
嵌入层通过将每个整数替换为一个向量来解决这个问题：一个浮点数列表，通常长度在50到1024之间。这些向量在训练过程中是可学习的，模型调整它们使得含义相似的词最终具有相似的向量。向量之间的相似性可以被衡量，这种衡量反映了词之间真实的语义相似性。
3.1 什么是查找表
嵌入层最简单的思维模型是一个电子表格。它在词汇表中每个 token 有一行，嵌入每个维度有一列。当模型看到 token 号42时，它查找电子表格的第42行并提取整行作为向量。该向量就是通过模型其余部分传播的内容。
该电子表格中的值从随机开始，在训练过程中被调整。这就是整个机制。
3.2 位置编码：第二个问题
这里有一个微妙的东西，大多数人第一次都会搞错：你将在下一节看到的注意力机制没有内在的顺序感。如果你向它输入 "the cat sat" 和 "sat the cat"，它会看到相同的 token，并且在没有额外信息的情况下产生相同的输出。但词序显然很重要。
位置编码修复了这一点。对于序列中的每个位置（位置0、位置1、位置2……），我们创建一个可学习的向量来表示该位置。在任何内容进入注意力机制之前，我们将此位置向量添加到词嵌入向量中。现在模型既看到词的含义，也看到它在序列中的位置。
class Embedding(nn.Module):
def __init__(self, vocab_size: int, embedding_dim: int):
"""
两个嵌入表在这里：
self.token_embedding - 每个词汇 token 一行，学习词的含义
self.position_embedding - 每个位置一行，学习词序
两个表有相同数量的列（embedding_dim），所以它们可以相加。
Args:
vocab_size (int): 词汇表中唯一 token 的总数。
embedding_dim (int): 表示每个 token 的向量大小。
更大 = 更有表达力，但更昂贵。
"""
super().__init__()
self.token_embedding = nn.Embedding(vocab_size, embedding_dim)
self.position_embedding = nn.Embedding(vocab_size, embedding_dim)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x: 形状为 (batch_size, sequence_length) 的 token 索引
Returns:
形状为 (batch_size, sequence_length, embedding_dim) 的组合嵌入
"""
# token 含义
word_vectors = self.token_embedding(x)
# 位置顺序 - x 同时作为位置索引
position_vectors = self.position_embedding(x)
# 相加为模型提供两个维度的信息
return word_vectors + position_vectors
4、注意力模块
阅读句子不仅仅是孤立地知道每个词的含义。它是关于理解哪些词与哪些其他词相关。在句子 "The bank by the river was steep" 中，单词 "bank" 的含义与 "The bank approved my loan" 中完全不同。模型需要查看周围的词来消歧。
注意力机制是语言模型查看周围词的方式。对于序列中的每个词，它问："在这个序列中，哪些其他词与理解当前这个词最相关？"然后它产生该词的新表示，融合来自相关其他词的信息。
4.1 Query、Key 和 Value
注意力机制使用三个概念，名称借用于信息检索。
Query 代表一个词在寻找什么。Key 代表每个词可以提供什么。Value 是在建立相关性后被混合进去的实际信息。
最简单的类比是搜索引擎。你输入搜索查询。引擎将你的查询与它知道的每个页面的 Key（页面标题、元数据）进行比较。具有最相关 Key 的页面将其 Value（实际内容）返回给你。注意力做同样的事情，但同时对序列中的每个词进行。
在实践中，Query、Key 和 Value 都通过将输入传递给三个独立的线性层来计算。它们从相同的数据开始，被投影到三个不同的"视图"。
4.2 掩码
还有一部分：掩码。在训练期间，模型一次看到完整的句子。但它被训练来预测下一个词，这意味着它只应该被允许查看之前的词，而不是未来的词。如果它能看到未来的词，它就会作弊——答案就在它的上下文中。
掩码通过将单词与任何未来单词之间的注意力分数设置为负无穷来强制执行这一点。在 softmax 步骤之后，负无穷变为零，这意味着这些位置对输出没有任何贡献。每个词只关注自身以及它之前的所有内容。
class AttentionBlock(nn.Module):
def __init__(self, embedding_dim: int, head_dim: int, context_size: int):
"""
一个注意力头。计算序列中每对 token 之间的相关性，
并使用该相关性构建丰富的表示。
Args:
embedding_dim (int): 输入向量的大小。
head_dim (int): Query、Key 和 Value 投影的大小。
在多头注意力中，这是 embedding_dim // num_heads。
context_size (int): 模型可以处理的最大序列长度。
确定因果掩码的大小。
"""
super().__init__()
self.query = nn.Linear(embedding_dim, head_dim, bias=False)
self.key = nn.Linear(embedding_dim, head_dim, bias=False)
self.value = nn.Linear(embedding_dim, head_dim, bias=False)
# 掩码不是可学习参数——它是固定的三角矩阵。
# register_buffer 使其成为模块的一部分（随模型移动到 GPU），
# 而不使其成为可训练权重。
ones = torch.ones(context_size, context_size)
self.register_buffer("mask", torch.tril(ones))
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x: 形状为 (batch_size, seq_len, embedding_dim) 的输入张量
Returns:
形状为 (batch_size, seq_len, head_dim) 的丰富表示
"""
B, T, C = x.shape
Q = self.query(x) # 每个 token 在寻找什么
K = self.key(x) # 每个 token 可以提供什么
V = self.value(x) # 要混合的实际内容
# 计算每对位置之间的相关性分数。
# 按 head_dim**-0.5 缩放防止分数变得太大，
# 导致 softmax 崩溃为 one-hot 分布（一个主导 token 淹没其他所有内容）。
scores = Q @ K.transpose(-2, -1) * (C ** -0.5)
# 应用因果掩码：未来位置变为 -inf，然后 softmax 后变为 0。
scores = scores.masked_fill(self.mask[:T, :T] == 0, float("-inf"))
# Softmax 将分数转换为总和为 1 的权重。
weights = F.softmax(scores, dim=-1)
# 根据这些权重混合 Value 向量。
return weights @ V
4.3 为什么这是模型的核心
Transformer 中的每个其他组件都在支持这一个。嵌入是为了给注意力提供有意义的比较对象。位置编码是为了给注意力提供顺序感。接下来的前馈层是为了处理注意力发现的内容。注意力是使语言模型与之前所有其他神经网络不同的机制，值得在继续之前花额外时间理解。
5、多头注意力
单个注意力头寻找 token 之间的一种关系类型。但语言有许多同时存在的关系类型。在 "She gave him the book she wrote" 中，一个关系在 "she" 和 "wrote" 之间（同一个人）。另一个关系在 "him" 和 "gave" 之间（接受者）。第三个在 "book" 和 "wrote" 之间（她写的东西）。
多头注意力并行运行多个注意力头，每个头有自己的 Query、Key 和 Value 权重。每个头专门寻找不同的模式。然后将它们的输出拼接并投影回原始维度。
class MultiHeadAttention(nn.Module):
def __init__(self, embedding_dim: int, num_heads: int, context_size: int):
"""
并行运行 num_heads 个注意力头并组合它们的输出。
每个头操作嵌入维度的切片（head_dim），
因此无论使用多少头，总计算量都与 embedding_dim 成比例。
Args:
embedding_dim (int): 输入和输出向量的大小。
num_heads (int): 并行注意力头的数量。
必须能整除 embedding_dim。
context_size (int): 最大序列长度。
"""
super().__init__()
assert embedding_dim % num_heads == 0, (
f"embedding_dim ({embedding_dim}) 必须能被 "
f"num_heads ({num_heads}) 整除"
)
head_dim = embedding_dim // num_heads
# 每个头是一个完整的 AttentionBlock，有自己的可学习权重。
# nn.ModuleList 正确注册它们，以便 PyTorch 跟踪它们的参数。
self.heads = nn.ModuleList([
AttentionBlock(embedding_dim, head_dim, context_size)
for _ in range(num_heads)
])
# 拼接后，投影回 embedding_dim。
# 这让模型可以跨头混合信息。
self.projection = nn.Linear(embedding_dim, embedding_dim)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
并行运行所有头，拼接它们的输出，投影。
Args:
x: 形状为 (batch_size, seq_len, embedding_dim) 的输入
Returns:
形状为 (batch_size, seq_len, embedding_dim) 的输出
"""
# 每个头产生 (batch_size, seq_len, head_dim)。
# 沿最后一个维度拼接给出 (batch_size, seq_len, embedding_dim)。
concatenated = torch.cat([head(x) for head in self.heads], dim=-1)
return self.projection(concatenated)
6、前馈层
注意力非常擅长在序列中的位置之间移动信息。它不擅长的是在信息被收集后处理它。前馈层就是进行这种处理的地方。
在注意力弄清楚"什么与什么相关"之后，前馈层对每个位置独立应用可学习的非线性变换。这是模型学习推理注意力揭示的关系的地方。
结构总是相同的：扩展维度，应用非线性，压缩回来。扩展（通常到嵌入维度的四倍）在推理步骤中为模型提供更大的工作空间。
class FeedForward(nn.Module):
def __init__(self, embedding_dim: int, ff_dim: int):
"""
一个两层网络，独立应用于每个位置。
扩展到 ff_dim（通常 4 * embedding_dim）为模型在压缩回来之前
表示复杂变换创造了空间。ReLU 引入了非线性，
使网络能学习纯线性模型无法学习的东西。
Args:
embedding_dim (int): 输入和输出维度。
ff_dim (int): 中间维度。通常 4 * embedding_dim。
"""
super().__init__()
self.expand = nn.Linear(embedding_dim, ff_dim)
self.activate = nn.ReLU()
self.compress = nn.Linear(ff_dim, embedding_dim)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.expand(x)
x = self.activate(x)
x = self.compress(x)
return x
7、Transformer 模块
Transformer 模块是模型的完整推理单元：多头注意力，后跟前馈层，并添加了两个重要内容：残差连接和层归一化。
残差连接（也称为跳跃连接）将模块的输入添加回其输出。如果注意力或前馈层学到了有用的东西，残差将其向前传递。如果它学到了有害的东西，原始信号被保留。这就是为什么可以堆叠许多 Transformer 模块而训练不会变得不可能。
层归一化稳定流经网络的值。没有它，值可能在许多层中不可控地增长或缩小。归一化使它们保持在可预测的范围内，使训练更快更稳定。标准做法是在将输入传递给每个子层之前进行归一化，这样注意力和前馈层总是接收到良好缩放的输入。
class TransformerBlock(nn.Module):
def __init__(
self,
embedding_dim: int,
num_heads: int,
context_size: int,
ff_dim: int
):
"""
一个完整的 Transformer 模块：堆叠 N 次形成语言模型主体的重复单元。
处理顺序（Pre-LN 变体，当今最常见）：
1. 归一化 → 注意力 → 添加残差
2. 归一化 → 前馈 → 添加残差
在子层之前归一化（Pre-LN）而不是之后，
在训练期间产生更稳定的梯度。
Args:
embedding_dim (int): 模型的核心向量维度。
num_heads (int): 注意力头的数量。
context_size (int): 最大序列长度。
ff_dim (int): 前馈中间维度。
"""
super().__init__()
self.norm_1 = nn.LayerNorm(embedding_dim)
self.attention = MultiHeadAttention(embedding_dim, num_heads, context_size)
self.norm_2 = nn.LayerNorm(embedding_dim)
self.feed_forward = FeedForward(embedding_dim, ff_dim)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# 步骤 1：带残差的注意力
# 先归一化，运行注意力，将结果添加回原始 x。
x = x + self.attention(self.norm_1(x))
# 步骤 2：带残差的前馈
# 先归一化，运行前馈，将结果添加回来。
x = x + self.feed_forward(self.norm_2(x))
return x
8、完整的语言模型
每个部分现在都已构建完成。语言模型按顺序堆叠它们：
嵌入层将 token 索引转换为带有位置信息的向量。一系列 Transformer 模块处理这些向量，每个模块通过关注上下文并对其进行推理来精炼表示。最终的层归一化稳定输出。线性投影（"语言模型头"）将嵌入维度映射到词汇表大小，为每个 token 产生一个分数。分数最高的 token 是模型的预测。
import torch
import torch.nn as nn
import torch.nn.functional as F
import requests
class LanguageModel(nn.Module):
@staticmethod
def load_shakespeare():
"""
下载 Tiny Shakespeare（约1MB）并返回原始文本字符串。
缓存在内存中——调用一次并传递结果。
"""
url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
response = requests.get(url)
response.raise_for_status()
print(f"数据集已加载：{len(response.text):,} 个字符")
return response.text
def __init__(
self,
vocab_size: int,
embedding_dim: int,
num_heads: int,
num_layers: int,
context_size: int,
ff_dim: int
):
"""
一个完整的仅解码器语言模型。
Args:
vocab_size (int): 唯一 token 的数量。确定嵌入表大小
和 LM 头的输出维度。
embedding_dim (int): 整个模型的核心向量维度。
所有层使用此大小的向量进行通信。
num_heads (int): 每个 Transformer 模块的注意力头。
必须能整除 embedding_dim。
num_layers (int): 堆叠的 Transformer 模块数量。
更多层 = 更深推理，更多计算。
context_size (int): 模型一次处理的最大 token 数。
注意力掩码由此值确定大小。
ff_dim (int): 每个模块的前馈中间维度。
通常 4 * embedding_dim。
"""
super().__init__()
self.embedding = Embedding(vocab_size, embedding_dim)
self.blocks = nn.Sequential(*[
TransformerBlock(embedding_dim, num_heads, context_size, ff_dim)
for _ in range(num_layers)
])
self.final_norm = nn.LayerNorm(embedding_dim)
self.lm_head = nn.Linear(embedding_dim, vocab_size)
def forward(
self,
x: torch.Tensor,
targets: torch.Tensor = None
) -> tuple[torch.Tensor, torch.Tensor | None]:
"""
前向传播。
Args:
x: 形状为 (batch_size, seq_len) 的 token 索引
targets: 可选的训练真实值下一个 token。
如果提供，计算并返回交叉熵损失。
Returns:
logits: 形状为 (batch_size, seq_len, vocab_size) 的原始分数。
loss: 如果提供 targets 则为交叉熵损失，否则为 None。
"""
x = self.embedding(x)
x = self.blocks(x)
x = self.final_norm(x)
logits = self.lm_head(x)
loss = None
if targets is not None:
B, T, V = logits.shape
loss = F.cross_entropy(
logits.view(B * T, V),
targets.view(B * T)
)
return logits, loss
@torch.no_grad()
def generate(
self,
prompt: torch.Tensor,
max_new_tokens: int,
context_size: int
) -> torch.Tensor:
"""
从提示自回归生成新 token。
Args:
prompt: 形状为 (1, seq_len) 的起始 token 索引
max_new_tokens: 要生成多少个 token
context_size: 模型的最大上下文窗口
Returns:
包含提示 + 生成 token 的完整序列
"""
self.eval()
for _ in range(max_new_tokens):
context = prompt[:, -context_size:]
logits, _ = self(context)
last_logits = logits[:, -1, :]
probs = F.softmax(last_logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
prompt = torch.cat([prompt, next_token], dim=1)
return prompt
9、训练模型
以下是一个完整、可运行的训练循环。每节都注释了它在做什么和为什么，而不仅仅是代码说了什么。
import torch
import torch.nn as nn
import torch.nn.functional as F
import requests
# ── 配置 ──────────────────────────────────────────────
EMBEDDING_DIM = 128 # 整个模型的向量大小
NUM_HEADS = 4 # 必须能整除 EMBEDDING_DIM
NUM_LAYERS = 4 # 堆叠的 Transformer 模块数量
CONTEXT_SIZE = 64 # 模型一次看到的最大 token 数
FF_DIM = 512 # 前馈中间大小（4 × EMBEDDING_DIM）
BATCH_SIZE = 32 # 每个训练步骤的序列数
LEARNING_RATE = 3e-4 # Transformer 的标准起点
MAX_STEPS = 5000 # 总训练步骤数
# ── 数据准备 ──────────────────────────────────────────
raw_text = LanguageModel.load_shakespeare()
vocab = Tokenizer.build_vocab(raw_text)
tokenizer = Tokenizer(vocab)
vocab_size = len(vocab)
data = torch.tensor(tokenizer.encode(raw_text), dtype=torch.long)
split = int(0.9 * len(data))
train_data = data[:split]
val_data = data[split:]
def get_batch(split: str) -> tuple[torch.Tensor, torch.Tensor]:
"""
从数据集中采样一批随机的 (input, target) 对。
对于批次中的每个序列：
input = 位置 [i : i + CONTEXT_SIZE] 的 token
target = 位置 [i+1 : i + CONTEXT_SIZE + 1] 的 token
target 是 input 向右移动一个位置。
这就是下一个 token 预测的设置方式：对于输入序列中的每个 token，
target 是它的下一个 token。
"""
source = train_data if split == "train" else val_data
indices = torch.randint(len(source) - CONTEXT_SIZE, (BATCH_SIZE,))
x = torch.stack([source[i : i + CONTEXT_SIZE ] for i in indices])
y = torch.stack([source[i + 1 : i + CONTEXT_SIZE + 1] for i in indices])
return x, y
# ── 模型和优化器 ──────────────────────────────────────
model = LanguageModel(
vocab_size = vocab_size,
embedding_dim = EMBEDDING_DIM,
num_heads = NUM_HEADS,
num_layers = NUM_LAYERS,
context_size = CONTEXT_SIZE,
ff_dim = FF_DIM,
)
optimiser = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
print(f"模型参数：{sum(p.numel() for p in model.parameters()):,}")
# ── 训练循环 ──────────────────────────────────────────
for step in range(MAX_STEPS):
model.train()
x_batch, y_batch = get_batch("train")
logits, loss = model(x_batch, targets=y_batch)
optimiser.zero_grad()
loss.backward()
optimiser.step()
if step % 500 == 0:
model.eval()
with torch.no_grad():
_, val_loss = model(*get_batch("val"))
print(f"步骤 {step:5d} | 训练损失: {loss.item():.4f} | 验证损失: {val_loss.item():.4f}")
# ── 生成 ──────────────────────────────────────────────
model.eval()
prompt_text = "T"
prompt_tokens = torch.tensor(tokenizer.encode(prompt_text)).unsqueeze(0)
generated = model.generate(
prompt = prompt_tokens,
max_new_tokens = 200,
context_size = CONTEXT_SIZE,
)
output_text = tokenizer.decode(generated[0].tolist())
print("\n生成的文本：")
print(output_text)
9.1 你刚刚构建了什么
让我们明确每个部分及其贡献，现在全貌已经可见。
分词器将文本转换为整数，以便模型可以数学方式处理，并将整数转换回文本，以便你可以阅读输出。没有它，模型根本无法处理语言。
嵌入层将这些整数转换为可以承载含义的向量。token 嵌入学习每个词的含义。位置嵌入学习每个词在序列中的位置。它们一起为模型提供内容和上下文。
注意力模块弄清楚序列中哪些 token 彼此相关。因果掩码确保模型只向后看，而不是向前看，这样它就不能在训练时通过偷看答案作弊。
多头注意力并行运行多个注意力头，每个头专门寻找不同类型的关系。这给了模型同时跟踪多种语言模式的能力。
前馈层处理注意力揭示的信息。注意力在位置之间移动信息；前馈层独立地在每个位置推理该信息。
Transformer 模块将注意力、前馈、残差连接和归一化组合成一个重复单元。堆叠这些模块赋予模型推理的深度。
语言模型头将最终表示映射回词汇空间，为每个可能的下一个 token 产生一个分数。分数最高者获胜。
训练循环通过反复采样序列、计算模型预测有多错、以及微调模型中的每个参数使其更正确来教模型做出更好的预测。
10、用 nn.Transformer 的相同模型
你现在已经手动构建了每个组件。你理解 Query 是什么，为什么掩码存在，残差连接保护什么，以及为什么层归一化属于子层之前而不是之后。这些知识使下一个版本可读。
PyTorch 附带一个名为 nn.Transformer 的内置类。它将多头注意力和前馈层打包成一个现成的模块。使用它，你可以用明显更少的行数构建相同的语言模型。权衡是内部细节被 API 隐藏了——这正是为什么这个版本排在第二位而不是第一位。在没有自己构建组件的情况下阅读它产生的是对代码形状的理解，而不是对机制的理解。
这个版本使用与从零开始版本完全相同的分词器、相同的嵌入层、相同的训练循环和相同的生成逻辑。唯一改变的是 Transformer 模块本身。
10.1 nn.Transformer 提供什么
nn.Transformer 给你一个完整的编码器-解码器 Transformer。对于仅解码器的语言模型，你只使用解码器端。PyTorch 还提供 nn.TransformerDecoderLayer 和 nn.TransformerDecoder 正是用于这个目的：一个单独的解码器模块和它们的堆叠。
在内部，nn.TransformerDecoderLayer 包含你构建的相同部分：多头注意力、前馈网络、层归一化和残差连接。配置参数直接映射到你已经知道的那些。
10.2 因果掩码
nn.Transformer 不会为你构建的一件事是因果掩码。你在前向传播时显式传入它。PyTorch 提供了一个实用函数来生成它：nn.Transformer.generate_square_subsequent_mask(size)。这产生与你的 AttentionBlock 注册为缓冲区的相同上三角负无穷值掩码——只是按需生成而不是存储在模块上。
# ── nn.Transformer 版本 ──────────────────────────────
class TransformerLanguageModel(nn.Module):
def __init__(
self,
vocab_size: int,
embedding_dim: int,
num_heads: int,
num_layers: int,
context_size: int,
ff_dim: int,
dropout: float = 0.1
):
"""
一个使用 PyTorch 的 nn.TransformerDecoder 构建的仅解码器语言模型。
与从零开始版本相比：
- 嵌入层完全相同
- nn.TransformerDecoderLayer 替换你的 TransformerBlock
- nn.TransformerDecoder 替换 nn.Sequential 模块堆叠
- LM 头和生成逻辑完全相同
- 因果掩码在前向时生成，不存储为缓冲区
Args:
vocab_size (int): 唯一 token 的总数。
embedding_dim (int): 核心向量维度。在 PyTorch 文档中称为 d_model。
num_heads (int): 每层的注意力头。必须整除 embedding_dim。
num_layers (int): 堆叠的解码器层数。
context_size (int): 最大序列长度。
ff_dim (int): 前馈中间维度（dim_feedforward）。
dropout (float): 应用于每个解码器层内的 dropout 率。
0.1 是标准起始值。
"""
super().__init__()
self.context_size = context_size
# 与从零开始版本完全相同
self.embedding = Embedding(vocab_size, embedding_dim)
# 一个解码器层：你的 TransformerBlock 的 nn.Transformer 等效物。
# batch_first=True 告诉 PyTorch 我们的张量是 (batch, seq, features)
# 而不是旧版默认的 (seq, batch, features)。
decoder_layer = nn.TransformerDecoderLayer(
d_model = embedding_dim,
nhead = num_heads,
dim_feedforward = ff_dim,
dropout = dropout,
batch_first = True,
)
# 堆叠 num_layers 个解码器层。
# nn.TransformerDecoder 是你的 nn.Sequential 模块堆叠的等效物。
self.transformer = nn.TransformerDecoder(
decoder_layer = decoder_layer,
num_layers = num_layers,
)
self.final_norm = nn.LayerNorm(embedding_dim)
self.lm_head = nn.Linear(embedding_dim, vocab_size)
def forward(
self,
x: torch.Tensor,
targets: torch.Tensor = None
) -> tuple[torch.Tensor, torch.Tensor | None]:
"""
Args:
x: 形状为 (batch_size, seq_len) 的 token 索引
targets: 可选的真实值下一个 token 用于损失计算
Returns:
logits: 形状 (batch_size, seq_len, vocab_size)
loss: 如果提供 targets 则为交叉熵损失，否则为 None
"""
B, T = x.shape
x = self.embedding(x) # (B, T, embedding_dim)
# 为当前序列长度生成因果掩码。
# 这与你的 AttentionBlock 中相同的三角掩码，
# 表示为加法负无穷值而不是布尔填充。PyTorch 的注意力实现
# 在 softmax 之前将此直接添加到注意力分数。
causal_mask = nn.Transformer.generate_square_subsequent_mask(
T,
device=x.device
)
# nn.TransformerDecoder 期望一个目标序列和一个 memory 序列
#（完整编码器-解码器模型中编码器的输出）。
# 对于仅解码器语言模型，我们对两者传递相同的张量。
# tgt_mask 在自注意力侧强制因果性。
x = self.transformer(
tgt = x,
memory = x,
tgt_mask = causal_mask,
)
x = self.final_norm(x)
logits = self.lm_head(x)
loss = None
if targets is not None:
B, T, V = logits.shape
loss = F.cross_entropy(
logits.view(B * T, V),
targets.view(B * T)
)
return logits, loss
@torch.no_grad()
def generate(
self,
prompt: torch.Tensor,
max_new_tokens: int,
context_size: int
) -> torch.Tensor:
"""
与从零开始版本完全相同的生成逻辑。
接口相同——替换模型，保留循环。
"""
self.eval()
for _ in range(max_new_tokens):
context = prompt[:, -context_size:]
logits, _ = self(context)
last_logits = logits[:, -1, :]
probs = F.softmax(last_logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
prompt = torch.cat([prompt, next_token], dim=1)
return prompt
10.3 替换它
第8部分的训练循环和生成代码需要一处更改：将 LanguageModel(...) 替换为 TransformerLanguageModel(...)。其他所有行保持不变，因为两个类共享相同的前向签名和相同的生成接口。
# 将第8部分训练设置中的这一行替换为：
model = LanguageModel(
vocab_size = vocab_size,
embedding_dim = EMBEDDING_DIM,
num_heads = NUM_HEADS,
num_layers = NUM_LAYERS,
context_size = CONTEXT_SIZE,
ff_dim = FF_DIM,
)
# 替换为：
model = TransformerLanguageModel(
vocab_size = vocab_size,
embedding_dim = EMBEDDING_DIM,
num_heads = NUM_HEADS,
num_layers = NUM_LAYERS,
context_size = CONTEXT_SIZE,
ff_dim = FF_DIM,
dropout = 0.1,
)
# 训练循环和生成块中的其他所有内容保持不变。
10.4 nn.Transformer 给你什么，不给你什么
两个版本产生相同的架构，并将在相同数据上收敛到相似的损失值。区别完全在于控制和透明度。
nn.TransformerDecoderLayer 是一个完整的、经过测试的、PyTorch 维护的实现。它在底层使用 PyTorch 优化的 F.scaled_dot_product_attention，这比你的 AttentionBlock 中的手动矩阵乘法明显更快，特别是在 GPU 上。如果你在训练真实的东西并且不需要修改注意力内部，nn.Transformer 版本是更快更实用的选择。
它不给你的是改变注意力机制本身的能力。如果你想尝试不同的掩码策略、修改的评分函数、不同的头组合方式，或者现代研究中出现的任何架构变体，你需要从零开始版本。nn.TransformerDecoderLayer 的内部设计不支持逐块覆盖。
这是诚实的权衡：便利性和速度与控制和透明度。大多数应用开发者想要前者。大多数研究者需要后者。知道如何构建两者意味着你可以选择而不是被迫。
11、结束语
这个模型按现代标准很小，而且是有意为之。架构是真实的。训练循环是真实的。生成逻辑是真实的。小的是配置：嵌入维度、层数、上下文大小、数据集。
扩大其中任何一个，你就踏上了产生你听说过的每个大语言模型的相同路径。架构并没有根本不同。成分是相同的。改变的是规模、数据和计算。
最有效的下一步实验，按工作量排序：
在更大更有趣的文本文件上训练。莎士比亚、Wikipedia 转储或特定领域数据集将产生明显不同的生成，并教你很多关于模型实际学到了什么。
增加层数和嵌入维度，然后观察验证损失如何变化以及训练需要多长时间。这给你对控制所有现代 AI 的计算与质量权衡的直接直觉。
在 Transformer 模块内添加 dropout 层用于正则化。当模型的参数多于训练数据来证明它们时，dropout 通过在训练期间随机零化激活来防止过拟合。
将字符级分词器替换为 tiktoken 库中的子词分词器。这在保持词汇表大小可控的同时大幅增加了实际词汇表覆盖率，并且它是每个生产语言模型使用的。
构建它。打破它。修复它。扩展它。这是从阅读文章到理解技术的唯一路径。
原文链接：Building an LLM from Scratch with Pytorch
汇智网翻译整理，转载请标明出处
