GTC 2023|Colossal-AI系統(tǒng)的數(shù)據(jù)并行優(yōu)化方案

Colossal-AI的改變：

從Colossal-AI的GitHub star數(shù)量中看出，這是衡量其開源社區(qū)的指標(biāo)。與其它非常流行的開源系統(tǒng)（如Spark、 MongoDB等）相比，在相同時間內(nèi)Colossal-AI使用增長速度是最快的。

世界上37個國家的用戶正在使用Colossal-AI，還有各種其他人工智能生態(tài)系統(tǒng)使用Colsssal-AI。

Colsssal-AI的重要功能——N-Dim并行系統(tǒng)

這可以將你的復(fù)雜模型可用硬件上，以便最大化使用并行同時最小化通信。

Colsssal-AI提供了完整的并行方法，包括數(shù)據(jù)并行、張量并行、流水線并行和序列并行，同時還有團隊提出的2維、2.5維和3維張量并行以及序列并行。

其中優(yōu)化了數(shù)據(jù)并行，可以讓你以最小的損失在相當(dāng)大的規(guī)換下訓(xùn)練模型。

大規(guī)模優(yōu)化器如何優(yōu)化數(shù)據(jù)并行

與單個GPU訓(xùn)練相比數(shù)據(jù)并行訓(xùn)練將整個數(shù)據(jù)集分為幾個分區(qū)。

每個GPU擁有1個數(shù)據(jù)集分區(qū)，模型在每個GPU上復(fù)制，開在其本地數(shù)據(jù)集分區(qū)進行訓(xùn)練，在反向傳播期間計算的梯度，被平均并用于更新模型副本，可以確保在不同GPU之間同步模型。

然而，簡單的增加批次大小，通常會導(dǎo)致泛化性能下降并降低計算精度。為了解決這些問題，研究出一種新的方法，可擴展的大規(guī)模優(yōu)化器LARS和LAMB，適合大批量設(shè)置。

在ViT模型和ImageNet數(shù)據(jù)集上對可擴展的大規(guī)模優(yōu)化器LARS和LAMB進行理論分析。他們可最大程度使用GPU資源，將ViT-B/32的訓(xùn)練時間從一個GPU的73小時減少到200個GPU的0.68小時，并實現(xiàn)幾乎線性增加的加速度，并保證收斂。

編輯：黃飛