news 2026/9/6 5:14:21

用PyTorch从零构建GPT:Transformer核心原理与实战

作者头像

张小明

前端开发工程师

1.2k 24
文章封面图
用PyTorch从零构建GPT:Transformer核心原理与实战

先问一个问题:你平时调用 GPT 接口、用 LangChain 做应用、或者浏览各种大模型的论文解读时,有没有想过 GPT 内部的 transformer 层到底是怎么写的?注意力机制里的 Q、K、V 矩阵是如何把一句话逐步编码成上下文向量的?当我们在说“大语言模型”时,它究竟是一堆什么样的张量运算和参数更新?

市面上关于大语言模型的文章很多,但大多是“怎么调用 API”“怎么写提示词”,真正从零开始用 PyTorch 搭建 GPT 结构的内容反而比较零散。本文将围绕“用 PyTorch 从零构建 GPT”这一主题,把 GPT 的核心原理、模型结构、训练流程、文本生成推理完整走一遍。

文章会以实战代码为主,适合这几类读者:

  • 已经会用 PyTorch 搭建简单神经网络,想进一步理解 Transformer 和 GPT 源码的开发者。
  • 准备阅读 GPT、LLaMA 等大模型源码,但被 attention、layer norm、位置编码绕晕的新手。
  • 需要在一个小规模数据集上自己做“类 GPT”实验的学生或算法工程师。

通过本文,你会掌握 GPT 的关键技术点,并拿到一份可以在本地训练和推理解的完整代码。

1. 背景与核心概念

1.1 从语言模型说起

“大语言模型”这个词听起来很庞大,但把它拆到最小单元,核心任务仍然是:给定一段已知文本,预测下一个词。比如输入“今天天气真”,模型要计算下一个字是“好”“差”“热”还是“冷”的概率分布。

传统语言模型基于 n-gram 统计,只能考虑附近几个词,无法处理长距离依赖。后来循环神经网络(RNN、LSTM)能建模序列,但存在训练速度慢、长期记忆衰减的问题,而且难以并行处理序列中的所有位置。

2017 年《Attention Is All You Need》提出 Transformer,彻底改变了序列建模的方式。Transformer 不再像 RNN 那样按时间步逐个读取输入,而是把所有 token 一次性送入网络,通过自注意力机制计算 token 之间的相关度。这种并行计算方式大幅提升了训练效率,也让长距离依赖建模成为可能。

GPT 全称是 Generative Pre-trained Transformer,属于 Transformer 的 decoder 分支,只保留自回归结构。它的训练思路分两阶段:

  • 第一阶段在海量语料上做无监督预训练,目标是预测下一个 token;
  • 第二阶段可以用带标签数据做有监督微调(SFT),让模型学会用户期望的回答方式。

1.2 GPT 与 Transformer 的关系

Transformer 原论文包含 Encoder 和 Decoder 两个部分。Encoder 负责把输入序列编码成语义向量,适合理解类任务;Decoder 负责逐 token 生成输出,适合生成类任务。

GPT 只使用 Decoder 的 masked self-attention 部分。它不需要 Encoder,因为语言建模的目标就是从历史 token 预测下一个 token,当前时刻只能看到过去的信息,不能看到未来的 token。这种限制由 attention mask 实现。

这里要重点区分三个概念:

  • Encoder:双向注意力,能同时看到上下文,适合 BERT 这类理解模型。
  • Decoder:单向注意力,只能看左侧历史 token,适合 GPT 这类生成模型。
  • Encoder-Decoder:完整 Transformer,适合机器翻译、文本摘要等“序列到序列”任务。

GPT 系列模型架构的核心是 Transformer decoder block,一个 block 内包含:

  • masked multi-head self-attention;
  • feed-forward network(FFN);
  • 每个子层都做残差连接和层归一化。

堆叠多个这样的 decoder block,就得到 GPT 的主干结构。模型规模增大时,主要的参数来自 embedding 层、注意力层和 FFN 层。

