大模型千卡訓練總結！難點到底在哪里？
本文經自動駕駛之心公眾號授權轉載，轉載請聯系出處。
最近看到知乎一個回答，把千卡訓練的難度吹上天了。但其實真正用過千卡就會發現也就那么幾個點。于是想寫一篇文章簡單講講。
本文將包括3個部分：首先我們將討論千卡訓練的難題，以及應該在什么時候使用千卡訓練；接著，我們將討論如何在一千張卡上開始訓練，如何讓他達到近乎線性的性能提升；最后我們將展開討論一些千卡訓練當中仍然懸而未決（至少對于開源社區來說）的問題。
為什么千卡訓練是困難的？
千卡訓練和八卡訓練的區別是—顯卡多了一百多倍。
這意味著什么呢？
- 通信時間增加
- 故障概率增加
這倆問題都很好理解。
時間上，PyTorch內部支持NCCL/Gloo/MPI三個通信后端（請務必使用NCCL。其中AllReduce操作會會根據具體硬件配置走Ring AllReduce和Tree AllReduce。Ring的時間復雜度是，Tree的時間復雜度是 。就算是理論上128節點也比單節點慢至少七倍，實踐當中跨節點通訊要遠比單節點慢得多。
故障上，一個節點出問題的概率是p，128個節點就是1-(1-p)^128。也就是說如果一個操作在一個訓練當中的出錯概率是1%，那么在128節點當中的出錯概率就是72.37%。
此外，隨著規模的增大，許多問題都會變得難以忍受。比如數據增強要花0.1s，一億條數據就是278個小時（當然這只是胡拆的一個數字，實際有各種機制所以不會有這么大影響。
因此，錢多燒手并不是使用千卡訓練的理由。閑得蛋疼可能是，但你得多蛋疼才能想出這么折磨自己的idea？
千卡訓練解決的問題是大模型&大數據問題。如果你的訓練時間沒有超過8192GPU日，那么你絕對不需要一千張顯卡。
看到這里，絕大多數人已經可以關掉這篇文章了。除非你的模型和數據都以B（十億）來作為計量單位。當然如果你正在廁所里手機沒電想看點兒東西解悶兒的話（雖然我很懷疑是否會有人把他打出來……那么可以繼續往下看
如何使用一千張卡訓練？
如何提高計算效率？
這件事情其實是一個case by case的事情。因為通信、計算速度啥的受硬件影響更多。而每一個集群的硬件拓撲都是不一樣的。同樣是A100集群，我全DGX節點，每一張A100都是SXM接口并配一塊兒專屬的IB網卡。你一個小破普惠服務器插8張PCI-E A100，IB卡一個節點只給一張。那咱倆遇到的問題就完全不是一個問題。
因此，要討論如何提高訓練效率、減少訓練耗時，我們首先要了解訓練耗時在哪里。那么，一個訓練步的耗時在哪里呢？需要謹記，沒有profile的優化是沒有意義的。
你可能會說，forward backward sync。很好，這說明你了解PyTorch的基本流程。不過現實當中要復雜得多。
- dataset讀取數據，構建輸出
- dataloader collate數據，進行數據預處理
- 模型forward計算輸出
- loss compute
- 模型backward計算梯度
- 模型sync梯度
- 優化器step更新權重
- 打印log
當然這是可以無限細分下去的，但一般這些就夠了。需要注意的是，除了4-7的耗時是真耗時，其他都需要通過異步操作來蓋掉。這也是我們的優化目標。
異步執行在PyTorch的dataloader、CUDA和分布式當中都存在。前者可以通過設置num_workers和prefetch_count為0來關閉，后兩者可以通過cuda.synchornize和dist.barrier來執行手動同步。在profile時，我們需要首先需要測整個step的時長。然后再在每次測量前執行手動同步來計算每個部分的時長。如果前者的總耗時等于后者4-7的耗時之和，那么通常不需要執行任何操作。但這種情況在千卡操作中幾乎不可能發生。
第6步通信往往需要耗費大量時間。因此，我們還需要進一步優化通信。
以下內容是對《PyTorch Distributed: Experiences on Accelerating Data Parallel Training》論文的概括，有感興趣的同學建議通讀并背誦全文。
Paper: https://arxiv.org/abs/2006.15704計算-通信重疊
在PyTorch當中，梯度的通信和反向傳播是交疊進行的。也就是說，每完成一層的梯度計算，都會立即觸發當前層的同步。實現起來也很簡單，每個進程在完成自己第k層的梯度計算后都會觸發一個鉤子來給計數器+1s。當計數器達到進程數是開火進行梯度通信。有很多同學在計算梯度過程中遇到過RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one.錯誤，這就是因為有的模塊沒有參與計算loss，導致梯度同步卡住了。需要注意，當find_unused_parameters=True時，PyTorch分布式使用nn.Module.__init__當中定義sub-module的反向順序來作為梯度桶的構建順序。因此，確保模塊定義和調用的順序一致對于高效訓練來說很重要。
梯度合桶
盡管理論上來說，同步發生的越及時，重合度越高，性能越好。但實際上每次發起通信都是有開銷的。因此，現實當中梯度同步并不是越多越好越快越好。為此，PyTorch引入了梯度合桶機制，通過把多個Tensor裝在一個桶里再通信桶來減少通信次數從而減少總耗時。合桶的Buffer Size等等參數往往需要針對硬件和模型來調整從而取得最好的通信效果。PyTorch的默認參數是從0.x時代祖傳下來的，這一參數通常都需要調節。
梯度累加
當你做完所有操作之后，驚喜的發現TMD怎么同步時間還是單節點的好幾倍。這其實是正常情況……實際上超過256卡的訓練想要把通信蓋掉就是一件不可能的事情。你說老師我看FB論文說他們256卡就是線性提升啊…那這里不得不提的一個策略就是梯度累加了。梯度累加會執行k次forward+backward之后再執行優化器步進。這有很多好處，首先對于大模型batch size通常不能開多大，梯度累加可以提升等效batch size。其次累加期間的backward不需要通信梯度，加快了訓練速度。
少即是快
Python是一種很慢的代碼。當然你說JIT trace+torch.compile有提升我也不反對，但對于最高效率來說，只有必須要存在的代碼和不存在的代碼兩種。
抱抱臉的Transformers就是一個反例。兩個sub-Module就能寫完的TransformerLayer他們硬是能寫出來一堆…偏偏他們還信奉Single Model File Policy……我尋思你這完全不考慮繼承的封這么多層是要搞雞毛啊？正例反而是PyTorch……（笑死，我竟然會夸臉書代碼寫得好。具體來說就是nn.functional當中的各種實現。你會發現他們第一行往往是handle_torch_func。熟悉Python裝飾器的小伙汁通常要問了，為啥這里不用個裝飾器統一一下？因為裝飾器會引入額外的函數調用，額外的函數調用就是額外的開銷。
因此，如果你想確保最高的效率，寫一個簡單的訓練代碼和模型代碼非常重要。畢竟，1%的效率提升，節省的可能是數百個GPU日。
如何平穩訓練
這一段當中中咱們只討論你能控制的問題。
捕捉不致命的異常
故障率高的問題其實很好解決。在訓練當中，大部分異常都是非致命異常，捉住他們就好了。https://danling.org/utils/decorators/#danling.utils.decorators.catch 是我之前寫的一個裝飾器，它的作用就是catch異常，然后調回調函數（默認當然就是把錯誤打印到log里）。所有你需要做的只是使用它來裝飾非fatal的操作。
在實際應用當中，我們遇到的最常見的問題是存ckpt寫滿了磁盤（不準笑，從商湯到深勢再到上海AI Lab，這個問題在哪兒都有出現。咱也不知道為啥肯買那么多顯卡但不肯多插點兒硬盤，咱也不敢問）。catch住所有保存操作，如果你有閑心可以在回調里刪一下之前的ckpt。沒嫌心的話…大不了重訓一次嘛（逃。第二常見的問題，你猜對了……存log寫滿了硬盤……所以所有logging操作也都是要catch的。這就是為啥我都用tmux然后開很長的緩存窗口，總是能搶救一些log出來的。
咳咳，說點兒正經的。任何聯網操作都是需要catch的，常見的聯網操作主要包括從ceph讀取數據和…寫log到遠程（逃。其他就沒啥了吧，我見過有大哥嘗試恢復OOM的，但效果似乎不是很好，至少我自己沒用過。簡單來說，唯一不應捕捉的錯誤是集群炸了。
那有的大兄弟就說了，集群沒爆炸，但是有兩張卡突然掉了咋辦。這個咱第三部分再討論。
過程也很重要
有用過[丹靈]http://danling.org的同學可能比較熟悉。丹靈其他地方都很輕量，唯獨實驗管理這里寫的很復雜。現代丹靈會將創建一個三個級別的實驗目錄，project/experiment-run/timestamp。其中project是用戶給出的，experiment和run分別是通過代碼版本和配置計算出來的，timestamp就是運行開始的時間。也就是說，如果代碼和配置是完全一樣的，丹靈就會認為這是同一個運行。在設置中打開auto_resum就會自動找最新的一個檢查點（這就是為啥最后一級要用時間戳）來加載。其實微軟用的amlt更好用，他甚至還會創建一個代碼的diff文件夾來幫助你回憶當初代碼修改了些啥。
收斂，收斂，收斂
模型訓著訓著發散了幾乎是每個訓大模型的人都會遇到的問題。輸出和loss只要有nan果斷丟掉。梯度先clip by value再clip by norm都是常規操作。哦對了，還有初始化……關于大模型收斂性的論文有一堆，此處不再贅述。
比更大，還更大，再更大
彈性訓練
實際上當你的訓練超過2048個GPU日時，在整個訓練過程當中發生單個GPU甚至單個節點下線是再正常不過的事情了。
PyTorch在1.10就引入了torchelastic彈性訓練機制，用過的都罵娘。等下，讓我先罵一遍，呸。ok咱們繼續吧。
我印象當中在微軟的最后一輪面試當中被問到了這個問題：如何設計一個彈性分布式系統。
我的回答很教科書。每k分鐘，系統會做一次AllReduce來統計存活進程數，然后選舉出一個主進程。主進程會計算好每個進程的rank和local rank進行broadcast。所有進程每次forward開始時向主進程發送一個心跳包來匯報狀態。主進程會根據心跳包來確定這一個step參與同步的機器有多少。但很可惜，2024年了。還是沒人去寫。
大小梯度同步
我一直認為梯度同步不應該以GPU/進程為單位。而應該分為大同步（節點間同步）和小同步（節點內同步）。小同步可以更高頻的進行，大同步則可以更慢的執行。這樣不僅能提高實際的梯度同步頻率，降低同步總耗時，并且還能天然的去結合小batch和大batch訓練的優點—節點內小batch關注個體，節點間大batch關注整體。
