本發(fā)明涉及大模型訓練，尤其涉及一種大模型并行訓練通信優(yōu)化方法。

背景技術(shù)：

1、傳統(tǒng)的數(shù)據(jù)并行性在所有數(shù)據(jù)并行進程中靜態(tài)地維護整個訓練過程中所需的所有模型狀態(tài)，導致內(nèi)存冗余。zero數(shù)據(jù)并行通過對模型狀態(tài)進行分區(qū)而不是復制來消除數(shù)據(jù)并行進程中的內(nèi)存冗余，這樣每個計算設備上的內(nèi)存消耗與數(shù)據(jù)并行維度成反比，zero數(shù)據(jù)并行也因此在學術(shù)界以及工業(yè)界流行起來。zero有主要的三個優(yōu)化階段，分別將優(yōu)化器狀態(tài)、梯度和模型參數(shù)分割存儲在多個計算設備上，減少了傳統(tǒng)數(shù)據(jù)并行中每個計算設備上存儲這些狀態(tài)所需的內(nèi)存。但是，這樣帶來了更多的訓練期間的通信開銷。

技術(shù)實現(xiàn)思路

1、為了解決現(xiàn)有技術(shù)中存在的問題，本申請實施例提供了一種大模型并行訓練通信優(yōu)化的方法、裝置、計算設備、計算機存儲介質(zhì)及包含計算機程序的產(chǎn)品，能夠降低通信開銷，提高并行訓練的效率。

2、第一方面，本申請實施例提供一種大模型并行訓練通信優(yōu)化方法，應用于多個計算設備組成的系統(tǒng)，每個計算設備配置有g(shù)pu，gpu包含至少一個進程，進程用于大模型并行訓練中的一個任務分區(qū)，方法用于對zero數(shù)據(jù)的zero-1階段的優(yōu)化，方法包括：執(zhí)行模型的向前計算以及反向計算，確定梯度信息；基于梯度信息，更新本地優(yōu)化器狀態(tài)中的動量；通過1-bit?allreduce通信對更新后的動量進行處理，并更新壓縮過程中產(chǎn)生的誤差；利用聚合后的動量以及方差，更新每個計算設備上存儲的模型參數(shù)分區(qū)。

3、在一些可能的實現(xiàn)方式中，1-bit?allreduce通信包括：使用all-to-all通信原語將壓縮的動量數(shù)據(jù)塊發(fā)送給其他計算設備，并接收來自其他計算設備的壓縮動量數(shù)據(jù)塊；對所有接收到的動量數(shù)據(jù)塊進行規(guī)約操作；使用allgather通信原語收集所有計算設備的動量數(shù)據(jù)塊，完成動量的全局聚合。

4、在一些可能的實現(xiàn)方式中，更新壓縮過程中產(chǎn)生的誤差，包括：記錄每次壓縮產(chǎn)生的誤差，并在下一次迭代中用于補償壓縮后的動量。

5、在一些可能的實現(xiàn)方式中，方法應用于對zero數(shù)據(jù)的zero-2階段的優(yōu)化，確定梯度信息之后，方法還包括：使用1-bit?reduce-scatter進行通信，確定聚合后的動量；基于聚合后的動量以及固定方差，完成模型參數(shù)的更新。

6、在一些可能的實現(xiàn)方式中，方法還包括：通過allgather對模型參數(shù)進行通信。

7、在一些可能的實現(xiàn)方式中，方法應用于對zero數(shù)據(jù)的zero-3階段的優(yōu)化，方法包括：每個進程執(zhí)行模型的前向計算，使用allgather通信操作收集所有計算設備上的完整模型信息；完成前向計算后，每個計算設備將不屬于當前參數(shù)分區(qū)的部分進行丟棄操作；執(zhí)行模型的反向計算，確定梯度信息，使用allgather通信操作收集所有計算設備上的梯度信息；進行1-bit?reduce-scatter通信以及更新壓縮誤差；計算設備再次丟棄非當前分區(qū)的參數(shù)信息，只保留當前分區(qū)必要的數(shù)據(jù)；利用聚合后的梯度信息和方差，計算每個進程獨立更新其分配的模型參數(shù)分區(qū)。