1.3 为什么要用 PyTorch 从零构建 GPT

现在有很多现成的大模型框架,比如 HuggingFace Transformers 里直接调用GPT2LMHeadModel就能加载预训练模型。那为什么还要自己手写一遍?

原因主要有三个。

第一,理解源码。GPT 的 huggingface 实现包含大量工程化逻辑,新手很难一眼看出核心张量流。自己写一遍,能把nn.Embeddingnn.Linearmasked attention这些组件的输入输出形状理清楚。

第二,方便实验。在做毕业设计、论文复现或者公司内部小规模预训练实验时,我们经常需要改动 attention 结构、换激活函数、改位置编码方式。如果只会调用现成接口,改动成本很高;自己写核心模型,改一行就能测试新想法。

第三,更深入理解大模型训练中的细节。比如为什么 GPT 需要独立的位置编码表?为什么推理时要缓存 KV?为什么采样时 temperature 会影响文本多样性?这些问题的答案都藏在模型结构的细节里。

2. 环境准备与版本说明

2.1 操作系统与 Python 环境

本文示例代码可以在 Windows、macOS 或 Linux 上运行。因为涉及矩阵运算和自动求导,建议先安装 Anaconda 或 Miniconda 来管理 Python 环境。

创建独立环境:

conda create -n gpt_from_scratch python=3.10 conda activate gpt_from_scratch

Python 版本建议 3.10 或更高。需要说明的是,PyTorch 对不同 Python 版本的支持范围不同,3.10 是一个非常稳妥的选择。

2.2 安装 PyTorch

PyTorch 的安装命令取决于你本机的 CUDA 版本。在开始之前,先用下面的命令确认显卡驱动和 CUDA 能力:

nvidia-smi

如果输出正常,说明有 NVIDIA GPU 可用;如果你的机器没有独立显卡,也不用担心,本文示例模型很小,用 CPU 也可以完成训练,只是速度会慢一些。

以 CUDA 12.x 环境为例,安装命令:

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

如果你只需要 CPU 版本:

pip install torch

版本说明:本文示例代码基于 PyTorch 2.x 编写,核心 API 在 1.13 及以上版本都能运行。如果你使用的是更早版本,建议升级到 2.x,因为 PyTorch 2.x 对编译优化和显存管理做得更好。

2.3 安装 tiktoken

训练 GPT 时,我们需要一个分词器(tokenizer)。tiktoken 是 OpenAI 开源的 BPE 分词库,可以让我们像 GPT 模型一样把文本切分成 token。

pip install tiktoken

安装完成后,用 Python 简单验证一下:

import tiktoken enc = tiktoken.get_encoding("gpt2") tokens = enc.encode("Hello, world!") print(tokens) print(enc.decode(tokens))

预期输出类似:

[15496, 11, 995, 0] Hello, world!

能看到文本被切分成 token ID 序列,并且可以解码回原文本,说明环境没问题。

3. GPT 核心原理拆解

在写代码之前,先把 GPT 的几个核心组件逐一拆开。不要小看这一部分,后面代码能不能看懂,关键就在这里。

3.1 Tokenizer:文本如何变成数字

计算机无法直接处理文本,神经网络也一样。我们需要把一段话变成数字序列,这个过程叫 tokenization。

一种最简单的方法是字符级切分:每个字符对应一个 ID。这种方法实现简单,但序列很长,且无法体现词和子词的语义信息。GPT 使用的是 BPE(Byte Pair Encoding,字节对编码),它先把文本按 UTF-8 编码成字节序列,然后迭代合并出现频率最高的字节对,生成一个子词词表。

使用 BPE 分词器后,一段文本会被切分成类似这样的 token:

"GPT is great" -> ["GPT", " is", " great"]

tiktokenget_encoding("gpt2")使用的就是 GPT-2 时代的 BPE 词表,词汇量 50257。在实际训练中,如果训练数据比较大,一般会重新训练一个 tokenizer;为了简化,我们这里直接使用现成的 gpt2 分词器。

