从零开始手写 LLM 模型架构与实现
大型语言模型(LLM)是一种以其实现通用语言理解和生成能力而闻名的语言模型。通过在计算密集型自监督和半监督训练过程中从文本文档中学习统计关系来获得这些能力。LLM 是遵循 Transformer 架构的人工神经网络。
这是一个基于神经网络的模型训练,采用了一些新旧技术:, , , , , , , 。
本文详细介绍了从零开始手写 LLM 模型的全过程。内容包括数据准备、Tokenization、词向量 Embedding、位置编码、Transformer 核心组件(QKV、多头注意力)、残差连接与层归一化,以及最终的预测与训练循环实现。文章通过 Python 和 PyTorch 代码示例,展示了如何构建基础的 Transformer 架构,帮助读者深入理解大模型的核心原理与实现细节。

大型语言模型(LLM)是一种以其实现通用语言理解和生成能力而闻名的语言模型。通过在计算密集型自监督和半监督训练过程中从文本文档中学习统计关系来获得这些能力。LLM 是遵循 Transformer 架构的人工神经网络。
这是一个基于神经网络的模型训练,采用了一些新旧技术:, , , , , , , 。
tokenizationembeddingposition encodingfeed-forwardnormalizationsoftmaxlinear transformationmulti-head attention下图阐述了整个模型,接下来以一个小文本来演练下。
![模型整体架构图]
安装如下工具包:
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
使用 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、词向量
# 把数据转成张量形式
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())
到此,我们的数据准备已经完成。开始进入模型领域。
将 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])。
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 模型所需要的全部输入。
包括 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
多头机制,为了把 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]
# 注意力机制
# 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())
继续 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
# 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)
定义 Wo 线性层进行投影。
# 定义 Wo
Wo=nn.Linear(d_model,d_model)
output=Wo(A)
此时多头注意力已经结束,需要做个残差链接,也很简单,只需要和原始数据求和即可。
# 残差链接
# 从多头出来和原始数据进行
output=output+x
# 层归一化
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
#再进行一层归一化
output=layer_norm(output)
最终一个线性变化,此时需要把所有的 token 作为最终预测,会得到一个分值,选择最大的那一个。
# 最终一个线性变化,此时需要把所有的 token 作为最终预测,会得到一个分值,选择最大的那一个
output=nn.Linear(d_model,max_token_value+1)(output)
output.shape
probabilities = torch.softmax(output, dim=-1)
pd.DataFrame(probabilities[0].detach().cpu().numpy())
获取预测值索引:
predicted_index = torch.argmax(output[0,0]).item()
encoding.decode([predicted_index])
#看看句子
encoding.decode(x_batch[0].tolist())
为了让模型真正具备学习能力,我们需要定义损失函数和优化器,并执行反向传播。
# 定义优化器和损失函数
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 的基本结构,并补充了训练循环的关键逻辑。理解这些底层原理有助于更好地掌握大模型的技术细节,为进一步研究和开发奠定基础。

微信公众号「极客日志」,在微信中扫描左侧二维码关注。展示文案:极客日志 zeeklog
使用加密算法(如AES、TripleDES、Rabbit或RC4)加密和解密文本明文。 在线工具,加密/解密文本在线工具,online
生成新的随机RSA私钥和公钥pem证书。 在线工具,RSA密钥对生成器在线工具,online
基于 Mermaid.js 实时预览流程图、时序图等图表,支持源码编辑与即时渲染。 在线工具,Mermaid 预览与可视化编辑在线工具,online
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具,curl 转代码在线工具,online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具,Base64 字符串编码/解码在线工具,online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具,Base64 文件转换器在线工具,online