ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

从零手写极简大语言模型:原理、训练与部署实战

从零手写极简大语言模型:原理、训练与部署实战 学完 Python 基础正在啃深度学习突然想搞明白“大语言模型到底是怎么被造出来的”——很多人都会卡在这个问题上。看论文看不懂细节看开源库又觉得是个黑盒最后只能停留在“会用 API”的层面。如果你想把这个认知缺口补上手写一个极简版大语言模型是性价比最高的路线。这篇文章对应深度学习 Python 教程的第 15、16 章核心任务是带着你从零构建一个小型 LLM。重点不是复现 GPT-4 或者 LLaMA 这种级别的模型而是把大模型的关键环节全部落到代码上数据准备、分词器、Transformer 架构、注意力机制、训练循环、推理采样最后再讨论精度管理和接口部署。走完整套流程你对 LLM 的认知会从“一个黑盒 API”变成“一条可以手动调试的训练链路”。文章不会绕弯子结构按环节拆开每一步都给出可运行的代码标注容易出现问题的位置最后还会补充模型部署、显存观察和批量任务的内容。适合已经掌握 Python 和 PyTorch 基础、想动手建模型的读者。1. 核心能力速览能力项说明项目类型从零构建 LLM 的教学实现不依赖 Hugging Face 高封装库核心功能文本数据预处理、Tokenizer、Transformer 训练、文本生成推理技术栈Python PyTorch运行环境支持 CPU 起步推荐使用 NVIDIA GPU 加速显存需求小型模型在 CPU / 小显存 GPU 上即可运行具体占用以实际模型参数为准启动方式脚本训练 推理文件可逐步运行无需一键包是否支持 API训练完成后可用 FastAPI 包装接口服务是否支持批量任务支持可在训练数据加载和推理阶段批量执行适合读者熟悉 Python正在学习深度学习想了解 LLM 内部机制的开发者2. 适用场景与学习路线这类“从零构建 LLM”的内容解决的核心问题不是“怎么调 API”而是“大模型内部发生了什么”。如果你工作中只需要调用现成模型这篇文章的很多细节可以当作背景知识但如果你需要做模型微调、私有化部署、定制 prompt 策略、调试推理性能那底层原理就不只是背景知识了。从学习路径看比较顺的顺序是先理解语言模型的目标给定前文预测下一个 token 的概率分布。再理解 tokenizer把文本拆成模型能处理的 token。然后搭 Transformer decoder 结构这是 GPT 系列模型的核心。接下来写因果注意力让模型只能看到过去的 token不能看到未来的 token。跑通训练循环用交叉熵损失更新参数。最后写推理函数用 temperature、top-k 等策略控制生成效果。这套流程走完你再看开源模型代码至少知道每一层在做什么而不是对着源码发愣。需要说明的是边界这里实现的是教学版小型 LLM不是可以直接商用的大规模模型。真实生产环境还需要考虑分布式训练、高质量数据清洗、对齐调优、推理加速等大量工程问题。版权方面训练数据如果来自公开网络要注意数据许可协议生成内容的用途也要注意合规边界不能把模型用于制造虚假信息、绕过安全限制等场景。3. 环境准备与前置条件动手之前先把环境确认好。从零构建 LLM 不是零依赖你需要准备好下面的东西。3.1 基础环境清单环境项建议操作系统Windows / Linux / macOS 均可推荐 LinuxPython3.9 或更高版本PyTorch2.0 或更高版本根据 CUDA 版本选择安装命令编辑器VS Code、PyCharm 均可GPU有 NVIDIA 显卡最好没有也能跑小模型只是训练慢磁盘空间训练数据 模型权重预留 10 GB 以上比较稳妥3.2 检查本地环境打开终端逐条执行下面的命令确认 Python 和 PyTorch 可用python --version python -c import torch; print(torch.__version__) python -c import torch; print(torch.cuda.is_available())如果 PyTorch 还没有安装可以根据自己的 CUDA 环境从官网选择安装命令。CPU 版本可以用下面的方式安装pip install torch有 NVIDIA GPU 的情况下建议先跑一条命令确认驱动可用再安装对应版本的 PyTorch。具体安装命令应以 PyTorch 官网选型为准不要直接照搬别人博客里的版本号CUDA 版本不匹配很容易踩坑。4. 从零构建小型 LLM 的核心步骤接下来进入正题。下面基于一个极简 GPT 风格模型把从零构建 LLM 的关键代码逐一拆开。这里的实现思路也是很多公开讲义共同采用的路线用 decoder-only Transformer配合因果注意力在字符级或者简单词级数据上训练。4.1 数据准备与预处理小模型和小数据是绝配。这里使用一份开源文本作为训练数据或者你自己准备一个 txt 文件也可以。注意数据来源要合规建议用无版权争议的文本比如经典公开语料。先把数据处理成模型可读的形式。最直接的方式是读入文本统计所有出现过的字符建立字符到索引、索引到字符的映射import torch import torch.nn as nn import torch.nn.functional as F # 读取训练文本替换为你的文件路径 with open(data.txt, r, encodingutf-8) as f: text f.read() chars sorted(list(set(text))) vocab_size len(chars) stoi {ch: i for i, ch in enumerate(chars)} itos {i: ch for i, ch in enumerate(chars)} def encode(s): return [stoi[c] for c in s] def decode(ids): return .join([itos[i] for i in ids]) data torch.tensor(encode(text), dtypetorch.long) print(f字符表大小: {vocab_size}) print(f数据总长度: {len(data)})这一步做的是“字符级分词”每个字符是一个 token。真实项目中会用 BPE 或 SentencePiece 做子词分词以减小词表并提高训练效率。但从零构建时字符级分词最容易理解也最容易跑通。4.2 构造训练样本上下文窗口与批量采样语言模型的输入是“前面一段 token”输出是“下一个 token”。所以要构造很多个(上下文, 目标)对。设定block_size表示上下文窗口长度batch_size表示一次训练取多少个样本block_size 64 batch_size 32 def get_batch(split): n int(0.9 * len(data)) train_data data[:n] val_data data[n:] source train_data if split train else val_data ix torch.randint(len(source) - block_size, (batch_size,)) x torch.stack([source[i:i block_size] for i in ix]) y torch.stack([source[i 1:i block_size 1] for i in ix]) return x, y这里x是输入 token 序列y是目标 token 序列两者整体右移一位。训练时模型看到x后要预测y。4.3 位置编码与 Token 嵌入Transformer 本身没有顺序概念所以需要加入位置信息。常用方式是把 token 嵌入和位置嵌入相加。这里用最简单的可学习位置嵌入class TokenEmbedding(nn.Module): def __init__(self, vocab_size, n_embd): super().__init__() self.token_embedding nn.Embedding(vocab_size, n_embd) self.position_embedding nn.Embedding(block_size, n_embd) def forward(self, idx): B, T idx.shape token_emb self.token_embedding(idx) pos torch.arange(T, deviceidx.device).unsqueeze(0) pos_emb self.position_embedding(pos) return token_emb pos_embn_embd是嵌入维度小型模型可以设成 128 或者 256。维度越大模型表达能力越强但训练时间也越长。4.4 单头注意力注意力机制是 LLM 的核心。它的作用是让每个 token 根据上下文中的其他 token 更新自己的表示。自注意力计算过程输入序列分别通过三个线性变换得到 Query、Key、Value。计算 Query 和 Key 的点积得到注意力分数。除以sqrt(d_k)做缩放防止点积过大。对当前 token 之前的注意力分数做遮蔽保证因果性。softmax 归一化后乘以 Value。实现如下class SelfAttention(nn.Module): def __init__(self, n_embd, head_size): super().__init__() self.key nn.Linear(n_embd, head_size, biasFalse) self.query nn.Linear(n_embd, head_size, biasFalse) self.value nn.Linear(n_embd, head_size, biasFalse) self.register_buffer(tril, torch.tril(torch.ones(block_size, block_size))) def forward(self, x): B, T, C x.shape k self.key(x) q self.query(x) v self.value(x) attn_scores q k.transpose(-2, -1) * (C ** -0.5) attn_scores attn_scores.masked_fill(self.tril[:T, :T] 0, float(-inf)) attn_weights F.softmax(attn_scores, dim-1) out attn_weights v return outregister_buffer把下三角矩阵注册到模型中它不会被训练但会跟着模型一起移动设备。4.5 多头注意力单头注意力只能捕捉一种注意力模式。真实大模型会用多头注意力让不同头关注不同关系。把多个头拼起来再接一个线性层class MultiHeadAttention(nn.Module): def __init__(self, n_embd, n_head, head_size): super().__init__() self.heads nn.ModuleList([ SelfAttention(n_embd, head_size) for _ in range(n_head) ]) self.proj nn.Linear(head_size * n_head, n_embd) def forward(self, x): out torch.cat([h(x) for h in self.heads], dim-1) return self.proj(out)实际项目里为了效率会把多头注意力合并到一次矩阵乘法中但手写版本用循环拼接更直观也更容易调试。4.6 前馈网络与残差连接注意力层之后通常接一个两层全连接网络做非线性变换。同时每个子层外面套一层残差连接再加 LayerNorm。残差连接让深层的梯度能顺利回传LayerNorm 让训练更稳定class FeedForward(nn.Module): def __init__(self, n_embd): super().__init__() self.net nn.Sequential( nn.Linear(n_embd, 4 * n_embd), nn.ReLU(), nn.Linear(4 * n_embd, n_embd) ) def forward(self, x): return self.net(x) class Block(nn.Module): def __init__(self, n_embd, n_head): super().__init__() head_size n_embd // n_head self.sa MultiHeadAttention(n_embd, n_head, head_size) self.ffwd FeedForward(n_embd) self.ln1 nn.LayerNorm(n_embd) self.ln2 nn.LayerNorm(n_embd) def forward(self, x): x x self.sa(self.ln1(x)) x x self.ffwd(self.ln2(x)) return xhead_size通常等于n_embd // n_head这样多头拼接后维度不变方便后续残差连接。4.7 组装 GPT 模型把上面所有模块组装起来。模型结构TokenEmbedding 得到输入表示。多层 Block 做特征提取。LayerNorm。Linear 映射到词表大小输出每个 token 的 logits。class GPT(nn.Module): def __init__(self, vocab_size, n_embd, n_head, n_layer): super().__init__() self.token_embedding TokenEmbedding(vocab_size, n_embd) self.blocks nn.Sequential(*[ Block(n_embd, n_head) for _ in range(n_layer) ]) self.ln_f nn.LayerNorm(n_embd) self.lm_head nn.Linear(n_embd, vocab_size) def forward(self, idx): B, T idx.shape x self.token_embedding(idx) x self.blocks(x) x self.ln_f(x) logits self.lm_head(x) return logits这里的n_layer是 Transformer 层数控制模型深度。小型实验可以从 2 层或 4 层开始。4.8 训练循环模型定义好了接下来写训练循环。优化器用 AdamW损失函数用交叉熵。这里直接把 logits 和目标 token 都展平后计算损失model GPT(vocab_sizevocab_size, n_embd128, n_head4, n_layer4) optimizer torch.optim.AdamW(model.parameters(), lr1e-3) torch.no_grad() def estimate_loss(): model.eval() out {} for split in [train, val]: losses [] for _ in range(50): x, y get_batch(split) logits model(x) loss F.cross_entropy(logits.view(-1, vocab_size), y.view(-1)) losses.append(loss.item()) out[split] sum(losses) / len(losses) model.train() return out for step in range(5000): x, y get_batch(train) logits model(x) loss F.cross_entropy(logits.view(-1, vocab_size), y.view(-1)) optimizer.zero_grad() loss.backward() optimizer.step() if step % 500 0: print(fstep {step}, loss {loss.item():.4f}) print(estimate_loss())训练过程中需要注意的是batch 里每个样本的 loss 是平均的view(-1, vocab_size)把[B, T, V]展平成[B*T, V]目标展平成[B*T]。这个写法在小型模型中很常见。4.9 文本生成与采样训练完成后写一个生成函数。采样时用temperature控制随机性值越小越保守值越大越多样。还可以加top-k过滤只保留概率最高的 k 个 tokentorch.no_grad() def generate(model, idx, max_new_tokens, temperature0.8, top_kNone): model.eval() for _ in range(max_new_tokens): idx_cond idx if idx.size(1) block_size else idx[:, -block_size:] logits model(idx_cond) logits logits[:, -1, :] / temperature if top_k is not None: v, _ torch.topk(logits, min(top_k, logits.size(-1))) logits[logits v[:, [-1]]] float(-inf) probs F.softmax(logits, dim-1) next_token torch.multinomial(probs, num_samples1) idx torch.cat([idx, next_token], dim-1) return idx context torch.tensor([encode(今天天气)], dtypetorch.long) print(decode(generate(model, context, 100).tolist()[0]))生成效果取决于训练数据和模型规模。字符级小模型生成的文本通常不能很好地模拟语义但已经能看出模型学到了训练数据中的一些模式。5. 训练稳定性的关键精度管理与设备选择从零构建 LLM 时很多同学把注意力放在模型结构上忽略了精度问题。实际上精度选择对训练稳定性和显存占用影响很大。5.1 fp32、fp16、bf16、tf32 是什么这是大模型训练中绕不开的一组概念精度格式位数指数位尾数位特点fp3232823最稳定精度高显存占用大fp1616510训练速度快但范围小容易溢出bf161687范围与 fp32 相同但精度低深模型训练常用tf3219810Tensor Core 专用截断格式介于 fp32 和 fp16 之间从数值稳定性角度看bf16 的指数范围和 fp32 一样不容易溢出所以在大模型训练中比 fp16 更常用。但 bf16 的尾数位更少精度略低对有些任务可能有影响。fp16 的问题在于数值范围小梯度容易溢出通常需要配合 loss scaling 使用。5.2 从零构建时怎么选如果你想让模型尽快训练稳跑起来优先用默认的 fp32。小型模型配 fp32 完全可接受也不会有太多溢出的坑。如果你要训练更大模型或者在 GPU 显存有限的情况下想扩大 batch size可以切换到混合精度from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for step in range(5000): x, y get_batch(train) optimizer.zero_grad() with autocast(): logits model(x) loss F.cross_entropy(logits.view(-1, vocab_size), y.view(-1)) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()这段代码在支持混合精度的 GPU 上有效。如果你用 CPU 训练或者显卡不支持混合精度直接跑 fp32 就行。5.3 显存占用观察方法训练过程中可以用下面两条命令实时观察显存占用nvidia-smi也可以看进程号对应的显存变化。更精细的做法是在训练脚本里轮询 PyTorch 的显存统计print(torch.cuda.memory_allocated() / 1024 ** 2, MB)需要注意显存占用主要受模型参数量、batch size、序列长度影响。想要降低显存占用可以减小 batch size、缩短 block_size、降低模型维度或者开启梯度累积。6. 训练后如何部署保存、加载与接口服务模型训练完成不部署起来意义少一半。从零构建 LLM 之后下一步是把模型权重保存下来然后加载权重对外提供服务。6.1 保存和加载模型建议保存状态字典而不是整个模型对象torch.save(model.state_dict(), tiny_llm.pt)需要重新加载时model GPT(vocab_sizevocab_size, n_embd128, n_head4, n_layer4) model.load_state_dict(torch.load(tiny_llm.pt, map_locationcpu)) model.eval()保存时最好把 tokenizer 的映射也一并保存否则推理时缺少字符映射表模型无法正确解码。6.2 用 FastAPI 包装成接口模型加载后可以用 FastAPI 快速暴露一个生成接口from fastapi import FastAPI from pydantic import BaseModel app FastAPI() class GenerateRequest(BaseModel): prompt: str max_new_tokens: int 100 temperature: float 0.8 top_k: int | None 40 app.post(/generate) def generate_endpoint(req: GenerateRequest): context torch.tensor([encode(req.prompt)], dtypetorch.long) output_ids generate(model, context, req.max_new_tokens, req.temperature, req.top_k) text decode(output_ids.tolist()[0]) return {text: text}启动服务uvicorn api_server:app --host 127.0.0.1 --port 8000然后可以用 curl 测试curl -X POST http://127.0.0.1:8000/generate \ -H Content-Type: application/json \ -d {prompt: 你好, max_new_tokens: 50}这里的encode和decode必须和训练时保持一致。接口服务上线时建议只监听127.0.0.1避免暴露到外网同时加上简单的鉴权或访问控制防止接口被恶意调用。6.3 批量任务推理阶段做批量生成时可以把多条 prompt 放在同一个 batch 里利用 GPU 并行能力。最简单的方式是在批处理脚本里循环调用更高效的方式是把输入拼接成同一个 tensor注意每个样本的上下文长度不能超过block_size。批量任务的架构建议{ input_dir: ./prompts, output_dir: ./outputs, batch_size: 8, max_new_tokens: 128, temperature: 0.8 }每次处理一批输入把结果写入独立文件同时记录失败原因和重试次数。7. 常见问题与排查方法这部分是实际动手最容易踩坑的地方整理成表格方便查阅。问题现象可能原因排查方式解决方案训练 loss 不下降学习率过大或过小打印每个 step 的 loss观察是否震荡调整学习率或者换用 AdamW 默认参数生成结果重复率太高temperature 太低或模型太小提高 temperature 到 0.8-1.0增加 top-k 采样扩展训练数据显存不足 OOMbatch_size 太大或 block_size 太长观察 nvidia-smi 显存占用减小 batch_size 或 block_size开启混合精度加载权重时报尺寸不匹配模型结构参数不一致对比保存时和加载时的 n_embd、n_head、n_layer重新创建相同结构的模型实例CUDA 不可用驱动或 PyTorch 版本不匹配torch.cuda.is_available()返回 False按官方网站重新安装对应 CUDA 版本的 PyTorch生成结果全是一个字符模型训练不充分或 vocab 构造错误检查训练 loss 是否仍然很高增加训练步数检查 encode/decode 是否对应推理速度慢模型在 CPU 上跑或序列长度很大观察推理时的 CPU/GPU 占用换 GPU 推理缩短生成长度减小模型层数接口返回超时生成长度太大测试不同 max_new_tokens 的耗时服务端设置合理超时或者加入任务队列生成的文本不连贯字符级模型表达力不足检查训练数据规模和模型层数使用 BPE 分词或换用预训练模型微调「生成的文本不连贯」这个问题在小模型训练中很常见不是代码 bug而是模型容量和数据规模的限制。如果期望生成自然流畅的长文本就需要更大更高质的数据集以及更复杂的子词分词方案。8. 最佳实践与学习建议以下建议来自常见的从零构建 LLM 的实践不是某个模型的特调方案适合所有读者参考。8.1 先把流程跑通再追求效果第一次动手时模型参数不要设太大。用 2 层 Transformer、128 维嵌入、几千条训练数据先跑通整个训练和生成流程。确认流程没问题后再逐步扩大模型规模。否则一上来就训练大模型遇到问题很难判断是架构写错还是工程配置问题。8.2 保留一份最小可运行基线把已经跑通的代码保存为一个最小可运行版本。之后改动模型架构、增加新功能时以这份基线做对比能快速定位是改代码引入的问题还是原本就存在的问题。8.3 目录管理要规范训练数据、模型权重、生成结果、日志文件要分开存放。一个简单的目录结构是这样的project/ ├── data/ │ └── train.txt ├── checkpoints/ │ └── tiny_llm.pt ├── outputs/ │ └── generated.txt └── src/ ├── model.py ├── train.py ├── generate.py └── api_server.py8.4 批量任务要加日志和重试批量生成或批量训练时别把输出一次性全部放在内存里。逐条写入文件记录每一条的处理状态遇到失败跳过并记录原因后续再重试。8.5 合规提醒文本生成模型如果训练数据来自网络内容要确认数据的版权许可。涉及人物肖像、个人声音、真实姓名等信息时必须获得授权。模型生成的内容不能用于制造虚假信息、传播不实内容或规避平台安全限制。发布模型、接口服务或生成结果之前要做内容审核和效果复核。9. 总结与下一步这篇文章带着你从数据准备、Tokenizer、Transformer 架构、注意力机制、训练循环、文本生成走到接口部署覆盖了从零构建 LLM 的完整主链路。核心收获是理解 GPT 风格模型的结构和训练流程能在本地跑通一个可生成文本的小型模型。最先应该验证的是训练 loss 是否下降这是代码是否正确的最重要信号。最容易踩的坑是模型构造参数和加载权重参数不一致以及 CUDA 版本不匹配导致 GPU 不可用。之后再去看 BPE 分词、RoPE 位置编码、KV Cache、混合精度训练这些工程优化会顺利很多。如果接下来想继续扩展可以从这几个方向入手用 BPE 或 SentencePiece 替换字符级分词提升模型对真实文本的处理能力加入更大的数据集训练更长时间观察模型生成质量的提升引入 LoRA 等微调方法在预训练模型上做参数高效微调学习向量化推理和批量调度把模型服务优化到适合多用户调用的状态。
返回列表