3.2 Token Embedding 与 Position Embedding

分词后,每个 token 对应一个 ID,例如15496。这个整数需要转换成向量才能进入神经网络,这一步由nn.Embedding完成。

nn.Embedding本质上是一个查找表,输入形状为[batch_size, seq_len]的 token ID 矩阵,输出形状为[batch_size, seq_len, n_embd]的向量序列。

但是 Transformer 模型本身不包含任何序列顺序信息。对于注意力机制来说,把句子“我打你”和“你打我”输入进去,如果不加位置信息,模型无法区分词的先后顺序。所以需要额外加入位置编码。

GPT 使用可学习的位置编码表(learned positional embedding),每个位置都有一个独立向量,维度与 token embedding 相同。位置编码表的形状是[block_size, n_embd],其中block_size是模型支持的最大序列长度。

在代码中,最终输入模型的特征是:

x = token_embedding(tokens) + position_embedding(positions)

3.3 多头自注意力机制

注意力机制是 GPT 的核心。它的作用是让每个 token 在计算特征表示时,关注序列中其他 token 的信息。

3.3.1 Self-Attention 的 Q、K、V

自注意力(Self-Attention)的输入是同一个序列经过线性变换得到的三组矩阵:

  • Query(查询向量):表示当前 token “想查找什么”;
  • Key(键向量):表示当前 token “能提供什么”;
  • Value(值向量):表示当前 token “实际携带的信息”。

对于序列中每个 token,我们用它的 Query 向量与所有 token 的 Key 向量做点积,得到注意力分数;注意力分数经过 softmax 归一化后,对 Value 向量加权求和。

公式:

Attention(Q, K, V) = softmax(Q * K^T / sqrt(d_k)) * V

其中d_k是 Key 向量的维度,除以sqrt(d_k)是为了防止点积值过大,导致 softmax 梯度消失。

Python 伪代码如下:

import torch import torch.nn as nn import torch.nn.functional as F def scaled_dot_product_attention(Q, K, V, mask=None): d_k = Q.size(-1) scores = torch.matmul(Q, K.transpose(-2, -1)) / d_k ** 0.5 if mask is not None: scores = scores.masked_fill(mask == 0, float("-inf")) attn_weights = F.softmax(scores, dim=-1) return torch.matmul(attn_weights, V), attn_weights
3.3.2 Masked Attention:为什么 GPT 只能看到过去

GPT 是自回归模型,生成第t个 token 时,只能看到前t-1个 token。如果在训练时让当前位置看到了未来的 token,模型就相当于照抄答案,学不到真正的预测能力。

实现方式很直接:构造一个上三角矩阵,把未来位置的注意力分数设置为负无穷,softmax 之后这些位置的权重就变成 0。

例如序列长度为 4,合法的 mask 矩阵为:

1 0 0 0 1 1 0 0 1 1 1 0 1 1 1 1

i行第j列为 1 表示位置i可以关注位置j;为 0 表示不允许关注。

3.3.3 多头注意力

多头注意力就是把注意力计算分成多个头并行做。假设特征维度是 768,使用 12 个头,那么每个头的维度是 64。每个头独立计算注意力,最后拼接起来再过一层线性变换。

多头注意力的好处是让模型从不同表示子空间捕捉信息。有的头可能关注语法关系,有的头可能关注远距离指代关系。

代码中可以用一个nn.Linear同时生成全部头的 Q、K、V,再通过 reshape 拆分到不同头,这样实现高效很多。

3.4 前馈神经网络与残差、LayerNorm

每个 Transformer block 在 attention 之后还会接一个前馈神经网络(FFN),通常包含两层全连接和一个激活函数:

FFN(x) = GELU(x W1 + b1) W2 + b2

GPT 中 FFN 的中间维度一般是模型维度的 4 倍,也就是说 768 维的输入会先升到 3072 维,再投影回 768 维。

