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

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 的基本结构，并补充了训练循环的关键逻辑。理解这些底层原理有助于更好地掌握大模型的技术细节，为进一步研究和开发奠定基础。