這次我們來看一項很多人學完 Transformer 但依然沒完全搞懂的事情不只是會調用nn.Transformer而是親手把 Linear Layer、Embedding Layer、參數初始化和反向傳播全部實現出來。本文對應的場景是 Stanford CS 336 3.3也就是在課程作業中要求用最樸素的 PyTorch 張量操作搭建語言模型組件并通過梯度檢查驗證手寫反向傳播的正確性。標題里的四個關鍵詞其實就對應了 Transformer 落地中四個最容易翻車的點參數初始化決定模型能不能收斂反向傳播決定你寫的網絡是不是真的“能學”Linear Layer 是所有前饋計算的地基Embedding Layer 則承擔著離散 token 到連續向量的映射。如果這幾個模塊只是靠nn.Linear和nn.Embedding一帶而過那遇到 loss 不降、梯度 NaN、維度對不上這類問題的時候往往只能瞎猜。自己手拼一遍之后很多報錯一眼就能看出問題。這篇文章會從零實現一個可訓練的最小 Transformer包含前向傳播、手寫 backward、參數初始化、梯度驗證最后在一個小模型上完成端到端訓練測試。全程使用 CPU 即可運行不依賴 GPU也不需要下載任何預訓練權重。適合正在啃 Transformer 源碼、準備算法面試、或者做課程作業時想把原理徹底吃透的讀者。1. 核心能力速覽先把這篇文章涉及的技術目標和運行約束講清楚。下面的表格不是項目 README而是你讀完這篇文章后應該具備的“能力驗收清單”。能力項說明項目類型從零手寫 Transformer 核心模塊的教學實現來源Stanford CS 336 Language Modeling 課程 3.3 小節主題核心模塊Linear Layer、Embedding Layer、LayerNorm、參數初始化、手寫反向傳播技術棧Python 3.9、PyTorch 2.x、NumPy 可選運行方式本地腳本運行支持 Jupyter Notebook硬件要求CPU 可跑顯存不是必需是否依賴 nn.Linear不依賴手動實現 forward/backward是否依賴 nn.Embedding不依賴手動實現查表與梯度回傳反向傳播方式手動推導并實現通過 torch.autograd.gradcheck 驗證驗證方式形狀測試、梯度對比、小模型端到端 loss 下降適合人群學習 Transformer 原理、面試準備、課程作業、源碼閱讀從表格可以看出這是一篇“先別用庫、先把原理跑通”的教程。它的核心產出不是一個大而全的訓練框架而是一套你可以自己掌控每個張量形狀和每個梯度去向的最小實現。2. 適用場景與使用邊界手寫 Transformer 并不是為了讓 PyTorch 用戶拋棄官方 API。相反它適合在以下幾個場景里發揮作用第一課程作業或實驗報告。CS 336 這類課程的核心目標就是讓學生脫離高級封裝理解語言模型內部到底發生了什么。自己實現 Linear、Embedding、LayerNorm 之后再去看nn.TransformerEncoder的實現會明顯感覺到“圖層”不再是黑盒。第二算法面試準備。很多公司面試會問“Transformer 參數初始化為什么用 Xavier”“Embedding 梯度怎么回傳”“softmax 反向傳播的雅可比是什么”。這些問題的答案不親手寫一遍很難講得清楚。第三源碼閱讀和二次開發。當你想修改某一層的行為比如給 Attention 換成線性注意力或者把 LayerNorm 換成 RMSNorm如果只依賴高層 API往往不知道改哪里但手寫過底層模塊之后你能很快定位到具體實現位置。同時也要明確使用邊界。這套手寫代碼主要用于教學和原理驗證不推薦直接用于大規模訓練。手寫循環、GPU 算子融合、FlashAttention 這些工程優化都不在范圍內。如果目標是在生產環境訓練大模型應該直接使用成熟框架PyTorch 的nn.Transformer、Hugging Face 的transformers或者 Megatron-LM 等。另一個必須強調的是合規邊界。課程代碼、作業要求、課件內容都有各自的版權和使用協議本文只講解通用技術原理和常見實踐不復制任何課程原始材料。訓練時不要使用未經授權的數據、人臉信息、版權文本或用戶隱私數據盡量在本地構造的玩具數據集上驗證。3. 環境準備與前置條件手寫層不需要復雜環境。建議使用 Python 3.9 以上版本和 PyTorch 2.x安裝命令pip install torch numpy tqdm如果沒有 GPU直接用 CPU 跑小模型即可。這里的關鍵不是訓練速度而是驗證前向輸出形狀與反向傳播梯度是否正確。建議的工程目錄結構transformer_from_scratch/ ├── layers/ │ ├── __init__.py │ ├── linear.py │ ├── embedding.py │ ├── layernorm.py │ └── attention.py ├── init.py ├── model.py ├── train.py └── tests/ └── test_gradcheck.py這種按層拆分的結構能讓你在調試時只關心單個模塊而不會一次面對整個網絡。動手之前還需要確認三塊數學基礎矩陣乘法求導公式、鏈式法則、softmax 的雅可比矩陣。不需要研究得很深但至少要知道對于一個線性變換y x W從上層傳回的梯度dy應該分別回傳到dx和dW公式是dx dy W.T、dW x.T dy。后面所有代碼都建立在這條基本規則上。4. 從零實現核心模塊4.1 手寫 Linear LayerLinear Layer 是最基礎的計算單元公式很簡單y x W b。但參數初始化和梯度回傳是兩個容易被忽略的細節。第一個細節是參數初始化。下面用 Xavier 均勻初始化它的上下界是bound sqrt(6 / (fan_in fan_out))。這種初始化適合線性層和 tanh 激活函數能讓每層輸出的方差保持在一個合理范圍內避免深層網絡時梯度消失或爆炸。第二個細節是反向傳播。除了dx要回傳給上一層還要計算dW和db供優化器更新。這里有一個常見錯誤把dW寫成了grad_output.T x正確寫法是x.T grad_output。推導方法是對損失函數求 W 的偏導保持張量形狀對齊即可。import torch import math class LinearLayer: def __init__(self, in_features, out_features, biasTrue): self.in_features in_features self.out_features out_features self.bias bias # Xavier 均勻初始化 bound math.sqrt(6.0 / (in_features out_features)) self.W torch.empty(in_features, out_features).uniform_(-bound, bound) self.b torch.zeros(out_features) if bias else None def forward(self, x): self.x x out x self.W if self.b is not None: out out self.b return out def backward(self, grad_output): dx grad_output self.W.T dW self.x.T grad_output db grad_output.sum(dim0) if self.b is not None else None return dx, dW, db這段代碼最值得注意的地方是forward里保存了self.x。這是手寫反向傳播的基礎backward 需要用到前向時輸入的原始值。很多新手寫的反向傳播出錯不是因為公式推導錯而是因為前向沒有保存必要的中間變量。4.2 手寫 Embedding LayerEmbedding Layer 本質上是一個查表操作。輸入是 token id 序列輸出是形狀為[batch_size, seq_len, embedding_dim]的向量序列。它的反向傳播很有特點并不是所有 embedding 行都會被更新只有當前 batch 中實際被查到的 token 行才有梯度。實現方式是用torch.index_add_把梯度累加到對應的行上。class EmbeddingLayer: def __init__(self, num_embeddings, embedding_dim, init_scale0.1): self.num_embeddings num_embeddings self.embedding_dim embedding_dim # 常用做法均勻分布 [-init_scale, init_scale] self.E torch.empty(num_embeddings, embedding_dim) self.E.uniform_(-init_scale, init_scale) def forward(self, tokens): self.tokens tokens return self.E[tokens] def backward(self, grad_output): dE torch.zeros_like(self.E) # 把梯度累加到被查到的行上 dE.index_add_(0, self.tokens.flatten(), grad_output.flatten(0, 1)) return dE這里有一個工程細節index_add_是原地操作調用前必須確保dE是全零矩陣。如果同一個 token 在一句話中出現多次它的梯度會自然累加這也符合 embedding 反向傳播的定義。你不需要為它手動去重。4.3 參數初始化決定 Transformer 能不能收斂的隱藏開關參數初始化在 PyTorch 使用里常常被忽略因為nn.Linear已經封裝好了默認初始化。但手寫模塊時每層初始化都必須自己處理。先看原始 Transformer 論文中的做法embedding 層使用均值 0、方差 1 的正態分布初始化之后乘以sqrt(d_model)。也就是說如果 embedding 矩陣初始化為N(0, 1)那么實際查表輸出要再乘一個sqrt(d_model)。這樣做的目的是讓輸入到 Attention 的向量范數不要太小避免 softmax 之前的所有點積都擠到極小值附近。GPT 風格的大模型通常采用更小的初始化方差例如N(0, 0.02)。原因很簡單模型層數越深殘差分支積累的隱狀態方差越大逐層用小方差初始化可以抑制這種累積。如果你在深層 Transformer 里發現訓練初期 logits 出現 NaN或者 loss 長時間不下降首先排查初始化方差是否過大。下面給出一個簡單的參數初始化工具函數def init_linear_weight(W, modexavier): if mode xavier: fan_in, fan_out W.shape bound math.sqrt(6.0 / (fan_in fan_out)) W.uniform_(-bound, bound) elif mode gpt: nn.init.normal_(W, std0.02) else: raise ValueError(fUnknown init mode: {mode}) def init_embedding(E, d_model, modenormal): if mode normal: E.normal_(mean0.0, std1.0) E.mul_(math.sqrt(d_model)) elif mode uniform: scale 0.1 E.uniform_(-scale, scale) else: raise ValueError(fUnknown init mode: {mode})從實踐角度看訓練一個小模型時建議 Linear 層優先用 Xavier 均勻初始化Embedding 層可以用uniform(-0.1, 0.1)這種比較保險的小范圍初始化。如果訓練不收斂再嘗試 GPT 風格的N(0, 0.02)并配合 warmup 學習率調度。4.4 手寫反向傳播并用 gradcheck 驗證手寫反向傳播最難的不是某一層的公式而是整個鏈路的梯度形狀。一個效率很高的驗證方式是用torch.autograd.gradcheck。它會用數值差分計算梯度再和你的手動 backward 梯度對比誤差在閾值內就說明反向傳播正確。from torch.autograd import gradcheck torch.manual_seed(0) linear LinearLayer(64, 128) x torch.randn(2, 32, 64, requires_gradTrue, dtypetorch.float64) # 注意gradcheck 需要高精度通常使用 float64 linear.W linear.W.double() linear.b linear.b.double() output linear.forward(x) grad_output torch.randn_like(output) dx_manual, dW_manual, db_manual linear.backward(grad_output) # 用 autograd 計算參考梯度 output_double x.double() linear.W.double() linear.b.double() dx_auto torch.autograd.grad(output_double.sum(), x)[0] print(dx 誤差:, (dx_manual - dx_auto).abs().max().item())gradcheck函數更嚴格可以直接傳一個自定義函數def linear_forward_wrapper(W, b, x): return x W b linear_layer LinearLayer(16, 32) linear_layer.W linear_layer.W.double() linear_layer.b linear_layer.b.double() x torch.randn(4, 16, dtypetorch.float64, requires_gradTrue) W linear_layer.W.clone().requires_grad_(True) b linear_layer.b.clone().requires_grad_(True) check gradcheck( lambda w, b_: linear_forward_wrapper(w, b_, x), (W, b), eps1e-6, atol1e-5, ) print(gradcheck passed:, check)如果你是在課程作業里驗證反向傳播務必把所有參數和輸入都轉成float64。數值差分在float32下誤差會很大導致明明公式正確gradcheck卻報錯。4.5 LayerNorm 與殘差Transformer 穩定訓練的第二道防線LayerNorm 在 Transformer 中負責把每一層的輸出拉回穩定的數值范圍。手寫時可以直接使用簡化形式y (x - mean) / sqrt(var eps)反向傳播用一個更緊湊的公式class LayerNorm: def __init__(self, dim, eps1e-5): self.gamma torch.ones(dim) self.beta torch.zeros(dim) self.eps eps def forward(self, x): self.x x self.mean x.mean(dim-1, keepdimTrue) self.var x.var(dim-1, unbiasedFalse, keepdimTrue) self.x_hat (x - self.mean) / torch.sqrt(self.var self.eps) return self.gamma * self.x_hat self.beta def backward(self, grad_output): N self.x.shape[-1] dy grad_output * self.gamma dx (dy - dy.mean(dim-1, keepdimTrue) - self.x_hat * (dy * self.x_hat).mean(dim-1, keepdimTrue) ) / torch.sqrt(self.var self.eps) dgamma (grad_output * self.x_hat).sum(dim(0, 1)) dbeta grad_output.sum(dim(0, 1)) return dx, dgamma, dbeta這個反向公式的推導過程相當經典。核心思路是歸一化操作同時依賴mean和var所以dx中必須包含dy.mean以及dy * x_hat的均值項。很多人手寫 LayerNorm 反向時少了一項導致gradcheck一直不過幾乎都是這個原因。殘差連接則是把輸入和子層輸出相加它的反向傳播不需要任何額外計算梯度在相加處直接“分叉”。實現時只要在 forward 里保存原始輸入即可。4.6 組裝一個最小 Transformer Block有了 Linear、Embedding、LayerNorm 和殘差可以組裝一個最小 Transformer Block。這里為了控制篇幅Attention 部分不手寫 backward而是用 PyTorch 基礎算子實現 forward反向交給 autograd。import torch import torch.nn.functional as F class MultiHeadAttention: def __init__(self, d_model, n_heads): self.d_model d_model self.n_heads n_heads self.head_dim d_model // n_heads self.Wq torch.empty(d_model, d_model) self.Wk torch.empty(d_model, d_model) self.Wv torch.empty(d_model, d_model) self.Wo torch.empty(d_model, d_model) for w in [self.Wq, self.Wk, self.Wv, self.Wo]: nn.init.xavier_uniform_(w) def forward(self, x): B, T, D x.shape q x self.Wq k x self.Wk v x self.Wv q q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) k k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) v v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) scores q k.transpose(-2, -1) / (self.head_dim ** 0.5) attn torch.softmax(scores, dim-1) out attn v out out.transpose(1, 2).reshape(B, T, D) return out self.Wo這里的關鍵知識點是縮放因子。Attention 公式里scores q k.T / sqrt(d_head)不能省略。如果不做縮放點積結果會隨著d_head增大而變大softmax 會進入飽和區梯度接近 0訓練會非常慢。5. 功能測試與效果驗證5.1 測試 Linear 前向形狀第一個測試很簡單確認輸出形狀正確。torch.manual_seed(0) linear LinearLayer(64, 128) x torch.randn(2, 32, 64) out linear.forward(x) print(Linear input shape:, x.shape) print(Linear output shape:, out.shape) # 期望輸出 (2, 32, 128)如果這里形狀不對問題幾乎一定出在W的定義方式上。W應該是(in_features, out_features)這樣x W才能把最后一維從 64 映射到 128。5.2 測試 Embedding 反向是否正確構造一個只有 0、1、2 三個 token 的序列驗證 embedding 各行的梯度。torch.manual_seed(0) embed EmbeddingLayer(num_embeddings3, embedding_dim8) tokens torch.tensor([[0, 1], [0, 2]]) out embed.forward(tokens) grad_output torch.randn_like(out) dE embed.backward(grad_output) print(dE shape:, dE.shape) print(dE row 0:, dE[0])判斷依據第 0 行梯度應該等于輸入中 token 0 對應位置梯度的累加。如果第 1 行或第 2 行也出現了不該有的梯度說明index_add_的索引寫錯了。5.3 用 gradcheck 驗證全部手寫層強烈建議給每個手寫層都寫一個 gradcheck。以 LinearLayer 為例用前文包裝函數方式驗證 W 和 b 的梯度EmbeddingLayer 則驗證 E 矩陣的梯度。def embedding_forward_wrapper(E, tokens): return E[tokens] E torch.randn(10, 16, dtypetorch.float64, requires_gradTrue) tokens torch.tensor([[1, 2, 3], [4, 5, 6]], dtypetorch.long) gradcheck_check gradcheck(embedding_forward_wrapper, (E, tokens)) print(Embedding gradcheck:, gradcheck_check)注意gradcheck要求輸入張量至少有一個是requires_gradTrue且所有浮點張量都建議使用float64。整數 token 張量不需要梯度。5.4 端到端小模型訓練觀察 loss 是否下降這是最后一個驗證能證明整個鏈路從參數初始化到反向傳播都沒有大問題。import torch import torch.optim as optim torch.manual_seed(0) vocab_size 64 d_model 32 n_heads 4 block MultiHeadAttention(d_modeld_model, n_headsn_heads) embed EmbeddingLayer(vocab_size, d_model) x torch.randint(0, vocab_size, (2, 8)) y torch.randint(0, vocab_size, (2, 8)) optim_params [ {params: [block.Wq, block.Wk, block.Wv, block.Wo]}, {params: [embed.E]}, ] optimizer optim.Adam(optim_params, lr1e-3) loss_fn torch.nn.CrossEntropyLoss() for step in range(20): h embed.forward(x) h block.forward(h) logits h embed.E.T loss loss_fn(logits.view(-1, vocab_size), y.view(-1)) optimizer.zero_grad() # 這里故意用 autograd 計算整條鏈路梯度因為 MultiHeadAttention 沒有手寫 backward loss.backward() optimizer.step() if step % 5 0: print(step, step, loss, loss.item())這段代碼用到了一個取巧但合理的設計embedding 矩陣既負責查表也作為輸出投影矩陣。這種“共享嵌入”的寫法在 GPT 系列中很常見。更好的實現是單獨定義一個輸出 Linear 層但從教學角度看共享嵌入能顯著減少參數。判斷標準如果 loss 在 20 步內明顯下降說明參數初始化、前向、反向、優化器調度全部工作正常如果 loss 原地不動優先檢查學習率是否太大或者太小其次檢查初始化方差。6. 從“手寫層”到“接口化調用”批量實驗與可復用設計這個項目沒有 HTTP API也不需要提供 Web 服務。但從工程化角度看你可以把自定義層設計成統一的“前向 反向”接口方便批量跑實驗。比如定義一個 BaseLayer 風格class BaseLayer: def forward(self, x): raise NotImplementedError def backward(self, grad_output): raise NotImplementedErrorLinearLayer、EmbeddingLayer、LayerNorm 都繼承這個接口這樣在訓練循環里可以統一調度。批量實驗最常見的場景是掃描不同初始化模式對收斂速度的影響。可以寫一個小腳本init_modes [xavier, gpt, uniform] learning_rates [1e-4, 1e-3, 1e-2] for init_mode in init_modes: for lr in learning_rates: model build_transformer(init_modeinit_mode) final_loss train_short(model, lrlr) print(init_mode, lr, final_loss)當你把“參數初始化”和“反向傳播”都變成可配置項之后就能系統性地觀察為什么 GPT 風格的小方差初始化在深層模型中表現更好為什么 Xavier 在淺層模型中已經足夠。這種實驗往往比單純閱讀論文更有價值。如果你確實想把模塊暴露成服務也可以額外包一層 FastAPI但這屬于教學項目之外的擴展不推薦在原理學習階段引入。7. 資源占用與性能觀察手寫模塊的最大問題是性能遠低于 PyTorch 原生算子。因為這里為了教學把 forward 和 backward 的中間張量都保存在內存中也沒有做算子融合。如果是大型模型訓練這種寫法會占用大量顯存。觀察資源使用可以這樣寫import time import psutil x torch.randn(4, 128, 256) linear LinearLayer(256, 256) # warmup linear.forward(x) t0 time.time() for _ in range(10): out linear.forward(x) _, _, _ linear.backward(torch.randn_like(out)) t1 time.time() print(平均耗時:, (t1 - t0) / 10) mem psutil.Process().memory_info().rss / 1024 / 1024 print(當前內存占用 MB:, mem)如果訓練時顯存不足有幾種通用降載方法減小 batch size減小序列長度降低 d_model 和 n_heads把不必要的中間變量從 forward 中移出例如 forward 結束后只保留反向需要用到的張量不保存完整 attention 矩陣對大模型訓練改用半精度或混合精度。從性能觀察的角度說手寫版的價值不在于快而在于你能清楚看到每個操作的時間和空間開銷。這也是理解后續 FlashAttention 為什么要合并算子、減少訪存的起點。8. 常見問題與排查方法問題現象可能原因排查方式解決方案gradcheck 一直報錯沒有開啟 float64檢查輸入張量和參數 dtype全部轉成torch.float64Linear 反向維度報錯dW寫成grad_output.T x打印x.shape、grad_output.shape改為x.T grad_outputEmbedding 梯度全部為 0查表后沒有保存 token 索引檢查forward是否保存self.tokens在 forward 中保存輸入 tokenloss 不下降學習率不合適或初始化方差太小打印梯度均值、方差調大 learning rate 或改用 GPT initloss 出現 NaN初始化方差太大或 Attention 沒有縮放打印 logits 是否存在 inf使用sqrt(d_head)縮放減小初始化方差LayerNorm 反向公式缺項忘記處理 mean 和 var 的依賴用 gradcheck 定位使用包含dy.mean與dy*x_hat項的簡化公式Attention 維度 reshape 錯誤Q、K、V 的 view/transpose 順序不對打印q.shape、k.shape、scores.shape先 view 成[B, T, n_heads, head_dim]再 transpose 成[B, n_heads, T, head_dim]CPU 訓練太慢batch size / 序列長度偏大用 time 統計 forward/backward 耗時減小 batch size 和 seq_len從上表能看出手寫 Transformer 的大部分報錯都集中在維度形狀和梯度公式上。我的建議是遇到問題不要直接搜答案先打印梯度的 shape 和數值范圍往往能更快定位。9. 最佳實踐與使用建議寫手寫 Transformer 項目時有幾條通用的工程建議值得遵守。第一按層拆分文件。不要把所有自定義層堆在一個 Python 文件里。按linear.py、embedding.py、layernorm.py、attention.py拆分每層只做一件事調試時直接定位對應文件。第二每寫一層就立即驗證一層。不要等到所有模塊寫完了再驗證。驗證工具就是torch.autograd.gradcheck它比任何肉眼觀察都可靠。每層驗證通過之后再進入下一步。第三先跑最小配置。第一個實驗建議使用vocab_size64、d_model32、n_heads4、batch_size2、seq_len8。這么小的配置可以在幾秒內完成完整訓練循環。如果最小配置都能跑通再逐步放大。第四保留一個 baseline 對照。用 PyTorch 官方nn.Linear、nn.Embedding或nn.TransformerEncoder實現同樣的網絡不斷對比前向輸出和梯度。這樣能快速排除“手寫層寫錯”和“訓練邏輯寫錯”兩類問題。第五模型文件、訓練代碼和測試代碼分目錄管理。避免把臨時調試代碼和正式代碼混在一起。如果你要批量跑初始化實驗建議加一個配置字典或者命令行參數。第六涉及真實數據時注意合規。本文教學項目只需要隨機 token 或自己構造的玩具數據不需要真實文本。如果后續擴展項目需要真實語料務必確認數據來源合法、授權完整不涉及個人隱私和版權內容。10. 總結與下一步這個手寫項目最值得驗證的功能就是“參數初始化 手寫反向傳播”的組合。你先跑通一層 Linear 的 gradcheck再試一個最小 Transformer 的 loss 下降就會理解為什么Transformer的每一層都需要精心設計梯度通路。最容易踩的坑有三個Linear 反向里dW的維度順序、LayerNorm 反向里缺少歸一化依賴項、Attention 里的縮放因子。這三個坑只要踩過一個以后再看到類似報錯就會很敏感。下一步建議做三件事第一給 MultiHeadAttention 也實現手寫 backward完整覆蓋所有模塊第二加入 causal mask讓 Attention 只能看到當前位置之前的 token變成一個真正的自回歸語言模型第三把學到的參數初始化經驗搬到一個小型文本數據集上觀察 loss 曲線在不同 init 和 learning rate 下的表現。做完這三步你對 Transformer 的掌握就不再停留在“會用層”的層面而是真正進入了“能寫層”的階段。 手寫Transformer核心模塊：從零實現Linear、Embedding與反向傳播
這次我們來看一項很多人學完 Transformer 但依然沒完全搞懂的事情不只是會調用nn.Transformer而是親手把 Linear Layer、Embedding Layer、參數初始化和反向傳播全部實現出來。本文對應的場景是 Stanford CS 336 3.3也就是在課程作業中要求用最樸素的 PyTorch 張量操作搭建語言模型組件并通過梯度檢查驗證手寫反向傳播的正確性。標題里的四個關鍵詞其實就對應了 Transformer 落地中四個最容易翻車的點參數初始化決定模型能不能收斂反向傳播決定你寫的網絡是不是真的“能學”Linear Layer 是所有前饋計算的地基Embedding Layer 則承擔著離散 token 到連續向量的映射。如果這幾個模塊只是靠nn.Linear和nn.Embedding一帶而過那遇到 loss 不降、梯度 NaN、維度對不上這類問題的時候往往只能瞎猜。自己手拼一遍之后很多報錯一眼就能看出問題。這篇文章會從零實現一個可訓練的最小 Transformer包含前向傳播、手寫 backward、參數初始化、梯度驗證最后在一個小模型上完成端到端訓練測試。全程使用 CPU 即可運行不依賴 GPU也不需要下載任何預訓練權重。適合正在啃 Transformer 源碼、準備算法面試、或者做課程作業時想把原理徹底吃透的讀者。1. 核心能力速覽先把這篇文章涉及的技術目標和運行約束講清楚。下面的表格不是項目 README而是你讀完這篇文章后應該具備的“能力驗收清單”。能力項說明項目類型從零手寫 Transformer 核心模塊的教學實現來源Stanford CS 336 Language Modeling 課程 3.3 小節主題核心模塊Linear Layer、Embedding Layer、LayerNorm、參數初始化、手寫反向傳播技術棧Python 3.9、PyTorch 2.x、NumPy 可選運行方式本地腳本運行支持 Jupyter Notebook硬件要求CPU 可跑顯存不是必需是否依賴 nn.Linear不依賴手動實現 forward/backward是否依賴 nn.Embedding不依賴手動實現查表與梯度回傳反向傳播方式手動推導并實現通過 torch.autograd.gradcheck 驗證驗證方式形狀測試、梯度對比、小模型端到端 loss 下降適合人群學習 Transformer 原理、面試準備、課程作業、源碼閱讀從表格可以看出這是一篇“先別用庫、先把原理跑通”的教程。它的核心產出不是一個大而全的訓練框架而是一套你可以自己掌控每個張量形狀和每個梯度去向的最小實現。2. 適用場景與使用邊界手寫 Transformer 并不是為了讓 PyTorch 用戶拋棄官方 API。相反它適合在以下幾個場景里發揮作用第一課程作業或實驗報告。CS 336 這類課程的核心目標就是讓學生脫離高級封裝理解語言模型內部到底發生了什么。自己實現 Linear、Embedding、LayerNorm 之后再去看nn.TransformerEncoder的實現會明顯感覺到“圖層”不再是黑盒。第二算法面試準備。很多公司面試會問“Transformer 參數初始化為什么用 Xavier”“Embedding 梯度怎么回傳”“softmax 反向傳播的雅可比是什么”。這些問題的答案不親手寫一遍很難講得清楚。第三源碼閱讀和二次開發。當你想修改某一層的行為比如給 Attention 換成線性注意力或者把 LayerNorm 換成 RMSNorm如果只依賴高層 API往往不知道改哪里但手寫過底層模塊之后你能很快定位到具體實現位置。同時也要明確使用邊界。這套手寫代碼主要用于教學和原理驗證不推薦直接用于大規模訓練。手寫循環、GPU 算子融合、FlashAttention 這些工程優化都不在范圍內。如果目標是在生產環境訓練大模型應該直接使用成熟框架PyTorch 的nn.Transformer、Hugging Face 的transformers或者 Megatron-LM 等。另一個必須強調的是合規邊界。課程代碼、作業要求、課件內容都有各自的版權和使用協議本文只講解通用技術原理和常見實踐不復制任何課程原始材料。訓練時不要使用未經授權的數據、人臉信息、版權文本或用戶隱私數據盡量在本地構造的玩具數據集上驗證。3. 環境準備與前置條件手寫層不需要復雜環境。建議使用 Python 3.9 以上版本和 PyTorch 2.x安裝命令pip install torch numpy tqdm如果沒有 GPU直接用 CPU 跑小模型即可。這里的關鍵不是訓練速度而是驗證前向輸出形狀與反向傳播梯度是否正確。建議的工程目錄結構transformer_from_scratch/ ├── layers/ │ ├── __init__.py │ ├── linear.py │ ├── embedding.py │ ├── layernorm.py │ └── attention.py ├── init.py ├── model.py ├── train.py └── tests/ └── test_gradcheck.py這種按層拆分的結構能讓你在調試時只關心單個模塊而不會一次面對整個網絡。動手之前還需要確認三塊數學基礎矩陣乘法求導公式、鏈式法則、softmax 的雅可比矩陣。不需要研究得很深但至少要知道對于一個線性變換y x W從上層傳回的梯度dy應該分別回傳到dx和dW公式是dx dy W.T、dW x.T dy。后面所有代碼都建立在這條基本規則上。4. 從零實現核心模塊4.1 手寫 Linear LayerLinear Layer 是最基礎的計算單元公式很簡單y x W b。但參數初始化和梯度回傳是兩個容易被忽略的細節。第一個細節是參數初始化。下面用 Xavier 均勻初始化它的上下界是bound sqrt(6 / (fan_in fan_out))。這種初始化適合線性層和 tanh 激活函數能讓每層輸出的方差保持在一個合理范圍內避免深層網絡時梯度消失或爆炸。第二個細節是反向傳播。除了dx要回傳給上一層還要計算dW和db供優化器更新。這里有一個常見錯誤把dW寫成了grad_output.T x正確寫法是x.T grad_output。推導方法是對損失函數求 W 的偏導保持張量形狀對齊即可。import torch import math class LinearLayer: def __init__(self, in_features, out_features, biasTrue): self.in_features in_features self.out_features out_features self.bias bias # Xavier 均勻初始化 bound math.sqrt(6.0 / (in_features out_features)) self.W torch.empty(in_features, out_features).uniform_(-bound, bound) self.b torch.zeros(out_features) if bias else None def forward(self, x): self.x x out x self.W if self.b is not None: out out self.b return out def backward(self, grad_output): dx grad_output self.W.T dW self.x.T grad_output db grad_output.sum(dim0) if self.b is not None else None return dx, dW, db這段代碼最值得注意的地方是forward里保存了self.x。這是手寫反向傳播的基礎backward 需要用到前向時輸入的原始值。很多新手寫的反向傳播出錯不是因為公式推導錯而是因為前向沒有保存必要的中間變量。4.2 手寫 Embedding LayerEmbedding Layer 本質上是一個查表操作。輸入是 token id 序列輸出是形狀為[batch_size, seq_len, embedding_dim]的向量序列。它的反向傳播很有特點并不是所有 embedding 行都會被更新只有當前 batch 中實際被查到的 token 行才有梯度。實現方式是用torch.index_add_把梯度累加到對應的行上。class EmbeddingLayer: def __init__(self, num_embeddings, embedding_dim, init_scale0.1): self.num_embeddings num_embeddings self.embedding_dim embedding_dim # 常用做法均勻分布 [-init_scale, init_scale] self.E torch.empty(num_embeddings, embedding_dim) self.E.uniform_(-init_scale, init_scale) def forward(self, tokens): self.tokens tokens return self.E[tokens] def backward(self, grad_output): dE torch.zeros_like(self.E) # 把梯度累加到被查到的行上 dE.index_add_(0, self.tokens.flatten(), grad_output.flatten(0, 1)) return dE這里有一個工程細節index_add_是原地操作調用前必須確保dE是全零矩陣。如果同一個 token 在一句話中出現多次它的梯度會自然累加這也符合 embedding 反向傳播的定義。你不需要為它手動去重。4.3 參數初始化決定 Transformer 能不能收斂的隱藏開關參數初始化在 PyTorch 使用里常常被忽略因為nn.Linear已經封裝好了默認初始化。但手寫模塊時每層初始化都必須自己處理。先看原始 Transformer 論文中的做法embedding 層使用均值 0、方差 1 的正態分布初始化之后乘以sqrt(d_model)。也就是說如果 embedding 矩陣初始化為N(0, 1)那么實際查表輸出要再乘一個sqrt(d_model)。這樣做的目的是讓輸入到 Attention 的向量范數不要太小避免 softmax 之前的所有點積都擠到極小值附近。GPT 風格的大模型通常采用更小的初始化方差例如N(0, 0.02)。原因很簡單模型層數越深殘差分支積累的隱狀態方差越大逐層用小方差初始化可以抑制這種累積。如果你在深層 Transformer 里發現訓練初期 logits 出現 NaN或者 loss 長時間不下降首先排查初始化方差是否過大。下面給出一個簡單的參數初始化工具函數def init_linear_weight(W, modexavier): if mode xavier: fan_in, fan_out W.shape bound math.sqrt(6.0 / (fan_in fan_out)) W.uniform_(-bound, bound) elif mode gpt: nn.init.normal_(W, std0.02) else: raise ValueError(fUnknown init mode: {mode}) def init_embedding(E, d_model, modenormal): if mode normal: E.normal_(mean0.0, std1.0) E.mul_(math.sqrt(d_model)) elif mode uniform: scale 0.1 E.uniform_(-scale, scale) else: raise ValueError(fUnknown init mode: {mode})從實踐角度看訓練一個小模型時建議 Linear 層優先用 Xavier 均勻初始化Embedding 層可以用uniform(-0.1, 0.1)這種比較保險的小范圍初始化。如果訓練不收斂再嘗試 GPT 風格的N(0, 0.02)并配合 warmup 學習率調度。4.4 手寫反向傳播并用 gradcheck 驗證手寫反向傳播最難的不是某一層的公式而是整個鏈路的梯度形狀。一個效率很高的驗證方式是用torch.autograd.gradcheck。它會用數值差分計算梯度再和你的手動 backward 梯度對比誤差在閾值內就說明反向傳播正確。from torch.autograd import gradcheck torch.manual_seed(0) linear LinearLayer(64, 128) x torch.randn(2, 32, 64, requires_gradTrue, dtypetorch.float64) # 注意gradcheck 需要高精度通常使用 float64 linear.W linear.W.double() linear.b linear.b.double() output linear.forward(x) grad_output torch.randn_like(output) dx_manual, dW_manual, db_manual linear.backward(grad_output) # 用 autograd 計算參考梯度 output_double x.double() linear.W.double() linear.b.double() dx_auto torch.autograd.grad(output_double.sum(), x)[0] print(dx 誤差:, (dx_manual - dx_auto).abs().max().item())gradcheck函數更嚴格可以直接傳一個自定義函數def linear_forward_wrapper(W, b, x): return x W b linear_layer LinearLayer(16, 32) linear_layer.W linear_layer.W.double() linear_layer.b linear_layer.b.double() x torch.randn(4, 16, dtypetorch.float64, requires_gradTrue) W linear_layer.W.clone().requires_grad_(True) b linear_layer.b.clone().requires_grad_(True) check gradcheck( lambda w, b_: linear_forward_wrapper(w, b_, x), (W, b), eps1e-6, atol1e-5, ) print(gradcheck passed:, check)如果你是在課程作業里驗證反向傳播務必把所有參數和輸入都轉成float64。數值差分在float32下誤差會很大導致明明公式正確gradcheck卻報錯。4.5 LayerNorm 與殘差Transformer 穩定訓練的第二道防線LayerNorm 在 Transformer 中負責把每一層的輸出拉回穩定的數值范圍。手寫時可以直接使用簡化形式y (x - mean) / sqrt(var eps)反向傳播用一個更緊湊的公式class LayerNorm: def __init__(self, dim, eps1e-5): self.gamma torch.ones(dim) self.beta torch.zeros(dim) self.eps eps def forward(self, x): self.x x self.mean x.mean(dim-1, keepdimTrue) self.var x.var(dim-1, unbiasedFalse, keepdimTrue) self.x_hat (x - self.mean) / torch.sqrt(self.var self.eps) return self.gamma * self.x_hat self.beta def backward(self, grad_output): N self.x.shape[-1] dy grad_output * self.gamma dx (dy - dy.mean(dim-1, keepdimTrue) - self.x_hat * (dy * self.x_hat).mean(dim-1, keepdimTrue) ) / torch.sqrt(self.var self.eps) dgamma (grad_output * self.x_hat).sum(dim(0, 1)) dbeta grad_output.sum(dim(0, 1)) return dx, dgamma, dbeta這個反向公式的推導過程相當經典。核心思路是歸一化操作同時依賴mean和var所以dx中必須包含dy.mean以及dy * x_hat的均值項。很多人手寫 LayerNorm 反向時少了一項導致gradcheck一直不過幾乎都是這個原因。殘差連接則是把輸入和子層輸出相加它的反向傳播不需要任何額外計算梯度在相加處直接“分叉”。實現時只要在 forward 里保存原始輸入即可。4.6 組裝一個最小 Transformer Block有了 Linear、Embedding、LayerNorm 和殘差可以組裝一個最小 Transformer Block。這里為了控制篇幅Attention 部分不手寫 backward而是用 PyTorch 基礎算子實現 forward反向交給 autograd。import torch import torch.nn.functional as F class MultiHeadAttention: def __init__(self, d_model, n_heads): self.d_model d_model self.n_heads n_heads self.head_dim d_model // n_heads self.Wq torch.empty(d_model, d_model) self.Wk torch.empty(d_model, d_model) self.Wv torch.empty(d_model, d_model) self.Wo torch.empty(d_model, d_model) for w in [self.Wq, self.Wk, self.Wv, self.Wo]: nn.init.xavier_uniform_(w) def forward(self, x): B, T, D x.shape q x self.Wq k x self.Wk v x self.Wv q q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) k k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) v v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) scores q k.transpose(-2, -1) / (self.head_dim ** 0.5) attn torch.softmax(scores, dim-1) out attn v out out.transpose(1, 2).reshape(B, T, D) return out self.Wo這里的關鍵知識點是縮放因子。Attention 公式里scores q k.T / sqrt(d_head)不能省略。如果不做縮放點積結果會隨著d_head增大而變大softmax 會進入飽和區梯度接近 0訓練會非常慢。5. 功能測試與效果驗證5.1 測試 Linear 前向形狀第一個測試很簡單確認輸出形狀正確。torch.manual_seed(0) linear LinearLayer(64, 128) x torch.randn(2, 32, 64) out linear.forward(x) print(Linear input shape:, x.shape) print(Linear output shape:, out.shape) # 期望輸出 (2, 32, 128)如果這里形狀不對問題幾乎一定出在W的定義方式上。W應該是(in_features, out_features)這樣x W才能把最后一維從 64 映射到 128。5.2 測試 Embedding 反向是否正確構造一個只有 0、1、2 三個 token 的序列驗證 embedding 各行的梯度。torch.manual_seed(0) embed EmbeddingLayer(num_embeddings3, embedding_dim8) tokens torch.tensor([[0, 1], [0, 2]]) out embed.forward(tokens) grad_output torch.randn_like(out) dE embed.backward(grad_output) print(dE shape:, dE.shape) print(dE row 0:, dE[0])判斷依據第 0 行梯度應該等于輸入中 token 0 對應位置梯度的累加。如果第 1 行或第 2 行也出現了不該有的梯度說明index_add_的索引寫錯了。5.3 用 gradcheck 驗證全部手寫層強烈建議給每個手寫層都寫一個 gradcheck。以 LinearLayer 為例用前文包裝函數方式驗證 W 和 b 的梯度EmbeddingLayer 則驗證 E 矩陣的梯度。def embedding_forward_wrapper(E, tokens): return E[tokens] E torch.randn(10, 16, dtypetorch.float64, requires_gradTrue) tokens torch.tensor([[1, 2, 3], [4, 5, 6]], dtypetorch.long) gradcheck_check gradcheck(embedding_forward_wrapper, (E, tokens)) print(Embedding gradcheck:, gradcheck_check)注意gradcheck要求輸入張量至少有一個是requires_gradTrue且所有浮點張量都建議使用float64。整數 token 張量不需要梯度。5.4 端到端小模型訓練觀察 loss 是否下降這是最后一個驗證能證明整個鏈路從參數初始化到反向傳播都沒有大問題。import torch import torch.optim as optim torch.manual_seed(0) vocab_size 64 d_model 32 n_heads 4 block MultiHeadAttention(d_modeld_model, n_headsn_heads) embed EmbeddingLayer(vocab_size, d_model) x torch.randint(0, vocab_size, (2, 8)) y torch.randint(0, vocab_size, (2, 8)) optim_params [ {params: [block.Wq, block.Wk, block.Wv, block.Wo]}, {params: [embed.E]}, ] optimizer optim.Adam(optim_params, lr1e-3) loss_fn torch.nn.CrossEntropyLoss() for step in range(20): h embed.forward(x) h block.forward(h) logits h embed.E.T loss loss_fn(logits.view(-1, vocab_size), y.view(-1)) optimizer.zero_grad() # 這里故意用 autograd 計算整條鏈路梯度因為 MultiHeadAttention 沒有手寫 backward loss.backward() optimizer.step() if step % 5 0: print(step, step, loss, loss.item())這段代碼用到了一個取巧但合理的設計embedding 矩陣既負責查表也作為輸出投影矩陣。這種“共享嵌入”的寫法在 GPT 系列中很常見。更好的實現是單獨定義一個輸出 Linear 層但從教學角度看共享嵌入能顯著減少參數。判斷標準如果 loss 在 20 步內明顯下降說明參數初始化、前向、反向、優化器調度全部工作正常如果 loss 原地不動優先檢查學習率是否太大或者太小其次檢查初始化方差。6. 從“手寫層”到“接口化調用”批量實驗與可復用設計這個項目沒有 HTTP API也不需要提供 Web 服務。但從工程化角度看你可以把自定義層設計成統一的“前向 反向”接口方便批量跑實驗。比如定義一個 BaseLayer 風格class BaseLayer: def forward(self, x): raise NotImplementedError def backward(self, grad_output): raise NotImplementedErrorLinearLayer、EmbeddingLayer、LayerNorm 都繼承這個接口這樣在訓練循環里可以統一調度。批量實驗最常見的場景是掃描不同初始化模式對收斂速度的影響。可以寫一個小腳本init_modes [xavier, gpt, uniform] learning_rates [1e-4, 1e-3, 1e-2] for init_mode in init_modes: for lr in learning_rates: model build_transformer(init_modeinit_mode) final_loss train_short(model, lrlr) print(init_mode, lr, final_loss)當你把“參數初始化”和“反向傳播”都變成可配置項之后就能系統性地觀察為什么 GPT 風格的小方差初始化在深層模型中表現更好為什么 Xavier 在淺層模型中已經足夠。這種實驗往往比單純閱讀論文更有價值。如果你確實想把模塊暴露成服務也可以額外包一層 FastAPI但這屬于教學項目之外的擴展不推薦在原理學習階段引入。7. 資源占用與性能觀察手寫模塊的最大問題是性能遠低于 PyTorch 原生算子。因為這里為了教學把 forward 和 backward 的中間張量都保存在內存中也沒有做算子融合。如果是大型模型訓練這種寫法會占用大量顯存。觀察資源使用可以這樣寫import time import psutil x torch.randn(4, 128, 256) linear LinearLayer(256, 256) # warmup linear.forward(x) t0 time.time() for _ in range(10): out linear.forward(x) _, _, _ linear.backward(torch.randn_like(out)) t1 time.time() print(平均耗時:, (t1 - t0) / 10) mem psutil.Process().memory_info().rss / 1024 / 1024 print(當前內存占用 MB:, mem)如果訓練時顯存不足有幾種通用降載方法減小 batch size減小序列長度降低 d_model 和 n_heads把不必要的中間變量從 forward 中移出例如 forward 結束后只保留反向需要用到的張量不保存完整 attention 矩陣對大模型訓練改用半精度或混合精度。從性能觀察的角度說手寫版的價值不在于快而在于你能清楚看到每個操作的時間和空間開銷。這也是理解后續 FlashAttention 為什么要合并算子、減少訪存的起點。8. 常見問題與排查方法問題現象可能原因排查方式解決方案gradcheck 一直報錯沒有開啟 float64檢查輸入張量和參數 dtype全部轉成torch.float64Linear 反向維度報錯dW寫成grad_output.T x打印x.shape、grad_output.shape改為x.T grad_outputEmbedding 梯度全部為 0查表后沒有保存 token 索引檢查forward是否保存self.tokens在 forward 中保存輸入 tokenloss 不下降學習率不合適或初始化方差太小打印梯度均值、方差調大 learning rate 或改用 GPT initloss 出現 NaN初始化方差太大或 Attention 沒有縮放打印 logits 是否存在 inf使用sqrt(d_head)縮放減小初始化方差LayerNorm 反向公式缺項忘記處理 mean 和 var 的依賴用 gradcheck 定位使用包含dy.mean與dy*x_hat項的簡化公式Attention 維度 reshape 錯誤Q、K、V 的 view/transpose 順序不對打印q.shape、k.shape、scores.shape先 view 成[B, T, n_heads, head_dim]再 transpose 成[B, n_heads, T, head_dim]CPU 訓練太慢batch size / 序列長度偏大用 time 統計 forward/backward 耗時減小 batch size 和 seq_len從上表能看出手寫 Transformer 的大部分報錯都集中在維度形狀和梯度公式上。我的建議是遇到問題不要直接搜答案先打印梯度的 shape 和數值范圍往往能更快定位。9. 最佳實踐與使用建議寫手寫 Transformer 項目時有幾條通用的工程建議值得遵守。第一按層拆分文件。不要把所有自定義層堆在一個 Python 文件里。按linear.py、embedding.py、layernorm.py、attention.py拆分每層只做一件事調試時直接定位對應文件。第二每寫一層就立即驗證一層。不要等到所有模塊寫完了再驗證。驗證工具就是torch.autograd.gradcheck它比任何肉眼觀察都可靠。每層驗證通過之后再進入下一步。第三先跑最小配置。第一個實驗建議使用vocab_size64、d_model32、n_heads4、batch_size2、seq_len8。這么小的配置可以在幾秒內完成完整訓練循環。如果最小配置都能跑通再逐步放大。第四保留一個 baseline 對照。用 PyTorch 官方nn.Linear、nn.Embedding或nn.TransformerEncoder實現同樣的網絡不斷對比前向輸出和梯度。這樣能快速排除“手寫層寫錯”和“訓練邏輯寫錯”兩類問題。第五模型文件、訓練代碼和測試代碼分目錄管理。避免把臨時調試代碼和正式代碼混在一起。如果你要批量跑初始化實驗建議加一個配置字典或者命令行參數。第六涉及真實數據時注意合規。本文教學項目只需要隨機 token 或自己構造的玩具數據不需要真實文本。如果后續擴展項目需要真實語料務必確認數據來源合法、授權完整不涉及個人隱私和版權內容。10. 總結與下一步這個手寫項目最值得驗證的功能就是“參數初始化 手寫反向傳播”的組合。你先跑通一層 Linear 的 gradcheck再試一個最小 Transformer 的 loss 下降就會理解為什么Transformer的每一層都需要精心設計梯度通路。最容易踩的坑有三個Linear 反向里dW的維度順序、LayerNorm 反向里缺少歸一化依賴項、Attention 里的縮放因子。這三個坑只要踩過一個以后再看到類似報錯就會很敏感。下一步建議做三件事第一給 MultiHeadAttention 也實現手寫 backward完整覆蓋所有模塊第二加入 causal mask讓 Attention 只能看到當前位置之前的 token變成一個真正的自回歸語言模型第三把學到的參數初始化經驗搬到一個小型文本數據集上觀察 loss 曲線在不同 init 和 learning rate 下的表現。做完這三步你對 Transformer 的掌握就不再停留在“會用層”的層面而是真正進入了“能寫層”的階段。