分散式訓練與高效推論:先找成本在哪裡|CMU 10-423 主題單元
模型放不下一張GPU,和模型已放得下卻生成太慢,是不同故障。分散式策略重新分配資料、參數與狀態;FlashAttention主要改善attention的記憶體IO;KV cache避免decode時重算歷史K/V。先找成本的位置,再選方法。
模型放不下一張GPU,和模型已放得下卻生成太慢,是不同故障。分散式策略重新分配資料、參數與狀態;FlashAttention主要改善attention的記憶體IO;KV cache避免decode時重算歷史K/V。先找成本的位置,再選方法。
11-868 第 14、15 講從 parameter server 講到 PyTorch DDP:先用 NCCL 的五個 collective(Broadcast、Reduce、AllReduce、ReduceScatter、AllGather)當積木,推出 ring 為什麼能讓廣播時間幾乎不隨 GPU 數增加,再把 AllReduce 拆成 ReduceScatter 加 AllGather。第二講拆 DDP 的兩個關鍵設計:把梯度分桶(預設 25 MB)、在反向傳播還沒結束時就開始同步。官方課表未列公開錄影連結,本文依投影片頁碼與 VLDB 2020 論文整理。
CMU 11-868 在 2026 春季用兩講回答「模型太大,一張 GPU 放不下」:L16 講切層的 pipeline parallelism(GPipe 的 micro-batch、1F1B、交錯 stage)和切矩陣的 tensor parallelism(Megatron-LM 對 FFN、attention、embedding 的切法),結論是節點內用 TP、跨節點用 PP、再外面疊 DP;L17 把 MoE 當成第三種切法——每張卡放不同的專家、其餘部分複製,代價換成 all-to-all 通訊與負載平衡,並以 GShard、DeepSpeed-MoE、DeepSeek-V3 的設計逐一說明。
L12 的主軸是「搬資料」。前段用 SambaNova SN40L 說明資料流架構與 metapipelining:同樣跑 Llama 3.1 8B,投影片說 RDU 每個 token 約 3 次 kernel 呼叫,GPU 約 800 次,差別來自能不能把整個 decoder 融合進一個 kernel。中段把尺度拉到資料中心:TP、PP、EP、DP 各自需要哪種集體通訊,以及計算與通訊重疊為什麼決定擴展效率。後段回到能源與 DRAM:搬一個 byte 比算一次貴得多,記憶體控制器、burst mode、HBM 都是在解同一個問題。這講沒有公開錄影,本文只依投影片。
每個核心都有自己的 cache,同一個位址就可能同時有好幾份副本,各核看到的值會不一樣。這不是加鎖能解決的問題,是硬體複製資料造成的。CS149 L14 先定義什麼叫 coherent,再拆解 snooping 的 MSI 協定:要寫就先廣播 BusRdX 讓別人作廢,MESI 多一個 E 狀態省掉「讀完再寫」的第二筆交易,directory 則把廣播改成點對點。對程式設計師最實際的後果是 false sharing:兩個 thread 寫不同變數,只因為落在同一條 cache line,就讓 cache line 在核心間來回彈,投影片的 demo 慢了三倍。
第 8 講要你換一個腦袋:不再想「每個 worker 做什麼」,而是把演算法寫成對序列的操作,例如 map、fold、scan、segmented scan、gather/scatter、sort、groupBy。這些原語都有高效的平行實作,能把不規則的平行變規則、把細粒度同步變粗粒度。代價是要多掃幾遍資料,所以很吃頻寬。
L9 開場就說:懂了 arithmetic intensity 與 roofline,你就幾乎懂了現代 AI 軟體面效能最佳化的一切。接著示範三件事:全連接層、卷積層、attention 最後都變成矩陣乘(GEMM);GEMM 要靠分塊(blocking)讓資料留在快取裡;層與層之間要融合(fusion),別把中間結果寫回 DRAM 再讀回來。softmax 可以分塊計算,這就是融合版 attention(FlashAttention 的核心想法)不必存下 N×N 矩陣的原因。
L13 問的是:寫快程式的專家太少,怎麼辦?投影片給三個答案。一是提高抽象層級:Halide 把「算什麼」(演算法)和「怎麼算」(排程)拆成兩種語言,同一段模糊濾波可以只改一行排程就換成分塊、向量化、多核版本。二是智慧搜尋:因為排程空間定義清楚,可以用搜尋加學出來的成本模型自動產生排程。三是新興的 LLM agent:讓模型寫 CUDA、執行、看 profiler、反思、再改,並用範例資料庫或 prompt 最佳化讓它自我改進。投影片最後把問題留給你:真正的價值在 DSL 設計,還是 LLM agent?
CS149 L16 分三段:先用 cache coherence 的眼光看鎖怎麼實作(test-and-set、test-and-test-and-set、ticket lock、CAS、LL/SC),再用一條排序 linked list 示範從一把大鎖改成 hand-over-hand 細粒度鎖,最後介紹 lock-free:單一生產者單一消費者佇列、用 CAS 寫的 stack、ABA 問題與 hazard pointer。投影片的結論很務實:在只有你的程式佔用機器的情境下,寫得好的鎖版本常常一樣快,而且好寫得多;lock-free 的價值在於執行緒可能被搶佔、page fault 的系統。
CUDA 的 grid、thread block、CUDA thread 是一套程式抽象;GPU 用 SM、warp 與硬體 block 排程器把它實作出來。這一講的核心是分清兩件事:thread block 之間系統可以任意排序,同一個 block 裡的 thread 則保證同時存在,所以 block 能用 shared memory 和 __syncthreads() 合作,也所以一個 SM 能塞幾個 block 由暫存器與 shared memory 的量決定。
L10 從一條式子出發:功耗固定時,效能只能靠「每個運算花多少焦耳」來換,而通用處理器把大部分能量花在取指令、解碼、搬資料,不是在算。投影片的經驗法則是 GPU 比 CPU 好約 10 倍 perf/watt,固定功能 ASIC 可達 100–1000 倍。接著用同一套標準(分塊 tensor、非同步計算與記憶體、運算單元直接互傳)檢視 H100 的 Tensor Core 與 TMA、Google TPU 的脈動陣列,以及可重組資料流架構。
L3 前半用高速公路與洗衣服的比喻分開延遲與頻寬,再算給你看:向量逐元素相乘在 V100 上效率不到 1%,因為記憶體送資料的速度追不上 ALU。後半的主題是「抽象 vs 實作」:ISPC 讓你用 SPMD 的方式思考(一群 program instance 各做一份),編譯器卻用 SIMD 指令實作。把這兩層混在一起,是這門課最常見的困惑來源。
CS149 L6 要你把「通訊」看得很廣:處理器和 cache、和記憶體、和另一台機器之間的資料搬移都算。現代平行處理器的算力遠高於頻寬,所以 arithmetic intensity(每搬一單位資料做多少計算)決定了你能不能把硬體餵飽。提高它的手段有三類:改分配方式減少必要通訊、用 blocking 與 loop fusion 減少 cache 造成的額外通訊、用分散與錯開存取降低 contention。
PA1 程式寫得不多,重點是分析:六個程式分別練 threads 的工作分配、SIMD 遮罩、ISPC gang 與 task、輸入資料如何影響 SIMD 效率、頻寬受限的 saxpy,以及用計時找 K-Means 的熱點。Written 1 則用紙筆題練峰值吞吐量、指令相依、管線、多執行緒藏延遲與 SIMD divergence。官方評分以 Stanford myth 機器為準,校外可以在自己的機器跑,但數字不能直接對照。
CS149 PA2 要你用 C++ 寫一個多核 CPU 上的任務執行庫,而且要寫四次:每次 run() 都開 thread、改成 spinning 的 thread pool、改成會睡覺的 thread pool,最後擴充成非同步、有依賴的 task graph。每一步都得是完全正確的系統,並跟官方參考實作比速度。官方評分機器是 AWS c7g.4xlarge;校外可以在自己的多核機器上跑,但數字不能直接和官方比。
PA4 把你丟到 AWS Trainium2 的一顆 NeuronCore 上。這裡沒有 cache 幫你決定什麼資料留在晶片上:SBUF(28 MiB)與 PSUM(2 MiB)都要你用 dma_copy 明確搬進搬出,partition 維度最多 128。Part 1 用 vector add 和 transpose 教你這些限制與 DMA 成本,Part 2 要你把 convolution 拆成一串 matmul、再和 max pool 融合成一個不落地到 HBM 的 kernel。Written 3 用 line buffer、兩次 box blur、softmax 硬體與 metapipelining 練同一件事:把中間結果留在晶片上。環境需要課程 private AMI 與付費 capacity block,校外實質只到 A2。
L4 給了一套平行化的思考流程:先分解(decomposition)找出彼此獨立的工作,再分配(assignment)給工作者,再協調(orchestration)通訊與同步,最後對應到硬體。Amdahl 定律提醒你循序部分決定 speedup 上限。貫穿全講的例子是 2D 網格解算器:原本的相依關係很難平行,改用紅黑著色換一種更新順序後,就能用資料平行或共享位址空間兩種模型寫出來。
L11 問的是:硬體為 AI 專用化之後,程式設計師要付出什麼?投影片以 H100 為例:要吃滿 Tensor Core,就得用 16×16 tile、讓 TMA 非同步搬資料、讓 producer 與 consumer warp 管線化,寫起來很複雜,所以有了 ThunderKittens 這類 DSL。另一條路是資料流架構(SambaNova SN40L):用 map/reduce/zip 等平行模式描述計算,編譯器做分塊、metapipelining 與佈局,投影片說它能把 Llama 3.1 8B 的整個 decoder 融成一個 kernel。
Coherence 只管單一位址;memory consistency 管的是不同位址之間的讀寫,在別的 thread 眼中以什麼順序生效。CS149 L15 用兩個 thread、兩個變數的例子說明:sequential consistency 下 r1 = r2 = 0 不可能,但每顆現代處理器都有的 write buffer 會讓讀跑到寫前面,於是它變成可能。TSO、PSO、weak ordering 依序放寬更多順序換取效能,fence 與同步原語再把需要的順序補回來。對應用程式設計師的結論很短:寫沒有 data race 的程式,用同步函式庫,C11、C++11、Java 5 就保證你看到 sequential consistency。
粗鎖好寫但慢,細鎖快但容易寫錯。Transactional memory 讓程式設計師只宣告 atomic { },由系統負責原子性與隔離。CS149 L17 講動機(failure atomicity、composability)與設計空間:資料版本管理分 eager(undo log)與 lazy(write buffer),衝突偵測分 pessimistic 與 optimistic。L18 拆 STM 的執行期資料結構與 McRT 演算法,再講 HTM 怎麼用 cache 的 R/W 位元加上 coherence 協定偵測衝突,最後是 Intel Haswell 的 RTM。Written 4 則把 MSI、LL/SC、鎖與記憶體順序、雙向 linked list 的細粒度鎖串成四題。
負載平衡的難處在於它和排程成本互相拉扯:任務切得越細越好平衡,但每次領任務都要付同步成本。CS149 L5 先把選項排成一條從 static 到 dynamic 的連續光譜,再拆解 Cilk 的執行期:每個 worker 一條 deque,spawn 時先跑 child、把 continuation 留給別人偷,閒置的 thread 從別人 deque 的頂端偷走最大塊的工作。
CS231N 第 11 講以 Llama3-405B 為貫穿範例,先講 GPU 硬體與叢集(H100、8 卡伺服器、24,576 張 GPU 的叢集),再把 Transformer activation 的四個維度對應到四種平行化:切 batch 是資料平行(再演進成 FSDP、HSDP),切序列是 context parallelism,切層是 pipeline parallelism,切通道是 tensor parallelism。中間插入 activation checkpointing(用重算換記憶體)和一份實用的擴展配方,並用 Model FLOPs Utilization(MFU)當調參目標:超過 30% 算好,超過 40% 算優秀。
GPT-3 的 fp16 權重就要 350GB,一張 80GB 的 A100 放不下,更別說梯度和 Adam 狀態。L19 講怎麼切:資料平行、ring all-reduce、ZeRO-1/2/3(每張卡能訓練的模型從 5B 推到 320B)、GPipe 把 pipeline 利用率從 25% 提到 57%、Megatron 式 tensor parallelism、Ulysses 與 Ring Attention 的 sequence parallelism。L20 講切完之後的通訊瓶頸:Alpa 自動搜平行策略、DGC 把梯度壓 277–608 倍還不掉精度、TernGrad 把梯度量化成三值,以及用延遲更新蓋掉網路延遲的 DGA。
演算法把模型縮小之後,系統層還能再榨多少?L11 用同一個矩陣乘法示範:loop reordering 快 12 倍、tiling 快 19 倍(Intel Xeon 4114),CUDA 版在 2080Ti 上端到端快 94 倍。後半講 TinyEngine 用的推論技巧:im2col、in-place depthwise 把峰值記憶體從 2×C×H×W 降到 (1+C)×H×W、pointwise 用 NHWC、depthwise 用 NCHW,以及少 2.25 倍乘法的 Winograd。
第七講不從 FSDP API 開始,而是用 broadcast、all-reduce、all-gather、reduce-scatter 與 all-to-all 建立通訊語言,再親手組出 data、tensor 與 pipeline parallelism。
第八講把平行化從原語提升到系統設計:ZeRO 逐階切 optimizer、gradient 與 parameter,TP/PP/SP/EP 再分別切 width、depth、sequence 與 experts;組合順序必須服從網路拓撲與動態 activation memory。