深度学习和 Python 结合最典型的落地场景,就是用代码从零构建一个大语言模型(LLM)。LLM 并不是一个神秘黑盒,它本质上是一个基于 Transformer 架构的深度神经网络,通过海量文本预测下一个词或字符,逐步学习语言规律。到第 15、16 章这个阶段,前面关于张量操作、反向传播、模型封装的知识基本已经齐备,真正要做的是把知识点串起来,完成从数据准备、模型定义、训练循环、权重保存到推理生成的全过程。
这篇文章会按“概念 -> 环境 -> 数据 -> 模型 -> 训练 -> 推理 -> 精度选型 -> 排查”的顺序,带你在 Python 中构建一个最小可运行的 GPT 风格 LLM。学完之后,你能理解 Transformer 各模块的作用,能自己跑通一个几百万参数的迷你模型,也能判断训练过程中 loss 不下降、显存不足、梯度溢出等问题应该从哪个方向查起。
1. 先拆解 LLM 的骨架:深度学习在其中做了什么
1.1 从“预测下一个词”这个目标出发
大语言模型的核心任务并不复杂:给定一段上文,预测下一个 token 是什么。例如输入“深度学习”,模型要预测出“的”或“模型”这类后续内容。这个看似简单的目标,在大量文本上反复训练后,模型会逐渐学到词与词之间的关系、语法结构、知识片段,甚至上下文推理能力。
在深度学习视角下,这个任务属于自监督学习。训练数据不需要人工标注,直接从原始文本中构造“上文 -> 下一个 token”的训练对即可。这也是 LLM 能够利用海量文本的根本原因。
1.2 Transformer 架构的三个关键模块
从零构建 LLM,核心是复现 Transformer 的 decoder-only 结构。理解以下三个模块就抓住了主线:
| 模块 | 作用 | 在 LLM 中的位置 |
|---|---|---|
| Token Embedding | 把离散的 token id 映射为连续向量 | 模型入口 |
| Self-Attention | 让每个位置能看到前文其他位置的信息 | 每层核心 |
| Feed-Forward Network | 对每个位置的向量做非线性变换 | 每层后半段 |
这三个模块之间通过残差连接和层归一化组合,形成 Transformer block。多个 block 叠加,再经过一个输出层,就构成了一个完整的 decoder-only LLM。
1.3 为什么第 15、16 章会集中讲 LLM 构建
基础深度学习课程前半部分通常讲线性回归、多层感知机、卷积网络、循环网络,这些内容更多处理固定长度输入。到 LLM 章节,输入变成了变长序列,模型需要处理上下文依赖,训练也从单样本预测变成批量序列预测。此时如果不理解张量形状、掩码机制、损失函数如何作用在序列上,很容易写出维度不匹配的代码。
因此后续内容的目标是:用最少代码实现一个可以训练和生成文本的迷你模型。理解了它,再去看 GPT、Llama 等开源实现的源码,思路会顺畅很多。
2. Python 环境、PyTorch 版本和硬件先对齐
2.1 环境清单
从零构建 LLM 最推荐的框架是 PyTorch,因为它的动态图和自动求导机制很适合研究型开发。学习阶段不需要一开始就上分布式训练,先在一台机器上把流程跑通即可。
推荐环境如下:
| 组件 | 推荐版本/配置 | 说明 |
|---|---|---|
| Python | 3.10 或 3.11 | 避免过旧版本导致依赖冲突 |
| PyTorch | 2.x | 对 AMP、Transformer 封装更友好 |
| CUDA | 11.8 或 12.x | 以 PyTorch 官方安装命令为准 |
| GPU | NVIDIA 显卡,显存 6GB 以上 | 学习入门可先小模型 CPU 跑通 |
| 虚拟环境工具 | venv 或 conda | 隔离项目依赖 |
在常见项目中,可以按这个顺序检查:
python --version pip list | grep torch nvidia-sminvidia-smi显示的是驱动支持的 CUDA 版本,不一定是 PyTorch 使用的版本。PyTorch 通过自己的 CUDA runtime 运行,只要驱动版本足够新即可。
2.2 安装 PyTorch 的正确方式
在常见项目中,推荐到 PyTorch 官网获取当前环境的 install 命令,避免手动安装后 CUDA 不可用。例如 Linux + CUDA 12.x 环境,通常使用:
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121安装后验证 CUDA 是否可用:
import torch print(torch.cuda.is_available()) print(torch.cuda.device_count()) print(torch.cuda.get_device_name(0))如果返回False,优先检查 PyTorch 版本和安装命令是否匹配,而不是先怀疑显卡驱动。
注意:
torch.cuda.is_available()返回True只代表 CUDA 可用,不代表显存足够训练模型。训练时还要及时观察显存占用。
2.3 最小项目结构
为了后续扩展,建议把代码拆成模块,而不是把全部逻辑写在一个脚本里。
llm-from-scratch/ |-- data/ | `-- prepare.py |-- model/ | |-- __init__.py | |-- gpt.py |-- train.py |-- generate.py |-- requirements.txt学习阶段目录少一点也能跑通,但保持这种拆分方式能让你在模型变大后更容易定位问题。
2.4 学习环境与生产环境的区别
学习环境中,数据集很小、模型只有几层,CPU 也能完成训练,只是速度慢。生产环境则还需要考虑:
- 模型参数外置化,训练配置使用配置文件管理。
- 日志、指标监控、checkpoint 定期保存与回滚。
- 混合精度训练减少显存占用。
- 数据清洗与去重,避免训练集混入重复或脏数据。
本文示例以学习为主,生产环境建议在示例基础上增加这些保障。
3. 数据准备:从原始文本到模型能学习的 token 序列
3.1 字符级与子词级的分词选择
要让模型处理文本,第一步是把文本转成数字。常用方法有两种:
- 字符级:把每个字符映射为一个 id。实现简单,适合入门,但序列长度长,学习效率低。
- 子词级:使用 BPE、WordPiece 等分词器,把常见词组合成 token,减少序列长度,是真实 LLM 的主流方案。
入门阶段建议先用字符级分词跑通训练流程,后续再替换为 Hugging Face 的tokenizers或transformers中的分词器。这样能把“数据问题”和“模型问题”分开排查。
3.2 构建训练样本的完整过程
假设原始文本是“深度学习从零构建大语言模型”,字符级分词后,每个字符有一个 id。接下来构造训练对:
- 输入:前 N 个字符的 id 序列。
- 目标:这 N 个字符向后移动一位得到的 id 序列。
例如context_length=8时:
输入: [深, 度, 学, 习, 从, 零, 构, 建] 目标: [度, 学, 习, 从, 零, 构, 建, 大]模型的任务是根据输入序列预测下一个 token,所以目标序列正好是输入序列右移一位。
3.3 用 PyTorch 的 Dataset 封装数据
下面代码展示如何把一段文本转换成训练所需的数据结构:
import torch from torch.utils.data import Dataset class CharDataset(Dataset): def __init__(self, text, context_length=64): self.chars = sorted(set(text)) self.stoi = {ch: i for i, ch in enumerate(self.chars)} self.itos = {i: ch for i, ch in enumerate(self.chars)} self.context_length = context_length self.data = torch.tensor([self.stoi[ch] for ch in text], dtype=torch.long) def __len__(self): return len(self.data) - self.context_length def __getitem__(self, idx): x = self.data[idx: idx + self.context_length] y = self.data[idx + 1: idx + self.context_length + 1] return x, y这段代码的关键点是:
stoi和itos完成字符到 id 的双向映射。__len__控制了最多能取多少个样本。x和y的形状都是(context_length,),训练时由 DataLoader 自动堆叠成 batch。
3.4 数据准备阶段最容易踩的坑
字符级分词的主要坑是训练和推理时使用不同的映射表。训练时模型学过的是stoi,推理时如果重新生成映射表,字符编号就变了,输出会完全错乱。
解决方式:训练结束后把chars、stoi和itos一起保存到 JSON 文件,推理时加载同一个映射表。
4. 核心实现:用 PyTorch 手写一个微型 GPT 风格 LLM
4.1 模型整体结构
本节实现一个简化版 GPT,它包含以下部分:
- token embedding:把 token id 映射为向量。
- position embedding:给每个位置加入位置信息。
- 多个 Transformer block。
- 最后的线性输出层。
为了保持可读性,这里先实现一个单层自注意力,再扩展为多头。
4.2 完整模型代码
import torch import torch.nn as nn import torch.nn.functional as F class SelfAttention(nn.Module): def __init__(self, embed_dim, head_size): super().__init__() self.key = nn.Linear(embed_dim, head_size, bias=False) self.query = nn.Linear(embed_dim, head_size, bias=False) self.value = nn.Linear(embed_dim, head_size, bias=False) self.register_buffer("tril", torch.tril(torch.ones(128, 128))) def forward(self, x): B, T, C = x.shape k = self.key(x) q = self.query(x) v = self.value(x) att = q @ k.transpose(-2, -1) * (C ** -0.5) att = att.masked_fill(self.tril[:T, :T] == 0, float("-inf")) att = F.softmax(att, dim=-1) y = att @ v return y自注意力做的事情是:每个 token 通过 query 去查询前文所有 token 的 key,得到注意力权重,再用权重加权 value。tril矩阵确保模型只能看到当前位置以及之前的位置,这是 decoder-only LLM 保持因果性的关键。
4.3 多头注意力与 Transformer Block
class MultiHeadAttention(nn.Module): def __init__(self, embed_dim, num_heads): super().__init__() self.heads = nn.ModuleList( [SelfAttention(embed_dim, embed_dim // num_heads) for _ in range(num_heads)] ) self.proj = nn.Linear(embed_dim, embed_dim) def forward(self, x): out = torch.cat([h(x) for h in self.heads], dim=-1) out = self.proj(out) return out多头注意力的意义是让模型在不同子空间里关注不同类型的依赖关系。例如一个头关注语法,另一个头关注指代关系。
class FeedForward(nn.Module): def __init__(self, embed_dim): super().__init__() self.net = nn.Sequential( nn.Linear(embed_dim, 4 * embed_dim), nn.ReLU(), nn.Linear(4 * embed_dim, embed_dim), ) def forward(self, x): return self.net(x) class TransformerBlock(nn.Module): def __init__(self, embed_dim, num_heads): super().__init__() self.ln1 = nn.LayerNorm(embed_dim) self.attn = MultiHeadAttention(embed_dim, num_heads) self.ln2 = nn.LayerNorm(embed_dim) self.ffn = FeedForward(embed_dim) def forward(self, x): x = x + self.attn(self.ln1(x)) x = x + self.ffn(self.ln2(x)) return x残差连接x = x + ...让梯度可以从深层直接回传,避免深层网络梯度消失。LayerNorm 则让每层输入分布稳定,加快收敛。
4.4 完整的 GPT 模型类
class GPT(nn.Module): def __init__(self, vocab_size, embed_dim, num_heads, num_layers, context_length): super().__init__() self.token_embedding = nn.Embedding(vocab_size, embed_dim) self.position_embedding = nn.Embedding(context_length, embed_dim) self.blocks = nn.Sequential( *[TransformerBlock(embed_dim, num_heads) for _ in range(num_layers)] ) self.ln_final = nn.LayerNorm(embed_dim) self.lm_head = nn.Linear(embed_dim, vocab_size) def forward(self, idx): B, T = idx.shape tok_emb = self.token_embedding(idx) pos = torch.arange(T, device=idx.device) pos_emb = self.position_embedding(pos) x = tok_emb + pos_emb x = self.blocks(x) x = self.ln_final(x) logits = self.lm_head(x) return logits这里logits的形状是(batch_size, context_length, vocab_size),第三维每个值表示对应 token 的得分。训练时,我们用这个分数和下一个 token 的真实 id 计算交叉熵损失。
4.5 模型参数速查表
| 参数 | 含义 | 学习阶段建议 | 规模增大时 |
|---|---|---|---|
vocab_size | 词表大小 | 字符集大小 | 几万 |
embed_dim | 向量维度 | 128 | 2048 以上 |
num_heads | 注意力头数 | 4 | 16 或 32 |
num_layers | Transformer 层数 | 2 到 4 | 几十 |
context_length | 最大上下文长度 | 64 或 128 | 4096 以上 |
参数不是越大越好。显存固定时,embed_dim和context_length直接决定显存占用,入门先用小参数确认流程正确。
5. 训练循环、精度选择与损失验证
5.1 训练循环的完整写法
训练一个 LLM,本质上是让模型在批量数据上预测下一个 token,并通过反向传播更新权重。
import torch from torch.utils.data import DataLoader def train(model, dataset, batch_size=32, lr=3e-4, max_steps=1000): model.train() optimizer = torch.optim.AdamW(model.parameters(), lr=lr) loader = DataLoader(dataset, batch_size=batch_size, shuffle=True) step = 0 epoch = 0 while step < max_steps: for x, y in loader: pred = model(x) loss = F.cross_entropy(pred.view(-1, pred.size(-1)), y.view(-1)) optimizer.zero_grad() loss.backward() optimizer.step() if step % 100 == 0: print(f"step {step}, loss {loss.item():.4f}") step += 1 if step >= max_steps: break epoch += 1这里pred.view(-1, pred.size(-1))把(batch, seq, vocab)展平成(batch * seq, vocab),y.view(-1)也展平成对应位置的真实 token id,从而让交叉熵损失直接作用在每一个时间步上。
5.2 如何判断训练是否正常
训练刚开始时,loss 应该接近log(vocab_size)的负值左右。例如词表大小为 65,随机初始化的模型损失大约在 4.17 附近。随着训练进行,loss 应逐渐下降。
如果 loss 一开始就非常低,或者完全不动,都要怀疑实现有问题:
| 现象 | 可能原因 |
|---|---|
loss 约为log(vocab_size)且不下降 | 数据加载、掩码或模型输出结构错误 |
| loss 快速降到 0 | 数据集太小,模型记住了训练文本 |
| loss 变为 NaN | 学习率过大或 fp16 梯度溢出 |
| loss 震荡很大 | batch 太小或学习率不合适 |
5.3 浮点精度选型:fp32、fp16、bf16、tf32
训练 LLM 时,浮点精度是一个绕不开的问题。默认情况下 PyTorch 全参数使用 fp32。fp32 足够稳定,但显存占用高、计算速度慢。为了提速和减少显存,需要了解四种常见精度。
| 精度 | 占用字节 | 指数位 | 尾数位 | 特点 | 适用场景 |
|---|---|---|---|---|---|
| fp32 | 4 | 8 | 23 | 精度高,范围大 | 默认精度、求 loss、保存权重 |
| fp16 | 2 | 5 | 10 | 速度较快,但小数值易溢出 | 混合精度训练的一部分 |
| bf16 | 2 | 8 | 7 | 范围与 fp32 相同,精度低 | 大模型预训练常用 |
| tf32 | 4 | 8 | 10 | NVIDIA Ampere 以上矩阵计算加速格式 | 单卡训练可尝试开启 |
fp16 的典型问题是当梯度值小于 6e-5 时可能下溢为 0,当损失值过大时又可能溢出为无穷大。推荐做法是使用 PyTorch 的自动混合精度 AMP,让前向计算用 fp16,梯度更新时仍维护一个 fp32 的权重副本。
from torch.cuda.amp import GradScaler, autocast scaler = GradScaler() with autocast(): pred = model(x) loss = F.cross_entropy(pred.view(-1, pred.size(-1)), y.view(-1)) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()这里GradScaler会在反向传播前把 loss 放大,完成反向传播后再把梯度还原,避免梯度过小被 fp16 吞掉。
5.4 保存 checkpoint
训练中断是常事,必须定期保存模型参数和优化器状态:
checkpoint = { "model_state": model.state_dict(), "optimizer_state": optimizer.state_dict(), "step": step, "loss": loss.item(), } torch.save(checkpoint, f"checkpoint_{step}.pt")恢复训练时,先加载模型结构,再加载model_state和optimizer_state。如果只保存模型参数而不保存优化器状态,恢复训练后可能因为学习率设置不同导致训练不稳定。
6. 文本生成:让训练完的模型真正输出内容
6.1 从 logits 到文本的转换
训练完成后,模型输出的是 logits。生成文本时,需要把 logits 转成概率分布,再从分布中采样 token:
def generate(model, idx, max_new_tokens=100): model.eval() for _ in range(max_new_tokens): idx_cond = idx[:, -context_length:] logits = model(idx_cond) logits = logits[:, -1, :] probs = F.softmax(logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) idx = torch.cat([idx, next_token], dim=1) return idx关键点:
idx_cond只取最近context_length个 token,防止超长输入。logits[:, -1, :]只关注最后一个位置的预测结果。torch.multinomial从概率分布中随机采样,增加文本多样性。
6.2 temperature 的影响
生成时一个常用参数是 temperature,它控制概率分布的平滑程度:
probs = F.softmax(logits / temperature, dim=-1)- temperature 越大,分布越平滑,文本越随机。
- temperature 越小,分布越尖锐,越倾向选择高概率 token。
- temperature 为 0 时变成贪心解码,每次选概率最高的 token。
| temperature | 效果 | 适用场景 |
|---|---|---|
| 0.2 | 输出稳定,重复率高 | 代码、结构化文本 |
| 0.8 | 平衡 | 日常对话 |
| 1.2 | 多样性强 | 创意写作 |
6.3 推理阶段最容易忽略的问题
训练时模型处于train模式,Dropout 和 BatchNorm 的行为与推理不同。从训练代码复用模型时,必须先调用model.eval(),否则结果不稳定。
注意:推理不要传整段超长文本给模型,而要在生成过程中不断裁剪到
context_length以内,否则位置编码会越界。
7. 常见问题排查链路
7.1 维度不匹配的检查顺序
从零构建 LLM 时,shape mismatch是最常见的报错。出现这类错误时,按以下顺序检查:
- 输入形状是否是
(batch, context_length)。 - token id 是否在
0 ~ vocab_size - 1之间。 position_embedding的输入是否只用了 0 到T-1。- 输出 logits 的最后一维是否等于
vocab_size。 - 损失计算时目标值是否在合法范围内。
7.2 loss 不下降的排查路径
loss 不下降需要从数据、模型、训练三个层面排查:
| 检查项 | 操作 |
|---|---|
| 数据是否正确 | 打印一批x和y,确认输入输出对齐 |
| 模型是否处于 train 模式 | 检查是否有 Dropout 影响 |
| 学习率是否合适 | 太大会震荡,太小几乎不动 |
| 是否做了正确的掩码 | 是 causal mask,不是 padding mask |
| 梯度是否正常 | 检查参数的grad是否存在,是否为 0 或 NaN |
7.3 显存不足的解决方向
显存不足(OOM)不代表代码错误,更多是资源与模型不匹配。处理顺序:
- 减小
batch_size,优先看是否仍 OOM。 - 减小
context_length。 - 减小
embed_dim或num_layers。 - 开启梯度累积,用小 batch 模拟大 batch。
- 使用混合精度
autocast减少 fp32 占用。
7.4 fp16 训练出现 NaN
fp16 训练出现 NaN 是一项高频问题。可能原因包括:
- 学习率过高,导致梯度过大。
- loss 本身过大,梯度溢出。
- 模型内部出现
sqrt或log的数值越界。 GradScaler没有正常调用。
可以先切回 fp32 训练,确认这是数据问题还是精度问题。如果 fp32 正常而 fp16 NaN,则优先缩小学习率,并检查是否在适当位置使用autocast。
8. 从零构建 LLM 的可复用清单与生产建议
8.1 学习阶段清单
每次从零构建 LLM,都可按这份清单自查:
- 文本转 token id 后,输入输出是否右移对齐。
- 位置编码长度是否大于等于
context_length。 - causal mask 是否只允许看到前文。
- loss 初始值是否接近
log(vocab_size)。 - 训练过程中 loss 是否稳步下降。
- 保存 checkpoint 时是否同时保存
stoi和itos。 - 生成时是否调用
model.eval()。 - 生成时输入是否裁剪到
context_length。
8.2 生产环境的额外保障
在真实项目中,模型还需要关注数据质量和运行稳定性:
- 训练数据必须清洗,避免重复文本、敏感内容、空文档。
- 日志中记录每次训练的 loss、学习率、显存占用、分词器版本。
- checkpoint 至少保留最近两份,方便回滚。
- 参数配置外置化,用 YAML 或 JSON 管理模型大小、训练步数等。
- 推理服务需要输入长度限制、超时控制和并发限制。
8.3 下一步扩展方向
跑通一个迷你 LLM 后,扩展方向很明确:
- 使用
tokenizers库实现 BPE,替代字符级分词。 - 引入 KV cache 加速推理。
- 使用分布式训练框架训练更大模型。
- 接入 Hugging Face
transformers,把自定义模型转成标准接口。 - 做模型微调实验,观察
instruction tuning对输出风格的影响。
从零构建 LLM 最大的价值,不在于写出一个能媲美商用大模型的系统,而在于让你真正理解每一层张量在做什么、每一处配置改变会影响什么。后续阅读任何开源模型源码时,对照这份实现去理解注意力、精度和损失函数,就会比单纯背诵 API 有效得多。