??【本文正在參與 AI.x社區(qū)AIGC創(chuàng)新先鋒者征文大賽】??
??http://www.scjtxx.cn/aigc/2223.html??

編者按： 怎樣在 10,000 個 H100 GPU 上訓(xùn)練大模型？如何充分利用每一塊 GPU 的算力？如何在這個復(fù)雜的 GPU 網(wǎng)絡(luò)中高效傳遞數(shù)據(jù)？當(dāng)不可避免的硬件故障發(fā)生時，又該如何快速恢復(fù)訓(xùn)練進(jìn)度？我們今天為大家?guī)淼奈恼轮?，作者為我們揭示了?yīng)對這些挑戰(zhàn)的關(guān)鍵策略。

作者 | Soumith Chintala

編譯 |?岳揚(yáng)

我的好友 Francois Fleuret 提出了上述問題。我迅速總結(jié)了一些在大規(guī)模訓(xùn)練領(lǐng)域中相當(dāng)普遍的知識，內(nèi)容分為三部分。

1. 首先，是如何將盡可能大的神經(jīng)網(wǎng)絡(luò)和 batch-size 適配到那 10000 張 H100s 上，這個步驟涉及到并行處理和使用節(jié)省內(nèi)存的各種技巧。
2. 其次，是如何在這些 GPU 之間盡可能高效地傳遞模型狀態(tài)信息（state）。
3. 最后，是如何在遇到硬件或軟件故障時，盡可能迅速地恢復(fù)系統(tǒng)。

01 如何將盡可能大的神經(jīng)網(wǎng)絡(luò)和 batch-size 適配到那 10000 張 H100s 上

1.1 并行策略

1. 在數(shù)據(jù)批次（batches）上進(jìn)行并行處理（數(shù)據(jù)并行（data parallel））
2. 在神經(jīng)網(wǎng)絡(luò)層上進(jìn)行并行處理（比如，將一層神經(jīng)網(wǎng)絡(luò)層分布到多個 GPU 上進(jìn)行計算）
3. 對神經(jīng)網(wǎng)絡(luò)的不同模型層進(jìn)行分割，以便它們能夠在不同的 GPU 上運(yùn)行（比如，前 N 層運(yùn)行在 GPU1 上，第 N+1 層到第 N+10 層運(yùn)行在 GPU2 上）

持續(xù)優(yōu)化并行策略，直到所有 GPU 都能被高效利用，達(dá)到最高利用率。

1.2 Checkpointing / Compute vs memorize

• 在執(zhí)行前向傳播時，需要保存一些中間結(jié)果以便后續(xù)計算反向傳播（save_for_backward）。然而，當(dāng)神經(jīng)網(wǎng)絡(luò)規(guī)模非常大時，為了處理更大的數(shù)據(jù)批次，更有效的方法是釋放這些中間結(jié)果，待到需要計算反向傳播時再重新計算。
• 類似 FSDP 這樣的技術(shù)，通過在單個 GPU 上只保留模型的分片來節(jié)省內(nèi)存。當(dāng)需要其他權(quán)重時，會從其他 GPU 聚合模型的完整權(quán)重。

02 盡可能高效地在 GPU 集群間傳遞模型狀態(tài)信息

2.1 Communication overlap 策略：

在需要 GPU 間通信時，應(yīng)盡可能早地啟動通信過程：

• 例如，當(dāng)?shù)?N 層完成反向傳播后，在第 N-1 層還在進(jìn)行反向傳播計算時，負(fù)責(zé)第 N 層的所有 GPU 可以同時開始執(zhí)行梯度全歸約操作。

2.2 探索并利用網(wǎng)絡(luò)底層拓?fù)浣Y(jié)構(gòu)：

在多個計算節(jié)點間傳遞大量模型狀態(tài)信息（如梯度、優(yōu)化器狀態(tài)信息）是一項復(fù)雜的任務(wù)。在使用 Sync SGD 時，需要盡可能快地集中傳輸這些狀態(tài)信息。

