今年我把大量業(yè)余時(shí)間投進(jìn)了一個(gè)叫ai-engineering-from-scratch的個(gè)人項(xiàng)目。簡(jiǎn)單說就是給自己立了條規(guī)矩凡是跟 AI 相關(guān)的環(huán)節(jié)能自己動(dòng)手實(shí)現(xiàn)的絕不直接調(diào)封裝好的接口。從手寫張量運(yùn)算開始到訓(xùn)練一個(gè)微型語言模型再到把模型量化部署成一個(gè)能用的服務(wù)整個(gè)過程走下來收獲最大的不是某個(gè)模型效果變好了而是終于弄明白了一個(gè)大語言模型從數(shù)據(jù)集到上線中間到底經(jīng)歷了什么出了問題時(shí)該往哪個(gè)方向查。這篇文章就圍繞這個(gè)項(xiàng)目把我踩過的坑、驗(yàn)證過確實(shí)有效的路線、以及現(xiàn)在還保留在收藏夾里的參考材料整理出來。不吹不黑直接上干貨。適合已經(jīng)會(huì)用 PyTorch、會(huì)調(diào) Transformers 庫但對(duì)底層實(shí)現(xiàn)還有點(diǎn)心里沒底的人也適合準(zhǔn)備進(jìn)入大模型應(yīng)用開發(fā)、想從底層建立手感的新手。看完你至少能少走我一半的彎路。1. 為什么我決定從零開始學(xué)AI工程先別急著調(diào)API1.1 框架用久了心里不踏實(shí)先說個(gè)真實(shí)感受。我最早做 NLP 項(xiàng)目時(shí)基本就是transformers庫一把梭加載模型、寫個(gè)訓(xùn)練循環(huán)、調(diào)參數(shù)、完事。模型跑不動(dòng)就換更大的顯卡效果不好就換更強(qiáng)的預(yù)訓(xùn)練權(quán)重。直到有一次模型在驗(yàn)證集上的損失怎么都降不下去我把學(xué)習(xí)率調(diào)低、把層數(shù)減少、把 dropout 拉滿全都沒用。當(dāng)時(shí)我能做的只有百度搜索關(guān)鍵詞然后挨個(gè)試網(wǎng)上流傳的玄學(xué)改法。那天之后我開始反思我不是在用 AI而是在調(diào)一個(gè)巨大的不透明的黑盒子。梯度怎么傳播的、學(xué)習(xí)率策略在內(nèi)部到底起了什么作用、顯存為什么爆、推理為什么慢這些問題我都答不上來。而這些問題恰恰是 AI 工程里最核心的問題。框架用久了AI 工程變成了一種配置工程這對(duì)工程師來說是非常危險(xiǎn)的。一旦遇到模型異常、數(shù)據(jù)泄漏、顯存瓶頸這類問題不懂底層就只能靠猜而靠猜的代價(jià)是極其昂貴的。所以我才決定做from scratch這件事。不是去重寫 PyTorch也不是去發(fā)明新一代算法而是把現(xiàn)代語言模型里最關(guān)鍵的模塊一個(gè)個(gè)抽出來用最樸素的方式實(shí)現(xiàn)一遍。注意力機(jī)制自己寫、分詞器自己寫、訓(xùn)練循環(huán)自己寫、量化腳本自己寫。寫完之后再回到框架里你會(huì)發(fā)現(xiàn)自己不是在使用框架而是在理解框架心態(tài)完全不一樣。1.2 從零手搓到底要學(xué)什么范圍怎么劃剛開始我也很迷茫因?yàn)锳I 工程這個(gè)范圍太大了。做 CV 的需要懂卷積、數(shù)據(jù)增強(qiáng)、分布式訓(xùn)練做推薦的需要懂特征工程、Embedding、在線學(xué)習(xí)做 NLP / LLM 的需要懂分詞、Transformer、微調(diào)、推理優(yōu)化。如果不能明確邊界這個(gè)項(xiàng)目很可能變成一個(gè)永遠(yuǎn)完不成的大坑。我的做法是先把全鏈路畫成一個(gè)流程圖再從里面選出必須親手做一遍的核心節(jié)點(diǎn)。我這里所謂的流程圖并不是畫給別人看的那種而是寫給自己確認(rèn)用的清單。我的清單大致是數(shù)據(jù)處理分詞器、批次構(gòu)建、數(shù)據(jù)采樣模型構(gòu)建Embedding、多頭注意力、前饋網(wǎng)絡(luò)、層歸一化訓(xùn)練系統(tǒng)損失函數(shù)、反向傳播、優(yōu)化器、學(xué)習(xí)率調(diào)度推理系統(tǒng)自回歸生成、KV Cache、采樣策略部署優(yōu)化量化、批處理、服務(wù)化接口模型進(jìn)階在基礎(chǔ)模型之上嘗試先思考、再回答的推理行為每個(gè)節(jié)點(diǎn)只追求能跑通、能解釋、能量化效果不做過度擴(kuò)展。比如我不會(huì)去手寫 CUDA kernel因?yàn)槟且呀?jīng)超出一般的 AI 工程范圍屬于系統(tǒng)底層優(yōu)化跟當(dāng)前目標(biāo)不對(duì)齊。先把上層的每一個(gè)工程環(huán)節(jié)摸透等將來遇到 perf 瓶頸時(shí)再下沉也不遲。這個(gè)范圍劃完之后我給自己定了三個(gè)里程碑第一用 NumPy 手寫一個(gè)能訓(xùn)練 MNIST 的迷你神經(jīng)網(wǎng)絡(luò)第二用 PyTorch 從零實(shí)現(xiàn)一個(gè)小型 GPT在幾十 MB 的數(shù)據(jù)上訓(xùn)練出能生成文本的模型第三對(duì)第二個(gè)模型做量化壓縮和推理加速部署成一個(gè) HTTP 服務(wù)。后面所有的工作都圍繞這三個(gè)里程碑展開。2. 我的學(xué)習(xí)路線把黑盒拆成白盒的六個(gè)階段2.1 六個(gè)階段一層層揭開隔層我給自己制定了遞進(jìn)式路線不是上來就寫 Transformer而是按照造一臺(tái)機(jī)器的順序去推進(jìn)。如果你也想照這條路走我建議你嚴(yán)格按順序來不要貪快。第一個(gè)階段是 Python 編程與數(shù)據(jù)操作。重點(diǎn)不是語法而是對(duì)數(shù)據(jù)到底長(zhǎng)什么樣有直覺。列表、字典、字符串處理、文件讀寫配合collections.Counter做詞頻統(tǒng)計(jì)這樣后面寫分詞器時(shí)不會(huì)發(fā)怵。第二個(gè)階段是手寫張量與自動(dòng)求導(dǎo)。現(xiàn)階段不建議直接上 PyTorch而是先用 NumPy 自己實(shí)現(xiàn)一個(gè)極簡(jiǎn)的張量類只支持標(biāo)量、向量、矩陣的加法乘法以及對(duì)某個(gè)參數(shù)的梯度記錄。這個(gè)階段的核心任務(wù)是理解反向傳播到底在做什么而不是把所有細(xì)節(jié)都搞完。我做了一個(gè)非常粗糙的自動(dòng)求導(dǎo)只有幾十行代碼但它讓我第一次看懂了鏈?zhǔn)椒▌t在程序里如何落地。第三個(gè)階段是手寫 MLP 和簡(jiǎn)單 CNN。用自己寫的自動(dòng)求導(dǎo)工具去訓(xùn)練一個(gè)識(shí)別手寫數(shù)字的模型把梯度下降、過擬合、正則化這些基礎(chǔ)概念體驗(yàn)一遍。很多人跳過這個(gè)階段直接學(xué)大模型但我強(qiáng)烈不建議。因?yàn)榇竽Ｐ鸵惨蕾囃瑯拥臄?shù)學(xué)機(jī)制如果連一個(gè)隱藏層的 MLP 都調(diào)不明白后面遇到 Transformer 的異常會(huì)更難定位。第四個(gè)階段是學(xué)習(xí)和實(shí)現(xiàn)注意力機(jī)制。注意力是任何現(xiàn)代語言模型的靈魂它做的事本質(zhì)上就是根據(jù)相關(guān)性從其他位置上取信息。先手寫單頭注意力再擴(kuò)展成多頭最后加上因果掩碼。第五個(gè)階段是把注意力堆成 Transformer實(shí)現(xiàn)一個(gè)微型 GPT。從 token 化文本開始訓(xùn)練一個(gè)幾百萬參數(shù)的小模型讓它生成看起來像模像樣的句子。第六個(gè)階段是推理優(yōu)化與部署。把訓(xùn)練好的權(quán)重量化成半精度或整數(shù)精度加上 KV Cache寫一個(gè)流式生成接口再封裝成一個(gè)所有人都能通過 HTTP 調(diào)用的服務(wù)。每個(gè)階段我都有對(duì)應(yīng)的完成標(biāo)準(zhǔn)。比如第四階段的完成標(biāo)準(zhǔn)不是理解了注意力公式而是寫出一個(gè)能跑、能反向傳播的多頭注意力模塊并在小數(shù)據(jù)集上訓(xùn)練收斂。這樣就不會(huì)陷入漫無目的的看論文狀態(tài)。2.2 參考書怎么用讀一遍不如改一遍在我動(dòng)手的過程中最常被問到的一句話是你到底用什么學(xué)的我的答案很直接《Build a Large Language Model From Scratch》這本書給了我很重要的路線參考。這本書從數(shù)據(jù)準(zhǔn)備開始一步步搭建類似 GPT 的結(jié)構(gòu)包含分詞、注意力、訓(xùn)練、微調(diào)等完整內(nèi)容。書里的代碼風(fēng)格非常樸素沒有花里胡哨的高級(jí)封裝很適合配合這個(gè)項(xiàng)目來理解每一個(gè)環(huán)節(jié)。但我還得提醒一點(diǎn)讀這本書千萬不要只停留在讀上。我第一次讀的時(shí)候以為看懂了結(jié)果合上書自己寫連masked_fill都寫錯(cuò)了位置。后來我換了個(gè)策略把書里的每段代碼都當(dāng)成參考實(shí)現(xiàn)然后關(guān)掉示例代碼自己重新實(shí)現(xiàn)一遍遇到卡殼再回去對(duì)照。這樣讀一遍書等于自己寫了三遍代碼效果完全不一樣。除了這本書我還會(huì)配合看一些原始論文和開源項(xiàng)目。書不會(huì)覆蓋所有細(xì)節(jié)比如混合精度訓(xùn)練、KV Cache 的顯存優(yōu)化、量化感知訓(xùn)練這些工程問題需要從零散的源碼和文章中補(bǔ)全。我的經(jīng)驗(yàn)是以書為主線建立骨架再以論文和源碼為枝葉填充細(xì)節(jié)最后用自己的代碼驗(yàn)證理解。3. 從零手搓一個(gè)微型語言模型的完整流程3.1 先解決詞的問題手寫 BPE 分詞器語言模型處理的是 token不是原始字符串。所以從零搭建模型第一步就是解決怎么把一段話切成 token。按空格切詞雖然簡(jiǎn)單但會(huì)得到一個(gè)巨大的詞表而且無法處理沒見過的新詞所以現(xiàn)代模型普遍使用字節(jié)對(duì)編碼BPE。BPE 的核心思想是先把文本看成 UTF-8 字節(jié)序列然后反復(fù)統(tǒng)計(jì)相鄰字節(jié)對(duì)的出現(xiàn)頻率把最高頻的對(duì)合并成一個(gè)新符號(hào)直到詞表達(dá)到目標(biāo)大小。手寫 BPE 時(shí)最關(guān)鍵的數(shù)據(jù)結(jié)構(gòu)是一個(gè)計(jì)數(shù)用的字典。首先是統(tǒng)計(jì)相鄰字節(jié)對(duì)頻率def get_stats(ids): counts {} for pair in zip(ids, ids[1:]): counts[pair] counts.get(pair, 0) 1 return counts然后在每一輪迭代中找到頻次最高的 pair把它的兩個(gè) id 合并成一個(gè)新的 iddef merge(ids, pair, new_id): result [] i 0 while i len(ids): if i len(ids) - 1 and ids[i] pair[0] and ids[i1] pair[1]: result.append(new_id) i 2 else: result.append(ids[i]) i 1 return result這個(gè)過程初看很簡(jiǎn)單但真正實(shí)現(xiàn)時(shí)有一個(gè)坑隨著合并的進(jìn)行id 長(zhǎng)度會(huì)變化新的高頻 pair 必須基于上一輪合并后的序列重新統(tǒng)計(jì)。如果理解錯(cuò)了合并順序就會(huì)亂掉分詞結(jié)果會(huì)有問題。我當(dāng)時(shí)就是在這里猶豫了很久后來意識(shí)到每次合并后全局重新統(tǒng)計(jì)才是標(biāo)準(zhǔn)做法而不是在一個(gè)固定列表上連續(xù)改。3.2 核心模塊手寫多頭注意力的幾個(gè)細(xì)節(jié)注意力模塊是整個(gè)模型中最容易寫錯(cuò)、也最值得親手實(shí)現(xiàn)的部分。常見的公式非常簡(jiǎn)單Q 和 K 做點(diǎn)積、除以根號(hào) d、加上掩碼、Softmax、再乘 V。但工程實(shí)現(xiàn)時(shí)有幾個(gè)細(xì)節(jié)特別值得注意。第一個(gè)是因果掩碼。語言模型生成時(shí)只能看前面的 token不能偷看后面的內(nèi)容。實(shí)現(xiàn)方式是把未來位置的值設(shè)成負(fù)無窮經(jīng)過 Softmax 之后概率趨近于零。我在第一次實(shí)現(xiàn)時(shí)用的掩碼矩陣是對(duì)的但維度 broadcast 沒對(duì)齊結(jié)果訓(xùn)練損失直接亂跳。排查了很久才發(fā)現(xiàn)是掩碼形狀錯(cuò)誤。手動(dòng)實(shí)現(xiàn)時(shí)建議把每一步的張量形狀打印出來別嫌麻煩。第二個(gè)是縮放因子的作用。除以sqrt(d_k)不只是為了讓數(shù)值范圍更好看而是防止點(diǎn)積結(jié)果過大導(dǎo)致 Softmax 進(jìn)入飽和區(qū)梯度變得極小。如果不縮放模型在小數(shù)據(jù)集上也能跑但訓(xùn)練會(huì)明顯變慢損失曲線的尾巴還會(huì)抖動(dòng)。第三點(diǎn)是多頭注意力中頭的意義。多頭不是簡(jiǎn)單地疊加多個(gè)注意力結(jié)果而是讓每個(gè)頭學(xué)到不同的相關(guān)性模式。有的頭可能關(guān)注前一個(gè)詞有的頭可能關(guān)注句法成分。雖然我們?cè)诖a層面只是把 d_model 切成幾段分別計(jì)算但正是這種參數(shù)獨(dú)立性讓模型表達(dá)力變強(qiáng)了。手寫時(shí)我建議先按循環(huán)每個(gè)頭的方式寫一版跑通后再改成矩陣并行版這樣對(duì)分塊的印象會(huì)特別深刻。第四點(diǎn)是殘差連接和層歸一化的位置。Transformer 每一層結(jié)構(gòu)是注意力 - 殘差 - 層歸一化 - 前饋網(wǎng)絡(luò) - 殘差 - 層歸一化。殘差連接讓梯度有一條高速公路可以直達(dá)底層層歸一化穩(wěn)定每一層的激活分布。順序如果搞錯(cuò)了模型的訓(xùn)練穩(wěn)定性和最終效果都有明顯差距。3.3 訓(xùn)練循環(huán)里的關(guān)鍵決策學(xué)習(xí)率、批次和損失模型搭完了訓(xùn)練循環(huán)也不像想象中那么簡(jiǎn)單。我最初寫的訓(xùn)練循環(huán)只是取數(shù)據(jù)、算 loss、反向傳播、更新結(jié)果模型一直不收斂。后來我逐項(xiàng)排查發(fā)現(xiàn)真正影響訓(xùn)練效果的是幾個(gè)細(xì)節(jié)隨機(jī)種子是否固定、數(shù)據(jù)是否被隨機(jī)打亂、學(xué)習(xí)率調(diào)度是否合理、梯度有沒有做裁剪。學(xué)習(xí)率我很推薦使用 warmup cosine 的調(diào)度方式。一開始用一個(gè)很小的學(xué)習(xí)率熱身讓參數(shù)更新不會(huì)太激進(jìn)然后再逐步降到接近零讓模型在后期做好精細(xì)調(diào)整。我用的配置大致是warmup 步數(shù)約占總步數(shù)的 3% 到 5%峰值學(xué)習(xí)率在1e-3到3e-4之間具體看模型規(guī)模和 batch size。批次大小也很關(guān)鍵。我的顯卡資源有限單批放不了太多樣本就用梯度累積來模擬更大的批次。注意梯度累積不是直接改batch_size而是在多個(gè)小批次上累計(jì)梯度然后再做一次優(yōu)化器更新。梯度累積的步數(shù)需要根據(jù)顯存實(shí)測(cè)來定我試過累積 4 步效果最穩(wěn)。數(shù)據(jù)處理上我加了一條保護(hù)線驗(yàn)證集絕對(duì)不參與訓(xùn)練而且驗(yàn)證集的數(shù)據(jù)順序每次都要固定。有一次我偷懶直接在數(shù)據(jù)類里給訓(xùn)練集和驗(yàn)證集用了同一個(gè) shuffler結(jié)果驗(yàn)證損失一直往下掉訓(xùn)練損失卻不降后來才發(fā)現(xiàn)是驗(yàn)證集里混進(jìn)了訓(xùn)練樣本數(shù)據(jù)泄漏的坑差點(diǎn)沒把我逼瘋。訓(xùn)練過程中我習(xí)慣每 20 個(gè) step 打印一次 loss 和當(dāng)前學(xué)習(xí)率每 200 個(gè) step 在驗(yàn)證集上算一次困惑度。損失曲線的形態(tài)非常有信息量如果訓(xùn)練 loss 下降但驗(yàn)證 loss 上升那是過擬合如果訓(xùn)練 loss 不降那多半是學(xué)習(xí)率太高、數(shù)據(jù)有泄漏或者模型結(jié)構(gòu)寫錯(cuò)了。3.4 采樣策略怎么讓模型說得像人話模型訓(xùn)練完之后生成階段同樣有很多坑。語言模型是逐步生成 token 的每次預(yù)測(cè)下一個(gè) token 的概率分布然后從分布中采樣。如果每次都選概率最大的 token就是貪心解碼雖然穩(wěn)定但容易重復(fù)、死板。如果每次都純隨機(jī)采樣句子會(huì)變得前言不搭后語。我實(shí)測(cè)下來配合使用temperature、top-k 和 top-p效果最好。Temperature 控制分布的尖銳程度溫度小于 1 的時(shí)候概率集中到高概率 token 上生成的文本更保守溫度大于 1 的時(shí)候分布更平緩文本更大膽。Top-k 是只從概率最高的 k 個(gè) token 中采樣避免給那些幾乎不可能的 token 太多機(jī)會(huì)。Top-p 則是動(dòng)態(tài)選擇累計(jì)概率超過閾值的 token 集合比 top-k 更靈活。我常用的一個(gè)組合是temperature0.8, top_k40, top_p0.9。這個(gè)組合在我訓(xùn)練的幾百萬參數(shù)小模型上生成結(jié)果既有一定的多樣性也不會(huì)很快陷入重復(fù)循環(huán)。如果你發(fā)現(xiàn)模型總是重復(fù)同樣的句子可以適當(dāng)把 temperature 調(diào)高一點(diǎn)或者增大 top-p 閾值如果發(fā)現(xiàn)上下文不連貫則降低 temperature 會(huì)更有效。4. 從生成到推理小模型也能學(xué)會(huì)思考鏈4.1 推理模型與傳統(tǒng)模型到底差在哪里當(dāng)我把基礎(chǔ)語言模型跑通之后我注意到行業(yè)里的討論熱點(diǎn)開始轉(zhuǎn)向另一個(gè)方向讓模型在回答之前先做一段內(nèi)部推理再給出最終答案。傳統(tǒng)語言模型是輸入問題立刻輸出答案而推理模型的生成結(jié)果是先輸出一系列思考過程再輸出最終結(jié)論。這個(gè)過程類似于人在紙上打草稿草稿越長(zhǎng)越有機(jī)會(huì)通過逐步演算推導(dǎo)出正確答案。從工程角度看兩者最主要的差別不在模型結(jié)構(gòu)而在訓(xùn)練方式。一個(gè)普通模型在訓(xùn)練時(shí)看到的是問題 - 答案的配對(duì)而一個(gè)具備推理能力的模型在訓(xùn)練時(shí)會(huì)看到問題 - 思考過程 - 答案的完整鏈路。思考過程被當(dāng)作可學(xué)習(xí)的文本參與訓(xùn)練模型因此學(xué)會(huì)了在回答前主動(dòng)進(jìn)行推理。另一個(gè)關(guān)鍵點(diǎn)是強(qiáng)化學(xué)習(xí)在這種訓(xùn)練中的應(yīng)用。為了讓模型學(xué)會(huì)找到更好的思考路徑社區(qū)里出現(xiàn)了很多從零構(gòu)建推理模型的項(xiàng)目。核心做法是讓模型對(duì)同一個(gè)問題生成多條不同的推理路徑和答案然后根據(jù)答案是否正確以及推理路徑是否合理給予不同強(qiáng)度的更新信號(hào)。做得好的路徑會(huì)得到正向強(qiáng)化做得差的路徑會(huì)被抑制。這個(gè)過程不依賴人工標(biāo)注的標(biāo)準(zhǔn)推理文本而是通過模型自身的探索來進(jìn)化成本低很多效果卻可能很好。我個(gè)人的立場(chǎng)是如果你想搞懂 modern LLM 的完整能力邊界從生成到推理這一步非常值得復(fù)現(xiàn)。它也是從會(huì)做 Next Token Prediction向會(huì)解決問題進(jìn)階的關(guān)鍵一步。4.2 我在小模型上復(fù)現(xiàn)推理能力的實(shí)操嘗試?yán)碚撉逦笪覜Q定在自己的小模型上進(jìn)行嘗試。我先準(zhǔn)備了一批帶思考鏈的訓(xùn)練數(shù)據(jù)。數(shù)據(jù)來源不是去爬什么高不可攀的內(nèi)容而是自己寫規(guī)則生成對(duì)于數(shù)學(xué)加減法、字符串反轉(zhuǎn)、簡(jiǎn)單邏輯題先生成問題再人工寫一段分步驟的思考文本最后給出答案。數(shù)據(jù)量控制在幾萬條以內(nèi)一方面是因?yàn)槲宜懔τ邢蘖硪环矫嬉沧阋则?yàn)證小模型是否能學(xué)到推理行為。訓(xùn)練時(shí)我用的是監(jiān)督微調(diào)在已經(jīng)訓(xùn)練好的基礎(chǔ)模型之上用這批帶思考鏈的數(shù)據(jù)繼續(xù)訓(xùn)練。關(guān)鍵技巧是思考過程的特殊 token 分隔。我在文本中插入[THINK]和[ANSWER]兩個(gè)特殊標(biāo)記讓模型明確知道哪里是思考區(qū)、哪里是答案區(qū)。推理時(shí)生成到[ANSWER]之前的部分就是模型的思考過程。實(shí)測(cè)結(jié)果很有意思。一個(gè)參數(shù)量不到兩百萬的小模型經(jīng)過幾千步有思考鏈數(shù)據(jù)的微調(diào)之后面對(duì)簡(jiǎn)單加法問題時(shí)確實(shí)會(huì)先輸出類似先把十位和個(gè)位分別相加然后處理進(jìn)位這類的過程文本再給出最終答案。雖然思考過程經(jīng)常有廢話但正確率比起微調(diào)前有明顯提升。不過也有明顯的局限性模型在陌生題型上基本不具備泛化推理能力它只是記住了遇到問題要先輸出一段過程的模式。這說明 from scratch 復(fù)現(xiàn)推理模型不是簡(jiǎn)單加數(shù)據(jù)就行可能還需要引入更多的強(qiáng)化學(xué)習(xí)策略。但作為工程入門走一遍這個(gè)流程已經(jīng)讓我很清楚地理解了思考鏈的本質(zhì)它不神秘只是把未顯式建模的中間計(jì)算過程變成可訓(xùn)練的文本序列。做這一步時(shí)還有一個(gè)非常重要的提醒數(shù)據(jù)合規(guī)和版權(quán)問題永遠(yuǎn)不能繞開。不要從不明渠道下載所謂的全套數(shù)據(jù)或者盜版資源自己寫規(guī)則生成、使用合規(guī)的開源數(shù)據(jù)才是能長(zhǎng)期復(fù)用的做法。模型能力可以慢慢提升合規(guī)底線是絕對(duì)不能突破的。5. 工程化落地模型能跑只是第一步5.1 量化、KV Cache 與推理加速在本地把模型訓(xùn)練出來之后真正讓人頭疼的是怎么讓它跑得更快、占用資源更少。我一開始直接用 PyTorch 逐 token 生成文本結(jié)果慢得離譜。后來才發(fā)現(xiàn)我連最基礎(chǔ)的 KV Cache 都沒有加。KV Cache 的核心思想非常樸素生成下一個(gè) token 時(shí)前面所有 token 的 Key 和 Value 矩陣其實(shí)已經(jīng)算過了沒必要重新計(jì)算。把這些結(jié)果緩存起來每次生成只需要算新 token 的相關(guān)部分推理速度能提升一個(gè)量級(jí)。我的小模型加上 KV Cache 之后生成 200 token 的耗時(shí)直接降為原來的三分之一。除了 KV Cache量化是另一個(gè)大頭。訓(xùn)練好的模型參數(shù)默認(rèn)是 32 位浮點(diǎn)數(shù)占內(nèi)存大、計(jì)算慢。把權(quán)重轉(zhuǎn)成 16 位、8 位甚至 4 位整數(shù)就能非常顯著地壓縮模型體積。我做的是一種比較簡(jiǎn)單的訓(xùn)練后量化方法先統(tǒng)計(jì)每一層權(quán)重的數(shù)值范圍然后把浮點(diǎn)數(shù)映射到整數(shù)范圍推理時(shí)再反量化回浮點(diǎn)數(shù)。對(duì)于我的小模型從 FP32 壓到 INT8 后模型體積減小了接近四倍生成結(jié)果的質(zhì)量幾乎沒有肉眼可見的下降。量化過程里最容易踩的坑是校準(zhǔn)。直接拿模型權(quán)重的 min/max 做映射有時(shí)候會(huì)被離群點(diǎn)帶偏導(dǎo)致量化后某些層輸出異常。我當(dāng)時(shí)用一批驗(yàn)證集數(shù)據(jù)觀察每層的激活分布選擇合理的截?cái)喾秶笤倭炕Ч⒖谭€(wěn)定了很多。簡(jiǎn)單說量化不是簡(jiǎn)單的數(shù)學(xué)映射它需要數(shù)據(jù)來輔助確定映射參數(shù)。5.2 微調(diào)階段的數(shù)據(jù)配比與防災(zāi)難性遺忘工程化之后我開始嘗試在自己的基礎(chǔ)模型上做微調(diào)。初始階段最容易犯的錯(cuò)誤是用單一任務(wù)的數(shù)據(jù)把整個(gè)模型反復(fù)訓(xùn)練結(jié)果只要訓(xùn)練步數(shù)稍微多一點(diǎn)模型原有的通用能力就被大幅破壞了。這個(gè)現(xiàn)象有一個(gè)專門術(shù)語叫災(zāi)難性遺忘我一開始天真地以為只要學(xué)習(xí)率夠小就沒事后來發(fā)現(xiàn)遠(yuǎn)遠(yuǎn)不夠。解決災(zāi)難性遺忘我驗(yàn)證有效的方法有三個(gè)。第一是混入通用數(shù)據(jù)微調(diào)數(shù)據(jù)中至少保留 30% 到 50% 的基礎(chǔ)語料讓模型一邊學(xué)習(xí)新任務(wù)一邊維持舊知識(shí)。第二是降低微調(diào)階段的學(xué)習(xí)率比預(yù)訓(xùn)練階段低一個(gè)數(shù)量級(jí)。第三是使用 LoRA 這類參數(shù)高效微調(diào)方法只訓(xùn)練一小部分附加參數(shù)原有參數(shù)盡量不動(dòng)。LoRA 的原理是把要更新的權(quán)重矩陣分解成兩個(gè)低秩矩陣訓(xùn)練時(shí)只更新這兩個(gè)小矩陣推理時(shí)再合并回原權(quán)重。這樣不僅顯存占用大幅降低而且對(duì)原有參數(shù)的擾動(dòng)很小。我在實(shí)驗(yàn)中發(fā)現(xiàn)LoRA 加上混合數(shù)據(jù)能把新任務(wù)的表現(xiàn)提升不少同時(shí)依然能流暢生成普通文本效果比全量微調(diào)穩(wěn)定得多。數(shù)據(jù)配比這件事沒有萬能公式但有一個(gè)通用的排查方法每次微調(diào)完成后在上一個(gè)版本模型的測(cè)試集上做一次回歸評(píng)測(cè)如果得分明顯下降就說明數(shù)據(jù)配比或?qū)W習(xí)率可能有問題。我前前后后調(diào)整了三次配比最終確定了一個(gè)能兼顧兩邊的比例這個(gè)經(jīng)驗(yàn)值只對(duì)我的場(chǎng)景有效但添加回歸驗(yàn)證集這個(gè)方法對(duì)任何人都有用。5.3 從單機(jī)實(shí)驗(yàn)到服務(wù)化部署模型調(diào)好之后要把模型給別人用就得做服務(wù)化部署。我沒有引入過于復(fù)雜的架構(gòu)而是用了一個(gè)非常簡(jiǎn)單的思路把模型加載到一個(gè)常駐進(jìn)程里通過 HTTP 接口接收文本請(qǐng)求生成結(jié)果后返回。這樣做的好處是模型只在啟動(dòng)時(shí)加載一次不用每次請(qǐng)求都重新初始化。這里有一個(gè)非常容易踩的坑服務(wù)進(jìn)程必須設(shè)置超時(shí)機(jī)制。我的模型雖然只有幾百萬參數(shù)但在無 GPU 的環(huán)境下生成一段長(zhǎng)文本仍然可能耗時(shí)幾秒。如果客戶端超時(shí)時(shí)間設(shè)得太短就會(huì)頻繁報(bào)錯(cuò)。解決辦法有兩個(gè)一是用流式返回模型每生成一個(gè) token 就推給客戶端一部分二是把生成請(qǐng)求丟到任務(wù)隊(duì)列里異步處理客戶端去輪詢結(jié)果。我實(shí)際項(xiàng)目中先用了流式返回交互體驗(yàn)好很多。另外多進(jìn)程服務(wù)一定要小心模型副本占用的顯存或內(nèi)存。如果使用gunicorn這類多進(jìn)程服務(wù)器每個(gè) worker 都會(huì)加載一份模型副本。我試過默認(rèn)配置結(jié)果服務(wù)啟動(dòng)后內(nèi)存直接翻了幾倍。解決辦法是根據(jù)實(shí)際內(nèi)存限制 worker 數(shù)量或者改用多線程。這個(gè)細(xì)節(jié)不親自部署一次是很難從文檔里學(xué)到的。6. 我踩過的坑問題排查與避坑速查表6.1 訓(xùn)練不收斂三個(gè)最常見的元兇我的模型第一次訓(xùn)練時(shí)loss 一路降不下去我很確定自己代碼沒有語法錯(cuò)誤就把鍋甩給了數(shù)據(jù)量太小。后來我系統(tǒng)排查了一遍發(fā)現(xiàn)根本原因和學(xué)習(xí)率有關(guān)我用了條很低的學(xué)習(xí)率導(dǎo)致模型幾乎在原地踏步。把學(xué)習(xí)率從1e-4提到1e-3之后loss 立刻開始下降。檢查學(xué)習(xí)率絕對(duì)是排查不收斂的第一動(dòng)作。第二個(gè)常見元兇是數(shù)據(jù)順序沒有打亂。如果每個(gè) epoch 內(nèi)樣本順序完全不變模型會(huì)學(xué)到順序上的偽特征訓(xùn)練 loss 也能降但驗(yàn)證 loss 非常不穩(wěn)定。打亂數(shù)據(jù)之后問題基本就消失了。第三個(gè)元兇是初始化。使用 PyTorch 默認(rèn)初始化的 Transformer在小模型上訓(xùn)練時(shí)有時(shí)會(huì)遇到嚴(yán)重的梯度異常。我后來把每個(gè)線性層的權(quán)重標(biāo)準(zhǔn)差按照隱藏層維度的倒數(shù)來縮放也就是非常流行的 small initialization 技巧訓(xùn)練穩(wěn)定性明顯提升。不要小看這幾行初始化代碼它對(duì)整個(gè)訓(xùn)練過程影響巨大。還有一個(gè)容易被忽略的元兇損失函數(shù)和標(biāo)簽的對(duì)齊。我在實(shí)現(xiàn)交叉熵?fù)p失時(shí)一開始把 labels 設(shè)置成了形狀不匹配的 tensorPyTorch 沒有直接報(bào)錯(cuò)而是靜默廣播了導(dǎo)致模型學(xué)到了一個(gè)非常奇怪的分布。發(fā)現(xiàn)這個(gè)問題后我在訓(xùn)練循環(huán)里加了一個(gè)assert檢查張量形狀遇到形狀不一致的情況立刻中斷。6.2 顯存爆炸與訓(xùn)練速度過慢的排查思路當(dāng)我把模型從幾百萬參數(shù)擴(kuò)到一兩千萬參數(shù)時(shí)顯存瞬間告急。第一個(gè)原因是注意力矩陣的平方復(fù)雜度序列長(zhǎng)度為 512 時(shí)注意力矩陣是 512 乘以 512顯存開銷隨序列長(zhǎng)度平方增長(zhǎng)。解決辦法是把最長(zhǎng)序列控制在 256或者改用梯度累積減少單批顯存峰值。第二個(gè)原因是優(yōu)化器狀態(tài)比想象中更占內(nèi)存。Adam 優(yōu)化器本身要保存每個(gè)參數(shù)的動(dòng)量變量和方差變量再加上模型參數(shù)和梯度實(shí)際顯存占用接近參數(shù)的 7 到 12 倍。所以不要看到模型只有幾百萬參數(shù)就覺得內(nèi)存肯定夠用?；旌暇扔?xùn)練是很好的緩解手段但要注意梯度縮放和 loss scaling否則很容易出現(xiàn)溢出導(dǎo)致 loss 變成 NaN。我在排查顯存問題時(shí)寫了一個(gè)簡(jiǎn)單腳本利用 PyTorch 的顯存統(tǒng)計(jì)工具在訓(xùn)練循環(huán)前后打印reserved和allocated內(nèi)存很快就定位出哪個(gè)模塊占用了大頭。這種可視化統(tǒng)計(jì)方法比肉眼猜測(cè)高效得多。如果你也遇到 OOM建議先用這個(gè)思路搞清楚內(nèi)存分配在哪里再?zèng)Q定是減小序列長(zhǎng)度、減小 batch還是換一種注意力實(shí)現(xiàn)。6.3 常見問題速查表問題現(xiàn)象可能原因排查方向解決方案訓(xùn)練 loss 不降學(xué)習(xí)率過低或過高打印每個(gè) step 的 loss 和學(xué)習(xí)率使用 warmup cosine 調(diào)度調(diào)峰值學(xué)習(xí)率訓(xùn)練 loss 下降但驗(yàn)證 loss 不降數(shù)據(jù)泄漏或過擬合檢查驗(yàn)證集來源觀察訓(xùn)練/驗(yàn)證差距重建驗(yàn)證集混入通用數(shù)據(jù)驗(yàn)證 loss 劇烈震蕩數(shù)據(jù)順序未打亂 / batch 太小檢查數(shù)據(jù)加載器是否 shuffle開啟 shuffle或增大批次推理速度極慢沒有 KV Cache打印每 token 生成耗時(shí)實(shí)現(xiàn) KV Cache復(fù)用歷史 Key/Value模型生成總是重復(fù)采樣策略太保守觀察生成文本的重復(fù)率調(diào)高 temperature禁用重復(fù) n-gram量化后效果暴跌校準(zhǔn)數(shù)據(jù)不足 / 離群點(diǎn)影響對(duì)比量化前后每層輸出分布用驗(yàn)證集做校準(zhǔn)設(shè)置合理的截?cái)喾秶?wù)啟動(dòng)后內(nèi)存翻倍多進(jìn)程加載多份模型副本查看進(jìn)程內(nèi)存占用限制 worker 數(shù)或改用線程這些坑都是我實(shí)際踩過的。第 4 條 KV Cache 的坑尤其隱蔽因?yàn)槟Ｐ托〉臅r(shí)候速度差異不明顯一旦把序列長(zhǎng)度拉長(zhǎng)差距立刻顯現(xiàn)出來。第 6 條量化校準(zhǔn)則直接關(guān)系到模型能不能在低資源環(huán)境下部署值得多花時(shí)間仔細(xì)調(diào)。最后再分享一個(gè)我現(xiàn)在仍然在用的習(xí)慣每跑一次實(shí)驗(yàn)之前先把隨機(jī)種子固定好然后把數(shù)據(jù)處理流程寫成可復(fù)現(xiàn)的腳本任何一步出錯(cuò)都能從緩存重放。這個(gè)習(xí)慣在 ai-engineering-from-scratch 項(xiàng)目的后半段幫了我大忙因?yàn)楦膭?dòng)一旦多了你根本無法判斷效果變化到底來自數(shù)據(jù)、模型版本還是訓(xùn)練參數(shù)。固定種子、緩存數(shù)據(jù)、記錄每次實(shí)驗(yàn)的配置這三件事堅(jiān)持做下來整個(gè)項(xiàng)目才真正算得上工程而不是一次性的代碼練習(xí)。如果你也打算從零開始走一遍 AI 工程我建議你從寫一個(gè)最簡(jiǎn)單的手寫數(shù)字分類器起步再一步步走到語言模型和部署。這條路不易但走完之后的底氣是任何框架封裝都給不了的。