从零开始手写 LLM 模型架构与实现

从零开始手写 LLM 模型架构与实现

大型语言模型（LLM）是一种以其实现通用语言理解和生成能力而闻名的语言模型。通过在计算密集型自监督和半监督训练过程中从文本文档中学习统计关系来获得这些能力。LLM 是遵循 Transformer 架构的人工神经网络。

这是一个基于神经网络的模型训练，采用了一些新旧技术：tokenization, embedding, position encoding, feed-forward, normalization, softmax, linear transformation, multi-head attention。

下图阐述了整个模型，接下来以一个小文本来演练下。

![模型整体架构图]

0、数据准备

安装如下工具包：

pip install numpy requests torch tiktoken matplotlib pandas

导入对应的包：

import os
import requests
import pandas as pd
import matplotlib.pyplot as plt
import math
import tiktoken
import torch
import torch.nn as nn

这里采用 sales_textbook.txt 的数据作为训练数据和验证数据。把文本数据加载进入内存：

with open('sales_textbook.txt', 'r', encoding='utf-8') as f:
    text = f.read()

一些全局设置：

# 设置一些参数
batch_size = 4  # 批次
context_length = 16  # token 长度，最长 16 个单词
d_model = 64  # 维度
num_layers = 8  # Number of transformer blocks
num_heads = 4  # 多头 # 我们的代码中通过 d_model / num_heads = 来获取 head_size

1、Tokenization

使用 tiktoken 工具库进行处理，该工具是 OpenAI 提供的又快又轻量级，基于原始单词。

encoding = tiktoken.get_encoding("cl100k_base")
tokenized_text = encoding.encode(text) # 整个文本的单词数量 77,919
vocab_size = len(set(tokenized_text)) # 单词数：相当于哈希表里面的字典 3,771
max_token_value = max(tokenized_text) # 每个单词对应一个数值，这里找到最大值，100069

print(f"文本单词数量：{len(tokenized_text)}")
print(f"词汇表数量：{vocab_size}")
print(f"最大的单词 token 对应的编码值：{max_token_value}")

2、词向量 Word Embedding

把数据转成张量形式，并划分训练集和验证集。

# 2、词向量 
# 把数据转成张量形式
tokenized_text=torch.tensor(tokenized_text, dtype=torch.long)

# 首先吧数据集分成训练数据和验证数据（8:2）
# Split train and validation
split_idx = int(len(tokenized_text) * 0.8) # 取出 80% 的位置
train_data = tokenized_text[:split_idx]  #训练数据
val_data = tokenized_text[split_idx:]    #验证数据

# 准备训练数据
data = train_data
#随机取 4 个数字，范围 0-(77919*0.8-16)，为了构建批量数据
idxs = torch.randint(low=0, high=len(data) - context_length, size=(batch_size,))  
print(idxs)
# x 是一个 4*16 的结构，因为 idxs 有 4 个，每个里面都是 16 的长度 
# y 比 x 移后一位 
x_batch = torch.stack([data[idx:idx + context_length] for idx in idxs])
y_batch = torch.stack([data[idx + 1:idx + context_length + 1] for idx in idxs])
print(x_batch.shape,y_batch.shape)

查看 x_batch 和 y_batch 的具体内容：

import pandas as pd
pd.DataFrame(x_batch[0].numpy())

解码对应的 token 看看是什么：

encoding.decode([15749])
encoding.decode(x_batch[0].numpy())

到此，我们的数据准备已经完成。开始进入模型领域。

3、把数字转成输入向量

将 16 个数字转换成 64 维度的向量。

# 把 16 个数字转换成 64 维度（列）
# 找到最大的 token 值对应的数字
max_token_value=tokenized_text.max().item()
max_token_value
# 需要构造一个 max_token_value*64 维度的矩阵
input_embedding_lookup_table = torch.nn.Embedding(max_token_value+1, d_model)
#看下初始化权重信息，这些数据是在不断更新的
input_embedding_lookup_table.weight.data

将 x_batch 与权重绑定：

x_batch_embedding = input_embedding_lookup_table(x_batch)
y_batch_embedding = input_embedding_lookup_table(y_batch)
# 打印看看
x_batch_embedding.shape

输出为 torch.Size([4, 16, 64])。

4、位置信息

Transformer 需要位置编码来区分 token 的顺序。

# 先构建一个全 0 编码 的 16*64 的矩阵，先把形状搞出来，16*64 方便与 x_batch_embedding 的形状进行@运算
position_encoding_lookup_table = torch.zeros(context_length, d_model)
# torch.arange 形成一维 16 个数，然后扩展第二位
position = torch.arange(0, context_length, dtype=torch.float).unsqueeze(1)  # unsqueeze 函数主要是对数据维度进行扩充。
# 计算正玄余玄
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
position_encoding_lookup_table[:, 0::2] = torch.sin(position * div_term) #偶数位
position_encoding_lookup_table[:, 1::2] = torch.cos(position * div_term) #奇数位

position_encoding_lookup_table = position_encoding_lookup_table.unsqueeze(0).expand(batch_size, -1, -1) #add batch to the first dimension
print("Position Encoding Look-up Table: ", position_encoding_lookup_table.shape)
pd.DataFrame(position_encoding_lookup_table[0].numpy())

最终输入需要输入与位置信息加和：

x=x_batch_embedding+position_encoding_lookup_table
y=y_batch_embedding+position_encoding_lookup_table
# 看看此时形状
x.shape,y.shape

目前拿到给 transform 模型所需要的全部输入。

5、Transform 核心区域

