从零(From Scratch)手写大语言模型
从0开始不依赖 HuggingFace、Ollama 等现成的大模型封装,使用 PyTorch 完成网络张量运算和自动求导,从零定义网络架构、初始化随机权重,并在自己的消费级显卡上训练出一个约百万参数的字符级语言模型实验。
一、 核心环境与硬件底座配置
大模型底层涉及大量的矩阵并行计算,必须打通完全兼容的硬件加速链条。
1. 软件环境基底
- Python 版本:本项目选择 Python 3.11.x。实际使用时,应以当前 PyTorch 官方安装页面列出的支持版本为准。
- 隔离虚拟环境:在项目根目录下通过
python -m venv .venv初始化沙盒环境,并通过.\.venv\Scripts\Activate.ps1激活,确保项目依赖不污染全局系统。
## 如果报错没找到python 刷新环境变量
$env:Path += ";C:\Users\zhoumao\AppData\Local\Programs\Python\Python311;C:\Users\zhoumao\AppData\Local\Programs\Python\Python311\Scripts"
2. Blackwell 架构显卡加速适配
RTX 5060 属于 Blackwell 架构,对应 NVIDIA 的计算能力 12.0(通常写作 sm_120)。PyTorch 的 CUDA 安装包必须支持当前显卡和驱动环境,但不应该把某一个 CUDA 小版本写成所有机器都必须使用的固定版本。应先查看 PyTorch 官方安装矩阵,再选择可用的 CUDA 构建版本:
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu130
上面的 cu130 只是本次实验使用的安装示例,不是 RTX 5060 的唯一要求。实际安装后,应该以 torch.cuda.is_available() 和最小张量运算测试为准。
3. 深度验证 CUDA 硬件加速是否成功
python -c "import torch; print('PyTorch 版本:', torch.__version__); print('CUDA 是否可用:', torch.cuda.is_available()); print('当前显卡:', torch.cuda.get_device_name(0) if torch.cuda.is_available() else '无')"
完美的输出结果应该是这样的:
- PyTorch 版本: 2.x.x+cu130 (或者你安装的其他可用 CUDA 构建版本)
- CUDA 是否可用: True
- 当前显卡: NVIDIA GeForce RTX 5060
二、先明确目标:我到底在训练什么
一开始,我给这个项目起的名字是“从0开始训练一个问答模型”。但真正开始写代码之后,必须先把目标拆开:
- 从零实现一个语言模型:理解 Token、Embedding、Attention、Transformer 和训练循环。
- 让模型记住几篇文章:把自己的文章作为训练语料。
- 做一个基于文章的问答工具:用户提问时,系统能够从文章中找到依据并回答。
这三个目标并不是一回事。
从零手写的模型,最适合用来学习大语言模型的基本原理;而“基于几篇文章进行可靠问答”,更适合使用检索增强生成(RAG,Retrieval-Augmented Generation)。如果只把几篇文章直接喂给一个很小的模型,模型通常只能记住文字表面和问答格式,并不能真正建立一个可靠的知识库。
所以,这个项目最终采用了两条并行路线:
train.py:从零训练一个教学性质的字符级 Transformer 语言模型。retriever.py和ask.py:建立一个基于文章原文的检索问答基础。
前者用于理解模型是如何训练的,后者用于逐步构建真正可用的文章问答工具。
三、从字符开始的 Tokenizer
1. 为什么先使用字符级分词
为了避免一开始就引入复杂的中文分词器或现成的大模型 Tokenizer,我先实现了一个最简单的字符级 Tokenizer。这里的“从零”指 Tokenizer 和模型结构由项目自己实现,并不意味着完全脱离 PyTorch。
它做的事情很直接:
人工智能
↓
[人, 工, 智, 能]
↓
[token_id_1, token_id_2, token_id_3, token_id_4]
Tokenizer 负责两件事:
encode():把文字转换成整数编号。decode():把整数编号还原成文字。
词表由训练语料中出现过的不同字符组成。对于中文短文本来说,字符级方式很容易实现,也方便观察模型到底在预测什么;但它的缺点也非常明显:词表虽然不大,序列却会变得很长,而且模型很难直接理解词语和句子的语义。
2. 处理未知字符和标点差异
最初的版本遇到训练语料中没有出现过的字符时,会直接跳过这个字符。这会导致一个危险的问题:
输入:人工智能一词首次被正式提出是在什么场合?
实际送入模型:人工智能一词首次被正式提出是在什么合?
用户并不知道输入已经被悄悄修改了。
后来我加入了四个特殊符号:
<PAD> 补齐序列时使用
<UNK> 未知字符
<BOS> 序列开始
<EOS> 序列结束
未知字符现在会被替换成 <UNK>,并且打印警告,让问题显式暴露出来,而不是静默丢失。
同时增加了 normalize_text(),在训练和推理时使用同一套文本归一化规则。例如把半角问号转换成中文全角问号,把连续空格压缩成一个空格。这样,下面两个问题在字符层面就不会被当成完全不同的输入:
AI是什么时候确立的?
AI是什么时候确立的?
下面是从 tokenizer.py 提取出的核心实现,省略了常量定义和部分注释;完整代码仍以项目文件为准:
SPECIAL_TOKENS = ["<PAD>", "<UNK>", "<BOS>", "<EOS>"]
class CharTokenizer:
def __init__(self, text: str, extra_chars: str = None):
chars = set(text)
if extra_chars:
chars |= set(extra_chars)
chars = sorted(chars)
vocab = SPECIAL_TOKENS + chars
self.stoi = {ch: i for i, ch in enumerate(vocab)}
self.itos = {i: ch for i, ch in enumerate(vocab)}
self.unk_id = self.stoi["<UNK>"]
self.vocab_size = len(vocab)
def encode(self, s: str, warn_unknown: bool = True):
ids = []
unknown_chars = set()
for c in s:
idx = self.stoi.get(c)
if idx is None:
idx = self.unk_id
unknown_chars.add(c)
ids.append(idx)
if warn_unknown and unknown_chars:
print(
f"[分词器警告] 发现未知字符: {sorted(unknown_chars)}"
)
return ids
def decode(self, ids):
return "".join(self.itos.get(i, "<UNK>") for i in ids)
标点归一化也使用独立函数,训练和推理调用的是同一个函数:
def normalize_text(s: str) -> str:
for half, full in _HALF_TO_FULL_PUNCT.items():
s = s.replace(half, full)
return re.sub(r"[ \t]+", " ", s)
因此,Tokenizer 的最小自测可以这样运行:
python tokenizer.py
遇到未登录字符时,输出会明确提示,而不是默默删除:
[分词器警告] 发现 1 个训练时没见过的字符: ['场']
四、手写一个类 GPT 的 Transformer
模型代码位于 model.py。它没有调用 nn.TransformerEncoder 或 nn.MultiheadAttention,而是自己实现了核心计算过程。
1. Token Embedding 和位置 Embedding
字符编号本身没有语义,模型需要把它转换成一个可以训练的向量:
tok_emb = self.token_embedding_table(idx)
但是注意力机制本身并不知道字符的先后顺序,因此还需要加入位置向量:
pos_emb = self.position_embedding_table(
torch.arange(T, device=idx.device)
)
x = tok_emb + pos_emb
最终模型看到的不是单纯的字符编号,而是“字符内容信息 + 所在位置信息”。
2. 因果自注意力
语言模型的训练目标是:给定前面的内容,预测下一个字符。
因此,第 i 个位置不能看到未来的字符,否则训练时模型就可以直接偷看答案。代码使用下三角矩阵实现因果掩码:
wei = wei.masked_fill(
self.tril[:T, :T] == 0,
float("-inf")
)
注意力的核心计算可以概括为:
Q = Query
K = Key
V = Value
attention(Q, K, V)
= softmax(QKᵀ / √d) V
在 model.py 中,单个注意力头的核心前向过程如下;这是与实际实现等价的整理版:
def forward(self, x):
B, T, C = x.shape
k = self.key(x)
q = self.query(x)
v = self.value(x)
wei = q @ k.transpose(-2, -1)
wei = wei * (k.shape[-1] ** -0.5)
# 只能看当前位置以及当前位置之前的内容
wei = wei.masked_fill(
self.tril[:T, :T] == 0,
float("-inf")
)
wei = F.softmax(wei, dim=-1)
wei = self.dropout(wei)
return wei @ v
这里的张量形状也很重要:
x: (batch_size, time, embedding_dim)
k: (batch_size, time, head_size)
q: (batch_size, time, head_size)
wei: (batch_size, time, time)
out: (batch_size, time, head_size)
多头注意力再把多个 Head 的结果拼接回来:
def forward(self, x):
out = torch.cat([h(x) for h in self.heads], dim=-1)
out = self.proj(out)
return self.dropout(out)
多个注意力头并行工作,每个头可以学习不同的关注模式;之后再把多个头拼接起来,通过线性层融合。
3. Transformer Block
每个 Block 包含两部分:
- 多头自注意力:让不同位置之间交换信息。
- 前馈网络:对每个位置获得的信息进行非线性加工。
同时使用 LayerNorm 和残差连接:
x = x + self.sa(self.ln1(x))
x = x + self.ffwd(self.ln2(x))
这就是一个简化版的 GPT 风格网络。模型最后通过 lm_head 把隐藏向量映射回整个字符词表,得到每一个字符作为“下一个字符”的概率。
完整模型的关键部分可以整理成下面这段;实际文件还包含类定义和模块初始化代码:
def forward(self, idx, targets=None):
B, T = idx.shape
tok_emb = self.token_embedding_table(idx)
pos_emb = self.position_embedding_table(
torch.arange(T, device=idx.device)
)
x = tok_emb + pos_emb
x = self.blocks(x)
x = self.ln_f(x)
logits = self.lm_head(x)
if targets is None:
loss = None
else:
B, T, C = logits.shape
logits = logits.view(B * T, C)
targets = targets.view(B * T)
loss = F.cross_entropy(logits, targets)
return logits, loss
这里的 logits 不是文字,而是每个位置对整个词表的打分。例如词表大小是 794,输入长度是 128,那么输出形状就是:
(batch_size, 128, 794)
模型在每一个位置都要从 794 个字符中预测下一个字符。
五、训练目标:预测下一个字符
train.py 会把所有文章读取后拼接成一段文本,然后随机截取长度为 BLOCK_SIZE 的片段。
假设输入是:
今天天气
对应的目标就是向右移动一个字符:
天天气很
也就是说,模型在每个位置都尝试预测输入序列中的下一个字符。训练过程包含四步:
- 从文章中随机取一个批次。
- 前向计算预测结果和交叉熵损失。
loss.backward()计算梯度。optimizer.step()更新模型参数。
项目中的批处理代码如下:
def get_batch(data, block_size, batch_size, device):
ix = torch.randint(
len(data) - block_size - 1,
(batch_size,)
)
x = torch.stack([
data[i:i + block_size] for i in ix
])
y = torch.stack([
data[i + 1:i + block_size + 1] for i in ix
])
return x.to(device), y.to(device)
训练循环没有交给高层训练框架,核心就是下面这段;反向传播本身由 PyTorch autograd 完成:
for step in range(MAX_ITERS):
xb, yb = get_batch(
train_data,
BLOCK_SIZE,
BATCH_SIZE,
DEVICE,
)
logits, loss = model(xb, yb)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
所以这里所谓的“训练”,实际就是不断重复取样、计算 loss、调用自动求导计算梯度,再更新参数。这个项目并没有手工推导或手写每个参数的反向传播公式。
当前实验使用的主要配置是:
BATCH_SIZE = 32
BLOCK_SIZE = 128
N_EMBD = 128
N_HEAD = 4
N_LAYER = 4
DROPOUT = 0.1
LEARNING_RATE = 3e-4
模型参数量大约为一百万。这个规模足够用来观察训练过程,但和真正的大语言模型之间仍然有数量级差距。
六、第一次训练:模型为什么会答对
最初的训练语料包括原文章和自动扩充的问答文件。训练结束时出现了这样的结果:
train loss: 0.1001
val loss: 8.3748
训练损失非常低,验证损失却越来越高。这是典型的过拟合。
模型之所以能够回答类似下面的问题:
问:人工智能一词首次被正式提出是在什么场合?
答:
并不一定表示它理解了“人工智能起源”这个抽象概念。更可能的情况是,它记住了训练文本中的局部字符模式:
问:人工智能这个概念是什么时候提出的?
答:1956年,在达特茅斯会议上首次正式提出。
当输入和训练数据足够相似时,字符级模型可以表现得像是在回答问题;但换一种问法,或者询问训练语料中没有的内容,就可能答非所问。
因此,训练损失下降并不等于问答能力可靠。必须同时观察验证损失,并使用没有直接出现在训练文本中的问题进行测试。
七、整理语料,而不是盲目增加重复数据
原来的 qa_augmented.txt 中包含大量相同问题和相同答案,只是改变了少量表达方式。例如同一个“深度学习为什么发展”的答案被重复了很多次。
这类数据短期内可能提高固定问题的命中率,但也会带来两个问题:
- 模型更加依赖固定模板。
- 重复答案在检索结果中占据多个位置。
所以我保留了原始文件用于回溯,但新增了 qa_curated.txt,只保留基于原文章整理的高信息密度问答。
整理后的问题不再只是简单改写,而是覆盖不同类型:
- 事实型:人工智能概念何时正式提出?
- 解释型:深度学习为什么快速发展?
- 对比型:深度神经网络和早期规则系统有什么不同?
- 条件型:为什么统一内存设备不一定更快?
- 局限型:大模型为什么不能回答所有问题?
- 数值型:文章给出的显存、带宽和参数数据是多少?
默认训练和检索现在使用:
sample_article.txt
qa_curated.txt
旧版 qa_augmented.txt 不再参与默认训练和检索,避免重复内容继续污染实验结果。
八、验证集和提前停止
第二次训练时,数据量减少到了约 9559 个字符,训练日志如下:
step 0 | train loss 6.8243 | val loss 6.8267
step 200 | train loss 3.2577 | val loss 5.8143
step 400 | train loss 2.1408 | val loss 6.0299
step 600 | train loss 1.4014 | val loss 6.4511
step 800 | train loss 0.8377 | val loss 6.9806
验证集连续 3 次变差,提前停止训练
可以看到,验证集最好的结果出现在 step 200。之后训练集损失继续下降,但验证集损失反而持续升高。
因此训练代码加入了两个保护:
- 记录验证集损失最低时的模型参数。
- 验证集连续多次变差时提前停止训练。
最终保存的是验证集表现最好的参数,而不是训练结束时的参数。这样可以避免把明显过拟合的最后一步模型保存下来。
对应代码如下:
best_val_loss = float("inf")
best_state = None
bad_evals = 0
if losses["val"] < best_val_loss:
best_val_loss = losses["val"]
best_state = {
key: value.detach().cpu().clone()
for key, value in model.state_dict().items()
}
bad_evals = 0
else:
bad_evals += 1
if bad_evals >= PATIENCE:
print("验证集连续变差,提前停止训练")
break
if best_state is not None:
model.load_state_dict(best_state)
不过,当前的验证集仍然只是按文本末尾切出的 10% 内容,并不是真正独立的问答测试集。更严格的做法应该是单独建立测试问题文件,确保测试问题和训练问题不重复。
九、为什么没有继续强行训练
使用当前模型生成下面的问题:
python generate.py --prompt "问:深度学习为什么发展这么快?`n答:"
在一次实际运行中得到过类似这样的结果:
答:文章一次内存。
答:参数里多知识可能在时间到生成答:RToken 5090B...
这不是简单的“训练步数不够”。当前模型同时受到几个因素限制:
- 训练语料太少。
- 字符级序列很长。
- 模型参数量有限。
- 问答数据之间存在格式和主题切换。
- 生成目标本质上只是预测下一个字符。
继续增加训练步数,通常只会让模型更牢地记住训练文本,并不一定让它更懂问题。对于当前项目,generate.py 更适合作为观察语言模型行为的实验入口,而不是可靠的问答接口。
十、从纯生成转向文章检索
到这里,项目开始从“训练一个小模型”转向“构建一个文章问答系统”。
真正基于几篇文章进行问答,最重要的第一步不是继续训练,而是先找到和问题相关的原文。
于是新增了三个文件:
retriever.py 文章切片和检索
build_index.py 构建索引
ask.py 命令行查询
1. 文章切片
文章先按空行拆分成段落,过长的段落再切成较短的片段,并保留一定重叠内容。
这样做的原因是:
- 整篇文章太长,不方便精确匹配。
- 片段太短,可能丢失上下文。
- 保留重叠区域,可以减少切分点带来的信息损失。
实际的切片入口会遍历 data/articles/*.txt,并记录来源文件。下面是从 retriever.py 提取出的核心逻辑:
def load_chunks(data_dir="data/articles", max_chars=500):
chunks = []
seen_text = set()
for path in sorted(Path(data_dir).glob("*.txt")):
if path.name in EXCLUDED_FILES:
continue
text = normalize_text(
path.read_text(encoding="utf-8")
)
paragraphs = re.split(r"\n\s*\n+", text)
for paragraph in paragraphs:
paragraph = paragraph.strip()
if not paragraph:
continue
for chunk in _split_long_paragraph(
paragraph,
max_chars=max_chars,
):
if chunk in seen_text:
continue
seen_text.add(chunk)
chunks.append({
"source": path.name,
"text": chunk,
})
return chunks
这里的 seen_text 是去重的关键。旧版 qa_augmented.txt 中有大量重复问答,因此当前代码还会默认排除它:
EXCLUDED_FILES = {"qa_augmented.txt"}
2. 字符 n-gram
中文没有直接使用第三方分词库,而是使用字符 unigram 和 bigram:
深度学习
↓
深、度、学、习、深度、度学、学习
这种方式实现简单,对中文短问题比较实用,也能处理没有登录词表的新词。
对应的特征提取代码只有几行:
def _features(text: str):
text = re.sub(
r"\s+",
"",
normalize_text(text),
).lower()
return (
list(text)
+ [
text[i:i + 2]
for i in range(len(text) - 1)
]
)
例如输入“深度学习”,会得到字符特征和相邻双字符特征。这样即使问题中出现了训练语料里没有完整收录的新词,也有机会通过字符片段产生匹配。
3. TF-IDF 检索
每个文章片段都会被转换成一个向量。查询时,把用户问题也转换成向量,再计算它和每个片段之间的相似度。
相似度越高,说明片段和问题共享的主题特征越多,排序就越靠前。
项目没有引入向量数据库,而是直接用 Python 字典保存稀疏向量。下面是片段向量计算的核心逻辑,实际代码还包括类初始化和索引保存:
counts = Counter(_features(chunk["text"]))
vector = {}
for feature, count in counts.items():
idf = math.log(
(1 + self.doc_count)
/ (1 + self.df[feature])
) + 1
vector[feature] = (
(1 + math.log(count)) * idf
)
norm = math.sqrt(
sum(value * value for value in vector.values())
) or 1.0
vector = {
key: value / norm
for key, value in vector.items()
}
查询时使用同样的公式生成问题向量,再计算余弦相似度:
score = sum(
query.get(key, 0.0) * value
for key, value in vector.items()
)
最后按得分排序:
scored.sort(
key=lambda item: item[0],
reverse=True,
)
运行索引构建:
python build_index.py
实际构建后得到:
已索引 154 个去重后的文章片段,保存到 index.json
4. 去重和最低相关度
由于旧版问答数据中有很多重复内容,索引阶段会过滤完全相同的片段,查询结果也会再次去重,避免同一段答案占满 top_k。
同时增加了最低相关度阈值。例如:
python ask.py "深度学习为什么发展这么快?"
可以得到:
[1] qa_curated.txt 相关度=0.3678
问:深度学习为什么在21世纪后快速发展?
答:计算能力提升和大数据积累为深度学习提供了条件,深度神经网络能够从大量数据中自动学习特征。
[2] sample_article.txt 相关度=0.2026
进入二十一世纪后,随着计算能力的提升和大数据的积累,机器学习尤其是深度学习技术迅速崛起。
第一条是整理后的直接问答,第二条是原文章依据。两者互相印证,说明检索结果是有意义的。
ask.py 并不是把所有片段无条件打印出来,而是使用最低相关度过滤:
results = retriever.search(
args.question,
args.top_k,
args.min_score,
)
在检索器内部,核心过滤逻辑可以简化表示为:
def search(self, question, top_k=3, min_score=0.12):
...
for score, chunk in scored:
if score < min_score:
break
results.append({
"score": round(score, 4),
**chunk,
})
if len(results) >= max(1, top_k):
break
return results
这个阈值不是“理解了问题”的证明,只是一个保守的工程过滤器:相似度太低时宁可不返回,也不把通用句子伪装成答案。
对于:
python ask.py "鸡蛋为什么快?"
系统会输出:
文章中没有找到足够相关的依据。
这一步非常重要。一个文章问答系统不应该为了“必须回答”而强行从文章里找一个看起来相近的句子。
在完善过滤前,类似“你的问题”这样的占位文本可能会因为共享“问题”等常见字符,命中“AI能够回答我们的问题”之类的句子。后来又加入了通用占位问题拦截:
你的问题
请输入问题
文章讲了什么
这些输入现在会要求用户提供更具体的问题,而不会伪装成有效检索结果。
十一、到检索为止,我得到了什么
到目前为止,这个项目已经完成了一个完整的实验闭环:
文章
↓
字符级 Tokenizer
↓
Embedding
↓
因果自注意力 Transformer
↓
预测下一个字符
↓
观察过拟合
↓
整理问答语料
↓
文章切片与去重
↓
TF-IDF 检索原文依据
其中,真正可以用于文章问答的部分已经从“让小模型自由生成”转向了“先检索原文依据”。
当前推荐的使用方式是:
python build_index.py
python ask.py "深度学习为什么发展这么快?"
这时系统返回的是文章中最相关的证据,而不是没有依据地自由编造答案。
ask.py 的主流程可以简化表示为:
if Path(args.index).exists():
retriever = ArticleRetriever.load(args.index)
else:
retriever = ArticleRetriever(load_chunks())
retriever.save(args.index)
results = retriever.search(
args.question,
args.top_k,
args.min_score,
)
if not results:
print("文章中没有找到足够相关的依据。")
return
for number, result in enumerate(results, 1):
print(
f"[{number}] {result['source']} "
f"相关度={result['score']}"
)
print(result["text"])
十二、如何复现实验
项目根目录下的实际操作顺序如下:
.\.venv\Scripts\Activate.ps1
python build_index.py
python ask.py "深度学习为什么发展这么快?"
del checkpoint.pt
python train.py
python generate.py --prompt "问:深度学习为什么发展这么快?`n答:"
其中,build_index.py 和 ask.py 不依赖已经训练好的 checkpoint.pt;它们使用文章文件直接建立索引。train.py 才会创建新的 checkpoint.pt,而 generate.py 依赖这个 checkpoint。
当前训练脚本默认排除旧版高重复语料:
EXCLUDED_FILES = {"qa_augmented.txt"}
因此默认训练和检索使用的是:
sample_article.txt
qa_curated.txt
如果增加或修改了文章,需要重新执行 python build_index.py;如果希望字符级模型也学习新文章,则还需要重新执行 python train.py。
十三、当前阶段的边界
当前系统已经具备:
- 从自己的文章构建本地索引。
- 对中文问题进行基础相关性检索。
- 过滤完全重复的片段。
- 设置最低相关度,减少无关命中。
- 在没有足够依据时拒答。
- 返回原文来源文件和相关度。
但它还没有完成:
- 把多个原文片段自动整理成自然语言答案。
- 处理复杂的跨段落推理。
- 使用真正的语义向量模型。
- 建立完整的独立问答测试集。
- 对答案进行事实一致性评估。
因此,现阶段最准确的描述是:
这是一个从零手写的字符级语言模型实验,以及一个已经完成基础检索能力的文章问答系统原型。
“检索”解决的是:应该参考文章中的哪些内容。
后续的“生成”才负责:如何把这些内容组织成简洁、自然的答案。
在没有接入可靠生成模型之前,先把检索依据做好,是比继续堆训练步数更稳妥的一步。
十四、项目源码压缩包
项目使用到的完整源码压缩包: