大規模訓練

張量平行(tensor parallelism)

有些權重矩陣大到單一裝置塞不下,或大到套用太慢。張量平行把單一運算子切到多個裝置上:它不複製矩陣乘法,而是切分權重矩陣本身,讓每個裝置算出結果的一部分。經典的 transformer 配方把第一個前饋矩陣按欄切分,使每個裝置產生隱藏激活的一片,再把第二個矩陣按列切分,於是只需一次 all-reduce 就能把各部分輸出加總回完整結果。

注意力則按頭切分——每個裝置擁有查詢、鍵、值投影的一個子集,獨立計算那些頭,並在輸出投影後做一次 all-reduce。如此安排是為了讓前向每個區塊只需兩次集合通訊(反向也兩次),把切面維持得很小。由於每個微步都要等一次完整激活張量的 all-reduce,張量平行需要極高頻寬與極低延遲,這也是為何它通常被限制在單節點的 NVLink 範圍內。

當瓶頸出在某一層——隱藏維度巨大,或每裝置激活必須縮小——張量平行就是你要拿出來用的槓桿。超過八路後它很少單獨使用:一旦越過 NVLink 邊界,all-reduce 成本就會主導,因此較大規模的訓練會把它巢狀嵌進管線與資料平行裡。

Y = \text{Dropout}\big(\,[\,XA_1,\,XA_2\,]\,B\,\big),\quad B=[B_1;B_2]

Megatron 的 MLP 切法:欄平行的 A(不需通訊),接著列平行的 B,其部分積由一次 all-reduce 加總。

張量平行切分的是每個運算子的數學,因此結果與未切分的層在位元上等價——和管線平行不同,它不改變相依結構,只改變乘法發生的地點。

又稱
TPintra-layer parallelism張量平行層內平行