news 2026/8/28 14:59:06

从零构建GPT风格LLM:Python与PyTorch实现Transformer大语言模型实战

作者头像

张小明

前端开发工程师

1.2k 24
文章封面图
从零构建GPT风格LLM:Python与PyTorch实现Transformer大语言模型实战

深度学习和 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-smi

nvidia-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 的tokenizerstransformers中的分词器。这样能把“数据问题”和“模型问题”分开排查。

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

这段代码的关键点是:

  • stoiitos完成字符到 id 的双向映射。
  • __len__控制了最多能取多少个样本。
  • xy的形状都是(context_length,),训练时由 DataLoader 自动堆叠成 batch。

3.4 数据准备阶段最容易踩的坑

字符级分词的主要坑是训练和推理时使用不同的映射表。训练时模型学过的是stoi,推理时如果重新生成映射表,字符编号就变了,输出会完全错乱。

解决方式:训练结束后把charsstoiitos一起保存到 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向量维度1282048 以上
num_heads注意力头数416 或 32
num_layersTransformer 层数2 到 4几十
context_length最大上下文长度64 或 1284096 以上

参数不是越大越好。显存固定时,embed_dimcontext_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 足够稳定,但显存占用高、计算速度慢。为了提速和减少显存,需要了解四种常见精度。

精度占用字节指数位尾数位特点适用场景
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, f"checkpoint_{step}.pt")

恢复训练时,先加载模型结构,再加载model_stateoptimizer_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是最常见的报错。出现这类错误时,按以下顺序检查:

  1. 输入形状是否是(batch, context_length)
  2. token id 是否在0 ~ vocab_size - 1之间。
  3. position_embedding的输入是否只用了 0 到T-1
  4. 输出 logits 的最后一维是否等于vocab_size
  5. 损失计算时目标值是否在合法范围内。

7.2 loss 不下降的排查路径

loss 不下降需要从数据、模型、训练三个层面排查:

检查项操作
数据是否正确打印一批xy,确认输入输出对齐
模型是否处于 train 模式检查是否有 Dropout 影响
学习率是否合适太大会震荡,太小几乎不动
是否做了正确的掩码是 causal mask,不是 padding mask
梯度是否正常检查参数的grad是否存在,是否为 0 或 NaN

7.3 显存不足的解决方向

显存不足(OOM)不代表代码错误,更多是资源与模型不匹配。处理顺序:

  1. 减小batch_size,优先看是否仍 OOM。
  2. 减小context_length
  3. 减小embed_dimnum_layers
  4. 开启梯度累积,用小 batch 模拟大 batch。
  5. 使用混合精度autocast减少 fp32 占用。

7.4 fp16 训练出现 NaN

fp16 训练出现 NaN 是一项高频问题。可能原因包括:

  • 学习率过高,导致梯度过大。
  • loss 本身过大,梯度溢出。
  • 模型内部出现sqrtlog的数值越界。
  • GradScaler没有正常调用。

可以先切回 fp32 训练,确认这是数据问题还是精度问题。如果 fp32 正常而 fp16 NaN,则优先缩小学习率,并检查是否在适当位置使用autocast

8. 从零构建 LLM 的可复用清单与生产建议

8.1 学习阶段清单

每次从零构建 LLM,都可按这份清单自查:

  • 文本转 token id 后,输入输出是否右移对齐。
  • 位置编码长度是否大于等于context_length
  • causal mask 是否只允许看到前文。
  • loss 初始值是否接近log(vocab_size)
  • 训练过程中 loss 是否稳步下降。
  • 保存 checkpoint 时是否同时保存stoiitos
  • 生成时是否调用model.eval()
  • 生成时输入是否裁剪到context_length

8.2 生产环境的额外保障

在真实项目中,模型还需要关注数据质量和运行稳定性:

  • 训练数据必须清洗,避免重复文本、敏感内容、空文档。
  • 日志中记录每次训练的 loss、学习率、显存占用、分词器版本。
  • checkpoint 至少保留最近两份,方便回滚。
  • 参数配置外置化,用 YAML 或 JSON 管理模型大小、训练步数等。
  • 推理服务需要输入长度限制、超时控制和并发限制。

8.3 下一步扩展方向

跑通一个迷你 LLM 后,扩展方向很明确:

  • 使用tokenizers库实现 BPE,替代字符级分词。
  • 引入 KV cache 加速推理。
  • 使用分布式训练框架训练更大模型。
  • 接入 Hugging Facetransformers,把自定义模型转成标准接口。
  • 做模型微调实验,观察instruction tuning对输出风格的影响。

从零构建 LLM 最大的价值,不在于写出一个能媲美商用大模型的系统,而在于让你真正理解每一层张量在做什么、每一处配置改变会影响什么。后续阅读任何开源模型源码时,对照这份实现去理解注意力、精度和损失函数,就会比单纯背诵 API 有效得多。

版权声明: 本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若内容造成侵权/违法违规/事实不符,请联系邮箱:809451989@qq.com进行投诉反馈,一经查实,立即删除!
网站建设 2026/8/28 14:57:54

C++函数模板:从泛型编程基础到高级实战应用

1. 项目概述&#xff1a;为什么我们需要函数模板&#xff1f;干了这么多年C&#xff0c;从学生时代到带团队做项目&#xff0c;我见过太多重复的代码。最典型的就是&#xff0c;为了处理不同类型的数据&#xff0c;程序员不得不写一堆功能几乎一模一样、只是参数类型不同的函数…

作者头像 李华
网站建设 2026/8/28 14:52:50

Linux 常用命令及权限管理练习指南

目录 1.使用Linux常用命令 a.启用计算机后用pwd查看当前所在目录 b.用ls列出此目录中的文件和目录 c.在当前目录创建测试目录test d.利用ls&#xff0c;确认创建成功 e.进入test目录&#xff0c;并利用pwd查看 f.利用touch创建空文件newfile g.用ll命令列出所有文件 2…

作者头像 李华
网站建设 2026/8/28 14:48:05

AI Agent 办公自动化实战:从豆包工作看飞书多维表格与机器人开发

豆包工作这类 Agent 产品的出现&#xff0c;正在把办公软件从一个“工具型平台”变成“智能执行平台”。本文会从字节跳动发布豆包工作、并与飞书深度打通这一产品动态出发&#xff0c;拆解 AI Agent 在办公协作场景中的技术定位&#xff0c;然后落到工程实践&#xff1a;如何基…

作者头像 李华