同时,每个子层都使用“残差连接 + 层归一化”的结构:

  • 残差连接(Residual Connection):让梯度能直接跨层传播,避免深层网络梯度消失;
  • 层归一化(LayerNorm):对每个样本的特征维做标准化,加速训练收敛。

注意 GPT 使用的是 pre-LayerNorm 结构,即先做 LayerNorm 再做 Attention/FFN。这与原始 Transformer 论文中的 post-LayerNorm 略有不同,但现代大模型几乎都采用 pre-LayerNorm,训练更稳定。

3.5 语言模型头与损失函数

经过多层 Transformer block 后,我们得到形状为[batch_size, seq_len, n_embd]的向量序列。要预测下一个 token,需要把向量映射回词表大小,这一层称为语言模型头(LM Head):

logits = nn.Linear(n_embd, vocab_size)(x)

得到每个位置上所有 token 的概率分布后,训练目标就是交叉熵损失,计算预测分布与真实下一个 token 的差距。

假设输入序列是[1, 2, 3, 4],模型会同时预测:

  • 给定[1],预测2
  • 给定[1, 2],预测3
  • 给定[1, 2, 3],预测4

这就是 GPT 的“并行 teacher forcing”训练方式:一次前向计算,得到整个序列所有位置的预测结果。

4. 完整实战:用 PyTorch 从零构建 GPT

现在进入本文的主体部分。我们将搭建一个 mini 版 GPT 模型,并在一个小型数据集上完成训练和文本生成。模型规模很小,完全可以在 CPU 上运行。

4.1 项目结构

建议创建如下目录:

gpt_from_scratch/ ├── data.py # 数据加载与 tokenizer 封装 ├── model.py # GPT 模型结构 ├── train.py # 训练脚本 └── generate.py # 文本生成脚本

4.2 数据准备

为了演示,我们使用一个英文小语料:莎士比亚的《哈姆雷特》部分片段,或者任何一本公开领域的英文书。这里为了代码简单,直接内置一段文本作为训练语料,你也可以用自己准备的.txt文件替换。

数据加载代码:

# 文件路径:data.py import tiktoken import torch from torch.utils.data import Dataset class TextDataset(Dataset): def __init__(self, text, block_size=128): self.enc = tiktoken.get_encoding("gpt2") self.tokens = self.enc.encode(text) self.block_size = block_size def __len__(self): return len(self.tokens) - self.block_size def __getitem__(self, idx): chunk = self.tokens[idx: idx + self.block_size + 1] x = torch.tensor(chunk[:-1], dtype=torch.long) y = torch.tensor(chunk[1:], dtype=torch.long) return x, y def load_text(): sample_text = """ To be, or not to be, that is the question: Whether 'tis nobler in the mind to suffer The slings and arrows of outrageous fortune, Or to take arms against a sea of troubles And by opposing end them. To die, to sleep; No more; and by a sleep to say we end The heart-ache and the thousand natural shocks That flesh is heir to: 'tis a consummation Devoutly to be wish'd. To die, to sleep; To sleep, perchance to dream; ay, there's the rub. """ return sample_text

这里的关键点:

  • 每个样本的输入x和标签y是错位一个 token 的序列,对应“预测下一个 token”的任务。
  • tiktoken编码后,文本中的空格和换行也会被保留为 token,不需要额外清理。

4.3 定义 GPT 模型

下面编写完整的 GPT 模型代码。代码尽量简洁,但保留 GPT 所有核心组件。