8、第二方面，本申請實施例提供一種大模型并行訓練通信優(yōu)化裝置，部署于多個計算設備組成的系統(tǒng)，每個計算設備配置有g(shù)pu，gpu包含至少一個進程，進程用于大模型并行訓練中的一個任務分區(qū)，裝置用于對zero數(shù)據(jù)的zero-1階段的優(yōu)化，裝置包括：梯度信息獲取模塊，用于執(zhí)行模型的向前計算以及反向計算，確定梯度信息；動量更新模塊，用于基于梯度信息，更新本地優(yōu)化器狀態(tài)中的動量；處理模塊，用于通過1-bit?allreduce通信對更新后的動量進行處理，并更新壓縮過程中產(chǎn)生的誤差；處理模塊，還用于利用聚合后的動量以及方差，更新每個計算設備上存儲的模型參數(shù)分區(qū)。

9、第三方面，本申請實施例提供一種計算機可讀存儲介質(zhì)，包括計算機可讀指令，當計算機讀取并執(zhí)行所述計算機可讀指令時，使得計算機執(zhí)行如第一方面任一項所述的方法。

10、第四方面，本申請實施例提供一種計算設備，包括處理器和存儲器，其中，所述存儲器中存儲有計算機程序指令，所述計算機程序指令被所述處理器運行時，執(zhí)行如第一方面任一項所述的方法。

11、第五方面，本申請實施例提供一種包含計算機程序的產(chǎn)品，當計算機程序產(chǎn)品在處理器上運行時，使得所述處理器執(zhí)行如第一方面任一項所述的方法。

技術(shù)特征：

1.一種大模型并行訓練通信優(yōu)化方法，其特征在于，應用于多個計算設備組成的系統(tǒng)，每個所述計算設備配置有g(shù)pu，所述gpu包含至少一個進程，所述進程用于大模型并行訓練中的一個任務分區(qū)，所述方法用于對zero數(shù)據(jù)的zero-1階段的優(yōu)化，所述方法包括：

2.根據(jù)權(quán)利要求1所述的方法，其特征在于，所述1-bit?allreduce通信包括：

3.根據(jù)權(quán)利要求1所述的方法，其特征在于，所述更新壓縮過程中產(chǎn)生的誤差，包括：記錄每次壓縮產(chǎn)生的誤差，并在下一次迭代中用于補償壓縮后的動量。

4.根據(jù)權(quán)利要求1所述的方法，其特征在于，所述方法應用于對zero數(shù)據(jù)的zero-2階段的優(yōu)化，所述確定梯度信息之后，所述方法還包括：

5.根據(jù)權(quán)利要求4所述的方法，其特征在于，所述方法還包括：通過allgather對模型參數(shù)進行通信。

6.根據(jù)權(quán)利要求1所述的方法，其特征在于，所述方法應用于對zero數(shù)據(jù)的zero-3階段的優(yōu)化，所述方法包括：

7.一種大模型并行訓練通信優(yōu)化裝置，其特征在于，部署于多個計算設備組成的系統(tǒng)，每個所述計算設備配置有g(shù)pu，所述gpu包含至少一個進程，所述進程用于大模型并行訓練中的一個任務分區(qū)，所述裝置用于對zero數(shù)據(jù)的zero-1階段的優(yōu)化，所述裝置包括：

技術(shù)總結(jié)
一種大模型并行訓練通信優(yōu)化方法，應用于多個計算設備組成的系統(tǒng)，每個計算設備配置有GPU，GPU包含至少一個進程，進程用于大模型并行訓練中的一個任務分區(qū)，方法用于對ZeRO數(shù)據(jù)的ZeRO?1階段的優(yōu)化，該方法包括：執(zhí)行模型的向前計算以及反向計算，確定梯度信息；基于梯度信息，更新本地優(yōu)化器狀態(tài)中的動量；通過1?bit?Allreduce通信對更新后的動量進行處理，并更新壓縮過程中產(chǎn)生的誤差；利用聚合后的動量以及方差，更新每個計算設備上存儲的模型參數(shù)分區(qū)。本方法可以降低通信開銷，提高并行訓練的效率。

技術(shù)研發(fā)人員：王彥棡,謝一浩,石君函,汪嘉恒,周純葆,王玨,劉芳
受保護的技術(shù)使用者：中國科學院計算機網(wǎng)絡信息中心
技術(shù)研發(fā)日：
技術(shù)公布日：2024/12/17