習(xí)筆記：從GPU顯存到分布式訓(xùn)練與推理部署的體系化認知)
1. 從一張白紙到體系化認知我為什么要啃下AI Infra這塊硬骨頭第一次聽到“AI Infra”這個詞的時候我正蹲在一個模型訓(xùn)練任務(wù)的日志前面盯著GPU利用率曲線像心電圖一樣忽高忽低。那時候我以為搞AI就是調(diào)參、煉丹、看loss直到有人問我一句“你這個模型從數(shù)據(jù)加載到梯度同步中間到底經(jīng)過了哪些環(huán)節(jié)”我愣住了。那一刻我才意識到自己一直在用黑盒的方式做事情而對支撐整個AI系統(tǒng)運轉(zhuǎn)的基礎(chǔ)設(shè)施層幾乎一無所知。AI Infra直譯過來就是“人工智能基礎(chǔ)設(shè)施”。它不是某一個具體的框架或者工具而是一整套讓AI模型能夠被高效訓(xùn)練、部署和運行的底層支撐體系。你可以把它理解成AI世界里的“水電煤”——沒有它再漂亮的算法也只是紙面上的公式。它涵蓋的范圍非常廣從最底層的芯片和驅(qū)動到中間的通信庫、訓(xùn)練框架、推理引擎再到上層的資源調(diào)度、監(jiān)控運維每一層都有大量的技術(shù)細節(jié)和工程取舍。我寫這份學(xué)習(xí)筆記的初衷很簡單市面上講AI Infra的資料要么太散要么太淺要么一上來就堆術(shù)語把人勸退。我希望把自己從零開始摸索的過程記錄下來形成一個有邏輯、有層次、能落地的知識全景。不管你是剛?cè)胄械乃惴üこ處熯€是想往Infra方向轉(zhuǎn)型的后端開發(fā)者或者只是對“模型到底怎么跑起來的”感到好奇的技術(shù)愛好者這份筆記都能給你一條清晰的路徑。接下來的內(nèi)容我會按照“計算底座—訓(xùn)練基礎(chǔ)設(shè)施—推理與部署—資源調(diào)度與運維”這條主線展開每一塊都會講清楚它解決什么問題、核心技術(shù)點是什么、實際工程中怎么選型和避坑。這些都是我在實際項目中踩過坑、翻過文檔、跟同行交流之后沉淀下來的東西不是教科書式的羅列。2. 計算底座GPU、顯存與算力單位的那些事2.1 為什么AI Infra繞不開GPU聊AI Infra第一個繞不開的就是GPU。很多人知道GPU比CPU適合做深度學(xué)習(xí)但說不清楚為什么。簡單來說CPU像是一個數(shù)學(xué)教授什么題都會做但一次只能做幾道GPU像是一萬個中學(xué)生每個人只會做簡單的加減乘除但一萬人同時動手總量就非?？捎^。深度學(xué)習(xí)里的矩陣乘法、卷積運算本質(zhì)上就是大量簡單計算的堆疊所以GPU的并行架構(gòu)天然適合。但這里有個容易被忽略的點GPU并不是越貴越好關(guān)鍵要看你的任務(wù)是什么類型。訓(xùn)練大模型需要高顯存帶寬和大的顯存容量因為模型參數(shù)、梯度、優(yōu)化器狀態(tài)都要放在顯存里而推理任務(wù)更看重吞吐量和延遲對顯存的要求相對低一些。我見過有人拿消費級顯卡去訓(xùn)大模型結(jié)果batch size只能設(shè)成1訓(xùn)練速度慢到懷疑人生。這不是顯卡不好是場景不匹配。2.2 顯存到底被誰吃掉了顯存是AI Infra里最稀缺的資源之一。很多人第一次訓(xùn)練模型時都會遇到“CUDA out of memory”的報錯然后一臉茫然我模型才幾百兆怎么24G顯存就不夠了這里我列一下顯存的主要消耗方模型參數(shù)這是最直觀的比如一個7B參數(shù)的模型如果用FP16存儲大約需要14GB。梯度訓(xùn)練時每個參數(shù)都要存對應(yīng)的梯度大小和參數(shù)一樣又是14GB。優(yōu)化器狀態(tài)如果用Adam每個參數(shù)要存一階矩和二階矩又是兩倍參數(shù)大小FP32下就是56GB。激活值前向傳播過程中每一層的輸出都要存下來供反向傳播使用這部分和batch size、序列長度強相關(guān)。臨時緩沖區(qū)通信、算子計算過程中的臨時空間。所以一個7B模型用Adam全量微調(diào)光模型相關(guān)的狀態(tài)就要80GB以上還沒算激活值。這就是為什么大模型訓(xùn)練要用分布式、要用ZeRO、要用混合精度。理解顯存的去向是做好AI Infra的第一步。2.3 算力單位的換算與實測差異買卡的時候銷售會跟你說“這張卡多少TFLOPS”但實際跑起來你會發(fā)現(xiàn)根本達不到。這里有幾個概念需要厘清精度類型說明典型用途FP32單精度浮點傳統(tǒng)訓(xùn)練、科學(xué)計算FP16半精度浮點混合精度訓(xùn)練BF16腦浮點大模型訓(xùn)練首選INT88位整數(shù)推理量化INT44位整數(shù)極致推理壓縮TFLOPS標稱值通常是在理想條件下測出來的實際利用率受限于顯存帶寬、通信開銷、算子效率等因素。我在實際項目中觀察到訓(xùn)練任務(wù)能達到標稱算力的30%到50%就已經(jīng)算不錯了。所以做容量規(guī)劃的時候不要按標稱值算打個對折甚至更低才靠譜。注意不同精度下的算力不能直接比較。一張卡FP16算力是FP32的幾倍但如果你用FP32跑就只能享受FP32的算力。選卡之前先確定你的訓(xùn)練精度策略。3. 訓(xùn)練基礎(chǔ)設(shè)施從單卡到千卡集群的跨越3.1 數(shù)據(jù)并行、模型并行與流水線并行當你只有一張卡的時候訓(xùn)練很簡單把數(shù)據(jù)喂進去算梯度更新參數(shù)。但當模型大到一張卡放不下或者數(shù)據(jù)多到一張卡算不完的時候就需要分布式訓(xùn)練。主流的并行策略有三種數(shù)據(jù)并行是最直觀的每張卡復(fù)制一份完整的模型把數(shù)據(jù)切分給不同的卡各自算梯度然后同步梯度。這種方式實現(xiàn)簡單但要求單卡能放下整個模型。PyTorch的DDP就是典型的數(shù)據(jù)并行。模型并行是把模型本身切開放到不同的卡上。比如Transformer的每一層放到不同的卡或者把一個大矩陣乘法拆開。這種方式能訓(xùn)更大的模型但通信開銷大實現(xiàn)復(fù)雜。流水線并行是模型并行的一種優(yōu)化把模型按層切成多個階段不同階段放在不同卡上數(shù)據(jù)像流水線一樣依次經(jīng)過各個階段。關(guān)鍵是做好微批次劃分讓各階段盡量不空閑。實際訓(xùn)練大模型時通常是三種并行混合使用業(yè)界叫“3D并行”。我剛開始接觸的時候覺得這套東西太復(fù)雜后來發(fā)現(xiàn)核心思想就一句話在顯存、計算、通信三者之間找平衡。顯存不夠就切模型計算不夠就加數(shù)據(jù)并行通信瓶頸就調(diào)整切分策略。3.2 通信原語AllReduce、AllGather與ReduceScatter分布式訓(xùn)練里卡和卡之間要交換信息這些交換操作叫通信原語。最核心的三個是AllReduce每張卡都有一份數(shù)據(jù)匯總后每張卡都得到全局結(jié)果。數(shù)據(jù)并行的梯度同步就用這個。AllGather每張卡有一份數(shù)據(jù)收集后每張卡都得到所有卡的數(shù)據(jù)拼接。ReduceScatter每張卡有一份數(shù)據(jù)匯總后把結(jié)果切分給各卡。這三個原語看起來簡單但組合起來能實現(xiàn)各種并行策略。比如ZeRO的核心思想就是用AllGather和ReduceScatter替代AllReduce把優(yōu)化器狀態(tài)、梯度、參數(shù)分片存儲從而大幅降低單卡顯存占用。通信的瓶頸通常不在帶寬而在延遲和同步。我踩過的一個坑是在跨節(jié)點通信時如果網(wǎng)絡(luò)拓撲沒配好AllReduce的時間會暴漲。后來查資料才知道NCCL會根據(jù)拓撲自動選擇ring或者tree算法但如果網(wǎng)卡和GPU的親和性沒設(shè)置好數(shù)據(jù)繞遠路性能就下來了。3.3 混合精度訓(xùn)練省顯存還能提速混合精度訓(xùn)練是現(xiàn)在大模型訓(xùn)練的標配。核心思路是前向和反向用FP16或BF16計算參數(shù)更新用FP32。這樣既能利用低精度的高算力又能保持數(shù)值穩(wěn)定性。但混合精度不是簡單地把數(shù)據(jù)類型一改就完事。有幾個關(guān)鍵點損失縮放FP16的表示范圍窄梯度太小會變成0。所以要在損失上乘一個縮放因子反向傳播后再除回來。主權(quán)重副本優(yōu)化器里要保留一份FP32的參數(shù)副本更新時用FP32前向時轉(zhuǎn)成FP16。算子兼容性不是所有算子都支持FP16有些必須用FP32框架會自動處理但你要知道哪些地方可能成為瓶頸。BF16比FP16更好用因為它的表示范圍和FP32一樣不需要損失縮放數(shù)值穩(wěn)定性更好?，F(xiàn)在新出的卡基本都支持BF16如果條件允許我建議優(yōu)先用BF16。3.4 Checkpoint與容錯訓(xùn)練中斷了怎么辦大模型訓(xùn)練動輒幾天甚至幾周中間難免遇到硬件故障、網(wǎng)絡(luò)抖動、任務(wù)被搶占。如果沒有Checkpoint機制一次中斷就前功盡棄。Checkpoint的設(shè)計有幾個考量頻率太頻繁影響訓(xùn)練速度太稀疏中斷后損失大。通常按時間間隔或者步數(shù)間隔來設(shè)。內(nèi)容至少要存模型參數(shù)、優(yōu)化器狀態(tài)、當前步數(shù)、數(shù)據(jù)加載器的狀態(tài)。少存一樣恢復(fù)后結(jié)果就可能不對。存儲大模型Checkpoint動輒幾百GB存到本地盤不現(xiàn)實一般要存到分布式文件系統(tǒng)或?qū)ο蟠鎯Α．惒奖４姹４鍯heckpoint時如果阻塞訓(xùn)練會浪費算力。好的框架支持異步保存先把狀態(tài)拷到CPU內(nèi)存再慢慢寫盤。我遇到過最坑的情況是Checkpoint存了但恢復(fù)后loss對不上。排查半天發(fā)現(xiàn)是數(shù)據(jù)加載器的隨機種子沒存導(dǎo)致恢復(fù)后數(shù)據(jù)順序變了。這種細節(jié)在文檔里往往不會寫只有踩過才知道。4. 推理與部署模型上線后的真實戰(zhàn)場4.1 推理和訓(xùn)練到底差在哪很多人覺得推理比訓(xùn)練簡單其實不然。訓(xùn)練關(guān)注的是吞吐量和收斂速度推理關(guān)注的是延遲、吞吐、成本三者的平衡。訓(xùn)練可以容忍偶爾的失敗重試推理要求7x24穩(wěn)定服務(wù)。訓(xùn)練可以用大batch提高利用率推理的請求是實時到來的batch大了延遲就高。推理場景大致分兩類在線推理要求低延遲比如對話系統(tǒng)、推薦排序離線推理追求高吞吐比如批量圖片處理、數(shù)據(jù)標注。這兩類的優(yōu)化方向完全不同。4.2 推理引擎的核心優(yōu)化手段現(xiàn)在主流的推理引擎有TensorRT、ONNX Runtime、Triton等它們做的事情大同小異算子融合是把多個小算子合并成一個大算子減少kernel啟動開銷和顯存讀寫。比如把卷積、偏置、激活融合成一個。這個優(yōu)化在GPU上效果特別明顯因為kernel啟動本身就有開銷。量化是把FP16或FP32的權(quán)重和激活值轉(zhuǎn)成INT8甚至INT4減少顯存占用和計算量。量化分訓(xùn)練后量化和量化感知訓(xùn)練前者簡單但精度損失大后者需要在訓(xùn)練時模擬量化誤差精度更好但麻煩。動態(tài)批處理是把多個請求攢成一個batch一起算提高GPU利用率。但攢多久是個問題攢太久延遲高攢太少吞吐低。通常設(shè)一個最大等待時間超時就發(fā)。KV Cache優(yōu)化是針對自回歸生成模型的。生成每個token時前面token的Key和Value可以緩存下來不用重復(fù)計算。但KV Cache會占顯存序列越長占得越多?，F(xiàn)在有PagedAttention之類的技術(shù)把KV Cache分頁管理減少碎片。4.3 部署形態(tài)從單機到服務(wù)化小規(guī)模部署可能就是一個Python腳本加載模型起個Flask服務(wù)。但生產(chǎn)環(huán)境要考慮的東西多得多模型版本管理新模型上線要能灰度、能回滾。自動擴縮容流量高峰自動加實例低谷自動減。健康檢查實例掛了要能自動摘除。監(jiān)控告警延遲、吞吐、錯誤率、GPU利用率都要盯著。這些需求催生了模型服務(wù)化平臺比如Triton Inference Server、TorchServe等。它們提供了統(tǒng)一的接口、多模型管理、動態(tài)批處理、指標暴露等功能。用這些平臺能省不少事但也要理解它們的限制比如對自定義算子的支持、對特殊輸入格式的處理。提示推理部署最容易忽略的是預(yù)處理和后處理。模型本身可能只占端到端延遲的30%剩下70%都在數(shù)據(jù)搬運、格式轉(zhuǎn)換、業(yè)務(wù)邏輯上。優(yōu)化推理之前先用profiler看清楚時間花在哪。5. 資源調(diào)度與集群運維讓算力不浪費5.1 調(diào)度系統(tǒng)的核心挑戰(zhàn)一個GPU集群少則幾十張卡多則上千張。怎么把訓(xùn)練任務(wù)、推理服務(wù)、開發(fā)調(diào)試等不同需求合理地分配到這些卡上就是調(diào)度系統(tǒng)要解決的問題。核心挑戰(zhàn)有幾個碎片化一個任務(wù)要8張卡但集群里只有7張空閑就得等。等多久不知道可能幾分鐘可能幾小時。親和性有些任務(wù)需要卡在同一個節(jié)點內(nèi)因為跨節(jié)點通信慢。有些任務(wù)需要特定的卡型號。優(yōu)先級緊急任務(wù)要能插隊低優(yōu)先級任務(wù)要能被搶占。公平性多個團隊共享集群時不能讓一個團隊占滿所有資源。Kubernetes是目前最流行的調(diào)度底座但原生K8s對GPU的支持比較基礎(chǔ)。所以有了各種擴展比如通過Device Plugin暴露GPU資源通過調(diào)度器擴展實現(xiàn)gang scheduling一組Pod要么都調(diào)度要么都不調(diào)度。5.2 利用率提升的常見手段GPU很貴閑置就是浪費。提升利用率的手段包括分時復(fù)用多個任務(wù)共享一張卡按時間片切換。適合開發(fā)調(diào)試場景不適合生產(chǎn)訓(xùn)練。顯存超分把多個小模型塞到一張卡上各自跑各自的。推理場景常用。彈性訓(xùn)練任務(wù)能動態(tài)增減卡數(shù)集群有空閑就加有高優(yōu)先級任務(wù)就減。這需要訓(xùn)練框架支持動態(tài)擴縮容。混部把推理和訓(xùn)練混在一起推理用不滿的卡給訓(xùn)練用。但要注意資源隔離別互相影響。我在實際運維中發(fā)現(xiàn)利用率低往往不是技術(shù)問題而是流程問題。比如團隊申請了卡但一直不用或者任務(wù)排隊策略不合理導(dǎo)致大任務(wù)餓死。技術(shù)手段能解決一部分但制度和流程的優(yōu)化同樣重要。5.3 監(jiān)控體系看得見才能管得住AI集群的監(jiān)控比普通服務(wù)復(fù)雜因為要監(jiān)控的層次多監(jiān)控層次關(guān)鍵指標常用工具硬件層GPU溫度、功耗、顯存、利用率DCGM、nvidia-smi通信層帶寬、延遲、重傳率NCCL日志、IB監(jiān)控框架層訓(xùn)練速度、loss曲線、梯度范數(shù)TensorBoard、WandB服務(wù)層請求延遲、QPS、錯誤率Prometheus、Grafana業(yè)務(wù)層模型效果、轉(zhuǎn)化率自定義埋點監(jiān)控的目的不只是發(fā)現(xiàn)問題還要能定位問題。比如訓(xùn)練變慢了是數(shù)據(jù)加載慢、通信慢、還是計算慢沒有細粒度的監(jiān)控就只能靠猜。6. 我踩過的那些坑與總結(jié)的經(jīng)驗6.1 環(huán)境配置版本兼容性是萬惡之源AI Infra最煩人的事情之一就是環(huán)境配置。CUDA版本、驅(qū)動版本、框架版本、通信庫版本任何一個不匹配都可能出問題。我遇到過CUDA版本和PyTorch不匹配導(dǎo)致無法使用GPU也遇到過NCCL版本和驅(qū)動不兼容導(dǎo)致多卡訓(xùn)練hang住。我的經(jīng)驗是用容器別在裸機上折騰。Docker加上NVIDIA Container Toolkit可以把環(huán)境固化下來。但容器也不是銀彈GPU直通、共享內(nèi)存、網(wǎng)絡(luò)模式都有坑。最好是用官方維護的鏡像作為基礎(chǔ)別自己從零搭。6.2 性能調(diào)優(yōu)先測量再優(yōu)化性能問題最忌諱拍腦袋。我見過有人一上來就改batch size、改學(xué)習(xí)率結(jié)果越調(diào)越差。正確的做法是先用profiler測量找到真正的瓶頸。PyTorch有torch.profilerNVIDIA有Nsight Systems這些都是好工具。測量之后你會發(fā)現(xiàn)瓶頸往往不在你以為的地方。比如訓(xùn)練慢可能不是GPU算得慢而是數(shù)據(jù)加載的IO慢或者CPU預(yù)處理慢。6.3 故障排查從日志和指標入手分布式訓(xùn)練的故障排查是最考驗人的。任務(wù)掛了日志可能幾千行關(guān)鍵信息就一兩行。我的習(xí)慣是先看有沒有明顯的報錯比如OOM、NCCL timeout?？锤鞴?jié)點的GPU利用率如果有的卡忙有的卡閑可能是負載不均?？赐ㄐ湃罩綨CCL會打印ring的建立過程如果卡在建ring可能是網(wǎng)絡(luò)問題?？碈heckpoint恢復(fù)是否正常有時候是恢復(fù)邏輯有bug。還有一個經(jīng)驗保留現(xiàn)場。任務(wù)掛了別急著重啟先把日志、nvidia-smi輸出、網(wǎng)絡(luò)狀態(tài)都存下來。重啟之后現(xiàn)場就沒了再想排查就難了。6.4 學(xué)習(xí)路徑建議如果你剛開始接觸AI Infra我建議按這個順序來先跑通單卡訓(xùn)練理解模型訓(xùn)練的基本流程。再跑通單機多卡理解數(shù)據(jù)并行和通信。然后嘗試多機多卡理解分布式訓(xùn)練的復(fù)雜性。接著學(xué)推理部署把訓(xùn)練好的模型服務(wù)化。最后學(xué)調(diào)度和運維理解集群級別的管理。每一步都要動手做光看文檔是學(xué)不會的。遇到問題先自己排查實在搞不定再查資料或問人。排查的過程本身就是最好的學(xué)習(xí)。這個領(lǐng)域變化很快新硬件、新框架、新優(yōu)化手段層出不窮。但底層的東西變化沒那么快計算、存儲、通信、調(diào)度這些核心問題十年二十年都不會過時。把基礎(chǔ)打牢上層的東西學(xué)起來就快。最后分享一個我自己的習(xí)慣我會維護一個“坑本”每次遇到問題解決了就記下來包括現(xiàn)象、排查過程、根因、解決方案。下次再遇到類似問題翻一下坑本就能快速定位。這個習(xí)慣幫我省了大量時間也讓我對系統(tǒng)的理解越來越深。 AI Infra學(xué)習(xí)筆記：從GPU顯存到分布式訓(xùn)練與推理部署的體系化認知
1. 從一張白紙到體系化認知我為什么要啃下AI Infra這塊硬骨頭第一次聽到“AI Infra”這個詞的時候我正蹲在一個模型訓(xùn)練任務(wù)的日志前面盯著GPU利用率曲線像心電圖一樣忽高忽低。那時候我以為搞AI就是調(diào)參、煉丹、看loss直到有人問我一句“你這個模型從數(shù)據(jù)加載到梯度同步中間到底經(jīng)過了哪些環(huán)節(jié)”我愣住了。那一刻我才意識到自己一直在用黑盒的方式做事情而對支撐整個AI系統(tǒng)運轉(zhuǎn)的基礎(chǔ)設(shè)施層幾乎一無所知。AI Infra直譯過來就是“人工智能基礎(chǔ)設(shè)施”。它不是某一個具體的框架或者工具而是一整套讓AI模型能夠被高效訓(xùn)練、部署和運行的底層支撐體系。你可以把它理解成AI世界里的“水電煤”——沒有它再漂亮的算法也只是紙面上的公式。它涵蓋的范圍非常廣從最底層的芯片和驅(qū)動到中間的通信庫、訓(xùn)練框架、推理引擎再到上層的資源調(diào)度、監(jiān)控運維每一層都有大量的技術(shù)細節(jié)和工程取舍。我寫這份學(xué)習(xí)筆記的初衷很簡單市面上講AI Infra的資料要么太散要么太淺要么一上來就堆術(shù)語把人勸退。我希望把自己從零開始摸索的過程記錄下來形成一個有邏輯、有層次、能落地的知識全景。不管你是剛?cè)胄械乃惴üこ處熯€是想往Infra方向轉(zhuǎn)型的后端開發(fā)者或者只是對“模型到底怎么跑起來的”感到好奇的技術(shù)愛好者這份筆記都能給你一條清晰的路徑。接下來的內(nèi)容我會按照“計算底座—訓(xùn)練基礎(chǔ)設(shè)施—推理與部署—資源調(diào)度與運維”這條主線展開每一塊都會講清楚它解決什么問題、核心技術(shù)點是什么、實際工程中怎么選型和避坑。這些都是我在實際項目中踩過坑、翻過文檔、跟同行交流之后沉淀下來的東西不是教科書式的羅列。2. 計算底座GPU、顯存與算力單位的那些事2.1 為什么AI Infra繞不開GPU聊AI Infra第一個繞不開的就是GPU。很多人知道GPU比CPU適合做深度學(xué)習(xí)但說不清楚為什么。簡單來說CPU像是一個數(shù)學(xué)教授什么題都會做但一次只能做幾道GPU像是一萬個中學(xué)生每個人只會做簡單的加減乘除但一萬人同時動手總量就非?？捎^。深度學(xué)習(xí)里的矩陣乘法、卷積運算本質(zhì)上就是大量簡單計算的堆疊所以GPU的并行架構(gòu)天然適合。但這里有個容易被忽略的點GPU并不是越貴越好關(guān)鍵要看你的任務(wù)是什么類型。訓(xùn)練大模型需要高顯存帶寬和大的顯存容量因為模型參數(shù)、梯度、優(yōu)化器狀態(tài)都要放在顯存里而推理任務(wù)更看重吞吐量和延遲對顯存的要求相對低一些。我見過有人拿消費級顯卡去訓(xùn)大模型結(jié)果batch size只能設(shè)成1訓(xùn)練速度慢到懷疑人生。這不是顯卡不好是場景不匹配。2.2 顯存到底被誰吃掉了顯存是AI Infra里最稀缺的資源之一。很多人第一次訓(xùn)練模型時都會遇到“CUDA out of memory”的報錯然后一臉茫然我模型才幾百兆怎么24G顯存就不夠了這里我列一下顯存的主要消耗方模型參數(shù)這是最直觀的比如一個7B參數(shù)的模型如果用FP16存儲大約需要14GB。梯度訓(xùn)練時每個參數(shù)都要存對應(yīng)的梯度大小和參數(shù)一樣又是14GB。優(yōu)化器狀態(tài)如果用Adam每個參數(shù)要存一階矩和二階矩又是兩倍參數(shù)大小FP32下就是56GB。激活值前向傳播過程中每一層的輸出都要存下來供反向傳播使用這部分和batch size、序列長度強相關(guān)。臨時緩沖區(qū)通信、算子計算過程中的臨時空間。所以一個7B模型用Adam全量微調(diào)光模型相關(guān)的狀態(tài)就要80GB以上還沒算激活值。這就是為什么大模型訓(xùn)練要用分布式、要用ZeRO、要用混合精度。理解顯存的去向是做好AI Infra的第一步。2.3 算力單位的換算與實測差異買卡的時候銷售會跟你說“這張卡多少TFLOPS”但實際跑起來你會發(fā)現(xiàn)根本達不到。這里有幾個概念需要厘清精度類型說明典型用途FP32單精度浮點傳統(tǒng)訓(xùn)練、科學(xué)計算FP16半精度浮點混合精度訓(xùn)練BF16腦浮點大模型訓(xùn)練首選INT88位整數(shù)推理量化INT44位整數(shù)極致推理壓縮TFLOPS標稱值通常是在理想條件下測出來的實際利用率受限于顯存帶寬、通信開銷、算子效率等因素。我在實際項目中觀察到訓(xùn)練任務(wù)能達到標稱算力的30%到50%就已經(jīng)算不錯了。所以做容量規(guī)劃的時候不要按標稱值算打個對折甚至更低才靠譜。注意不同精度下的算力不能直接比較。一張卡FP16算力是FP32的幾倍但如果你用FP32跑就只能享受FP32的算力。選卡之前先確定你的訓(xùn)練精度策略。3. 訓(xùn)練基礎(chǔ)設(shè)施從單卡到千卡集群的跨越3.1 數(shù)據(jù)并行、模型并行與流水線并行當你只有一張卡的時候訓(xùn)練很簡單把數(shù)據(jù)喂進去算梯度更新參數(shù)。但當模型大到一張卡放不下或者數(shù)據(jù)多到一張卡算不完的時候就需要分布式訓(xùn)練。主流的并行策略有三種數(shù)據(jù)并行是最直觀的每張卡復(fù)制一份完整的模型把數(shù)據(jù)切分給不同的卡各自算梯度然后同步梯度。這種方式實現(xiàn)簡單但要求單卡能放下整個模型。PyTorch的DDP就是典型的數(shù)據(jù)并行。模型并行是把模型本身切開放到不同的卡上。比如Transformer的每一層放到不同的卡或者把一個大矩陣乘法拆開。這種方式能訓(xùn)更大的模型但通信開銷大實現(xiàn)復(fù)雜。流水線并行是模型并行的一種優(yōu)化把模型按層切成多個階段不同階段放在不同卡上數(shù)據(jù)像流水線一樣依次經(jīng)過各個階段。關(guān)鍵是做好微批次劃分讓各階段盡量不空閑。實際訓(xùn)練大模型時通常是三種并行混合使用業(yè)界叫“3D并行”。我剛開始接觸的時候覺得這套東西太復(fù)雜后來發(fā)現(xiàn)核心思想就一句話在顯存、計算、通信三者之間找平衡。顯存不夠就切模型計算不夠就加數(shù)據(jù)并行通信瓶頸就調(diào)整切分策略。3.2 通信原語AllReduce、AllGather與ReduceScatter分布式訓(xùn)練里卡和卡之間要交換信息這些交換操作叫通信原語。最核心的三個是AllReduce每張卡都有一份數(shù)據(jù)匯總后每張卡都得到全局結(jié)果。數(shù)據(jù)并行的梯度同步就用這個。AllGather每張卡有一份數(shù)據(jù)收集后每張卡都得到所有卡的數(shù)據(jù)拼接。ReduceScatter每張卡有一份數(shù)據(jù)匯總后把結(jié)果切分給各卡。這三個原語看起來簡單但組合起來能實現(xiàn)各種并行策略。比如ZeRO的核心思想就是用AllGather和ReduceScatter替代AllReduce把優(yōu)化器狀態(tài)、梯度、參數(shù)分片存儲從而大幅降低單卡顯存占用。通信的瓶頸通常不在帶寬而在延遲和同步。我踩過的一個坑是在跨節(jié)點通信時如果網(wǎng)絡(luò)拓撲沒配好AllReduce的時間會暴漲。后來查資料才知道NCCL會根據(jù)拓撲自動選擇ring或者tree算法但如果網(wǎng)卡和GPU的親和性沒設(shè)置好數(shù)據(jù)繞遠路性能就下來了。3.3 混合精度訓(xùn)練省顯存還能提速混合精度訓(xùn)練是現(xiàn)在大模型訓(xùn)練的標配。核心思路是前向和反向用FP16或BF16計算參數(shù)更新用FP32。這樣既能利用低精度的高算力又能保持數(shù)值穩(wěn)定性。但混合精度不是簡單地把數(shù)據(jù)類型一改就完事。有幾個關(guān)鍵點損失縮放FP16的表示范圍窄梯度太小會變成0。所以要在損失上乘一個縮放因子反向傳播后再除回來。主權(quán)重副本優(yōu)化器里要保留一份FP32的參數(shù)副本更新時用FP32前向時轉(zhuǎn)成FP16。算子兼容性不是所有算子都支持FP16有些必須用FP32框架會自動處理但你要知道哪些地方可能成為瓶頸。BF16比FP16更好用因為它的表示范圍和FP32一樣不需要損失縮放數(shù)值穩(wěn)定性更好?，F(xiàn)在新出的卡基本都支持BF16如果條件允許我建議優(yōu)先用BF16。3.4 Checkpoint與容錯訓(xùn)練中斷了怎么辦大模型訓(xùn)練動輒幾天甚至幾周中間難免遇到硬件故障、網(wǎng)絡(luò)抖動、任務(wù)被搶占。如果沒有Checkpoint機制一次中斷就前功盡棄。Checkpoint的設(shè)計有幾個考量頻率太頻繁影響訓(xùn)練速度太稀疏中斷后損失大。通常按時間間隔或者步數(shù)間隔來設(shè)。內(nèi)容至少要存模型參數(shù)、優(yōu)化器狀態(tài)、當前步數(shù)、數(shù)據(jù)加載器的狀態(tài)。少存一樣恢復(fù)后結(jié)果就可能不對。存儲大模型Checkpoint動輒幾百GB存到本地盤不現(xiàn)實一般要存到分布式文件系統(tǒng)或?qū)ο蟠鎯Α．惒奖４姹４鍯heckpoint時如果阻塞訓(xùn)練會浪費算力。好的框架支持異步保存先把狀態(tài)拷到CPU內(nèi)存再慢慢寫盤。我遇到過最坑的情況是Checkpoint存了但恢復(fù)后loss對不上。排查半天發(fā)現(xiàn)是數(shù)據(jù)加載器的隨機種子沒存導(dǎo)致恢復(fù)后數(shù)據(jù)順序變了。這種細節(jié)在文檔里往往不會寫只有踩過才知道。4. 推理與部署模型上線后的真實戰(zhàn)場4.1 推理和訓(xùn)練到底差在哪很多人覺得推理比訓(xùn)練簡單其實不然。訓(xùn)練關(guān)注的是吞吐量和收斂速度推理關(guān)注的是延遲、吞吐、成本三者的平衡。訓(xùn)練可以容忍偶爾的失敗重試推理要求7x24穩(wěn)定服務(wù)。訓(xùn)練可以用大batch提高利用率推理的請求是實時到來的batch大了延遲就高。推理場景大致分兩類在線推理要求低延遲比如對話系統(tǒng)、推薦排序離線推理追求高吞吐比如批量圖片處理、數(shù)據(jù)標注。這兩類的優(yōu)化方向完全不同。4.2 推理引擎的核心優(yōu)化手段現(xiàn)在主流的推理引擎有TensorRT、ONNX Runtime、Triton等它們做的事情大同小異算子融合是把多個小算子合并成一個大算子減少kernel啟動開銷和顯存讀寫。比如把卷積、偏置、激活融合成一個。這個優(yōu)化在GPU上效果特別明顯因為kernel啟動本身就有開銷。量化是把FP16或FP32的權(quán)重和激活值轉(zhuǎn)成INT8甚至INT4減少顯存占用和計算量。量化分訓(xùn)練后量化和量化感知訓(xùn)練前者簡單但精度損失大后者需要在訓(xùn)練時模擬量化誤差精度更好但麻煩。動態(tài)批處理是把多個請求攢成一個batch一起算提高GPU利用率。但攢多久是個問題攢太久延遲高攢太少吞吐低。通常設(shè)一個最大等待時間超時就發(fā)。KV Cache優(yōu)化是針對自回歸生成模型的。生成每個token時前面token的Key和Value可以緩存下來不用重復(fù)計算。但KV Cache會占顯存序列越長占得越多?，F(xiàn)在有PagedAttention之類的技術(shù)把KV Cache分頁管理減少碎片。4.3 部署形態(tài)從單機到服務(wù)化小規(guī)模部署可能就是一個Python腳本加載模型起個Flask服務(wù)。但生產(chǎn)環(huán)境要考慮的東西多得多模型版本管理新模型上線要能灰度、能回滾。自動擴縮容流量高峰自動加實例低谷自動減。健康檢查實例掛了要能自動摘除。監(jiān)控告警延遲、吞吐、錯誤率、GPU利用率都要盯著。這些需求催生了模型服務(wù)化平臺比如Triton Inference Server、TorchServe等。它們提供了統(tǒng)一的接口、多模型管理、動態(tài)批處理、指標暴露等功能。用這些平臺能省不少事但也要理解它們的限制比如對自定義算子的支持、對特殊輸入格式的處理。提示推理部署最容易忽略的是預(yù)處理和后處理。模型本身可能只占端到端延遲的30%剩下70%都在數(shù)據(jù)搬運、格式轉(zhuǎn)換、業(yè)務(wù)邏輯上。優(yōu)化推理之前先用profiler看清楚時間花在哪。5. 資源調(diào)度與集群運維讓算力不浪費5.1 調(diào)度系統(tǒng)的核心挑戰(zhàn)一個GPU集群少則幾十張卡多則上千張。怎么把訓(xùn)練任務(wù)、推理服務(wù)、開發(fā)調(diào)試等不同需求合理地分配到這些卡上就是調(diào)度系統(tǒng)要解決的問題。核心挑戰(zhàn)有幾個碎片化一個任務(wù)要8張卡但集群里只有7張空閑就得等。等多久不知道可能幾分鐘可能幾小時。親和性有些任務(wù)需要卡在同一個節(jié)點內(nèi)因為跨節(jié)點通信慢。有些任務(wù)需要特定的卡型號。優(yōu)先級緊急任務(wù)要能插隊低優(yōu)先級任務(wù)要能被搶占。公平性多個團隊共享集群時不能讓一個團隊占滿所有資源。Kubernetes是目前最流行的調(diào)度底座但原生K8s對GPU的支持比較基礎(chǔ)。所以有了各種擴展比如通過Device Plugin暴露GPU資源通過調(diào)度器擴展實現(xiàn)gang scheduling一組Pod要么都調(diào)度要么都不調(diào)度。5.2 利用率提升的常見手段GPU很貴閑置就是浪費。提升利用率的手段包括分時復(fù)用多個任務(wù)共享一張卡按時間片切換。適合開發(fā)調(diào)試場景不適合生產(chǎn)訓(xùn)練。顯存超分把多個小模型塞到一張卡上各自跑各自的。推理場景常用。彈性訓(xùn)練任務(wù)能動態(tài)增減卡數(shù)集群有空閑就加有高優(yōu)先級任務(wù)就減。這需要訓(xùn)練框架支持動態(tài)擴縮容。混部把推理和訓(xùn)練混在一起推理用不滿的卡給訓(xùn)練用。但要注意資源隔離別互相影響。我在實際運維中發(fā)現(xiàn)利用率低往往不是技術(shù)問題而是流程問題。比如團隊申請了卡但一直不用或者任務(wù)排隊策略不合理導(dǎo)致大任務(wù)餓死。技術(shù)手段能解決一部分但制度和流程的優(yōu)化同樣重要。5.3 監(jiān)控體系看得見才能管得住AI集群的監(jiān)控比普通服務(wù)復(fù)雜因為要監(jiān)控的層次多監(jiān)控層次關(guān)鍵指標常用工具硬件層GPU溫度、功耗、顯存、利用率DCGM、nvidia-smi通信層帶寬、延遲、重傳率NCCL日志、IB監(jiān)控框架層訓(xùn)練速度、loss曲線、梯度范數(shù)TensorBoard、WandB服務(wù)層請求延遲、QPS、錯誤率Prometheus、Grafana業(yè)務(wù)層模型效果、轉(zhuǎn)化率自定義埋點監(jiān)控的目的不只是發(fā)現(xiàn)問題還要能定位問題。比如訓(xùn)練變慢了是數(shù)據(jù)加載慢、通信慢、還是計算慢沒有細粒度的監(jiān)控就只能靠猜。6. 我踩過的那些坑與總結(jié)的經(jīng)驗6.1 環(huán)境配置版本兼容性是萬惡之源AI Infra最煩人的事情之一就是環(huán)境配置。CUDA版本、驅(qū)動版本、框架版本、通信庫版本任何一個不匹配都可能出問題。我遇到過CUDA版本和PyTorch不匹配導(dǎo)致無法使用GPU也遇到過NCCL版本和驅(qū)動不兼容導(dǎo)致多卡訓(xùn)練hang住。我的經(jīng)驗是用容器別在裸機上折騰。Docker加上NVIDIA Container Toolkit可以把環(huán)境固化下來。但容器也不是銀彈GPU直通、共享內(nèi)存、網(wǎng)絡(luò)模式都有坑。最好是用官方維護的鏡像作為基礎(chǔ)別自己從零搭。6.2 性能調(diào)優(yōu)先測量再優(yōu)化性能問題最忌諱拍腦袋。我見過有人一上來就改batch size、改學(xué)習(xí)率結(jié)果越調(diào)越差。正確的做法是先用profiler測量找到真正的瓶頸。PyTorch有torch.profilerNVIDIA有Nsight Systems這些都是好工具。測量之后你會發(fā)現(xiàn)瓶頸往往不在你以為的地方。比如訓(xùn)練慢可能不是GPU算得慢而是數(shù)據(jù)加載的IO慢或者CPU預(yù)處理慢。6.3 故障排查從日志和指標入手分布式訓(xùn)練的故障排查是最考驗人的。任務(wù)掛了日志可能幾千行關(guān)鍵信息就一兩行。我的習(xí)慣是先看有沒有明顯的報錯比如OOM、NCCL timeout?？锤鞴?jié)點的GPU利用率如果有的卡忙有的卡閑可能是負載不均?？赐ㄐ湃罩綨CCL會打印ring的建立過程如果卡在建ring可能是網(wǎng)絡(luò)問題?？碈heckpoint恢復(fù)是否正常有時候是恢復(fù)邏輯有bug。還有一個經(jīng)驗保留現(xiàn)場。任務(wù)掛了別急著重啟先把日志、nvidia-smi輸出、網(wǎng)絡(luò)狀態(tài)都存下來。重啟之后現(xiàn)場就沒了再想排查就難了。6.4 學(xué)習(xí)路徑建議如果你剛開始接觸AI Infra我建議按這個順序來先跑通單卡訓(xùn)練理解模型訓(xùn)練的基本流程。再跑通單機多卡理解數(shù)據(jù)并行和通信。然后嘗試多機多卡理解分布式訓(xùn)練的復(fù)雜性。接著學(xué)推理部署把訓(xùn)練好的模型服務(wù)化。最后學(xué)調(diào)度和運維理解集群級別的管理。每一步都要動手做光看文檔是學(xué)不會的。遇到問題先自己排查實在搞不定再查資料或問人。排查的過程本身就是最好的學(xué)習(xí)。這個領(lǐng)域變化很快新硬件、新框架、新優(yōu)化手段層出不窮。但底層的東西變化沒那么快計算、存儲、通信、調(diào)度這些核心問題十年二十年都不會過時。把基礎(chǔ)打牢上層的東西學(xué)起來就快。最后分享一個我自己的習(xí)慣我會維護一個“坑本”每次遇到問題解決了就記下來包括現(xiàn)象、排查過程、根因、解決方案。下次再遇到類似問題翻一下坑本就能快速定位。這個習(xí)慣幫我省了大量時間也讓我對系統(tǒng)的理解越來越深。