# 文件路径:model.py import torch import torch.nn as nn import torch.nn.functional as F class LayerNorm(nn.Module): def __init__(self, n_embd, eps=1e-5): super().__init__() self.eps = eps self.gamma = nn.Parameter(torch.ones(n_embd)) self.beta = nn.Parameter(torch.zeros(n_embd)) def forward(self, x): mean = x.mean(-1, keepdim=True) var = x.var(-1, keepdim=True, unbiased=False) x = (x - mean) / torch.sqrt(var + self.eps) return x * self.gamma + self.beta class CausalSelfAttention(nn.Module): def __init__(self, n_embd, n_head, block_size, dropout=0.0): super().__init__() assert n_embd % n_head == 0 self.n_head = n_head self.n_embd = n_embd self.head_dim = n_embd // n_head self.c_attn = nn.Linear(n_embd, 3 * n_embd, bias=False) self.c_proj = nn.Linear(n_embd, n_embd, bias=False) self.dropout = nn.Dropout(dropout) # 上三角掩码矩阵,用于屏蔽未来位置 self.register_buffer( "mask", torch.tril(torch.ones(block_size, block_size)).view( 1, 1, block_size, block_size ), ) def forward(self, x): B, T, C = x.shape # 一次性生成 Q、K、V qkv = self.c_attn(x) # [B, T, 3*C] q, k, v = torch.split(qkv, self.n_embd, dim=-1) # 拆分成多头: [B, n_head, T, head_dim] q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) # 缩放点积注意力 att = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5) att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf")) att = F.softmax(att, dim=-1) att = self.dropout(att) y = att @ v # [B, n_head, T, head_dim] y = y.transpose(1, 2).contiguous().view(B, T, C) y = self.c_proj(y) return y class MLP(nn.Module): def __init__(self, n_embd, dropout=0.0): super().__init__() self.fc1 = nn.Linear(n_embd, 4 * n_embd) self.fc2 = nn.Linear(4 * n_embd, n_embd) self.dropout = nn.Dropout(dropout) def forward(self, x): x = self.fc1(x) x = F.gelu(x) x = self.fc2(x) x = self.dropout(x) return x class TransformerBlock(nn.Module): def __init__(self, n_embd, n_head, block_size, dropout=0.0): super().__init__() self.ln1 = LayerNorm(n_embd) self.attn = CausalSelfAttention(n_embd, n_head, block_size, dropout) self.ln2 = LayerNorm(n_embd) self.mlp = MLP(n_embd, dropout) def forward(self, x): x = x + self.attn(self.ln1(x)) x = x + self.mlp(self.ln2(x)) return x class GPT(nn.Module): def __init__( self, vocab_size=50257, n_embd=128, n_head=4, n_layer=4, block_size=128, dropout=0.0, ): super().__init__() self.token_embedding = nn.Embedding(vocab_size, n_embd) self.position_embedding = nn.Embedding(block_size, n_embd) self.blocks = nn.Sequential( *[ TransformerBlock(n_embd, n_head, block_size, dropout) for _ in range(n_layer) ] ) self.ln_f = LayerNorm(n_embd) self.lm_head = nn.Linear(n_embd, vocab_size, bias=False) # 权重初始化 self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) def forward(self, idx, targets=None): B, T = idx.shape assert T <= self.position_embedding.weight.shape[0], "序列长度超出位置编码范围" pos = torch.arange(0, T, dtype=torch.long, device=idx.device) pos = pos.unsqueeze(0) # [1, T] x = self.token_embedding(idx) + self.position_embedding(pos) x = self.blocks(x) x = self.ln_f(x) logits = self.lm_head(x) loss = None if targets is not None: B, T, C = logits.shape logits = logits.view(B * T, C) targets = targets.view(B * T) loss = F.cross_entropy(logits, targets) return logits, loss def generate(self, idx, max_new_tokens=100, temperature=1.0): """自回归生成新 token""" self.eval() for _ in range(max_new_tokens): idx_cond = idx[:, -self.position_embedding.weight.shape[0]:] logits, _ = self(idx_cond) logits = logits[:, -1, :] / temperature probs = F.softmax(logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) idx = torch.cat((idx, next_token), dim=1) return idx

这个模型需要注意几个地方:

  • register_buffer注册的 mask 会随模型一起移动到 GPU/CPU,不需要手动管理设备。
  • c_attn一次性输出 3 倍维度的 Q、K、V,比分别定义三个线性层效率更高。
  • generate方法里只取最后一个位置的 logits 做采样,这是因为自回归模型当前位置的预测只依赖之前所有 token。
  • 采样时除以 temperature,让概率分布更尖锐或更平滑,从而控制生成文本的随机性。

