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

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

作者 | Soumith Chintala

編譯 |?岳揚(yáng)

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

1. 首先，是如何將盡可能大的神經(jīng)網(wǎng)絡(luò)和 batch-size 適配到那 10000 張 H100s 上，這個(gè)步驟涉及到并行處理和使用節(jié)省內(nèi)存的各種技巧。
2. 其次，是如何在這些 GPU 之間盡可能高效地傳遞模型狀態(tài)信息（state）。
3. 最后，是如何在遇到硬件或軟件故障時(shí)，盡可能迅速地恢復(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ò)層分布到多個(gè) GPU 上進(jìn)行計(jì)算）
3. 對(duì)神經(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í)行前向傳播時(shí)，需要保存一些中間結(jié)果以便后續(xù)計(jì)算反向傳播（save_for_backward）。然而，當(dāng)神經(jīng)網(wǎng)絡(luò)規(guī)模非常大時(shí)，為了處理更大的數(shù)據(jù)批次，更有效的方法是釋放這些中間結(jié)果，待到需要計(jì)算反向傳播時(shí)再重新計(jì)算。
• 類(lèi)似 FSDP 這樣的技術(shù)，通過(guò)在單個(gè) GPU 上只保留模型的分片來(lái)節(jié)省內(nèi)存。當(dāng)需要其他權(quán)重時(shí)，會(huì)從其他 GPU 聚合模型的完整權(quán)重。

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

2.1 Communication overlap 策略：

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

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

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

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

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

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

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

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

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

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

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

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

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

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

原文鏈接：

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