深度学习和 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因为它的动态图和自动求导机制很适合研究型开发。学习阶段不需要一开始就上分布式训练先在一台机器上把流程跑通即可。推荐环境如下组件推荐版本/配置说明Python3.10 或 3.11避免过旧版本导致依赖冲突PyTorch2.x对 AMP、Transformer 封装更友好CUDA11.8 或 12.x以 PyTorch 官方安装命令为准GPUNVIDIA 显卡显存 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_length8时输入: [深, 度, 学, 习, 从, 零, 构, 建] 目标: [度, 学, 习, 从, 零, 构, 建, 大]模型的任务是根据输入序列预测下一个 token所以目标序列正好是输入序列右移一位。3.3 用 PyTorch 的 Dataset 封装数据下面代码展示如何把一段文本转换成训练所需的数据结构import torch from torch.utils.data import Dataset class CharDataset(Dataset): def __init__(self, text, context_length64): 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], dtypetorch.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 风格 LLM4.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, biasFalse) self.query nn.Linear(embed_dim, head_size, biasFalse) self.value nn.Linear(embed_dim, head_size, biasFalse) 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 Blockclass 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, deviceidx.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向量维度1282048 以上num_heads注意力头数416 或 32num_layersTransformer 层数2 到 4几十context_length最大上下文长度64 或 1284096 以上参数不是越大越好。显存固定时embed_dim和context_length直接决定显存占用入门先用小参数确认流程正确。5. 训练循环、精度选择与损失验证5.1 训练循环的完整写法训练一个 LLM本质上是让模型在批量数据上预测下一个 token并通过反向传播更新权重。import torch from torch.utils.data import DataLoader def train(model, dataset, batch_size32, lr3e-4, max_steps1000): model.train() optimizer torch.optim.AdamW(model.parameters(), lrlr) loader DataLoader(dataset, batch_sizebatch_size, shuffleTrue) 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(fstep {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 足够稳定但显存占用高、计算速度慢。为了提速和减少显存需要了解四种常见精度。精度占用字节指数位尾数位特点适用场景fp324823精度高范围大默认精度、求 loss、保存权重fp162510速度较快但小数值易溢出混合精度训练的一部分bf16287范围与 fp32 相同精度低大模型预训练常用tf324810NVIDIA 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, fcheckpoint_{step}.pt)恢复训练时先加载模型结构再加载model_state和optimizer_state。如果只保存模型参数而不保存优化器状态恢复训练后可能因为学习率设置不同导致训练不稳定。6. 文本生成让训练完的模型真正输出内容6.1 从 logits 到文本的转换训练完成后模型输出的是 logits。生成文本时需要把 logits 转成概率分布再从分布中采样 tokendef generate(model, idx, max_new_tokens100): 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_samples1) idx torch.cat([idx, next_token], dim1) 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 或 NaN7.3 显存不足的解决方向显存不足OOM不代表代码错误更多是资源与模型不匹配。处理顺序减小batch_size优先看是否仍 OOM。减小context_length。减小embed_dim或num_layers。开启梯度累积用小 batch 模拟大 batch。使用混合精度autocast减少 fp32 占用。7.4 fp16 训练出现 NaNfp16 训练出现 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 Facetransformers把自定义模型转成标准接口。做模型微调实验观察instruction tuning对输出风格的影响。从零构建 LLM 最大的价值不在于写出一个能媲美商用大模型的系统而在于让你真正理解每一层张量在做什么、每一处配置改变会影响什么。后续阅读任何开源模型源码时对照这份实现去理解注意力、精度和损失函数就会比单纯背诵 API 有效得多。