4.4 编写训练脚本

训练脚本中,我们使用 AdamW 优化器。AdamW 是 GPT 训练的标配优化器,因为它对权重衰减的处理更合理,能提升泛化性能。

# 文件路径:train.py import torch from torch.utils.data import DataLoader from data import TextDataset, load_text from model import GPT def main(): # 超参数 batch_size = 8 block_size = 128 n_embd = 128 n_head = 4 n_layer = 4 learning_rate = 3e-4 epochs = 50 device = "cuda" if torch.cuda.is_available() else "cpu" text = load_text() dataset = TextDataset(text, block_size=block_size) loader = DataLoader(dataset, batch_size=batch_size, shuffle=True) model = GPT( vocab_size=50257, n_embd=n_embd, n_head=n_head, n_layer=n_layer, block_size=block_size, dropout=0.1, ).to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.01) print(f"数据集 token 数: {len(dataset.tokens)}") print(f"使用设备: {device}") for epoch in range(epochs): model.train() total_loss = 0 for x, y in loader: x, y = x.to(device), y.to(device) optimizer.zero_grad() logits, loss = model(x, y) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() total_loss += loss.item() avg_loss = total_loss / len(loader) print(f"Epoch {epoch+1:3d}/{epochs}, loss = {avg_loss:.4f}") torch.save(model.state_dict(), "gpt_model.pt") print("训练完成,模型已保存到 gpt_model.pt") if __name__ == "__main__": main()

训练过程中值得关注的点:

  • clip_grad_norm_是梯度裁剪,防止梯度爆炸。对语言模型来说,这是稳定训练的重要技巧。
  • 数据量很小时,loss 下降会比较快,但生成的文本可能仍然不通顺。实际项目中需要更大规模的语料和更多的训练步数。
  • 这里模型参数量大约 800 万,远小于真实的 GPT 模型,所以 CPU 完全可以跑。

4.5 编写文本生成脚本

训练完成后,加载模型权重并生成文本。

# 文件路径:generate.py import torch import tiktoken from model import GPT def main(): device = "cuda" if torch.cuda.is_available() else "cpu" model = GPT( vocab_size=50257, n_embd=128, n_head=4, n_layer=4, block_size=128, dropout=0.0, ).to(device) model.load_state_dict(torch.load("gpt_model.pt", map_location=device)) model.eval() enc = tiktoken.get_encoding("gpt2") prompt = "To be, or not to be" input_ids = enc.encode(prompt) input_ids = torch.tensor([input_ids], dtype=torch.long, device=device) output_ids = model.generate(input_ids, max_new_tokens=50, temperature=0.8) output_text = enc.decode(output_ids[0].tolist()) print("生成结果:") print(output_text) if __name__ == "__main__": main()

4.6 运行与预期结果

依次运行:

python train.py python generate.py

训练时如果数据量较小,loss 可能在 0.1 以内;如果你替换成了自己的语料,loss 会在几十个 epoch 后逐渐下降。

生成的部分文本大概率不完全通顺,因为我们的训练数据只有一小段莎士比亚原文。但你能明显看出模型学到了一些英文单词拼写和常见词搭配,这说明模型结构本身是有效的。换用更大的语料、更长的训练时间后,生成质量会显著改善。

5. 常见问题与排查思路

5.1 常见报错汇总

下表整理了手写 GPT 过程中最常见的几个问题:

问题现象常见原因解决思路
训练时 loss 没有下降学习率过大/过小,或 tokenizer 词表与模型 vocab_size 不匹配检查学习率,确认tiktoken.get_encoding("gpt2")的词汇量与模型vocab_size一致
运行时提示“序列长度超出位置编码范围”输入 token 长度超过block_size对输入做截断,或增大block_size
生成的文本全是乱码或重复字符采样温度过低,或模型训练不充分调高 temperature,增加训练数据量
CUDA out of memory模型过大或 batch_size 过大降低 batch_size,减少 block_size,或使用梯度累积
注意力分数出现 NaN数据中存在过大的数值,网络未做梯度裁剪增加 LayerNorm 稳定性,添加梯度裁剪

