无库无捷径，PyTorch 手写完整 Transformer 大语言模型 LLM 调用**font stylecolor:rgb(0, 206, 185);background-color:rgb(252, 252, 252);from transformers import AutoModel/font**只要一行代码但真正理解 Transformer 的人少之又少。多数开发者对 Attention 的认知停留在Q、K、V 三个矩阵的口号层面——面试时背得出公式动手时写不出一个完整的 Block。本文的目标很纯粹不用任何高层封装库只用 PyTorch 原生算子从零实现一个完整的 Decoder-only Transformer 语言模型GPT 架构也是 GPT-4 / LLaMA / Qwen 的基础架构。每一行代码对应论文中的一个公式。一、架构对齐Decoder-only Transformer输入 Token IDs │ ▼ Token Embedding RoPE 位置编码 │ ▼ ┌─────────────────────────────┐ │ Causal Self-Attention │ │ Residual LayerNorm │ × N 层 ├─────────────────────────────┤ │ Feed-Forward (MLP) │ │ Residual LayerNorm │ └─────────────────────────────┘ │ ▼ LayerNorm → Linear Head → Logits核心实现清单Token Embedding → RoPE 位置编码 → Multi-Head Attention Causal Mask → Feed-Forward → Pre-Norm 残差 → 自回归生成。二、基础配置import torch import torch.nn as nn import torch.nn.functional as F import math class Config: vocab_size 10000 # 词表大小 d_model 512 # 模型维度 n_heads 8 # 注意力头数 n_layers 6 # Block 层数 d_ff 2048 # FFN 中间维度 max_seq_len 512 # 最大序列长度 dropout 0.1 pad_id 0 config Config() device cuda if torch.cuda.is_available() else cpu三、逐层实现3.1 Token Embeddingclass TokenEmbedding(nn.Module): def __init__(self, vocab_size, d_model): super().__init__() self.embedding nn.Embedding(vocab_size, d_model) self.d_model d_model def forward(self, x): # 论文要求embedding 乘以 sqrt(d_model) 做尺度对齐 return self.embedding(x) * math.sqrt(self.d_model)3.2 旋转位置编码RoPERoPE 是 LLaMA / Qwen 等主流 LLM 的标准位置编码方案。它将位置信息编码为 Q/K 向量的旋转支持相对位置外推。class RotaryPositionalEmbedding(nn.Module): def __init__(self, d_model, max_seq_len): super().__init__() d_half d_model // 2 # 频率θ_i 10000^(-2i/d) freqs 1.0 / (10000 ** (torch.arange(0, d_half).float() / d_half)) positions torch.arange(max_seq_len).float() angles torch.outer(positions, freqs) # [seq_len, d_half] self.register_buffer(cos, angles.cos()) self.register_buffer(sin, angles.sin()) def forward(self, x): seq_len x.size(1) cos self.cos[:seq_len].unsqueeze(0) sin self.sin[:seq_len].unsqueeze(0) d_half x.size(-1) // 2 x1, x2 x[..., :d_half], x[..., d_half:] # 旋转(x1 i*x2) * (cos i*sin) return torch.cat([x1 * cos - x2 * sin, x1 * sin x2 * cos], dim-1)3.3 Multi-Head Causal Self-Attention论文公式Attention(Q,K,V)softmax(QKTdk)VAttention(Q,K,V)softmax(dkQKT)Vclass MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads, dropout0.1): super().__init__() assert d_model % n_heads 0 self.n_heads n_heads self.d_head d_model // n_heads self.W_q nn.Linear(d_model, d_model, biasFalse) self.W_k nn.Linear(d_model, d_model, biasFalse) self.W_v nn.Linear(d_model, d_model, biasFalse) self.W_o nn.Linear(d_model, d_model, biasFalse) self.dropout nn.Dropout(dropout) def forward(self, x, ropeNone, maskNone): B, S, _ x.shape # 1. 线性投影 → 切分多头 Q self.W_q(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2) K self.W_k(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2) V self.W_v(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2) # Q/K/V: [B, n_heads, S, d_head] # 2. RoPE 注入位置信息仅作用于 Q 和 K if rope is not None: Q rope(Q) K rope(K) # 3. 注意力分数 Q K^T / sqrt(d_head) scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_head) # 4. 因果掩码上三角置 -infsoftmax 后概率为 0 if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) # 5. Softmax → 加权 V → 合并多头 → 输出投影 attn self.dropout(F.softmax(scores, dim-1)) out torch.matmul(attn, V) # [B, n_heads, S, d_head] out out.transpose(1, 2).contiguous().view(B, S, -1) return self.W_o(out)3.4 Feed-Forward Networkclass FeedForward(nn.Module): def __init__(self, d_model, d_ff, dropout0.1): super().__init__() self.w1 nn.Linear(d_model, d_ff) self.w2 nn.Linear(d_ff, d_model) self.dropout nn.Dropout(dropout) def forward(self, x): return self.w2(self.dropout(F.relu(self.w1(x))))3.5 Transformer BlockPre-Norm 残差结构现代 LLM 统一采用 Pre-Norm先 Norm 再 Attention训练稳定性远优于原始论文的 Post-Norm。class TransformerBlock(nn.Module): def __init__(self, config): super().__init__() self.attn MultiHeadAttention(config.d_model, config.n_heads, config.dropout) self.ffn FeedForward(config.d_model, config.d_ff, config.dropout) self.norm1 nn.LayerNorm(config.d_model) self.norm2 nn.LayerNorm(config.d_model) self.dropout nn.Dropout(config.dropout) def forward(self, x, ropeNone, maskNone): # Sub-layer 1: Attention 残差 x x self.dropout(self.attn(self.norm1(x), roperope, maskmask)) # Sub-layer 2: FFN 残差 x x self.dropout(self.ffn(self.norm2(x))) return x3.6 完整模型class TransformerLLM(nn.Module): def __init__(self, config): super().__init__() self.config config self.token_embedding TokenEmbedding(config.vocab_size, config.d_model) self.rope RotaryPositionalEmbedding(config.d_model, config.max_seq_len) self.blocks nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)]) self.final_norm nn.LayerNorm(config.d_model) self.lm_head nn.Linear(config.d_model, config.vocab_size, biasFalse) # 权重共享Embedding 和 LM Head 用同一套权重 self.lm_head.weight self.token_embedding.embedding.weight self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean0.0, std0.02) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean0.0, std0.02) def forward(self, input_ids): B, S input_ids.shape # 因果掩码下三角矩阵 mask torch.tril(torch.ones(S, S, deviceinput_ids.device)).unsqueeze(0).unsqueeze(0) x self.token_embedding(input_ids) for block in self.blocks: x block(x, ropeself.rope, maskmask) x self.final_norm(x) return self.lm_head(x) # [B, S, vocab_size]3.7 验证前向传播model TransformerLLM(config).to(device) n_params sum(p.numel() for p in model.parameters()) print(f参数量: {n_params/1e6:.1f}M) dummy torch.randint(1, config.vocab_size, (2, 64)).to(device) logits model(dummy) print(f输出: {logits.shape}) # [2, 64, 10000]四、训练与推理4.1 训练单步def train_step(model, batch, optimizer, config): model.train() input_ids batch[:, :-1].to(device) # 输入去掉最后一个 token targets batch[:, 1:].to(device) # 目标左移一位 logits model(input_ids) loss F.cross_entropy( logits.reshape(-1, config.vocab_size), targets.reshape(-1), ignore_indexconfig.pad_id ) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) # 梯度裁剪 optimizer.step() return loss.item()4.2 自回归生成torch.no_grad() def generate(model, prompt_ids, max_new_tokens100, temperature0.8, top_k50): model.eval() for _ in range(max_new_tokens): context prompt_ids[:, -config.max_seq_len:] logits model(context)[:, -1, :] / temperature # Top-K 采样 v, _ torch.topk(logits, top_k) logits[logits v[:, [-1]]] float(-inf) probs F.softmax(logits, dim-1) next_token torch.multinomial(probs, num_samples1) prompt_ids torch.cat([prompt_ids, next_token], dim1) return prompt_ids五、模型规模与训练数据工程调整**font stylecolor:rgb(0, 206, 185);background-color:rgb(252, 252, 252);Config/font**即可缩放模型规模配置d_modeln_headsn_layers参数量对标教学版51286~38M—GPT-2 Small7681212~85MGPT-2GPT-2 Medium10241624~350MGPT-2 Medium工业级训练的真正瓶颈往往不在模型代码而在数据吞吐——高质量语料的采集和清洗需要完整的数据管线支撑。以亿牛云的数据采集场景为例从外部站点批量采集训练语料时IP 限流和并发频率是首要约束import requests # 亿牛云 API 代理提取出口 IP api_url http://ip.16yun.cn:817/myip/pl/ORDER_ID/?sSIGNuUSERformatjson proxy_info requests.get(api_url, timeout10).json()[0] proxy fhttp://{proxy_info[ip]}:{proxy_info[port]} # 批量采集语料 resp requests.get(https://corpus.example.com/batch, proxies{http: proxy, https: proxy}, timeout30)采集返回403需检查白名单配置返回429需降低请求频率。稳定的数据采集是 LLM 训练的第一道基础设施保障。六、踩坑速查陷阱根因解法Loss 变 NaN梯度爆炸梯度裁剪**font stylecolor:rgb(0, 206, 185);max_norm1.0/font** Pre-NormLoss 不下降学习率不匹配前 200 步线性预热再用余弦衰减生成重复采样策略单一Top-K 温度缩放显存溢出batch/seq 过大减小 batch_size 或用梯度累积维度报错**font stylecolor:rgb(0, 206, 185);d_model % n_heads ≠ 0/font**确保**font stylecolor:rgb(0, 206, 185);d_model/font**能被**font stylecolor:rgb(0, 206, 185);n_heads/font**整除七、总结手写一遍 Transformer收获的是对每个组件的深层理解Embedding离散 token → 连续空间乘 √d_model 做尺度对齐RoPE旋转编码位置信息到 Q/K支持相对位置外推AttentionQ·Kᵀ 算相似度 → softmax 归一化 → 加权 V计算复杂度 O(n²·d)Causal Mask上三角填 -inf实现自回归生成Pre-Norm 残差信息直通通道让深层网络可训练权重共享Embedding 和 LM Head 共享权重参数减半当面试官问Attention 的计算复杂度为什么是 O(n²)时你能立刻回答因为 Q·Kᵀ 那一步矩阵乘法的维度是**font stylecolor:rgb(0, 206, 185);background-color:rgb(252, 252, 252);[seq, d] × [d, seq]/font**。这种理解层次是调用**font stylecolor:rgb(0, 206, 185);background-color:rgb(252, 252, 252);from_pretrained()/font**永远无法获得的。