??【本文正在參與 AI.x社區AIGC創新先鋒者征文大賽】??
??http://www.jxzklqfsx.com/aigc/2223.html??

編者按： 怎樣在 10,000 個 H100 GPU 上訓練大模型？如何充分利用每一塊 GPU 的算力？如何在這個復雜的 GPU 網絡中高效傳遞數據？當不可避免的硬件故障發生時，又該如何快速恢復訓練進度？我們今天為大家帶來的文章中，作者為我們揭示了應對這些挑戰的關鍵策略。

作者 | Soumith Chintala

編譯 |?岳揚

我的好友 Francois Fleuret 提出了上述問題。我迅速總結了一些在大規模訓練領域中相當普遍的知識，內容分為三部分。

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

01 如何將盡可能大的神經網絡和 batch-size 適配到那 10000 張 H100s 上

1.1 并行策略

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

持續優化并行策略，直到所有 GPU 都能被高效利用，達到最高利用率。

1.2 Checkpointing / Compute vs memorize

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

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

2.1 Communication overlap 策略：

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

• 例如，當第 N 層完成反向傳播后，在第 N-1 層還在進行反向傳播計算時，負責第 N 層的所有 GPU 可以同時開始執行梯度全歸約操作。

2.2 探索并利用網絡底層拓撲結構：

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

網絡中可能包含多層交換機，并具備 RDMA 能力（可以直接將 GPU 內存中的數據復制到網卡，完全繞過 CPU 內存），同時擁有前端和后端網卡（前端網卡連接到如 NFS 之類的存儲系統，后端網卡則將 GPU 連接到集群中的其他 GPU）。

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

像 NCCL 這樣的庫能夠智能地識別底層網絡拓撲，并在執行 all-reduce 和其他通信操作時加以利用。

在這樣的大規模計算中，我們還必須調整交換機和網卡中的數據包路由算法，以實現有效的負載均衡。交換機也需要大量的 HBM 內存（不僅僅是 GPU 需要），因為當數據包排隊等待時，需要在某個地方排隊而不會被丟棄——這就是交換機級別的 HBM 內存。

03 如何在遇到硬件或軟件故障時，盡可能迅速地恢復系統？

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

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

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

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

本期互動內容 ??

?還記得你第一次配置分布式訓練環境時的經歷嗎？有什么想對新手說的建議？

原文鏈接：

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