5.2 关键排查步骤

如果你训练出来的模型效果很差,建议按下面顺序排查:

第一步:确认数据预处理。打印一个样本的xy,确认y确实是x的下一个 token。如果错位逻辑写错,模型永远学不对。

第二步:确认 mask 是否正确。把 mask 打印出来,看是否上三角全为 0。如果 mask 没生效,模型在训练时会偷看未来 token,导致训练 loss 很低但生成效果很差。

第三步:确认 embedding 层维度。GPT 中 token embedding 和 position embedding 矩阵都是随机初始化的,如果初始化不合理,训练可能非常慢。建议使用均值为 0、标准差 0.02 的正态分布初始化,这也是 GPT-2 原始代码的做法。

第四步:确认优化器和学习率。语言模型常用的学习率区间是1e-43e-4,过大会导致 loss 震荡。同时不要忘记梯度裁剪。

6. 最佳实践与工程建议

6.1 从小规模实验出发

不要一上来就训练一个 7B 参数的大模型。工程实践中,先用一个几十 MB 的小语料、小模型(几百万参数)把代码流程跑通,确认 loss 能正常下降、文本能正常生成,然后逐步扩大规模。这样做能节省大量调试时间。

6.2 使用梯度累积扩展 batch size

真实的大语言模型训练往往需要很大的 batch size,但单张显卡的显存有限。梯度累积的思路是:先计算多个 mini-batch 的梯度,然后将梯度累加,最后做一次参数更新,效果等价于更大的 batch size。

PyTorch 中的实现思路如下:

accumulation_steps = 4 optimizer.zero_grad() for i, (x, y) in enumerate(loader): x, y = x.to(device), y.to(device) loss = model(x, y)[1] / accumulation_steps loss.backward() if (i + 1) % accumulation_steps == 0: optimizer.step() optimizer.zero_grad()

注意 loss 除以accumulation_steps,是为了让梯度累加后的量级与正常 batch size 训练保持一致。

6.3 混合精度训练

PyTorch 2.x 中可以使用torch.autocastGradScaler做混合精度训练。在支持 CUDA 的环境下,混合精度能显著减少显存占用,并提速训练。

scaler = torch.amp.GradScaler("cuda") for x, y in loader: x, y = x.to(device), y.to(device) optimizer.zero_grad() with torch.amp.autocast("cuda"): _, loss = model(x, y) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

混合精度并不是必须的,但对大规模训练非常有帮助。CPU 环境下不需要使用。

6.4 推理阶段的 KV Cache

在文本生成时,模型需要逐个 token 生成。如果不做任何缓存,每生成一个新 token,都要重新计算之前所有 token 的 Key 和 Value,计算量很大。

真实 GPT 推理(例如 HuggingFace 的generate)会使用 KV Cache:把之前 step 计算出来的 K、V 矩阵缓存下来,下一轮只计算新 token 的 Q、K、V,与缓存拼接后再做注意力计算。

KV Cache 是理解大模型推理性能的核心概念。本文为了保持代码简洁没有实现,但在深入研究大模型部署、vLLM、TensorRT-LLM 等推理框架时,这个机制必须掌握。

6.5 权重保存与加载

训练完成后保存权重时,最好同时保存模型配置和优化器状态,方便断点续训。

# 保存 torch.save({ "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "epoch": epoch, "config": { "n_embd": n_embd, "n_head": n_head, "n_layer": n_layer, "block_size": block_size, "vocab_size": 50257, } }, "checkpoint.pt") # 恢复 checkpoint = torch.load("checkpoint.pt") model.load_state_dict(checkpoint["model_state_dict"]) optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

