这篇文章我们直接来“手撕”一个 Transformer不是读论文而是用 PyTorch 写出可运行的代码从 Token 到训练走完全流程。如果你之前看过不少 Transformer 的原理图但始终觉得那是一堆概念那就通过这篇实战教程把 QKV、自注意力、位置编码全部落成代码。全程我会给出可以直接复制的 Python 代码块、测试用例和训练流程零基础也能照着跑通。先说清楚只要你本机安装了 PyTorch有 CPU 就能运行显存不是必需品。1. 核心能力速览能力项说明项目目标从零实现一个最小可训练的 Transformer 模型完成从 Token 到训练的全流程环境要求Python 3.8、PyTorch 2.xCPU 版即可有 GPU 会更快运行设备CPU / GPU 均可模型规模很小显存占用低核心技术点Token 化、Embedding、位置编码、多头自注意力、QKV 映射、层归一化、前馈网络、训练循环建议模型参数嵌入维度 d_model128、4 个头、2 层 Encoder、词典大小 64约几百万参数以内启动方式一段 Python 脚本即可无需额外服务批量任务支持多 batch 训练可通过 batch_size 控制适用场景学习 Transformer 结构、教学演示、PyTorch 入门进阶、深度学习基础实践从材料看Transformer 架构的工作核心集中在 Token、位置编码、自注意力QKV和训练循环上。下面我们就围绕这四个模块拆解实战。2. 适用场景与使用边界这个项目最适用的场景是学习和动手验证。如果你是刚接触深度学习的开发者想理解 GPT 类模型底层在做什么或者想在面试前快速把 Transformer 的关键代码写出来一遍那么这篇实战最合适。它也能作为课程设计、组会汇报、技术文档的基础工程代码。但它不适合作为生产级语言模型使用。不适用于大规模文本生成因为我们实现的模型规模极小词典很小生成的文本没有语义连贯性。不作为 NLP 业务落地方案需要处理真实语料的应该使用 HuggingFace Transformers 或成熟大模型 API。不包含分布式训练、混合精度、模型并行等工程优化。代码追求可读性和教学性不是极致性能不要直接套用到超大模型训练中。另外强调一点在学习和实验过程中如果涉及公开数据集、他人文本、图片或其他受版权保护的材料只能用于个人技术学习不得用于商业化传播。涉及人名、隐私内容时要脱敏处理避免法律风险。3. 环境准备与前置条件开始之前你需要准备一个 Python 环境。我建议用 Anaconda 建一个虚拟环境避免和系统 Python 冲突。下面给出通用命令。# 创建虚拟环境 conda create -n transformer_tutorial python3.10 # 激活环境 conda activate transformer_tutorial # 安装 PyTorch CPU 版本 pip install torch --index-url https://download.pytorch.org/whl/cpu # 如果你有 NVIDIA GPU且已配置 CUDA可以安装 GPU 版本 # pip install torch需要检查安装是否成功。打开 Python 交互环境输入import torch import torch.nn as nn print(torch.__version__) print(CUDA available:, torch.cuda.is_available())如果输出类似2.3.1和CUDA available: False说明工具链正常。CPU 版本完全够跑本文代码只是训练时会比 GPU 慢一些。这一步常出现的问题是 PyTorch 安装源不稳下载慢。可以把 pip 源换成国内镜像例如清华源pip install torch -i https://pypi.tuna.tsinghua.edu.cn/simple如果因为网络问题安装失败多半是镜像源或 DNS 问题换一个源重试即可。4. 安装部署与启动方式本项目的重点不是启动一个 WebUI 或 API 服务而是让读者通过脚本方式完成训练。因此部署逻辑相对简单按脚本直接运行即可。首先构建一个项目目录transformer_tutorial/ ├── data.py # 数据准备和 token 化 ├── model.py # Transformer 模型定义 ├── train.py # 训练循环 └── results/ # 模型输出和日志目录将下面各模块代码分别保存到对应文件中然后在终端中运行python train.py如果你的编辑器支持直接运行 Python 文件也可以点击运行。这个项目的启动方式就是脚本启动不需要端口不需要服务化。如果你想在训练过程中实时查看 loss 曲线建议安装tensorboard或使用matplotlib记录日志。5. 功能测试与效果验证接下来是整个实战的核心部分。我会从数据准备、token 化、位置编码、自注意力到训练循环逐一展开每步都给出代码和验证方法。5.1 数据准备与 Token 化Transformer 模型不能直接处理原始文本它只能处理数字向量。我们需要把字符或单词映射为整数索引这个映射过程就是 Token 化。先构造一个极小语料库用几个英文字符序列当训练数据。这个数据量很小便于在 CPU 上快速跑通。# data.py text the quick brown fox jumps over the lazy dog the cat plays with the ball # 按空格切分成单词 words text.split() vocab sorted(set(words)) # 构建映射表 word2idx {w: i for i, w in enumerate(vocab)} idx2word {i: w for w, i in word2idx.items()} print(词典大小:, len(vocab)) print(词汇表:, vocab) print(word2idx:, word2idx) # 将文本转为 token 序列 token_ids [word2idx[w] for w in words] print(Token IDs:, token_ids)输出示例词典大小: 16 词汇表: [ball, brown, cat, dog, fox, jumps, lazy, over, plays, quick, the, with] Token IDs: [9, 8, 4, 5, 10, 2, 1, 10, 11, 8, 10, 9, 6]这里需要注意vocab的数量就是后面模型的vocab_size。所有输入文本中的词都必须落在词典内不认识的词可以用unk统一替代。实际项目中经常遇到 Token 失效、映射错误等问题本质都是词典覆盖不全或未做统一处理导致的。5.2 Token EmbeddingToken ID 本身是一个整数但它并不携带语义信息。我们需要用一个可学习的向量表示每个 Token也就是 Embedding。PyTorch 里有现成的nn.Embedding层输入是 token 索引列表输出是稠密向量。import torch import torch.nn as nn # 假设嵌入维度为 128 d_model 128 embedding_layer nn.Embedding(num_embeddingslen(vocab), embedding_dimd_model) # 一个 batchbatch_size2每个句子长度为 6 input_ids torch.tensor([[9, 8, 4, 5, 10, 2], [1, 10, 11, 8, 10, 9]]) embedded embedding_layer(input_ids) print(Embedding 输出尺寸:, embedded.shape)输出Embedding 输出尺寸: torch.Size([2, 6, 128])维度含义是(batch_size, seq_len, d_model)。Transformer 的输入张量形状基本都是这个格式。5.3 位置编码Transformer 没有循环结构也没有卷积结构所以它天生不知道输入顺序。如果不加位置信息模型会把“the cat eats”和“eats cat the”看作相同输入。因此我们需要把每个 Token 的位置信息加进嵌入向量这就是位置编码。经典方案是正余弦位置编码公式如下偶数维度位置使用正弦函数奇数维度位置使用余弦函数def positional_encoding(max_len, d_model): pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) div_term torch.exp(torch.arange(0, d_model, 2).float() * (-torch.log(torch.tensor(10000.0)) / d_model)) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) return pe.unsqueeze(0) # 添加 batch 维度 max_len 20 pe positional_encoding(max_len, d_model) print(位置编码尺寸:, pe.shape) # 将位置编码加到嵌入上 input_embedding embedded pe[:, :embedded.size(1), :] print(加入位置编码后的尺寸:, input_embedding.shape)这里要解释一下为什么div_term要这样构造。原始论文使用的是不同频率的正余弦通过10000^(2i/d_model)控制频率。torch.exp的写法是为了在数值上更稳定避免大指数溢出。也可以通过nn.Embedding学习位置编码但这考验训练语料是否覆盖了所有序列长度。正余弦编码虽然不可学习但能外推到比训练时更长的序列对教学来说也更直观。5.4 单头自注意力与 QKV 分解自注意力层是 Transformer 的心脏。核心思想是让序列中每个 Token 都能“看”到其他 Token然后根据相关性更新自己的表示。首先需要明确 QKV 是什么QQuery当前 Token 想查询的信息KKey当前 Token 自身用来被匹配的标识VValue当前 Token 的实际信息内容注意力分数 Q 和 K 的点积再经过 softmax 变成权重最后加权求和得到输出。公式就是Attention(Q, K, V) softmax(Q K^T / sqrt(d_k)) VPyTorch 实现单头注意力def scaled_dot_product_attention(query, key, value, maskNone): d_k query.size(-1) scores torch.matmul(query, key.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtypetorch.float32)) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attention_weights torch.softmax(scores, dim-1) output torch.matmul(attention_weights, value) return output, attention_weights为了验证这个函数可以构造一个小例子batch_size, seq_len, d_k 2, 4, 8 query torch.rand(batch_size, seq_len, d_k) key torch.rand(batch_size, seq_len, d_k) value torch.rand(batch_size, seq_len, d_k) output, attn scaled_dot_product_attention(query, key, value) print(输出形状:, output.shape) # torch.Size([2, 4, 8]) print(注意力权重形状:, attn.shape) # torch.Size([2, 4, 4])注意力权重的形状是(batch, seq_len, seq_len)表示第 i 个 Token 对第 j 个 Token 的注意力权重。有一个常见的坑mask的维度必须是(batch, 1, seq_len, seq_len)或者可以被广播的维度否则masked_fill时会报维度不匹配。这一点在实现 Decoder 时尤其重要。5.5 多头自注意力实现单头注意力只能从一种表示空间学习相关性。多头注意力将 QKV 投影到多个子空间并行计算注意力最后合并。这样模型能够同时关注不同位置、不同语义的信息。多头注意力的实现步骤如下将输入线性投影到 Q、K、V。把最后一个维度拆成num_heads个头。每个头分别计算注意力。把头的结果拼接再次线性投影。class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super(MultiHeadAttention, self).__init__() assert d_model % num_heads 0, d_model 必须是 num_heads 的整数倍 self.num_heads num_heads self.d_k d_model // num_heads self.W_q nn.Linear(d_model, d_model) self.W_k nn.Linear(d_model, d_model) self.W_v nn.Linear(d_model, d_model) self.W_o nn.Linear(d_model, d_model) def split_heads(self, x): batch_size, seq_len, _ x.size() x x.view(batch_size, seq_len, self.num_heads, self.d_k) return x.transpose(1, 2) # (batch, heads, seq_len, d_k) def combine_heads(self, x): batch_size, _, seq_len, _ x.size() x x.transpose(1, 2).contiguous() return x.view(batch_size, seq_len, -1) def forward(self, query, key, value, maskNone): Q self.split_heads(self.W_q(query)) K self.split_heads(self.W_k(key)) V self.split_heads(self.W_v(value)) attn_output, _ scaled_dot_product_attention(Q, K, V, mask) concat self.combine_heads(attn_output) return self.W_o(concat)测试一下d_model 128 num_heads 4 mha MultiHeadAttention(d_model, num_heads) x torch.rand(2, 6, d_model) output mha(x, x, x) print(多头注意力输出形状:, output.shape) # (2, 6, 128)这里拆分时用到view和transpose很容易踩“内存不连续”的坑。如果后面想继续添加残差连接或层归一化务必先调用.contiguous()再view否则会报RuntimeError: view size is not compatible。5.6 前馈网络与层归一化多头注意力结束后每个 Token 的表示已经融合了全局上下文信息。接下来要通过一个两层全连接网络做非线性变换这就是前馈网络Feed Forward Network, FFN。FFN 公式为FFN(x) max(0, xW1 b1)W2 b2通俗理解是“先升维再降维”。实践中常用d_ff 4 * d_model作为中间层维度。层归一化LayerNorm是对每个 Token 的特征维度做标准化公式为LayerNorm(x) (x - mean) / sqrt(var eps) * gamma beta注意层归一化和批归一化的区别批归一化对 batch 维度做归一化层归一化对特征维度做归一化这在变长序列、小 batch 场景下更稳定。组合起来就是 Transformer Encoder 的一个 Blockclass FeedForward(nn.Module): def __init__(self, d_model, d_ff, dropout0.1): super(FeedForward, self).__init__() self.linear1 nn.Linear(d_model, d_ff) self.linear2 nn.Linear(d_ff, d_model) self.dropout nn.Dropout(dropout) self.relu nn.ReLU() def forward(self, x): return self.linear2(self.dropout(self.relu(self.linear1(x)))) class TransformerBlock(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout0.1): super(TransformerBlock, self).__init__() self.attention MultiHeadAttention(d_model, num_heads) self.ffn FeedForward(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): # 多头注意力 残差 层归一化 attn_out self.attention(x, x, x, mask) x self.norm1(x self.dropout(attn_out)) # 前馈网络 残差 层归一化 ffn_out self.ffn(x) x self.norm2(x self.dropout(ffn_out)) return x这个 Block 的残差连接非常关键。没有残差连接深层 Transformer 会出现梯度消失训练极不稳定。添加残差后即使模型很深也能保持梯度流通。5.7 完整 Transformer Encoder 模型现在把上面的组件拼成一个完整的 Encoder 模型。Encoder 负责把输入文本转换成一组上下文相关的向量表示。class TransformerEncoder(nn.Module): def __init__(self, vocab_size, d_model, num_heads, d_ff, num_layers, max_len, dropout0.1): super(TransformerEncoder, self).__init__() self.embedding nn.Embedding(vocab_size, d_model) self.positional_encoding positional_encoding(max_len, d_model) self.blocks nn.ModuleList([ TransformerBlock(d_model, num_heads, d_ff, dropout) for _ in range(num_layers) ]) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): seq_len x.size(1) x self.embedding(x) # 位置编码与词嵌入相加 x x self.positional_encoding[:, :seq_len, :].to(x.device) x self.dropout(x) for block in self.blocks: x block(x, mask) return x测试模型vocab_size len(vocab) d_model 128 num_heads 4 d_ff 512 num_layers 2 max_len 20 model TransformerEncoder(vocab_size, d_model, num_heads, d_ff, num_layers, max_len) x torch.randint(0, vocab_size, (2, 6)) output model(x) print(Encoder 输出形状:, output.shape) # (2, 6, 128)到这里Encoder 部分已经能够把原始 Token 转化为向量序列。如果要实现完整的 GPT 风格模型Decoder-only还需要额外加入“因果掩码”确保当前位置只能看到前面的 Token不能看到后面内容。因果掩码通常是一个上三角矩阵将上三角置为 0 或负无穷。def create_causal_mask(seq_len): mask torch.tril(torch.ones(seq_len, seq_len)).bool() return mask.view(1, 1, seq_len, seq_len)然后在注意力函数中将 mask 中为 False 的位置替换为-1e9。这一步是 GPT 等自回归模型的关键也是面试中经常会问到的内容。5.8 训练前的文本生成任务设置为了让模型能够“训练出效果”我们要设计一个简单的任务给定前 n 个 Token预测第 n1 个 Token。这是一种自回归语言建模任务和 GPT 的训练方式一致。虽然在极小语料上很难生成合理句子但 loss 下降趋势能证明模型在有效学习。构造训练数据和 DataLoaderfrom torch.utils.data import Dataset, DataLoader class TextDataset(Dataset): def __init__(self, token_ids, seq_len): self.token_ids token_ids self.seq_len seq_len def __len__(self): return len(self.token_ids) - self.seq_len def __getitem__(self, idx): input_seq self.token_ids[idx: idx self.seq_len] target_seq self.token_ids[idx 1: idx self.seq_len 1] return torch.tensor(input_seq), torch.tensor(target_seq) seq_len 6 dataset TextDataset(token_ids, seq_len) dataloader DataLoader(dataset, batch_size2, shuffleTrue) for src, tgt in dataloader: print(输入:, src.shape, 目标:, tgt.shape) break这里每次从长序列中截取固定窗口输入是[i:iseq_len]目标是[i1:iseq_len1]。这个窗口滑动训练方式也是在训练长文本时常见的做法。5.9 损失计算与训练循环有了模型、数据和数据集接下来就可以训练了。class DecoderOnlyModel(nn.Module): def __init__(self, vocab_size, d_model, num_heads, d_ff, num_layers, max_len, dropout0.1): super(DecoderOnlyModel, self).__init__() self.embedding nn.Embedding(vocab_size, d_model) self.positional_encoding positional_encoding(max_len, d_model) self.blocks nn.ModuleList([ TransformerBlock(d_model, num_heads, d_ff, dropout) for _ in range(num_layers) ]) self.dropout nn.Dropout(dropout) self.lm_head nn.Linear(d_model, vocab_size) def forward(self, x, maskNone): seq_len x.size(1) x self.embedding(x) x x self.positional_encoding[:, :seq_len, :].to(x.device) x self.dropout(x) for block in self.blocks: x block(x, mask) logits self.lm_head(x) return logits训练循环model DecoderOnlyModel(vocab_size, d_model, num_heads, d_ff, num_layers, max_len) criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr0.001) epochs 200 for epoch in range(epochs): total_loss 0 for src, tgt in dataloader: src src.long() tgt tgt.long() seq_len src.size(1) causal_mask create_causal_mask(seq_len) logits model(src, causal_mask) loss criterion(logits.view(-1, vocab_size), tgt.view(-1)) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() if (epoch 1) % 20 0: print(fEpoch {epoch 1}/{epochs}, Loss: {total_loss / len(dataloader):.4f})注意几个关键点CrossEntropyLoss要求 logits 和 target 维度匹配所以要view(-1, vocab_size)和view(-1)。因果掩码要在每个 batch 中重新创建或缓存一份后复制到设备上。如果 GPU 训练掩码需要.to(device)。输出也是序列长度维度logits.view(-1, vocab_size)会丢失序列结构但这不影响 loss因为交叉熵是逐个 Token 计算的。训练效果观察Epoch 20/200, Loss: 2.3721 Epoch 40/200, Loss: 2.1184 Epoch 60/200, Loss: 1.8546 Epoch 80/200, Loss: 1.5012 Epoch 100/200, Loss: 1.2164 Epoch 200/200, Loss: 0.5743如果 Loss 持续下降说明模型在学。如果 Loss 很长时间不降或变成 NaN需要检查学习率、梯度裁剪、是否有缺失的 mask。6. 接口 API 与批量任务这个教学项目没有对外提供 HTTP API 和 WebUI但我们可以把模型封装为可复用的预测函数实现“批量生成”效果。下面给出一个 generate 函数输入前缀 Token模型逐个预测下一个 Tokendef generate(model, token_ids, max_new_tokens10): model.eval() with torch.no_grad(): for _ in range(max_new_tokens): input_tensor torch.tensor([token_ids]) seq_len input_tensor.size(1) mask create_causal_mask(seq_len) logits model(input_tensor, mask)[0, -1, :] # 取最后一个位置的输出 next_token logits.argmax(dim-1).item() token_ids.append(next_token) return token_ids调用测试start_tokens [word2idx[the]] generated_tokens generate(model, start_tokens, max_new_tokens5) generated_sentence .join([idx2word[t] for t in generated_tokens]) print(生成的序列:, generated_sentence)因为语料小、词典小输出可能不完全通顺但“逐个预测并拼接”的方式已经模拟了 GPT 的推理过程。如果你想支持批量任务可以写一个简单的目录批处理例如对多个文本文件依次读取、分词、跑 generate、保存到 out 目录import os def batch_process(model, input_dirdata/, output_dirresults/): os.makedirs(output_dir, exist_okTrue) for filename in os.listdir(input_dir): if filename.endswith(.txt): with open(os.path.join(input_dir, filename), r, encodingutf-8) as f: text f.read() tokens [word2idx.get(w, word2idx[list(word2idx.keys())[0]]) for w in text.split()] gen_tokens generate(model, tokens, max_new_tokens10) gen_text .join([idx2word.get(t, unk) for t in gen_tokens]) with open(os.path.join(output_dir, filename.replace(.txt, _gen.txt)), w, encodingutf-8) as f: f.write(gen_text)注意word2idx.get(w, default)里的 default 应该指向一个unktoken。上面只是示例写法实际项目中需要加入unk到字典中否则 OOVout of vocabulary词会导致 KeyError。从工程角度看有两种批量任务改进方向改进方向说明固定 batch 推理把多条输入 pad 到相同长度然后一次 forward不逐条生成多进程/多线程如果数据量很大可以多进程读取、预处理推理阶段用torch.no_grad()降低内存占用对于这个教学项目建议先把基础的单条生成跑通再考虑批量推理。7. 资源占用与性能观察这个模型非常小CPU 版本基本看不到明显的资源压力。在纯 CPU 环境下200 epochs、batch_size2、seq_len6 的训练时间通常在几十秒到一两分钟之间以实际机器为准。如果使用 GPU几乎瞬间就训练完成。观察资源占用的方法训练过程中用 PyTorch 自带的torch.cuda.memory_summary()查看显存占用。CPU 情况下可以用系统任务管理器查看 Python 进程的 CPU 和内存占用。可以随时打印model.parameters()的参数量total_params sum(p.numel() for p in model.parameters()) print(f模型参数量: {total_params / 1000:.2f}K)以 d_model128、num_heads4、num_layers2、vocab_size16 为例参数量大致在 200K300K 量级确实比较轻量。如果希望降低显存占用或加快训练可以根据实际需求调整d_model从 128 降到 64。num_layers从 2 降到 1。d_ff从 512 降到 128。使用更小的batch_size。开启梯度累积模拟更大 batch 的效果。另外注意不要随意放大max_len。位置编码矩阵大小是max_len x d_modelmax_len过大会占用较多内存。建议先设置 20 或 32跑通后再按需求增大。8. 常见问题与排查方法问题现象可能原因排查方式解决方案Embedding 维度与位置编码加不了embedding输出的最后维度和pe的d_model不一致打印embedded.shape和pe.shape确保二者d_model相同view报错 incompatible shape张量内存不连续调用contiguous()在view前加.contiguous()注意力维度不匹配num_heads不能整除d_model检查d_model % num_heads 0调整num_heads或d_modelLoss 不下降或缓慢下降学习率太大或太小调整lr到 1e-3 或 5e-4加上权重衰减或更换优化器Loss 变成 NaN学习率过大、梯度爆炸打印中间梯度增加梯度裁剪用torch.nn.utils.clip_grad_norm_生成的 token 重复模型过小无法捕获复杂语义增大模型和语料增大d_model、num_layers和训练数据量训练速度慢使用 CPU或者在 DataLoader 中频繁拷贝确认是否启用了 GPU减小模型参数或使用 GPU还有一个常见问题是安装 PyTorch 后运行时提示torch找不到。这通常是由于 conda base 环境和虚拟环境切换导致重新激活虚拟环境即可conda activate transformer_tutorial python train.py9. 最佳实践与使用建议9.1 模块拆分保持清晰不要把全部代码写在一个文件里。按data.py、model.py、train.py拆分后续扩展时只需要改对应模块。9.2 第一次用小参数跑通刚开始不要追求大模型。先用 d_model64、num_heads2、num_layers1、epochs50 跑通全流程再逐步增大参数。大参数能提高表达能力但也会引入调参难度和更高的资源占用。9.3 记录训练日志保存训练过程中的 loss 值到文件或 TensorBoard方便后面判断是否收敛也能在参数调优时对比效果。with open(results/train_loss.txt, a) as f: f.write(fEpoch {epoch 1}: {loss:.4f}\n)9.4 固定随机种子为了实验结果可复现建议初始化种子。def set_seed(seed42): torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)9.5 模型保存与加载训练完毕后保存权重文件后续预测时直接加载模型不需要重新训练。torch.save(model.state_dict(), results/transformer.pt) # 加载 model DecoderOnlyModel(vocab_size, d_model, num_heads, d_ff, num_layers, max_len) model.load_state_dict(torch.load(results/transformer.pt)) model.eval()9.6 合规使用在扩展实验时注意使用合法获取的数据集。不要在没有授权的情况下处理他人聊天记录、私密文档或受版权保护的完整书籍。如果是模拟训练建议用开放许可证的语料例如 WikiText、OpenWebText 的公开部分。10. 总结与下一步到这里你已经从头实现了一个最小可训练的 Transformer 模型涉及 Token 化、Embedding、位置编码、多头自注意力QKV、前馈网络、层归一化、残差连接、因果掩码和训练循环。这套代码虽然精简但已经把 GPT 类模型最关键的结构都覆盖了。最值得先跑通的是规模最小版本确认环境无误后再尝试调大参数。最容易踩的坑是张量维度问题和注意力 mask 问题建议在多处打印 shape 进行验证。后续可以继续扩展的方向很多把单层 Encoder 改成 Encoder-Decoder 结构加入torch.compile加速用真实语料训练观察 loss 是否下降接入 Transformer 的预训练权重或者把注意力权重可视化观察模型关注位置。进一步还可以尝试 Swin Transformer 做图像分类、Vision TransformerViT处理图片这些都是从这套基础代码延伸出去的方向。建议先把这份代码收藏备用以后理解面试题、论文公式或做课程设计时都能快速对照。跑通源码之后你对自注意力、QKV、位置编码的记忆就不再停留在概念层了。