基于 PyTorch 从零构建与训练大型语言模型入门指南
LLM（Large Language Model）是如今大多数 AI 聊天机器人的核心基础，例如 ChatGPT、Gemini、MetaAI、Mistral AI 等。这些 LLM 背后的核心架构通常是 Transformer。
本文介绍如何一步步使用 PyTorch 从零开始构建和训练一个大型语言模型（LLM）。该模型以 Transformer 架构为基础，实现英文到马来语的翻译功能，同时也适用于其他语言翻译任务。
步骤 1：加载数据集
为了让 LLM 模型能够执行从英文到马来语的翻译任务，需要使用含有英马双语对照的数据集。
这里选择 Huggingface 提供的'Helsinki-NLP/opus-100'数据集。它包含百万级的英文 - 马来语对照句对，足以确保模型训练的准确性。此外，该数据集还包含了 2000 条验证和测试数据，且已经预先完成了分割工作，省去了手动分割的繁琐步骤。
from datasets import load_dataset
# 加载 Helsinki-NLP/opus-100 数据集
dataset = load_dataset("Helsinki-NLP/opus-100", "en-my")
# 查看数据集结构
print(dataset)
步骤 2：创建分词器
Transformer 模型不处理原始文本，只处理数字。因此，需要将原始文本转换为数字格式。
这里使用名为 BPE（Byte Pair Encoding）的流行分词器来完成这一转换过程。这是一种子词级别的分词技术，已在 GPT-3 等先进模型中得到应用。
通过训练数据集来训练这个 BPE 分词器，生成英马双语的词汇表，这些词汇表是从语料中提取的独特标记的集合。
分词器的作用是将原始文本中的每个单词或子词映射到词汇表中的相应标记，并为这些标记分配唯一的索引或位置 ID。
这种子词分词方法的优势在于，它能有效解决 OOV 问题，即词汇表外单词的处理难题。通过这种方式，我们能够确保模型在处理翻译任务时，无论是常见词汇还是生僻词汇，都能准确无误地进行编码，为后续的嵌入表示打下坚实基础。
from tokenizers import Tokenizer, models, trainers
# 初始化 BPE 分词器
tokenizer = Tokenizer(models.BPE())
# 配置训练参数
trainer = trainers.BpeTrainer(
vocab_size=8000,
special_tokens=["<pad>", "<unk>", "<s>", "</s>"]
)
# 准备文件列表进行训练
files = ["train_en.txt", "train_my.txt"]
tokenizer.train(files, trainer)
步骤 3：准备数据集和数据加载器
在构建模型的第三步，着手准备数据集及其加载器。这一阶段的目标是为源语言（英语）和目标语言（马来语）的数据集做好训练与验证的准备。
为此，需要编写一个类，能够接收原始数据集，并利用英语和马来语的分词器（分别为 tokenizer_en 和 tokenizer_my）对文本进行编码处理。编码后的数据会通过数据加载器进行管理，该加载器将按照设定的批次大小（本例中为 10）来迭代处理数据集。
如有需要，还可以根据数据量和计算资源的实际情况，对批次大小进行调整。
import torch
from torch.utils.data import Dataset, DataLoader
class TranslationDataset(Dataset):
def __init__(self, dataset, tokenizer_en, tokenizer_my):
self.dataset = dataset
self.tokenizer_en = tokenizer_en
self.tokenizer_my = tokenizer_my
def __len__(self):
return len(self.dataset)
def __getitem__(self, idx):
source_text = self.dataset[idx]["translation"]["en"]
target_text = self.dataset[idx]["translation"]["my"]
# 编码
enc_en = self.tokenizer_en.encode(source_text).ids
enc_my = self.tokenizer_my.encode(target_text).ids
return {
"source_ids": torch.tensor(enc_en),
"target_ids": torch.tensor(enc_my)
}
# 实例化数据集和加载器
dataset_obj = TranslationDataset(dataset["train"], tokenizer_en, tokenizer_my)
dataloader = DataLoader(dataset_obj, batch_size=10, shuffle=True)
步骤 4：输入嵌入和位置编码
这一步进行输入嵌入和位置编码的处理。
首先，输入嵌入层负责将步骤 2 生成的标记 ID 序列转换为词汇表中的索引，并为每个标记生成一个 512 维的嵌入向量。
这个向量能够捕捉标记的深层语义特征。在多维空间中，相似的实体如狗和猫的向量会彼此接近，而与学校、家等不相似实体的向量则相隔较远。
其次，位置编码解决了 Transformer 架构在并行处理序列时可能忽略词序的问题。通过给每个标记的 512 维嵌入向量添加位置信息，保证模型能够理解词序对句子含义的影响。
具体来说，采用正弦和余弦函数对每个维度进行编码，其中正弦应用于偶数维度，余弦应用于奇数维度。这样，每个标记的嵌入向量不仅包含了其语义信息，还包含了其在句子中的位置信息。
import math
import torch.nn as nn
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=5000):
super().__init__()
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0)
self.register_buffer('pe', pe)
def forward(self, x):
return x + self.pe[:, :x.size(1)]
class EmbeddingLayer(nn.Module):
def __init__(self, vocab_size, d_model):
super().__init__()
self.embedding = nn.Embedding(vocab_size, d_model)
self.pos_encoding = PositionalEncoding(d_model)
def forward(self, x):
x = self.embedding(x) * math.sqrt(self.embedding.embedding_dim)
return self.pos_encoding(x)
步骤 5：多头注意力块
Transformer 模型的精髓在于自注意力机制，它赋予模型动态理解上下文的能力。而多头自注意力则进一步将这一能力细分，让模型能够同时从多个角度捕捉信息，从而更全面地理解句子。
如果熟悉矩阵乘法，掌握多头注意力机制其实相当简单。
首先，我们会从步骤 4 得到的编码输入创建三份副本：Q（查询）、K（键）、V（值）。这些副本将作为自注意力计算的基础。
随后，将 Q、K、V 分别与各自的权重矩阵进行矩阵乘法，这些权重矩阵将初始化为随机值，并在训练过程中不断更新。这一步骤引入了可学习的参数，帮助模型更好地捕捉信息。
按照论文中的设定，我们将使用 8 个头来进行多头注意力的计算。这意味着，每个经过矩阵乘法得到的查询、键、值向量都将被分割成 8 份，每份的维度为 d_k = d_model / num_heads。
接下来，每个查询向量将与序列中所有键向量的转置进行点积运算，得到注意力分数，这些分数反映了标记间的相似度。为避免模型过度关注高分数或忽略低分数，我们通过除以 d_k 的平方根来规范化这些分数。
在应用 softmax 函数之前，如果存在编码掩码，我们会将其与注意力分数结合，确保模型不会受到未来时间步的影响。Softmax 函数将这些分数转换为概率分布，然后这些概率将与相应的值向量相乘，得到每个头的输出。
最终，我们将 8 个注意力头的输出合并，并通过输出权重矩阵 W_o 进一步处理，得到多头自注意力的最终结果。这个结果能够综合考虑单词在句子中的不同上下文含义。
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, num_heads):
super().__init__()
assert d_model % num_heads == 0
self.d_model = d_model
self.num_heads = num_heads
self.head_dim = d_model // num_heads
self.W_q = nn.Linear(d_model, d_model)
self.W_k = nn.Linear(d_model, d_model)
self.W_v = nn.Linear(d_model, d_model)
self.W_o = nn.Linear(d_model, d_model)
def split_heads(self, x, batch_size):
x = x.view(batch_size, -1, self.num_heads, self.head_dim)
return x.permute(0, 2, 1, 3)
def forward(self, q, k, v, mask=None):
batch_size = q.shape[0]
Q = self.split_heads(self.W_q(q), batch_size)
K = self.split_heads(self.W_k(k), batch_size)
V = self.split_heads(self.W_v(v), batch_size)
score = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)
if mask is not None:
score = score.masked_fill(mask == 0, -1e9)
attention = torch.softmax(score, dim=-1)
output = torch.matmul(attention, V)
output = output.permute(0, 2, 1, 3).contiguous().view(batch_size, -1, self.d_model)
return self.W_o(output)
以上步骤构成了 Transformer 模型的核心组件。在实际训练中，还需要结合损失函数（如 CrossEntropyLoss）和优化器（如 AdamW）进行端到端的训练。后续章节将详细介绍模型的整体搭建与训练流程。