6.6 数据质量决定模型上限

大语言模型的训练遵循一个规律:模型结构决定能力的下限,数据质量决定能力的上限。如果你在自己的项目里微调 GPT,不要只关心模型参数量,更要投入时间清洗训练数据,过滤重复文本、去掉格式错乱内容、尽量保证语料的多样性。

7. 总结与下一步学习建议

这篇文章从零实现了一个 mini 版 GPT,核心代码只有不到 300 行。你已经掌握了 Token Embedding、Position Embedding、多头因果自注意力、LayerNorm、残差连接、FFN 和语言模型头的完整实现;也能用自己的脚本训练一个小模型并生成文本。

下一步可以从这几个方向继续深入:

  • 把字符级或者 BPE 分词器替换成 SentencePiece,训练支持中文的 tokenizer;
  • 在模型中加入 KV Cache,对比推理速度的提升;
  • 参考 GPT-2/LLaMA 的源码,把 RoPE(旋转位置编码)、SwiGLU 激活函数等现代改进加入模型;
  • 尝试用更大的自己的语料(例如维基百科导出数据)训练一个 1 亿参数左右的模型,感受训练 loss 和生成质量的真实变化。

如果本文对你有帮助,建议收藏备用;如果你在复现代码时遇到问题,可以把报错信息和模型配置发在评论区,我们一起排查。动手跑一遍代码,比看十篇文章的印象都深刻。

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

大模型+财务智能化落地:DeepSeek选型、部署与场景实践

/* MD / 富文本中的 .toc(含博客园搬家等嵌套结构);.toc-box 在侧栏,不受影响 */#content_views .toc,/* 编辑器常在目录前后插入空 p(:empty 仍占 20px),一并去掉避免顶空隙 */#content_views.markdown_views > p:empty:has(+ .toc),#content_views.markdown_views …

作者头像 李华
网站建设 2026/9/6 5:07:23

[特殊字符] ZorvAI 动态 UI 组件:让 AI 的回答「看得见、用得上」

&#x1f31f; 项目简介 ZorvAI 动态 UI&#xff08;quro-ui&#xff09; 是一套基于 Jetpack Compose 原生构建的可交互界面渲染框架。它让 AI 不再局限于「文字 代码块」的输出形态&#xff0c;而是能够主动地生成卡片、表单、列表、播放器、浏览器、富媒体等完整的交互式界…

作者头像 李华
网站建设 2026/9/6 5:06:15

PaddleOCR-VL 在 Intel Arc A770 上的部署与调优指南

/* MD / 富文本中的 .toc(含博客园搬家等嵌套结构);.toc-box 在侧栏,不受影响 */#content_views .toc,/* 编辑器常在目录前后插入空 p(:empty 仍占 20px),一并去掉避免顶空隙 */#content_views.markdown_views > p:empty:has(+ .toc),#content_views.markdown_views …

作者头像 李华
网站建设 2026/9/6 4:59:05

【无标题】AI获客工具包:5套即用提示词,帮自由职业者快速写高转化文案

这个工具包把“获客文案”拆成 5 个可以直接替换使用的提示词模板&#xff1a;冷邮件、LinkedIn 私信、跟进序列、提案、异议处理。适合自由职业者、小型代理机构和独立创业者。你只需要把服务、目标客户、成果案例填进去&#xff0c;就能快速生成高转化沟通内容。适用 ChatGPT…

作者头像 李华
网站建设 2026/9/6 4:58:18

三款AI论文写作软件亲测:从大纲到降重怎么选才不踩坑?

写论文这事&#xff0c;最怕的不是写不出来&#xff0c;而是写得心里没底。 题目改了七八版还怕选重了&#xff0c;文献下载了两百篇越读越乱&#xff0c;参考文献格式调到崩溃&#xff0c;交稿前还得担心重复率和AIGC检测。今年开学季一到&#xff0c;又有一波人在搜“AI论文工…

作者头像 李华