包括 Q(query)、K(key)、V(Value)，多头机制等，Wq、Wk、Wv 是一个 64*64 的矩阵，与我们输入的 x 进行相乘，得到 Q、K、V。

# 准备 Wq、Wk、Wv，也是权重矩阵
Wq=nn.Linear(d_model,d_model)  # 64*64
Wk=nn.Linear(d_model,d_model)
Wv=nn.Linear(d_model,d_model)

# 此处相乘是 x 的后两维与 Wq 等相乘   16*64  X   64*64  =  16*64
Q=Wq(x)  # Wq*x
K=Wk(x)
V=Wv(x)

#得到 4*16*64 的 QKV
Q.shape

6、多头机制

多头机制，为了把 64 个维度进行拆分多份，每一个头里面有一部分维度，这里是 4 个头，分别进行部分计算，然后合并。

#先拆 Q，里面第一维度第二维度不变，分别是 4，,16，增加第三维度 4，第四维度从之前的第三维 64 变成 16；这样就变成（4，,16，4，,16）；

Q=Q.reshape(batch_size,context_length,num_heads,d_model//num_heads)
K=K.reshape(batch_size,context_length,num_heads,d_model//num_heads)
V=V.reshape(batch_size,context_length,num_heads,d_model//num_heads)

# 在把二三维度转换，变成（4，,4，,16，,16）因为接下来需要做注意力机制，关注的是 token，而不是多头
Q = Q.transpose(1, 2) # [4, 4, 16, 16]
K = K.transpose(1, 2) # [4, 4, 16, 16]
V = V.transpose(1, 2) # [4, 4, 16, 16]

7、注意力机制

# 注意力机制
# K 转置
output=Q @ K.transpose(-2,-1)/math.sqrt(d_model//num_heads)

#做 softmax 之前做 mask 部分，一句话比如中国人民，在第三个字的时候，他只知道中国人，不知道第四位是民这个字，所以需要处理下，把未来的设置成 0
#形成一个三角区域
mask=torch.triu(torch.ones((context_length,context_length)),diagonal=1).bool()
output=output.masked_fill(mask,float('-inf'))
pd.DataFrame(output[0][0].detach().numpy())

8、Softmax

继续 softmax，把概率数值转换成 0-1 之间的百分比。

attention_score=torch.softmax(output,dim=-1) #[4, 4, 16, 16] [batch_size, num_heads, context_length, context_length]
attention_score.shape

9、Attention 积分

# attention_score @ V 的操作
A=attention_score@V
#torch.Size([4, 4, 16, 16])
print(A.shape)
#合并多头，之前把 64 拆成 4*16，改变了位置，这边需要换回来
A=A.permute(0,2,1,3).reshape(batch_size,context_length,d_model)
print(A.shape)

10、Wo 操作

定义 Wo 线性层进行投影。

# 定义 Wo
Wo=nn.Linear(d_model,d_model)
output=Wo(A)

11、残差链接

此时多头注意力已经结束，需要做个残差链接，也很简单，只需要和原始数据求和即可。

# 残差链接
# 从多头出来和原始数据进行
output=output+x

12、进入层归一化

# 层归一化
layer_norm=nn.LayerNorm(d_model)
layer_norm_output=layer_norm(output)

#前馈网络：先把维度放大，然后做激活函数，然后再维度缩回
output=nn.Linear(d_model,d_model*4)(layer_norm_output )
output=nn.ReLU()(output)
output=nn.Linear(d_model*4,d_model)(output)

output=output+layer_norm_output

13、再进入层归一化

#再进行一层归一化
output=layer_norm(output)

14、线性变化

最终一个线性变化，此时需要把所有的 token 作为最终预测，会得到一个分值，选择最大的那一个。

# 最终一个线性变化，此时需要把所有的 token 作为最终预测，会得到一个分值，选择最大的那一个
output=nn.Linear(d_model,max_token_value+1)(output)
output.shape

15、最后的 Softmax

probabilities = torch.softmax(output, dim=-1)
pd.DataFrame(probabilities[0].detach().cpu().numpy())

16、Test

获取预测值索引：

predicted_index = torch.argmax(output[0,0]).item()
encoding.decode([predicted_index])
#看看句子
encoding.decode(x_batch[0].tolist())

17、训练循环

为了让模型真正具备学习能力，我们需要定义损失函数和优化器，并执行反向传播。

# 定义优化器和损失函数
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# 模拟训练步骤
for epoch in range(10): # 假设训练 10 个 epoch
    optimizer.zero_grad()
    # 前向传播
    output = model(x_batch)
    # 计算损失 (注意 y_batch 需要 reshape 以匹配 CrossEntropyLoss 的要求)
    loss = criterion(output.view(-1, output.size(-1)), y_batch.view(-1))
    # 反向传播
    loss.backward()
    optimizer.step()
    print(f"Epoch {epoch}, Loss: {loss.item()}")

在实际应用中，通常会将上述模块封装成一个完整的 nn.Module 类，并在 DataLoader 上迭代多个 epoch 进行训练。通过这种方式，模型能够逐步学习到语言的统计规律，从而提升生成质量。

总结

本文详细演示了从零开始构建一个简化版 LLM 模型的过程，涵盖了数据预处理、Embedding、位置编码、多头注意力机制、前馈网络以及残差连接等核心组件。通过 PyTorch 框架实现了 Transformer 的基本结构，并补充了训练循环的关键逻辑。理解这些底层原理有助于更好地掌握大模型的技术细节，为进一步研究和开发奠定基础。