ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

用PyTorch从零构建GPT:核心组件与工程实践

用PyTorch从零构建GPT:核心组件与工程实践 今天文章的主题是“用 PyTorch 从零构建 GPT”。开头不绕弯子先从我一次带新人跑代码的真实场景说起因为那个场景几乎浓缩了所有初学者会遇到的问题。大概是一年前有个刚入 NLP 方向的同学问我他想自己写一个 GPT但照着开源代码敲了一遍前向传播还是跑不通。他把报错信息发给我我发现问题根本不在模型结构上而是他用的 PyTorch 版本和代码里的F.scaled_dot_product_attention不兼容。他折腾了两天最后把 PyTorch 降级到 2.0 以上问题就消失了。这种经历其实很常见。很多人以为从零构建 GPT 的难点在“注意力机制”“Transformer 架构”“训练策略”这些听起来很高级的概念上。但真正动手以后你会发现拦住你的往往不是算法理解而是环境、版本、数据类型、设备分配这些看起来特别不起眼的问题。这篇文章我想讲的不是把 GPT 的源码逐行抄一遍给你看。我更想和你聊清楚如果你真的想用 PyTorch 从零构建一个 GPT你最应该把精力花在什么地方哪些坑是可以提前避开的什么样才算真正“跑通”了以及从“能跑通”到“能用”之间还差哪些关键拼图。这个主判断先放在这里从零构建 GPT 这件事真正锻炼你的不是 Transformer 理论的背诵而是一套工程化拆解能力——把论文细节变成模型代码把模型代码变成可训练流程再把可训练流程变成可控的实验。任何一个环节掉链子你都会卡在原地。1. 先搞清楚“从零构建 GPT”到底是在构建什么很多教程喜欢把 GPT 描述成一栋宏伟的建筑但实际动手时你是在一块一块地砌砖。你不需要真的去实现一个 GPU 矩阵乘法库也不需要自己写 CUDA 内核。你要做的是用 PyTorch 这种深度学习框架提供的“积木”把论文《Attention Is All You Need》和 GPT 系列的技术报告中描述的结构一块一块还原出来。1.1 它不是一个“大”模型而是一套有边界的结构如果你搜索“大语言模型”看到的都是几百亿参数、几千张 GPU 之类的新闻。但你自己动手用 PyTorch 构建 GPT 时完全不需要追求“大”。你完全可以先写一个参数只有几百万的小模型放在 CPU 上也能跑甚至在普通的笔记本上就能完成前向传播和过拟合测试。这里有一个非常重要的认知转变GPT 并不是“因为大才叫 GPT”而是它采用了 GPT 这种架构。架构决定了模型如何组织信息流参数规模决定了它能装下多少知识。你从零构建时核心目标是把 GPT 架构中的几个关键组件准确实现出来Tokenizer分词器把文本切分成 token再把 token 映射成整数 ID。Token Embedding Position Embedding把 token ID 变成向量并把位置信息编码进去。Transformer Decoder Block这是重中之重包含多头自注意力机制、交叉注意力严格来说 GPT 的 Decoder Block 里没有交叉注意力这一点后文会细说、前馈神经网络和层归一化。输出投影层把最后一个隐藏状态映射到词表大小的概率分布上。很多人第一次接触会搞混一个点GPT 用的 Transformer Decoder和原始的 Transformer Decoder 并不完全一样。原始 Transformer 的 Decoder 里有交叉注意力cross-attention用来访问 Encoder 的输出而 GPT 里的 Decoder Block 只有 masked self-attention 和一个前馈网络。这个区别会导致你对结构的理解完全不同。注意写代码前先想清楚你构建的 GPT 是哪种形态。如果参考“GPT 风格的 decoder-only”结构就不要把 cross-attention 硬塞进去。1.2 结构看清了才能知道每一步代码在做什么当我拿到一份从零构建 GPT 的代码时我通常会按顺序检查四个模块是否存在且正确输入处理模块原始文本能否正确变成[batch_size, seq_len]的整数张量。位置信息注入是用可学习的 Positional Embedding还是用 Sinusoidal 编码还是用 RoPE旋转位置编码。注意力掩码GPT 是自回归模型必须保证当前位置的注意力只能看到前面位置不能看到未来。这个掩码做错了整个训练过程会崩掉但不报错因为损失依然会下降只是模型变成了“作弊”模型。分类头和损失函数输出维度是否等于词表大小损失是否在忽略填充位置的前提下计算。如果你能把这个 4 个点讲清楚写成注释和文档那这篇文章的真实价值就出来了。因为大多数教程只会给你看完整代码而不会告诉你“为什么这段代码要这样组织”。2. 环境搭建不是简单的 pip install而是版本对齐工程刚才我提到有过一个同学卡在 PyTorch 版本上两天没进展。这其实不是个例。从热搜词里能看到大量人都在搜“pytorch安装”“pytorch环境搭建”“anaconda配置pytorch环境”“ubuntu系统下载pytorch教程”“pytorch cu130”这类关键词。这说明环境搭建本身就是从零构建 GPT 的第一道坎而且拦住了很多人。2.1 为什么版本不对模型代码会出各种怪问题PyTorch 的 API 变化速度极快。举个例子torch.nn.functional.scaled_dot_product_attention是在 PyTorch 2.0 引入的高效注意力实现。如果你拿到的参考代码用了这个函数但你的 PyTorch 还是 1.x就会报一个不那么明显的错误甚至有时候只是慢不是直接崩。再比如torch.compile它在 PyTorch 2.0 被引入可以明显加速训练但它对 Python 版本、CUDA 版本和 GPU 型号都有要求。如果你是在 Windows 上用 pip 装的 CPU 版 PyTorch很多加速能力都用不上。所以在开始构建模型之前先确定几个版本而不是拿到最新版就装Python 版本建议 3.9 到 3.11太新有时候会遇到 wheel 不匹配PyTorch 版本建议 2.x除非你刻意学习 1.x 的旧代码CUDA 版本取决于你的驱动支持不要盲目选 cu130依赖包版本transformers、tiktoken、numpy、datasets 等2.2 我的建议流程先固定依赖再写代码这里我用一个通用流程来说明具体版本号会变化但顺序是稳定的# 1. 创建独立虚拟环境避免污染全局 Python conda create -n gpt-from-scratch python3.10 -y conda activate gpt-from-scratch # 2. 安装 PyTorch根据你的机器选择 CPU 版或 CUDA 版 # CPU 版先保证能跑通代码 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu # 3. 安装训练和分词相关依赖 pip install numpy tiktoken datasets matplotlib如果你的机器有 NVIDIA 显卡并且已经安装了合适的驱动你可以去 PyTorch 官网选择对应 CUDA 版本的安装命令。但这里有一个经验第一次用 PyTorch 从零构建 GPT不要直接上 CUDA 版。先用 CPU 版把完整流程跑通确认模型代码没有逻辑错误再安装 GPU 版加速训练。为什么因为 CPU 和 GPU 的报错信息不一样把两个变量混在一起排查你会浪费大量时间。环境搭建的核心逻辑是把你跑通的实验环境记录下来形成确定性的依赖清单pip freeze requirements.txt这样你未来重跑、换机器、给别人复现都能回到同一起点。这一步必须在第一行模型代码之前完成。3. 从零写一个最小 GPT 结构核心组件逐个拆解当环境稳定下来就可以开始写模型代码了。这个阶段我的建议是不要照抄完整开源项目而是用最少的代码把核心结构写出来。你可以不用考虑训练效率不用优化显存甚至可以不用考虑 GPU。你只需要一个能在小规模数据上过拟合的最小模型。3.1 最简结构的“零件清单”在开始写代码前先像列购物清单一样把零件列出来。下面这个清单是从零构建 GPT 时至少需要涉及的部分模块作用常见坑点分词器文本和 token ID 互相转换词表大小不一致导致 embedding 维度错位词嵌入层将 token ID 映射成向量忘记设置padding_idx或词表大小参数位置编码给模型提供序列顺序信息序列长度和位置编码长度不一致多头自注意力提取上下文关系mask 形状错误、attn_mask广播问题前馈网络非线性变换残差连接维度不匹配层归一化稳定训练归一化的维度和形状对不对输出投影层映射到词表概率分布权重与嵌入层是否共享不共享也没错这些零件组合起来从结构上看GPT 就是一个把“嵌入层 N 个 Transformer Decoder Block 输出层”串起来的网络。Block 里最核心的组件就是 masked multi-head self-attention。3.2 自注意力机制为什么 GPT 必须要 mask自注意力的公式看起来简单Attention(Q, K, V) softmax(Q K^T / sqrt(d_k)) V但在 GPT 这样的自回归模型里Q K^T之后不能直接 softmax。因为在生成第 t 个 token 时模型不能看到第 t1、t2 等未来 token。这需要在计算注意力分数后、softmax 之前把未来位置的值替换成一个极大的负数比如-inf这样 softmax 之后这些位置的权重就会变成 0。如果你用了F.scaled_dot_product_attention它可以通过is_causalTrue参数直接实现因果掩码非常方便。但为了理解机制我建议你至少手写一次完整的注意力函数哪怕它慢一点import torch import torch.nn as nn import torch.nn.functional as F class CausalSelfAttention(nn.Module): def __init__(self, embed_dim, num_heads, dropout0.0): super().__init__() assert embed_dim % num_heads 0 self.embed_dim embed_dim self.num_heads num_heads self.head_dim embed_dim // num_heads self.qkv nn.Linear(embed_dim, 3 * embed_dim) self.out_proj nn.Linear(embed_dim, embed_dim) self.dropout nn.Dropout(dropout) def forward(self, x): B, T, C x.size() qkv self.qkv(x) # [B, T, 3*C] q, k, v qkv.chunk(3, dim-1) # 拆分多头 q q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) k k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) v v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) # 缩放点积注意力 attn q k.transpose(-2, -1) * (self.head_dim ** -0.5) # 因果掩码下三角矩阵为 0上三角为 -inf causal_mask torch.tril(torch.ones(T, T, devicex.device, dtypetorch.bool)) attn attn.masked_fill(~causal_mask, float(-inf)) attn F.softmax(attn, dim-1) attn self.dropout(attn) y attn v # [B, num_heads, T, head_dim] y y.transpose(1, 2).contiguous().view(B, T, C) return self.out_proj(y)这段代码是“示例结构”不是某个开源项目里的原版。但它包含了你理解 GPT 必需的核心逻辑QKV 投影、多头拆分、缩放、因果掩码、softmax、attention 加权、输出投影。这里有一个很容易踩的坑causal_mask只在 T 等于当前序列长度时是对的。如果训练时输入序列长度固定这没问题但推理时如果想生成长度大于 T 的序列位置编码和掩码都需要外推这时就会出问题。这也是为什么很多 GPT 实现会使用相对位置编码或 RoPE。3.3 从一个小模型跑通到过拟合这才是真正的“从零构建”写完结构以后不要直接去找大型语料。先用一个很小的文本文件比如几千个字的莎士比亚作品片段或者你自己写的一段日志把它切成 token然后用一个小模型去过拟合这一段数据。你可以用一个非常小的配置config { vocab_size: 256, # 比如用字节级分词词表就是 0-255 block_size: 128, # 最大上下文长度 n_embd: 64, # 嵌入维度 n_head: 4, # 注意力头数 n_layer: 2, # Transformer Block 数量 batch_size: 4, learning_rate: 3e-4, }这种规模的模型在 CPU 上也能跑。关键目标是让训练 loss 持续下降最后在训练集上过拟合到非常小的数值。如果过拟合都做不到说明你的模型代码或数据管道有问题。这是最好的调试方式。一个常见经验是如果你连一小段文本都无法过拟合问题大概率出在输入 ID 和标签的对齐上。检查一下你生成的训练样本中输入序列和标签序列是不是错位了一个 token。GPT 训练时通常输入是第 i 到 T-1 个 token标签是第 i1 到 T 个 token。4. 训练一个迷你 GPT数据、损失与采样流程模型代码跑通后就要进入训练阶段。训练过程就像把零件组装好的汽车第一次点火。很多新手在这一步会发现模型不是“不能跑”而是“跑得非常诡异”loss 不下降、loss 变成 NaN、模型生成的内容全是重复 token。这些问题的原因往往不只在模型结构上而是在数据准备和训练配置上。4.1 大语言模型下载下来是什么先弄清数据和权重的区别热搜词里有一句“大语言模型下载下来是什么”非常能代表新手的困惑。很多人以为下载一个模型就像下载一个软件双击就能打开。但实际上你下载的是一个权重文件里面全是浮点数它必须配合模型结构代码和分词器文件才能运行。从零构建 GPT 时你也会碰到同样的问题。你的模型训练完之后需要把权重保存下来。通常做法是torch.save(model.state_dict(), mini_gpt.pt)这不是一个完整的模型包。它只是把模型参数以字典形式保存了。如果别人拿到这个.pt文件但没有你的model.py他根本无法把权重加载回去。所以如果你的最终目标是分享一个别人能用的模型不仅要保存权重还要保存配置文件、分词器文件和一份加载脚本。在 PyTorch 生态里更推荐的是用.pt或.pth保存完整的模型 checkpoint包括优化器状态、训练步数、loss 曲线等。如果你只是保存权重那就把它视为“推理权重”加载时要注意构建完全相同的模型结构。4.2 训练过程中的 loss 曲线就是你的仪表盘GPT 训练时你关注的核心指标通常有三个训练集 loss反映模型对训练数据的拟合程度。验证集 loss反映模型的泛化能力。生成效果语言模型特有的评估方式直接看模型在给定前缀后能生成什么。如果训练集 loss 一直不下降先检查学习率是否合适再检查输入数据是否存在大量噪声。如果是训练集 loss 下降但验证集 loss 升高那就是过拟合可以考虑增加数据、减小模型或者加 dropout。如果生成结果是乱码或无限重复可能是分词器配置错误、重复惩罚设置不当或者模型还没训练充分。我一般会建议用一个小工具来记录 loss 曲线比如matplotlib或在CSV文件里记录。这一个动作虽然简单但它能让你在训练过程中快速判断模型状态而不是等几个小时后才发现模型早就跑偏了。4.3 从“训练完成”到“模型能生成文字”还有最后一步写到这里很多教程会到此为止——损失下降了模型保存了任务“完成”了。但真正让 GPT 有可用价值的是生成过程。生成过程也是自回归的把当前 token 序列输入模型取最后一个位置的输出用 softmax 得到下一个 token 的概率分布然后从中采样一个 token接入序列继续循环。最基础的采样方法是贪心解码def generate(model, idx, max_new_tokens, block_size): model.eval() for _ in range(max_new_tokens): idx_cond idx[:, -block_size:] # 只保留最近的 block_size 个 token logits model(idx_cond) logits logits[:, -1, :] # 只取最后一个位置 probs F.softmax(logits, dim-1) next_token torch.argmax(probs, dim-1, keepdimTrue) idx torch.cat([idx, next_token], dim1) return idx贪心解码每次选概率最高的 token结果往往非常平淡而且容易重复。更自然的方法是引入温度参数和 top-k 采样。这些策略在真正使用模型时几乎必须掌握。但在从零构建阶段先跑通贪心解码就好不要一开始就上复杂采样否则你会分不清问题是出在模型没训练好还是采样策略不对。注意生成时序列长度超过block_size会导致问题。最常见的是位置编码越界或因果掩码维度错误。这个坑非常隐蔽很多人调了半天最后发现输入序列超过了训练时的最大上下文长度。5. 从“能跑”到“能用来做事”你离生产还差几块拼图当你成功训练完一个小型 GPT并且它能根据前缀生成看起来有点意思的文本之后你很容易产生一个错觉我已经会构建大语言模型了。从学习和锻炼工程能力的角度看确实有很大收获但距离真正“可用”还有很长的路。5.1 工程化缺失项数据质量、训练效率、服务化一个真实的大语言模型从训练到上线涉及的东西远不止模型结构本身。热搜词里很多人关注的“本地部署大语言模型”“大语言模型界面”就对应了这条链路。从工程视角看把训练代码跑通只是一小步。后续你还需要处理数据清洗与配比从原始语料到高质量训练数据需要去重、过滤垃圾文本、控制各类数据比例。这一步对模型最终效果的影响往往比模型结构和训练参数还大。训练加速与显存优化梯度累积、混合精度、分布式训练、模型并行、ZeRO 优化等。你训练的小模型用不上但这是从小模型走向大模型的必经之路。评估体系不能只看 loss 下降还要有评估集评测、人类反馈、对齐等。推理服务把模型部署成 API 接口做显存管理、并发控制、请求排队、流式输出。这些内容看起来和“从零构建 GPT”无关但其实是同一个命题。因为一旦你确定以后要长期在这条路上走从第一天就应该清楚模型代码只是一张门票完整的大语言模型工程还有全景图和路线图。越早意识到这一点你就越不会把“训练了一个小模型”当成最终的终点。5.2 如何继续深入一个可复用的进阶路线如果你已经完成了文章前面所有步骤下一个阶段可以按下面的路径继续往前探索。这也是我比较推荐的一个“三步延伸法”复现一个开源的中型 GPT 结构比如 1 亿参数左右的模型。找一个公开的教程或代码库自己照着思路写一遍先不追求完全一致重点是理解每个组件为什么这样设计。自己构造一个真实任务不要只用莎士比亚文本做学习玩具找一个你能拿到的、有明确领域结构的数据集比如代码补全、GitHub Issue 摘要、技术问答等训练一个针对特定任务的 GPT。从单机训练切到分布式训练哪怕你只有一台机器、两张显卡也可以体验DistributedDataParallelDDP。这会让你对“大语言模型为什么需要多卡”和“数据并行”有真正的直觉。这三步看起来每一步都不复杂但它能推动你从一个“会用 PyTorch 写模型的人”变成一个“能判断技术方案、能做实验设计、能排查训练链路问题”的人。这才是从零构建 GPT 真正值得训练的能力。6. 常见问题排查遇到错误先修思想和环境再扣代码最后整理一份适合放在手边的排查思路。不针对某一段代码而是对整个从零构建 GPT 的过程。如果你在实验过程中卡住按照下面的顺序逐层排查大概率能比盲目搜报错信息更快定位问题。先确认报错来自哪一层是数据管道、模型构造、前向传播、loss 计算、梯度回传还是采样生成不同层的报错完全指向不同原因。再确认环境是否对齐Python 版本、PyTorch 版本、CUDA 版本、transformers版本、tiktoken版本。版本不一致导致的怪问题靠读代码是找不出来的。接着检查输入输出形状对每个关键张量在调试模式下把.shape打印出来或者用assert检查形状。形状不一致是最容易发现也最容易被忽略的错误。然后核对掩码和数据对齐因果掩码的维度对不对训练样本中输入和标签是否错位block_size是否覆盖了所有序列长度。这一步出问题很多情况下模型也能训练但效果极差。最后检查数值稳定性loss 是否出现nan权重是否变成nan学习率是不是太大梯度是否爆掉。如果是优先降低学习率或者加入梯度裁剪。这部分如果你已经顺利跑到生成阶段大概率不会遇到灾难级的错误。但未来你写完更复杂的代码、迁移到其他框架或更高版本的 PyTorch 时这套排查链路可以反复复用。回到最开始的那个判断。用 PyTorch 从零构建 GPT本质上是一个工程训练而不是理论背诵。你可以在一天内看完注意力机制的讲解但要自己写出一段既能训练又能生成文本的最小 GPT通常需要几天甚至更久而且大部分时间会花在环境修复、形状调整、掩码错误和数据对齐上。这个时间花得值。因为只有真正写过一遍你才会理解为什么大语言模型需要海量数据、为什么自回归结构能生成自然语言、为什么一个看似简单的 Transformer Block 能装下那么多知识。这些理解比单纯跑一个开源代码能带来的收获要深得多。所以如果你正在准备用 PyTorch 构建自己的第一个 GPT我的建议只有一句话先装好环境再把最小模型跑通让它在小数据上过拟合一次。这一步做到了你的从零构建之路就已经成功了三分之一。剩下的三分之二都是在学会面对不确定性时不断修正自己判断的过程。
返回列表