現(xiàn)Linear、Embedding與反向傳播)
這次我們來看一項(xiàng)很多人學(xué)完 Transformer 但依然沒完全搞懂的事情不只是會(huì)調(diào)用nn.Transformer而是親手把 Linear Layer、Embedding Layer、參數(shù)初始化和反向傳播全部實(shí)現(xiàn)出來。本文對應(yīng)的場景是 Stanford CS 336 3.3也就是在課程作業(yè)中要求用最樸素的 PyTorch 張量操作搭建語言模型組件并通過梯度檢查驗(yàn)證手寫反向傳播的正確性。標(biāo)題里的四個(gè)關(guān)鍵詞其實(shí)就對應(yīng)了 Transformer 落地中四個(gè)最容易翻車的點(diǎn)參數(shù)初始化決定模型能不能收斂反向傳播決定你寫的網(wǎng)絡(luò)是不是真的“能學(xué)”Linear Layer 是所有前饋計(jì)算的地基Embedding Layer 則承擔(dān)著離散 token 到連續(xù)向量的映射。如果這幾個(gè)模塊只是靠nn.Linear和nn.Embedding一帶而過那遇到 loss 不降、梯度 NaN、維度對不上這類問題的時(shí)候往往只能瞎猜。自己手拼一遍之后很多報(bào)錯(cuò)一眼就能看出問題。這篇文章會(huì)從零實(shí)現(xiàn)一個(gè)可訓(xùn)練的最小 Transformer包含前向傳播、手寫 backward、參數(shù)初始化、梯度驗(yàn)證最后在一個(gè)小模型上完成端到端訓(xùn)練測試。全程使用 CPU 即可運(yùn)行不依賴 GPU也不需要下載任何預(yù)訓(xùn)練權(quán)重。適合正在啃 Transformer 源碼、準(zhǔn)備算法面試、或者做課程作業(yè)時(shí)想把原理徹底吃透的讀者。1. 核心能力速覽先把這篇文章涉及的技術(shù)目標(biāo)和運(yùn)行約束講清楚。下面的表格不是項(xiàng)目 README而是你讀完這篇文章后應(yīng)該具備的“能力驗(yàn)收清單”。能力項(xiàng)說明項(xiàng)目類型從零手寫 Transformer 核心模塊的教學(xué)實(shí)現(xiàn)來源Stanford CS 336 Language Modeling 課程 3.3 小節(jié)主題核心模塊Linear Layer、Embedding Layer、LayerNorm、參數(shù)初始化、手寫反向傳播技術(shù)棧Python 3.9、PyTorch 2.x、NumPy 可選運(yùn)行方式本地腳本運(yùn)行支持 Jupyter Notebook硬件要求CPU 可跑顯存不是必需是否依賴 nn.Linear不依賴手動(dòng)實(shí)現(xiàn) forward/backward是否依賴 nn.Embedding不依賴手動(dòng)實(shí)現(xiàn)查表與梯度回傳反向傳播方式手動(dòng)推導(dǎo)并實(shí)現(xiàn)通過 torch.autograd.gradcheck 驗(yàn)證驗(yàn)證方式形狀測試、梯度對比、小模型端到端 loss 下降適合人群學(xué)習(xí) Transformer 原理、面試準(zhǔn)備、課程作業(yè)、源碼閱讀從表格可以看出這是一篇“先別用庫、先把原理跑通”的教程。它的核心產(chǎn)出不是一個(gè)大而全的訓(xùn)練框架而是一套你可以自己掌控每個(gè)張量形狀和每個(gè)梯度去向的最小實(shí)現(xiàn)。2. 適用場景與使用邊界手寫 Transformer 并不是為了讓 PyTorch 用戶拋棄官方 API。相反它適合在以下幾個(gè)場景里發(fā)揮作用第一課程作業(yè)或?qū)嶒?yàn)報(bào)告。CS 336 這類課程的核心目標(biāo)就是讓學(xué)生脫離高級封裝理解語言模型內(nèi)部到底發(fā)生了什么。自己實(shí)現(xiàn) Linear、Embedding、LayerNorm 之后再去看nn.TransformerEncoder的實(shí)現(xiàn)會(huì)明顯感覺到“圖層”不再是黑盒。第二算法面試準(zhǔn)備。很多公司面試會(huì)問“Transformer 參數(shù)初始化為什么用 Xavier”“Embedding 梯度怎么回傳”“softmax 反向傳播的雅可比是什么”。這些問題的答案不親手寫一遍很難講得清楚。第三源碼閱讀和二次開發(fā)。當(dāng)你想修改某一層的行為比如給 Attention 換成線性注意力或者把 LayerNorm 換成 RMSNorm如果只依賴高層 API往往不知道改哪里但手寫過底層模塊之后你能很快定位到具體實(shí)現(xiàn)位置。同時(shí)也要明確使用邊界。這套手寫代碼主要用于教學(xué)和原理驗(yàn)證不推薦直接用于大規(guī)模訓(xùn)練。手寫循環(huán)、GPU 算子融合、FlashAttention 這些工程優(yōu)化都不在范圍內(nèi)。如果目標(biāo)是在生產(chǎn)環(huán)境訓(xùn)練大模型應(yīng)該直接使用成熟框架PyTorch 的nn.Transformer、Hugging Face 的transformers或者 Megatron-LM 等。另一個(gè)必須強(qiáng)調(diào)的是合規(guī)邊界。課程代碼、作業(yè)要求、課件內(nèi)容都有各自的版權(quán)和使用協(xié)議本文只講解通用技術(shù)原理和常見實(shí)踐不復(fù)制任何課程原始材料。訓(xùn)練時(shí)不要使用未經(jīng)授權(quán)的數(shù)據(jù)、人臉信息、版權(quán)文本或用戶隱私數(shù)據(jù)盡量在本地構(gòu)造的玩具數(shù)據(jù)集上驗(yàn)證。3. 環(huán)境準(zhǔn)備與前置條件手寫層不需要復(fù)雜環(huán)境。建議使用 Python 3.9 以上版本和 PyTorch 2.x安裝命令pip install torch numpy tqdm如果沒有 GPU直接用 CPU 跑小模型即可。這里的關(guān)鍵不是訓(xùn)練速度而是驗(yàn)證前向輸出形狀與反向傳播梯度是否正確。建議的工程目錄結(jié)構(gòu)transformer_from_scratch/ ├── layers/ │ ├── __init__.py │ ├── linear.py │ ├── embedding.py │ ├── layernorm.py │ └── attention.py ├── init.py ├── model.py ├── train.py └── tests/ └── test_gradcheck.py這種按層拆分的結(jié)構(gòu)能讓你在調(diào)試時(shí)只關(guān)心單個(gè)模塊而不會(huì)一次面對整個(gè)網(wǎng)絡(luò)。動(dòng)手之前還需要確認(rèn)三塊數(shù)學(xué)基礎(chǔ)矩陣乘法求導(dǎo)公式、鏈?zhǔn)椒▌t、softmax 的雅可比矩陣。不需要研究得很深但至少要知道對于一個(gè)線性變換y x W從上層傳回的梯度dy應(yīng)該分別回傳到dx和dW公式是dx dy W.T、dW x.T dy。后面所有代碼都建立在這條基本規(guī)則上。4. 從零實(shí)現(xiàn)核心模塊4.1 手寫 Linear LayerLinear Layer 是最基礎(chǔ)的計(jì)算單元公式很簡單y x W b。但參數(shù)初始化和梯度回傳是兩個(gè)容易被忽略的細(xì)節(jié)。第一個(gè)細(xì)節(jié)是參數(shù)初始化。下面用 Xavier 均勻初始化它的上下界是bound sqrt(6 / (fan_in fan_out))。這種初始化適合線性層和 tanh 激活函數(shù)能讓每層輸出的方差保持在一個(gè)合理范圍內(nèi)避免深層網(wǎng)絡(luò)時(shí)梯度消失或爆炸。第二個(gè)細(xì)節(jié)是反向傳播。除了dx要回傳給上一層還要計(jì)算dW和db供優(yōu)化器更新。這里有一個(gè)常見錯(cuò)誤把dW寫成了grad_output.T x正確寫法是x.T grad_output。推導(dǎo)方法是對損失函數(shù)求 W 的偏導(dǎo)保持張量形狀對齊即可。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。這是手寫反向傳播的基礎(chǔ)backward 需要用到前向時(shí)輸入的原始值。很多新手寫的反向傳播出錯(cuò)不是因?yàn)楣酵茖?dǎo)錯(cuò)而是因?yàn)榍跋驔]有保存必要的中間變量。4.2 手寫 Embedding LayerEmbedding Layer 本質(zhì)上是一個(gè)查表操作。輸入是 token id 序列輸出是形狀為[batch_size, seq_len, embedding_dim]的向量序列。它的反向傳播很有特點(diǎn)并不是所有 embedding 行都會(huì)被更新只有當(dāng)前 batch 中實(shí)際被查到的 token 行才有梯度。實(shí)現(xiàn)方式是用torch.index_add_把梯度累加到對應(yīng)的行上。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這里有一個(gè)工程細(xì)節(jié)index_add_是原地操作調(diào)用前必須確保dE是全零矩陣。如果同一個(gè) token 在一句話中出現(xiàn)多次它的梯度會(huì)自然累加這也符合 embedding 反向傳播的定義。你不需要為它手動(dòng)去重。4.3 參數(shù)初始化決定 Transformer 能不能收斂的隱藏開關(guān)參數(shù)初始化在 PyTorch 使用里常常被忽略因?yàn)閚n.Linear已經(jīng)封裝好了默認(rèn)初始化。但手寫模塊時(shí)每層初始化都必須自己處理。先看原始 Transformer 論文中的做法embedding 層使用均值 0、方差 1 的正態(tài)分布初始化之后乘以sqrt(d_model)。也就是說如果 embedding 矩陣初始化為N(0, 1)那么實(shí)際查表輸出要再乘一個(gè)sqrt(d_model)。這樣做的目的是讓輸入到 Attention 的向量范數(shù)不要太小避免 softmax 之前的所有點(diǎn)積都擠到極小值附近。GPT 風(fēng)格的大模型通常采用更小的初始化方差例如N(0, 0.02)。原因很簡單模型層數(shù)越深殘差分支積累的隱狀態(tài)方差越大逐層用小方差初始化可以抑制這種累積。如果你在深層 Transformer 里發(fā)現(xiàn)訓(xùn)練初期 logits 出現(xiàn) NaN或者 loss 長時(shí)間不下降首先排查初始化方差是否過大。下面給出一個(gè)簡單的參數(shù)初始化工具函數(shù)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})從實(shí)踐角度看訓(xùn)練一個(gè)小模型時(shí)建議 Linear 層優(yōu)先用 Xavier 均勻初始化Embedding 層可以用uniform(-0.1, 0.1)這種比較保險(xiǎn)的小范圍初始化。如果訓(xùn)練不收斂再嘗試 GPT 風(fēng)格的N(0, 0.02)并配合 warmup 學(xué)習(xí)率調(diào)度。4.4 手寫反向傳播并用 gradcheck 驗(yàn)證手寫反向傳播最難的不是某一層的公式而是整個(gè)鏈路的梯度形狀。一個(gè)效率很高的驗(yàn)證方式是用torch.autograd.gradcheck。它會(huì)用數(shù)值差分計(jì)算梯度再和你的手動(dòng) backward 梯度對比誤差在閾值內(nèi)就說明反向傳播正確。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 計(jì)算參考梯度 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函數(shù)更嚴(yán)格可以直接傳一個(gè)自定義函數(shù)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)如果你是在課程作業(yè)里驗(yàn)證反向傳播務(wù)必把所有參數(shù)和輸入都轉(zhuǎn)成float64。數(shù)值差分在float32下誤差會(huì)很大導(dǎo)致明明公式正確gradcheck卻報(bào)錯(cuò)。4.5 LayerNorm 與殘差Transformer 穩(wěn)定訓(xùn)練的第二道防線LayerNorm 在 Transformer 中負(fù)責(zé)把每一層的輸出拉回穩(wěn)定的數(shù)值范圍。手寫時(shí)可以直接使用簡化形式y(tǒng) (x - mean) / sqrt(var eps)反向傳播用一個(gè)更緊湊的公式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這個(gè)反向公式的推導(dǎo)過程相當(dāng)經(jīng)典。核心思路是歸一化操作同時(shí)依賴mean和var所以dx中必須包含dy.mean以及dy * x_hat的均值項(xiàng)。很多人手寫 LayerNorm 反向時(shí)少了一項(xiàng)導(dǎo)致gradcheck一直不過幾乎都是這個(gè)原因。殘差連接則是把輸入和子層輸出相加它的反向傳播不需要任何額外計(jì)算梯度在相加處直接“分叉”。實(shí)現(xiàn)時(shí)只要在 forward 里保存原始輸入即可。4.6 組裝一個(gè)最小 Transformer Block有了 Linear、Embedding、LayerNorm 和殘差可以組裝一個(gè)最小 Transformer Block。這里為了控制篇幅Attention 部分不手寫 backward而是用 PyTorch 基礎(chǔ)算子實(shí)現(xiàn) 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這里的關(guān)鍵知識(shí)點(diǎn)是縮放因子。Attention 公式里scores q k.T / sqrt(d_head)不能省略。如果不做縮放點(diǎn)積結(jié)果會(huì)隨著d_head增大而變大softmax 會(huì)進(jìn)入飽和區(qū)梯度接近 0訓(xùn)練會(huì)非常慢。5. 功能測試與效果驗(yàn)證5.1 測試 Linear 前向形狀第一個(gè)測試很簡單確認(rèn)輸出形狀正確。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應(yīng)該是(in_features, out_features)這樣x W才能把最后一維從 64 映射到 128。5.2 測試 Embedding 反向是否正確構(gòu)造一個(gè)只有 0、1、2 三個(gè) token 的序列驗(yàn)證 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])判斷依據(jù)第 0 行梯度應(yīng)該等于輸入中 token 0 對應(yīng)位置梯度的累加。如果第 1 行或第 2 行也出現(xiàn)了不該有的梯度說明index_add_的索引寫錯(cuò)了。5.3 用 gradcheck 驗(yàn)證全部手寫層強(qiáng)烈建議給每個(gè)手寫層都寫一個(gè) gradcheck。以 LinearLayer 為例用前文包裝函數(shù)方式驗(yàn)證 W 和 b 的梯度EmbeddingLayer 則驗(yàn)證 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要求輸入張量至少有一個(gè)是requires_gradTrue且所有浮點(diǎn)張量都建議使用float64。整數(shù) token 張量不需要梯度。5.4 端到端小模型訓(xùn)練觀察 loss 是否下降這是最后一個(gè)驗(yàn)證能證明整個(gè)鏈路從參數(shù)初始化到反向傳播都沒有大問題。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 計(jì)算整條鏈路梯度因?yàn)?MultiHeadAttention 沒有手寫 backward loss.backward() optimizer.step() if step % 5 0: print(step, step, loss, loss.item())這段代碼用到了一個(gè)取巧但合理的設(shè)計(jì)embedding 矩陣既負(fù)責(zé)查表也作為輸出投影矩陣。這種“共享嵌入”的寫法在 GPT 系列中很常見。更好的實(shí)現(xiàn)是單獨(dú)定義一個(gè)輸出 Linear 層但從教學(xué)角度看共享嵌入能顯著減少參數(shù)。判斷標(biāo)準(zhǔn)如果 loss 在 20 步內(nèi)明顯下降說明參數(shù)初始化、前向、反向、優(yōu)化器調(diào)度全部工作正常如果 loss 原地不動(dòng)優(yōu)先檢查學(xué)習(xí)率是否太大或者太小其次檢查初始化方差。6. 從“手寫層”到“接口化調(diào)用”批量實(shí)驗(yàn)與可復(fù)用設(shè)計(jì)這個(gè)項(xiàng)目沒有 HTTP API也不需要提供 Web 服務(wù)。但從工程化角度看你可以把自定義層設(shè)計(jì)成統(tǒng)一的“前向 反向”接口方便批量跑實(shí)驗(yàn)。比如定義一個(gè) BaseLayer 風(fēng)格class BaseLayer: def forward(self, x): raise NotImplementedError def backward(self, grad_output): raise NotImplementedErrorLinearLayer、EmbeddingLayer、LayerNorm 都繼承這個(gè)接口這樣在訓(xùn)練循環(huán)里可以統(tǒng)一調(diào)度。批量實(shí)驗(yàn)最常見的場景是掃描不同初始化模式對收斂速度的影響。可以寫一個(gè)小腳本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)當(dāng)你把“參數(shù)初始化”和“反向傳播”都變成可配置項(xiàng)之后就能系統(tǒng)性地觀察為什么 GPT 風(fēng)格的小方差初始化在深層模型中表現(xiàn)更好為什么 Xavier 在淺層模型中已經(jīng)足夠。這種實(shí)驗(yàn)往往比單純閱讀論文更有價(jià)值。如果你確實(shí)想把模塊暴露成服務(wù)也可以額外包一層 FastAPI但這屬于教學(xué)項(xiàng)目之外的擴(kuò)展不推薦在原理學(xué)習(xí)階段引入。7. 資源占用與性能觀察手寫模塊的最大問題是性能遠(yuǎn)低于 PyTorch 原生算子。因?yàn)檫@里為了教學(xué)把 forward 和 backward 的中間張量都保存在內(nèi)存中也沒有做算子融合。如果是大型模型訓(xùn)練這種寫法會(huì)占用大量顯存。觀察資源使用可以這樣寫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(平均耗時(shí):, (t1 - t0) / 10) mem psutil.Process().memory_info().rss / 1024 / 1024 print(當(dāng)前內(nèi)存占用 MB:, mem)如果訓(xùn)練時(shí)顯存不足有幾種通用降載方法減小 batch size減小序列長度降低 d_model 和 n_heads把不必要的中間變量從 forward 中移出例如 forward 結(jié)束后只保留反向需要用到的張量不保存完整 attention 矩陣對大模型訓(xùn)練改用半精度或混合精度。從性能觀察的角度說手寫版的價(jià)值不在于快而在于你能清楚看到每個(gè)操作的時(shí)間和空間開銷。這也是理解后續(xù) FlashAttention 為什么要合并算子、減少訪存的起點(diǎn)。8. 常見問題與排查方法問題現(xiàn)象可能原因排查方式解決方案gradcheck 一直報(bào)錯(cuò)沒有開啟 float64檢查輸入張量和參數(shù) dtype全部轉(zhuǎn)成torch.float64Linear 反向維度報(bào)錯(cuò)dW寫成grad_output.T x打印x.shape、grad_output.shape改為x.T grad_outputEmbedding 梯度全部為 0查表后沒有保存 token 索引檢查forward是否保存self.tokens在 forward 中保存輸入 tokenloss 不下降學(xué)習(xí)率不合適或初始化方差太小打印梯度均值、方差調(diào)大 learning rate 或改用 GPT initloss 出現(xiàn) NaN初始化方差太大或 Attention 沒有縮放打印 logits 是否存在 inf使用sqrt(d_head)縮放減小初始化方差LayerNorm 反向公式缺項(xiàng)忘記處理 mean 和 var 的依賴用 gradcheck 定位使用包含dy.mean與dy*x_hat項(xiàng)的簡化公式Attention 維度 reshape 錯(cuò)誤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 訓(xùn)練太慢batch size / 序列長度偏大用 time 統(tǒng)計(jì) forward/backward 耗時(shí)減小 batch size 和 seq_len從上表能看出手寫 Transformer 的大部分報(bào)錯(cuò)都集中在維度形狀和梯度公式上。我的建議是遇到問題不要直接搜答案先打印梯度的 shape 和數(shù)值范圍往往能更快定位。9. 最佳實(shí)踐與使用建議寫手寫 Transformer 項(xiàng)目時(shí)有幾條通用的工程建議值得遵守。第一按層拆分文件。不要把所有自定義層堆在一個(gè) Python 文件里。按linear.py、embedding.py、layernorm.py、attention.py拆分每層只做一件事調(diào)試時(shí)直接定位對應(yīng)文件。第二每寫一層就立即驗(yàn)證一層。不要等到所有模塊寫完了再驗(yàn)證。驗(yàn)證工具就是torch.autograd.gradcheck它比任何肉眼觀察都可靠。每層驗(yàn)證通過之后再進(jìn)入下一步。第三先跑最小配置。第一個(gè)實(shí)驗(yàn)建議使用vocab_size64、d_model32、n_heads4、batch_size2、seq_len8。這么小的配置可以在幾秒內(nèi)完成完整訓(xùn)練循環(huán)。如果最小配置都能跑通再逐步放大。第四保留一個(gè) baseline 對照。用 PyTorch 官方nn.Linear、nn.Embedding或nn.TransformerEncoder實(shí)現(xiàn)同樣的網(wǎng)絡(luò)不斷對比前向輸出和梯度。這樣能快速排除“手寫層寫錯(cuò)”和“訓(xùn)練邏輯寫錯(cuò)”兩類問題。第五模型文件、訓(xùn)練代碼和測試代碼分目錄管理。避免把臨時(shí)調(diào)試代碼和正式代碼混在一起。如果你要批量跑初始化實(shí)驗(yàn)建議加一個(gè)配置字典或者命令行參數(shù)。第六涉及真實(shí)數(shù)據(jù)時(shí)注意合規(guī)。本文教學(xué)項(xiàng)目只需要隨機(jī) token 或自己構(gòu)造的玩具數(shù)據(jù)不需要真實(shí)文本。如果后續(xù)擴(kuò)展項(xiàng)目需要真實(shí)語料務(wù)必確認(rèn)數(shù)據(jù)來源合法、授權(quán)完整不涉及個(gè)人隱私和版權(quán)內(nèi)容。10. 總結(jié)與下一步這個(gè)手寫項(xiàng)目最值得驗(yàn)證的功能就是“參數(shù)初始化 手寫反向傳播”的組合。你先跑通一層 Linear 的 gradcheck再試一個(gè)最小 Transformer 的 loss 下降就會(huì)理解為什么Transformer的每一層都需要精心設(shè)計(jì)梯度通路。最容易踩的坑有三個(gè)Linear 反向里dW的維度順序、LayerNorm 反向里缺少歸一化依賴項(xiàng)、Attention 里的縮放因子。這三個(gè)坑只要踩過一個(gè)以后再看到類似報(bào)錯(cuò)就會(huì)很敏感。下一步建議做三件事第一給 MultiHeadAttention 也實(shí)現(xiàn)手寫 backward完整覆蓋所有模塊第二加入 causal mask讓 Attention 只能看到當(dāng)前位置之前的 token變成一個(gè)真正的自回歸語言模型第三把學(xué)到的參數(shù)初始化經(jīng)驗(yàn)搬到一個(gè)小型文本數(shù)據(jù)集上觀察 loss 曲線在不同 init 和 learning rate 下的表現(xiàn)。做完這三步你對 Transformer 的掌握就不再停留在“會(huì)用層”的層面而是真正進(jìn)入了“能寫層”的階段。 手寫Transformer核心模塊：從零實(shí)現(xiàn)Linear、Embedding與反向傳播
這次我們來看一項(xiàng)很多人學(xué)完 Transformer 但依然沒完全搞懂的事情不只是會(huì)調(diào)用nn.Transformer而是親手把 Linear Layer、Embedding Layer、參數(shù)初始化和反向傳播全部實(shí)現(xiàn)出來。本文對應(yīng)的場景是 Stanford CS 336 3.3也就是在課程作業(yè)中要求用最樸素的 PyTorch 張量操作搭建語言模型組件并通過梯度檢查驗(yàn)證手寫反向傳播的正確性。標(biāo)題里的四個(gè)關(guān)鍵詞其實(shí)就對應(yīng)了 Transformer 落地中四個(gè)最容易翻車的點(diǎn)參數(shù)初始化決定模型能不能收斂反向傳播決定你寫的網(wǎng)絡(luò)是不是真的“能學(xué)”Linear Layer 是所有前饋計(jì)算的地基Embedding Layer 則承擔(dān)著離散 token 到連續(xù)向量的映射。如果這幾個(gè)模塊只是靠nn.Linear和nn.Embedding一帶而過那遇到 loss 不降、梯度 NaN、維度對不上這類問題的時(shí)候往往只能瞎猜。自己手拼一遍之后很多報(bào)錯(cuò)一眼就能看出問題。這篇文章會(huì)從零實(shí)現(xiàn)一個(gè)可訓(xùn)練的最小 Transformer包含前向傳播、手寫 backward、參數(shù)初始化、梯度驗(yàn)證最后在一個(gè)小模型上完成端到端訓(xùn)練測試。全程使用 CPU 即可運(yùn)行不依賴 GPU也不需要下載任何預(yù)訓(xùn)練權(quán)重。適合正在啃 Transformer 源碼、準(zhǔn)備算法面試、或者做課程作業(yè)時(shí)想把原理徹底吃透的讀者。1. 核心能力速覽先把這篇文章涉及的技術(shù)目標(biāo)和運(yùn)行約束講清楚。下面的表格不是項(xiàng)目 README而是你讀完這篇文章后應(yīng)該具備的“能力驗(yàn)收清單”。能力項(xiàng)說明項(xiàng)目類型從零手寫 Transformer 核心模塊的教學(xué)實(shí)現(xiàn)來源Stanford CS 336 Language Modeling 課程 3.3 小節(jié)主題核心模塊Linear Layer、Embedding Layer、LayerNorm、參數(shù)初始化、手寫反向傳播技術(shù)棧Python 3.9、PyTorch 2.x、NumPy 可選運(yùn)行方式本地腳本運(yùn)行支持 Jupyter Notebook硬件要求CPU 可跑顯存不是必需是否依賴 nn.Linear不依賴手動(dòng)實(shí)現(xiàn) forward/backward是否依賴 nn.Embedding不依賴手動(dòng)實(shí)現(xiàn)查表與梯度回傳反向傳播方式手動(dòng)推導(dǎo)并實(shí)現(xiàn)通過 torch.autograd.gradcheck 驗(yàn)證驗(yàn)證方式形狀測試、梯度對比、小模型端到端 loss 下降適合人群學(xué)習(xí) Transformer 原理、面試準(zhǔn)備、課程作業(yè)、源碼閱讀從表格可以看出這是一篇“先別用庫、先把原理跑通”的教程。它的核心產(chǎn)出不是一個(gè)大而全的訓(xùn)練框架而是一套你可以自己掌控每個(gè)張量形狀和每個(gè)梯度去向的最小實(shí)現(xiàn)。2. 適用場景與使用邊界手寫 Transformer 并不是為了讓 PyTorch 用戶拋棄官方 API。相反它適合在以下幾個(gè)場景里發(fā)揮作用第一課程作業(yè)或?qū)嶒?yàn)報(bào)告。CS 336 這類課程的核心目標(biāo)就是讓學(xué)生脫離高級封裝理解語言模型內(nèi)部到底發(fā)生了什么。自己實(shí)現(xiàn) Linear、Embedding、LayerNorm 之后再去看nn.TransformerEncoder的實(shí)現(xiàn)會(huì)明顯感覺到“圖層”不再是黑盒。第二算法面試準(zhǔn)備。很多公司面試會(huì)問“Transformer 參數(shù)初始化為什么用 Xavier”“Embedding 梯度怎么回傳”“softmax 反向傳播的雅可比是什么”。這些問題的答案不親手寫一遍很難講得清楚。第三源碼閱讀和二次開發(fā)。當(dāng)你想修改某一層的行為比如給 Attention 換成線性注意力或者把 LayerNorm 換成 RMSNorm如果只依賴高層 API往往不知道改哪里但手寫過底層模塊之后你能很快定位到具體實(shí)現(xiàn)位置。同時(shí)也要明確使用邊界。這套手寫代碼主要用于教學(xué)和原理驗(yàn)證不推薦直接用于大規(guī)模訓(xùn)練。手寫循環(huán)、GPU 算子融合、FlashAttention 這些工程優(yōu)化都不在范圍內(nèi)。如果目標(biāo)是在生產(chǎn)環(huán)境訓(xùn)練大模型應(yīng)該直接使用成熟框架PyTorch 的nn.Transformer、Hugging Face 的transformers或者 Megatron-LM 等。另一個(gè)必須強(qiáng)調(diào)的是合規(guī)邊界。課程代碼、作業(yè)要求、課件內(nèi)容都有各自的版權(quán)和使用協(xié)議本文只講解通用技術(shù)原理和常見實(shí)踐不復(fù)制任何課程原始材料。訓(xùn)練時(shí)不要使用未經(jīng)授權(quán)的數(shù)據(jù)、人臉信息、版權(quán)文本或用戶隱私數(shù)據(jù)盡量在本地構(gòu)造的玩具數(shù)據(jù)集上驗(yàn)證。3. 環(huán)境準(zhǔn)備與前置條件手寫層不需要復(fù)雜環(huán)境。建議使用 Python 3.9 以上版本和 PyTorch 2.x安裝命令pip install torch numpy tqdm如果沒有 GPU直接用 CPU 跑小模型即可。這里的關(guān)鍵不是訓(xùn)練速度而是驗(yàn)證前向輸出形狀與反向傳播梯度是否正確。建議的工程目錄結(jié)構(gòu)transformer_from_scratch/ ├── layers/ │ ├── __init__.py │ ├── linear.py │ ├── embedding.py │ ├── layernorm.py │ └── attention.py ├── init.py ├── model.py ├── train.py └── tests/ └── test_gradcheck.py這種按層拆分的結(jié)構(gòu)能讓你在調(diào)試時(shí)只關(guān)心單個(gè)模塊而不會(huì)一次面對整個(gè)網(wǎng)絡(luò)。動(dòng)手之前還需要確認(rèn)三塊數(shù)學(xué)基礎(chǔ)矩陣乘法求導(dǎo)公式、鏈?zhǔn)椒▌t、softmax 的雅可比矩陣。不需要研究得很深但至少要知道對于一個(gè)線性變換y x W從上層傳回的梯度dy應(yīng)該分別回傳到dx和dW公式是dx dy W.T、dW x.T dy。后面所有代碼都建立在這條基本規(guī)則上。4. 從零實(shí)現(xiàn)核心模塊4.1 手寫 Linear LayerLinear Layer 是最基礎(chǔ)的計(jì)算單元公式很簡單y x W b。但參數(shù)初始化和梯度回傳是兩個(gè)容易被忽略的細(xì)節(jié)。第一個(gè)細(xì)節(jié)是參數(shù)初始化。下面用 Xavier 均勻初始化它的上下界是bound sqrt(6 / (fan_in fan_out))。這種初始化適合線性層和 tanh 激活函數(shù)能讓每層輸出的方差保持在一個(gè)合理范圍內(nèi)避免深層網(wǎng)絡(luò)時(shí)梯度消失或爆炸。第二個(gè)細(xì)節(jié)是反向傳播。除了dx要回傳給上一層還要計(jì)算dW和db供優(yōu)化器更新。這里有一個(gè)常見錯(cuò)誤把dW寫成了grad_output.T x正確寫法是x.T grad_output。推導(dǎo)方法是對損失函數(shù)求 W 的偏導(dǎo)保持張量形狀對齊即可。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。這是手寫反向傳播的基礎(chǔ)backward 需要用到前向時(shí)輸入的原始值。很多新手寫的反向傳播出錯(cuò)不是因?yàn)楣酵茖?dǎo)錯(cuò)而是因?yàn)榍跋驔]有保存必要的中間變量。4.2 手寫 Embedding LayerEmbedding Layer 本質(zhì)上是一個(gè)查表操作。輸入是 token id 序列輸出是形狀為[batch_size, seq_len, embedding_dim]的向量序列。它的反向傳播很有特點(diǎn)并不是所有 embedding 行都會(huì)被更新只有當(dāng)前 batch 中實(shí)際被查到的 token 行才有梯度。實(shí)現(xiàn)方式是用torch.index_add_把梯度累加到對應(yīng)的行上。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這里有一個(gè)工程細(xì)節(jié)index_add_是原地操作調(diào)用前必須確保dE是全零矩陣。如果同一個(gè) token 在一句話中出現(xiàn)多次它的梯度會(huì)自然累加這也符合 embedding 反向傳播的定義。你不需要為它手動(dòng)去重。4.3 參數(shù)初始化決定 Transformer 能不能收斂的隱藏開關(guān)參數(shù)初始化在 PyTorch 使用里常常被忽略因?yàn)閚n.Linear已經(jīng)封裝好了默認(rèn)初始化。但手寫模塊時(shí)每層初始化都必須自己處理。先看原始 Transformer 論文中的做法embedding 層使用均值 0、方差 1 的正態(tài)分布初始化之后乘以sqrt(d_model)。也就是說如果 embedding 矩陣初始化為N(0, 1)那么實(shí)際查表輸出要再乘一個(gè)sqrt(d_model)。這樣做的目的是讓輸入到 Attention 的向量范數(shù)不要太小避免 softmax 之前的所有點(diǎn)積都擠到極小值附近。GPT 風(fēng)格的大模型通常采用更小的初始化方差例如N(0, 0.02)。原因很簡單模型層數(shù)越深殘差分支積累的隱狀態(tài)方差越大逐層用小方差初始化可以抑制這種累積。如果你在深層 Transformer 里發(fā)現(xiàn)訓(xùn)練初期 logits 出現(xiàn) NaN或者 loss 長時(shí)間不下降首先排查初始化方差是否過大。下面給出一個(gè)簡單的參數(shù)初始化工具函數(shù)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})從實(shí)踐角度看訓(xùn)練一個(gè)小模型時(shí)建議 Linear 層優(yōu)先用 Xavier 均勻初始化Embedding 層可以用uniform(-0.1, 0.1)這種比較保險(xiǎn)的小范圍初始化。如果訓(xùn)練不收斂再嘗試 GPT 風(fēng)格的N(0, 0.02)并配合 warmup 學(xué)習(xí)率調(diào)度。4.4 手寫反向傳播并用 gradcheck 驗(yàn)證手寫反向傳播最難的不是某一層的公式而是整個(gè)鏈路的梯度形狀。一個(gè)效率很高的驗(yàn)證方式是用torch.autograd.gradcheck。它會(huì)用數(shù)值差分計(jì)算梯度再和你的手動(dòng) backward 梯度對比誤差在閾值內(nèi)就說明反向傳播正確。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 計(jì)算參考梯度 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函數(shù)更嚴(yán)格可以直接傳一個(gè)自定義函數(shù)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)如果你是在課程作業(yè)里驗(yàn)證反向傳播務(wù)必把所有參數(shù)和輸入都轉(zhuǎn)成float64。數(shù)值差分在float32下誤差會(huì)很大導(dǎo)致明明公式正確gradcheck卻報(bào)錯(cuò)。4.5 LayerNorm 與殘差Transformer 穩(wěn)定訓(xùn)練的第二道防線LayerNorm 在 Transformer 中負(fù)責(zé)把每一層的輸出拉回穩(wěn)定的數(shù)值范圍。手寫時(shí)可以直接使用簡化形式y(tǒng) (x - mean) / sqrt(var eps)反向傳播用一個(gè)更緊湊的公式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這個(gè)反向公式的推導(dǎo)過程相當(dāng)經(jīng)典。核心思路是歸一化操作同時(shí)依賴mean和var所以dx中必須包含dy.mean以及dy * x_hat的均值項(xiàng)。很多人手寫 LayerNorm 反向時(shí)少了一項(xiàng)導(dǎo)致gradcheck一直不過幾乎都是這個(gè)原因。殘差連接則是把輸入和子層輸出相加它的反向傳播不需要任何額外計(jì)算梯度在相加處直接“分叉”。實(shí)現(xiàn)時(shí)只要在 forward 里保存原始輸入即可。4.6 組裝一個(gè)最小 Transformer Block有了 Linear、Embedding、LayerNorm 和殘差可以組裝一個(gè)最小 Transformer Block。這里為了控制篇幅Attention 部分不手寫 backward而是用 PyTorch 基礎(chǔ)算子實(shí)現(xiàn) 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這里的關(guān)鍵知識(shí)點(diǎn)是縮放因子。Attention 公式里scores q k.T / sqrt(d_head)不能省略。如果不做縮放點(diǎn)積結(jié)果會(huì)隨著d_head增大而變大softmax 會(huì)進(jìn)入飽和區(qū)梯度接近 0訓(xùn)練會(huì)非常慢。5. 功能測試與效果驗(yàn)證5.1 測試 Linear 前向形狀第一個(gè)測試很簡單確認(rèn)輸出形狀正確。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應(yīng)該是(in_features, out_features)這樣x W才能把最后一維從 64 映射到 128。5.2 測試 Embedding 反向是否正確構(gòu)造一個(gè)只有 0、1、2 三個(gè) token 的序列驗(yàn)證 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])判斷依據(jù)第 0 行梯度應(yīng)該等于輸入中 token 0 對應(yīng)位置梯度的累加。如果第 1 行或第 2 行也出現(xiàn)了不該有的梯度說明index_add_的索引寫錯(cuò)了。5.3 用 gradcheck 驗(yàn)證全部手寫層強(qiáng)烈建議給每個(gè)手寫層都寫一個(gè) gradcheck。以 LinearLayer 為例用前文包裝函數(shù)方式驗(yàn)證 W 和 b 的梯度EmbeddingLayer 則驗(yàn)證 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要求輸入張量至少有一個(gè)是requires_gradTrue且所有浮點(diǎn)張量都建議使用float64。整數(shù) token 張量不需要梯度。5.4 端到端小模型訓(xùn)練觀察 loss 是否下降這是最后一個(gè)驗(yàn)證能證明整個(gè)鏈路從參數(shù)初始化到反向傳播都沒有大問題。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 計(jì)算整條鏈路梯度因?yàn)?MultiHeadAttention 沒有手寫 backward loss.backward() optimizer.step() if step % 5 0: print(step, step, loss, loss.item())這段代碼用到了一個(gè)取巧但合理的設(shè)計(jì)embedding 矩陣既負(fù)責(zé)查表也作為輸出投影矩陣。這種“共享嵌入”的寫法在 GPT 系列中很常見。更好的實(shí)現(xiàn)是單獨(dú)定義一個(gè)輸出 Linear 層但從教學(xué)角度看共享嵌入能顯著減少參數(shù)。判斷標(biāo)準(zhǔn)如果 loss 在 20 步內(nèi)明顯下降說明參數(shù)初始化、前向、反向、優(yōu)化器調(diào)度全部工作正常如果 loss 原地不動(dòng)優(yōu)先檢查學(xué)習(xí)率是否太大或者太小其次檢查初始化方差。6. 從“手寫層”到“接口化調(diào)用”批量實(shí)驗(yàn)與可復(fù)用設(shè)計(jì)這個(gè)項(xiàng)目沒有 HTTP API也不需要提供 Web 服務(wù)。但從工程化角度看你可以把自定義層設(shè)計(jì)成統(tǒng)一的“前向 反向”接口方便批量跑實(shí)驗(yàn)。比如定義一個(gè) BaseLayer 風(fēng)格class BaseLayer: def forward(self, x): raise NotImplementedError def backward(self, grad_output): raise NotImplementedErrorLinearLayer、EmbeddingLayer、LayerNorm 都繼承這個(gè)接口這樣在訓(xùn)練循環(huán)里可以統(tǒng)一調(diào)度。批量實(shí)驗(yàn)最常見的場景是掃描不同初始化模式對收斂速度的影響。可以寫一個(gè)小腳本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)當(dāng)你把“參數(shù)初始化”和“反向傳播”都變成可配置項(xiàng)之后就能系統(tǒng)性地觀察為什么 GPT 風(fēng)格的小方差初始化在深層模型中表現(xiàn)更好為什么 Xavier 在淺層模型中已經(jīng)足夠。這種實(shí)驗(yàn)往往比單純閱讀論文更有價(jià)值。如果你確實(shí)想把模塊暴露成服務(wù)也可以額外包一層 FastAPI但這屬于教學(xué)項(xiàng)目之外的擴(kuò)展不推薦在原理學(xué)習(xí)階段引入。7. 資源占用與性能觀察手寫模塊的最大問題是性能遠(yuǎn)低于 PyTorch 原生算子。因?yàn)檫@里為了教學(xué)把 forward 和 backward 的中間張量都保存在內(nèi)存中也沒有做算子融合。如果是大型模型訓(xùn)練這種寫法會(huì)占用大量顯存。觀察資源使用可以這樣寫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(平均耗時(shí):, (t1 - t0) / 10) mem psutil.Process().memory_info().rss / 1024 / 1024 print(當(dāng)前內(nèi)存占用 MB:, mem)如果訓(xùn)練時(shí)顯存不足有幾種通用降載方法減小 batch size減小序列長度降低 d_model 和 n_heads把不必要的中間變量從 forward 中移出例如 forward 結(jié)束后只保留反向需要用到的張量不保存完整 attention 矩陣對大模型訓(xùn)練改用半精度或混合精度。從性能觀察的角度說手寫版的價(jià)值不在于快而在于你能清楚看到每個(gè)操作的時(shí)間和空間開銷。這也是理解后續(xù) FlashAttention 為什么要合并算子、減少訪存的起點(diǎn)。8. 常見問題與排查方法問題現(xiàn)象可能原因排查方式解決方案gradcheck 一直報(bào)錯(cuò)沒有開啟 float64檢查輸入張量和參數(shù) dtype全部轉(zhuǎn)成torch.float64Linear 反向維度報(bào)錯(cuò)dW寫成grad_output.T x打印x.shape、grad_output.shape改為x.T grad_outputEmbedding 梯度全部為 0查表后沒有保存 token 索引檢查forward是否保存self.tokens在 forward 中保存輸入 tokenloss 不下降學(xué)習(xí)率不合適或初始化方差太小打印梯度均值、方差調(diào)大 learning rate 或改用 GPT initloss 出現(xiàn) NaN初始化方差太大或 Attention 沒有縮放打印 logits 是否存在 inf使用sqrt(d_head)縮放減小初始化方差LayerNorm 反向公式缺項(xiàng)忘記處理 mean 和 var 的依賴用 gradcheck 定位使用包含dy.mean與dy*x_hat項(xiàng)的簡化公式Attention 維度 reshape 錯(cuò)誤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 訓(xùn)練太慢batch size / 序列長度偏大用 time 統(tǒng)計(jì) forward/backward 耗時(shí)減小 batch size 和 seq_len從上表能看出手寫 Transformer 的大部分報(bào)錯(cuò)都集中在維度形狀和梯度公式上。我的建議是遇到問題不要直接搜答案先打印梯度的 shape 和數(shù)值范圍往往能更快定位。9. 最佳實(shí)踐與使用建議寫手寫 Transformer 項(xiàng)目時(shí)有幾條通用的工程建議值得遵守。第一按層拆分文件。不要把所有自定義層堆在一個(gè) Python 文件里。按linear.py、embedding.py、layernorm.py、attention.py拆分每層只做一件事調(diào)試時(shí)直接定位對應(yīng)文件。第二每寫一層就立即驗(yàn)證一層。不要等到所有模塊寫完了再驗(yàn)證。驗(yàn)證工具就是torch.autograd.gradcheck它比任何肉眼觀察都可靠。每層驗(yàn)證通過之后再進(jìn)入下一步。第三先跑最小配置。第一個(gè)實(shí)驗(yàn)建議使用vocab_size64、d_model32、n_heads4、batch_size2、seq_len8。這么小的配置可以在幾秒內(nèi)完成完整訓(xùn)練循環(huán)。如果最小配置都能跑通再逐步放大。第四保留一個(gè) baseline 對照。用 PyTorch 官方nn.Linear、nn.Embedding或nn.TransformerEncoder實(shí)現(xiàn)同樣的網(wǎng)絡(luò)不斷對比前向輸出和梯度。這樣能快速排除“手寫層寫錯(cuò)”和“訓(xùn)練邏輯寫錯(cuò)”兩類問題。第五模型文件、訓(xùn)練代碼和測試代碼分目錄管理。避免把臨時(shí)調(diào)試代碼和正式代碼混在一起。如果你要批量跑初始化實(shí)驗(yàn)建議加一個(gè)配置字典或者命令行參數(shù)。第六涉及真實(shí)數(shù)據(jù)時(shí)注意合規(guī)。本文教學(xué)項(xiàng)目只需要隨機(jī) token 或自己構(gòu)造的玩具數(shù)據(jù)不需要真實(shí)文本。如果后續(xù)擴(kuò)展項(xiàng)目需要真實(shí)語料務(wù)必確認(rèn)數(shù)據(jù)來源合法、授權(quán)完整不涉及個(gè)人隱私和版權(quán)內(nèi)容。10. 總結(jié)與下一步這個(gè)手寫項(xiàng)目最值得驗(yàn)證的功能就是“參數(shù)初始化 手寫反向傳播”的組合。你先跑通一層 Linear 的 gradcheck再試一個(gè)最小 Transformer 的 loss 下降就會(huì)理解為什么Transformer的每一層都需要精心設(shè)計(jì)梯度通路。最容易踩的坑有三個(gè)Linear 反向里dW的維度順序、LayerNorm 反向里缺少歸一化依賴項(xiàng)、Attention 里的縮放因子。這三個(gè)坑只要踩過一個(gè)以后再看到類似報(bào)錯(cuò)就會(huì)很敏感。下一步建議做三件事第一給 MultiHeadAttention 也實(shí)現(xiàn)手寫 backward完整覆蓋所有模塊第二加入 causal mask讓 Attention 只能看到當(dāng)前位置之前的 token變成一個(gè)真正的自回歸語言模型第三把學(xué)到的參數(shù)初始化經(jīng)驗(yàn)搬到一個(gè)小型文本數(shù)據(jù)集上觀察 loss 曲線在不同 init 和 learning rate 下的表現(xiàn)。做完這三步你對 Transformer 的掌握就不再停留在“會(huì)用層”的層面而是真正進(jìn)入了“能寫層”的階段。