網(wǎng)絡(luò)中可能包含多層交換機(jī)，并具備 RDMA 能力（可以直接將 GPU 內(nèi)存中的數(shù)據(jù)復(fù)制到網(wǎng)卡，完全繞過 CPU 內(nèi)存），同時擁有前端和后端網(wǎng)卡（前端網(wǎng)卡連接到如 NFS 之類的存儲系統(tǒng)，后端網(wǎng)卡則將 GPU 連接到集群中的其他 GPU）。

因此，在執(zhí)行 all-reduce 或 scatter/gather 等通信操作時，充分利用這些網(wǎng)絡(luò)信息至關(guān)重要。例如，通過樹形歸約算法（tree-reduce），all-reduce 操作的時間復(fù)雜度可以降低到O(log(n))；同時，網(wǎng)絡(luò)光纖連接節(jié)點間的不同類型光纖對常數(shù)因子的影響，對于減少整體延遲時間也是非常重要的。

像 NCCL 這樣的庫能夠智能地識別底層網(wǎng)絡(luò)拓?fù)洌⒃趫?zhí)行 all-reduce 和其他通信操作時加以利用。

在這樣的大規(guī)模計算中，我們還必須調(diào)整交換機(jī)和網(wǎng)卡中的數(shù)據(jù)包路由算法，以實現(xiàn)有效的負(fù)載均衡。交換機(jī)也需要大量的 HBM 內(nèi)存（不僅僅是 GPU 需要），因為當(dāng)數(shù)據(jù)包排隊等待時，需要在某個地方排隊而不會被丟棄——這就是交換機(jī)級別的 HBM 內(nèi)存。

03 如何在遇到硬件或軟件故障時，盡可能迅速地恢復(fù)系統(tǒng)？

故障是不可避免的，涉及GPU、網(wǎng)卡、電纜等多種硬件。有些故障能夠迅速被發(fā)現(xiàn)，而有些則可能因為某個節(jié)點沒有按時響應(yīng)（比如 NCCL 的 all-reduce 操作卡住了）才被察覺。我們開發(fā)了多種工具來監(jiān)控機(jī)群的健康狀況，并盡可能快地將故障節(jié)點從機(jī)群中移除。這可不是一件容易的事。

在這種規(guī)模下，內(nèi)存位隨機(jī)翻轉(zhuǎn)導(dǎo)致的隱性數(shù)據(jù)損壞概率增加，可能導(dǎo)致訓(xùn)練 loss 值異常升高。雖然這種問題在小規(guī)模系統(tǒng)中很少見，但在大規(guī)模系統(tǒng)中則可能頻繁發(fā)生。在軟件層面提前檢測這種問題非常困難。一些硬件設(shè)備配備了內(nèi)置校驗和的電路，可以在計算后進(jìn)行校驗 —— 這樣，一旦發(fā)生位翻轉(zhuǎn)，硬件就能觸發(fā)中斷。但 H100 和之前的 NVIDIA GPU 都不具備這一功能。

為了應(yīng)對這些故障，我們需要盡可能頻繁且迅速地保存模型狀態(tài)信息；一旦發(fā)生故障，我們也要能夠迅速恢復(fù)并繼續(xù)訓(xùn)練。通常，我們會迅速將模型狀態(tài)信息另存到 CPU 內(nèi)存的一個獨立線程中，并在后臺將數(shù)據(jù)從 CPU 內(nèi)存寫入到磁盤或遠(yuǎn)程存儲系統(tǒng)。我們還以分片的形式保存模型狀態(tài)信息（利用了 torch.distributed 的 checkpointing 功能），也就是說，不是每個 GPU 都需要保存完整的模型權(quán)重；每個 GPU 只需保存一部分權(quán)重 —— 其余部分可以通過其他 GPU 的分片 checkpoints 來恢復(fù)。

Thanks for reading!

Hope you have enjoyed and learned new things from this blog!

About the authors

Soumith Chintala

Cofounded and lead?@PyTorch?at Meta. Also dabble in robotics at NYU. AI is delicious when it is accessible and open-source.

END

本期互動內(nèi)容 ??

?還記得你第一次配置分布式訓(xùn)練環(huán)境時的經(jīng)歷嗎？有什么想對新手說的建議？

原文鏈接：

https://soumith.ch/blog/2024-10-02-training-10k-scale.md.html