從BERT 到ChatGPT：大模型訓(xùn)練中的存儲系統(tǒng)挑戰(zhàn)與技術(shù)發(fā)展

馮楊洋 汪 慶 謝旻暉 舒繼武

（清華大學(xué)計算機(jī)科學(xué)與技術(shù)系 北京 100084）

（fyy21@mails.tsinghua.edu.cn）

OpenA? 推出的人工智能聊天機(jī)器人ChatGPT 在發(fā)布后的2 個月內(nèi)用戶數(shù)量就突破了1 億，成為歷史上用戶增長最快速的應(yīng)用軟件.其對文本卓越的理解和生成能力引起了工業(yè)界和學(xué)術(shù)界的廣泛關(guān)注.ChatGPT 使用了基于Transformer[1]結(jié)構(gòu)的大規(guī)模深度學(xué)習(xí)模型，即大模型.相較于傳統(tǒng)的深度學(xué)習(xí)模型結(jié)構(gòu)，Transformer 結(jié)構(gòu)可以將模型的參數(shù)量推向更高的數(shù)量級.如圖1 所示，2020 年由谷歌提出的BERT[2]模型的參數(shù)量為億級，而ChatGPT 所基于的GPT-3.5 的參數(shù)量已達(dá)到千億級，其下一代的GPT-4[3]的參數(shù)量甚至達(dá)到萬億級.

Fig.1 ?ncrease of model size圖1 模型規(guī)模增長

隨著大模型的參數(shù)量持續(xù)增長，其訓(xùn)練過程中的數(shù)據(jù)量不斷增加，對存儲的需求也在日益增加.但是，以往為互聯(lián)網(wǎng)、大數(shù)據(jù)等應(yīng)用設(shè)計的存儲技術(shù)不適用于大模型訓(xùn)練.在存儲性能方面，這些傳統(tǒng)的存儲技術(shù)并未針對大模型訓(xùn)練的計算模式、訪存模式和數(shù)據(jù)特征設(shè)計，無法在訓(xùn)練過程中充分發(fā)揮硬件的性能，影響了訓(xùn)練效率.在存儲容錯方面，大模型訓(xùn)練過程中的參數(shù)量龐大且更新頻繁，參數(shù)數(shù)據(jù)之間存在依賴關(guān)系，傳統(tǒng)的存儲容錯技術(shù)會為大模型訓(xùn)練帶來巨大的容錯開銷.

在存儲性能方面，現(xiàn)有工作提出了針對大模型訓(xùn)練的存儲加速技術(shù).這些技術(shù)可以總結(jié)為3 類：基于大模型計算模式的分布式顯存管理技術(shù)、大模型訓(xùn)練訪存感知的異構(gòu)存儲技術(shù)和大模型數(shù)據(jù)縮減技術(shù).基于大模型計算模式的分布式顯存管理技術(shù)依據(jù)訓(xùn)練中計算任務(wù)的劃分和計算任務(wù)之間的依賴關(guān)系，設(shè)計模型數(shù)據(jù)在GPU 集群中的分布和傳輸策略，減少數(shù)據(jù)在分布式集群中傳輸帶來的通信開銷.大模型訓(xùn)練訪存感知的異構(gòu)存儲技術(shù)利用模型訓(xùn)練過程中的訪存模式可預(yù)測的特性，設(shè)計了數(shù)據(jù)的預(yù)取和卸載策略，以充分利用計算掩蓋數(shù)據(jù)在各種存儲介質(zhì)間移動帶來的開銷.大模型數(shù)據(jù)縮減技術(shù)依據(jù)模型的數(shù)據(jù)特征，通過修改模型訓(xùn)練中的流程或降低數(shù)據(jù)精度，減少了模型訓(xùn)練中對存儲空間的需求.

在存儲容錯方面，現(xiàn)有工作針對大模型訓(xùn)練的存算模式提出了模型參數(shù)檢查點技術(shù)和冗余計算技術(shù).模型參數(shù)檢查點技術(shù)定期將數(shù)據(jù)保存至持久化介質(zhì)中以保證模型數(shù)據(jù)在設(shè)備故障后不丟失；冗余計算技術(shù)則在多張GPU 中重復(fù)計算相同版本的參數(shù)，以對模型訓(xùn)練數(shù)據(jù)容錯.

1 大模型訓(xùn)練的存儲挑戰(zhàn)

大模型訓(xùn)練的存儲容量需求高.大模型不僅參數(shù)量龐大，在訓(xùn)練的每一步還需要保存前向傳播過程中產(chǎn)生的激活量和用于參數(shù)更新的優(yōu)化器.

1）模型參數(shù)（model parameters）.近年來，大規(guī)模深度學(xué)習(xí)模型的參數(shù)量呈爆炸式增長，且這種趨勢仍在持續(xù)，如谷歌于2020 年提出了億級參數(shù)量的BERT，OpenA? 于2021 年和2023 年分別提出了千萬億級參數(shù)量的GPT-3 和萬億級參數(shù)量的GPT-4.海量的模型參數(shù)需要大量的存儲空間以支持大模型訓(xùn)練.以訓(xùn)練GPT-3 為例，其至少需要約800 GB 的顯存空間，約等于10 張H100 GPU 的顯存總和.

2）激活量（activation）.根據(jù)反向傳播算法，前向傳播階段產(chǎn)生的激活量需要先被保存，然后在反向傳播時用于計算梯度信息，最后在梯度計算結(jié)束后被釋放.在大模型訓(xùn)練過程中，激活量所占的存儲空間龐大.根據(jù)文獻(xiàn)[4]，激活量在訓(xùn)練過程中的存儲開銷占總存儲開銷的70%.這主要由3 個方面導(dǎo)致：①隨著模型層數(shù)的增多，所需要保存的激活量也隨之增加；②在大模型中單層特征的維度較大，這導(dǎo)致激活量的特征維度增大；③在模型訓(xùn)練中，每一步會使用更多的樣本以提高GPU 的計算利用率，這會導(dǎo)致激活量的樣本維度增加.

3）優(yōu)化器（optimizer）.計算得到的梯度被傳輸?shù)絻?yōu)化器中以更新得到新版本的模型參數(shù).為讓模型參數(shù)更新時能盡可能逼近最優(yōu)值，許多不同的優(yōu)化器被提出，如SGD[5]和Adam[6].這些優(yōu)化器需要保存額外的信息，如Adam 需要保存與梯度規(guī)模一致的一階動量和二階動量，其大小是模型參數(shù)的2 倍.

此外，大模型訓(xùn)練過程中的容錯需求高.大模型訓(xùn)練會使用大量GPU，這提高了故障的可能性[7].例如，Meta 團(tuán)隊在訓(xùn)練OPT-175B 的時候出現(xiàn)了百余次故障[8]；微軟的集群中平均45 min 就會出現(xiàn)1 次DNN 訓(xùn)練作業(yè)故障[9].

因此，大模型訓(xùn)練需要支持大規(guī)模數(shù)據(jù)的存儲系統(tǒng).傳統(tǒng)支持大規(guī)模數(shù)據(jù)存儲系統(tǒng)所使用的存儲技術(shù)可以分為3 類[10-12]：基于同構(gòu)存儲介質(zhì)的分布式存儲技術(shù)、基于異構(gòu)存儲介質(zhì)的異構(gòu)存儲技術(shù)和數(shù)據(jù)縮減技術(shù).傳統(tǒng)的大規(guī)模數(shù)據(jù)存儲系統(tǒng)還會采用日志、多副本以及糾刪碼等方式對數(shù)據(jù)進(jìn)行容錯.但是，這些傳統(tǒng)的大規(guī)模數(shù)據(jù)存儲技術(shù)是針對互聯(lián)網(wǎng)、大數(shù)據(jù)等應(yīng)用而設(shè)計的.這些技術(shù)一方面未充分利用大模型訓(xùn)練中的計算模式、訪存模式和數(shù)據(jù)特征，另一方面不適用于大模型數(shù)據(jù)更新數(shù)據(jù)量大、更新頻繁的特點，嚴(yán)重影響大模型訓(xùn)練的效率，其具體表現(xiàn)有4 個方面：

1）傳統(tǒng)的分布式存儲技術(shù)不適用于大模型訓(xùn)練的計算模式.大模型常使用多機(jī)多卡的配置進(jìn)行分布式訓(xùn)練，因此其數(shù)據(jù)也需要被存儲到多張GPU 顯存中.但傳統(tǒng)的分布式存儲技術(shù)并不適用于大模型訓(xùn)練的計算模式.一方面，大模型訓(xùn)練常使用的GPU具有計算資源和存儲資源強(qiáng)耦合的特點.系統(tǒng)在對其上的計算任務(wù)進(jìn)行劃分時，需要綜合考慮計算任務(wù)與存儲之間的依賴關(guān)系.但傳統(tǒng)的分布式存儲技術(shù)并未針對大模型訓(xùn)練的這一特征進(jìn)行設(shè)計，這可能導(dǎo)致數(shù)據(jù)不能及時傳輸至負(fù)責(zé)計算的GPU 上，影響訓(xùn)練效率.另一方面，傳統(tǒng)的分布式存儲技術(shù)未利用大模型訓(xùn)練中各個任務(wù)間的數(shù)據(jù)依賴關(guān)系進(jìn)行優(yōu)化，這可能導(dǎo)致相鄰任務(wù)間的數(shù)據(jù)傳輸方案非最優(yōu)，增加任務(wù)間通信的開銷.

2）傳統(tǒng)的異構(gòu)存儲技術(shù)對大模型訓(xùn)練中的訪存模式不感知.在傳統(tǒng)異構(gòu)存儲系統(tǒng)中，由于數(shù)據(jù)鏈路帶寬低、存儲介質(zhì)帶寬差異大，數(shù)據(jù)在不同存儲介質(zhì)間的移動開銷高.但是，大模型訓(xùn)練中的訪存模式可預(yù)測.系統(tǒng)可以利用數(shù)據(jù)的訪存信息對模型數(shù)據(jù)進(jìn)行預(yù)取，以減少數(shù)據(jù)在不同介質(zhì)間的移動開銷對計算產(chǎn)生的影響.傳統(tǒng)的異構(gòu)存儲系統(tǒng)并未利用這些訪存模式設(shè)計數(shù)據(jù)的預(yù)取和傳輸策略，因而無法達(dá)到訓(xùn)練的最佳性能.

3）傳統(tǒng)的存儲縮減技術(shù)不適用于大模型訓(xùn)練中的數(shù)據(jù)特征.在大模型訓(xùn)練過程中，參數(shù)和中間變量常采用全精度浮點數(shù)或半精度浮點數(shù)表示，以保證訓(xùn)練的精度，單位數(shù)據(jù)的存儲占用高.而且這些數(shù)據(jù)的稠密度高，難以通過傳統(tǒng)的壓縮方法緩解存儲壓力.

4）傳統(tǒng)的存儲容錯技術(shù)在大模型訓(xùn)練場景下容錯開銷大.大模型訓(xùn)練中的數(shù)據(jù)量龐大，并且數(shù)據(jù)更新頻繁.每一輪訓(xùn)練都會更新整個模型的參數(shù)以及部分優(yōu)化器的參數(shù).若系統(tǒng)使用傳統(tǒng)的數(shù)據(jù)容錯技術(shù)對此類數(shù)據(jù)進(jìn)行容錯，會導(dǎo)致大量并且頻繁的讀寫操作，降低了模型訓(xùn)練的效率.

近些年來，大量面向大模型訓(xùn)練的存儲技術(shù)被提出.其中一類工作著眼于大模型訓(xùn)練中的存儲性能.這類工作通過統(tǒng)一管理分布式顯存、利用異構(gòu)存儲介質(zhì)以及壓縮訓(xùn)練數(shù)據(jù)的方式滿足大規(guī)模模型訓(xùn)練中對存儲容量的需求.另一類工作著眼于大模型訓(xùn)練中的存儲容錯.這類工作通過持久化存儲模型訓(xùn)練的參數(shù)、在多GPU 上冗余訓(xùn)練模型參數(shù)以實現(xiàn)對模型訓(xùn)練數(shù)據(jù)的容錯.

2 面向大模型的存儲加速技術(shù)

2.1 基于大模型計算模式的分布式顯存管理技術(shù)

大模型訓(xùn)練中的計算模式可以分為層間并行和層內(nèi)并行.這2 種計算模式對應(yīng)不同的任務(wù)劃分方式和任務(wù)依賴關(guān)系.其中，層間并行以張量為粒度將模型劃分到多張GPU 上，單個計算任務(wù)僅依賴于負(fù)責(zé)計算相鄰張量的任務(wù)；層內(nèi)并行則是以張量的某一維度為粒度，將單個張量的計算拆分到多張GPU 上，單個計算任務(wù)依賴于多個負(fù)責(zé)計算相鄰張量分片的任務(wù).

若僅簡單地將模型數(shù)據(jù)劃分到多張GPU 上，而不考慮模型訓(xùn)練任務(wù)劃分和任務(wù)間的依賴關(guān)系，會從兩方面增加訓(xùn)練過程中的通信開銷：一方面，數(shù)據(jù)存儲的位置可能和計算任務(wù)所在的位置不一致，導(dǎo)致計算前需要先從遠(yuǎn)端GPU 傳輸所需的數(shù)據(jù)；另一方面，相鄰任務(wù)間的數(shù)據(jù)傳輸所采用的通信鏈路或者通信路徑非最優(yōu).

因此，現(xiàn)有工作依據(jù)大模型訓(xùn)練的2 種計算模式的任務(wù)劃分方式和任務(wù)間的依賴關(guān)系，對模型數(shù)據(jù)進(jìn)行劃分.其中：模型層間分片根據(jù)層間并行的計算模式，以張量為粒度對模型數(shù)據(jù)進(jìn)行拆分；模型層內(nèi)分片根據(jù)層內(nèi)并行的計算模式，以張量中的某一維度為粒度對模型數(shù)據(jù)進(jìn)行拆分.

2.1.1 基于模型層間并行的分布式顯存管理技術(shù)

模型層間分片方式是以張量為粒度，將模型數(shù)據(jù)劃分成多個小片，并分布式地存儲到多張GPU 中.此類工作在訓(xùn)練過程中傳輸?shù)臄?shù)據(jù)類型可以分為2 類：模型數(shù)據(jù)和激活量.在傳輸模型數(shù)據(jù)的工作中，每張GPU 負(fù)責(zé)所有模型分片的訓(xùn)練；在訓(xùn)練前，所需的模型數(shù)據(jù)被傳輸?shù)截?fù)責(zé)計算的GPU 顯存中；在訓(xùn)練后，GPU 將釋放不由此GPU 存儲的模型分片所占的顯存空間.在傳輸激活量的工作中，每張GPU 僅負(fù)責(zé)訓(xùn)練顯存中存儲的模型分片，并傳輸相鄰分片的激活量.

ZeRO-DP[13]是微軟公司提出的大模型訓(xùn)練系統(tǒng)，它采用數(shù)據(jù)并行的方式對模型進(jìn)行訓(xùn)練.為滿足模型訓(xùn)練過程中對存儲空間的需求，ZeRO-DP 采用層間分片的方式存儲模型，并訓(xùn)練過程中采用在GPU間傳輸模型數(shù)據(jù)的方案.

如圖2 所示，ZeRO-DP 將模型數(shù)據(jù)分為優(yōu)化器（Pos）、梯度（Pg）和模型參數(shù)（Pp）.對于優(yōu)化器，ZeRODP 將它按層均等地劃分成多份.每張GPU 僅存儲某一份優(yōu)化器分片，并僅負(fù)責(zé)更新其存儲的優(yōu)化器分片對應(yīng)的模型參數(shù).在每張GPU 完成模型參數(shù)更新之后，ZeRO-DP 將更新得到的模型參數(shù)傳輸?shù)狡渌鸊PU 中，覆蓋舊版本模型參數(shù).對于梯度，ZeRO-DP將每一份梯度分片聚合到存儲有對應(yīng)參數(shù)分片的GPU 中，從而減少全局的梯度數(shù)據(jù)存儲.對于模型參數(shù)，ZeRO-DP 將每一份參數(shù)分片存放在帶有對應(yīng)的優(yōu)化器分片的GPU 上.在訓(xùn)練過程中，若所需的模型分片不在負(fù)責(zé)計算的GPU 上，該GPU 則需從其他GPU 拉取模型分片，并在計算結(jié)束之后釋放對應(yīng)的顯存空間.

Fig.2 ?llustration of model partition in ZeRO-DP [13]圖2 ZeRO-DP 模型劃分示意圖[13]

用戶可以選擇是否對模型參數(shù)進(jìn)行劃分.若不對模型參數(shù)進(jìn)行劃分，僅對優(yōu)化器和梯度進(jìn)行劃分（Pos+g），則訓(xùn)練過程中的通信操作僅有歸約（All-Reduce）梯度和廣播（Broadcast）更新后的參數(shù)，但系統(tǒng)可訓(xùn)練的模型規(guī)模受限于單張GPU 的顯存大小.如需要擴(kuò)大訓(xùn)練的模型規(guī)模，則要對模型參數(shù)進(jìn)行劃分（Pos+g+p），但每次前向傳播和反向傳播前均需對所有參數(shù)進(jìn)行廣播，相比于Pos+g增加了1 次廣播通信.用戶可以根據(jù)GPU 通信和顯存的大小選擇所需的劃分組合.

流水線并行訓(xùn)練模式使用了層間分片的方式存儲模型.在訓(xùn)練過程中，流水線中的GPU 間僅傳輸相鄰分片的激活量.在流水線并行訓(xùn)練中，模型首先被劃分成多個子模型，每張GPU 負(fù)責(zé)計算某一個子模型.在訓(xùn)練過程中，一批樣本（minibatch）被拆分成多份子樣本（microbatch），系統(tǒng)以流水線的方式同時使用多份子樣本進(jìn)行訓(xùn)練，每張GPU 僅向存儲相鄰子模型的GPU 傳輸連接2 個子模型的激活量.

GPipe[14]是谷歌提出的基于流水線并行的模型訓(xùn)練框架.如圖3 所示，它首先使用一個樣本中的所有子樣本進(jìn)行前向傳播，接著對所有子樣本進(jìn)行反向傳播，最后根據(jù)此樣本訓(xùn)練得到的梯度進(jìn)行參數(shù)更新.由于前向傳播和反向傳播的順序不同，并且需要在使用一批樣本完成訓(xùn)練后暫停模型訓(xùn)練以對參數(shù)進(jìn)行更新，GPipe 產(chǎn)生了較多的計算空閑時間（即氣泡）.

Fig.3 ?llustration of GPipe[14] and PipeDream[15]圖3 GPipe[14]和PipeDream[15]示意圖

PipeDream[15]是英偉達(dá)公司提出的基于流水線并行的模型訓(xùn)練模式，它在流水線中將前向傳播和反向傳播交替進(jìn)行（1 forward 1 backward，1F1B），解決了GPipe 中的氣泡問題.在PipeDream 中，某一子樣本完成前向傳播傳播后會立即開始此子樣本的反向傳播.當(dāng)PipeDream 啟動并穩(wěn)定后，流水線中將不存在氣泡.但是，PipeDream 中需要存放多個版本的參數(shù)，增加了額外的存儲開銷.在1F1B 中，某一子樣本反向傳輸結(jié)束之后，子模型的參數(shù)立即被更新.此時，流水線中仍存在著使用原版本參數(shù)進(jìn)行前向傳播但未完成反向傳播的子樣本.為保證單個樣本前向傳播和反向傳播所使用的子模型參數(shù)版本一致，需要在GPU 中保存多個版本的模型數(shù)據(jù)，這增大了訓(xùn)練過程中的存儲開銷.

此外，PipeDream 的多版本參數(shù)還會影響收斂性.不同于GPipe 中所有子模型在完成一輪前向反向傳播后同步更新參數(shù)，PipeDream 中的子模型在完成一個子樣本的反向傳播后立即更新參數(shù).這會導(dǎo)致訓(xùn)練中單個子樣本在不同子模型使用的參數(shù)版本不一致，從而影響模型收斂.

表1 所示總結(jié)了GPipe，PipeDream 以及其他流水線在顯存開銷、氣泡和收斂性上的差異.PipeDream-2BW[17]沿用1F1B 的計算模式，同時通過梯度累積減少了參數(shù)更新頻率，進(jìn)而降低了所需存儲的模型參數(shù)版本數(shù)量.DAPPLE[18]則在PipeDream 中加入了類似于GPipe 的梯度同步更新，綜合了GPipe 和Pipe-Dream 兩者的優(yōu)勢.Chimera[19]則提出了雙頭流水線的訓(xùn)練模式，通過增加一定量的存儲開銷減少啟動流水線時帶來的氣泡.

Table 1 Comparison of Large Model Training Schemes Based on Different Pipeline Modes表1 基于不同流水線模式的大模型訓(xùn)練方案對比

2.1.2 基于模型層內(nèi)并行的分布式顯存管理技術(shù)

張量并行采用模型層內(nèi)分片的策略，該策略以張量的某一維度為粒度對模型進(jìn)行切分，并將切分得到的分片存儲在多張GPU 中.相比于模型的層間分片以層為粒度存儲數(shù)據(jù)，模型的層內(nèi)分片將某一層的數(shù)據(jù)以更細(xì)的粒度拆分.因此，模型的層內(nèi)分片可將單層的數(shù)據(jù)存儲到更多GPU 中，從而支持單層數(shù)據(jù)量更大的模型[20].

如圖4 所示，張量可以從行和列2 個維度被切分.以列對張量A切分時，每一個張量的分片的計算需要完整的輸入X，最后需要將所有分片計算得到的結(jié)果拼接以得到完整的結(jié)果.以行為粒度對張量A切分時，輸入X也需對應(yīng)按列進(jìn)行切分，最后需要將所有分片計算得到的結(jié)果相加以得到完整的計算結(jié)果.

Fig.4 Tensor splitting[20]圖4 張量切分[20]

根據(jù)是否切分輸入的張量和張量存放的方式，張量并行分為4 種不同的模式.其中每種張量并行的存儲開銷和通信開銷的對比如表2 所示.

Table 2 Comparison of Four Types of Tensor Parallelism表2 4 種張量并行的對比

1）1D 張量并行

1D 張量并行僅按行和列兩者中的某一維度對張量進(jìn)行切分.以式（1）的矩陣乘法（GEMM）為例，假設(shè)有P張GPU，且P=2，其中的參數(shù)A若按列被切分，如式（2）所示，每一參數(shù)切片與輸入相乘得到的子結(jié)果如式（3）所示.

下一個與之相鄰的線性GEMM 層參數(shù)B則按列進(jìn)行切分.如式（4）所示，每一個參數(shù)切片與上一層的子結(jié)果相乘，得到的結(jié)果通過All-Reduce 以得到最終的計算結(jié)果.

在1D 張量切分方式中，單張GPU 存儲1/P份參數(shù)張量和1 份完整的輸入張量.在計算過程中需要傳輸P-1 份的輸入張量.

2）2D 張量并行

1D 張量并行雖然將參數(shù)分布式存儲在每一張GPU 上，但每一張GPU 需要冗余保存一份完整的輸入.2D 張量并行[21]將參數(shù)和輸入同時按照行和列進(jìn)行切分，并分布式保存在多張GPU 上.以GEMM 為例，假設(shè)有P=q×q張GPU，其中q=2，GEMM 中的參數(shù)A同時按列和行進(jìn)行切分，與之對應(yīng)的輸入X也按照行和列進(jìn)行切分，切分結(jié)果如式（6）（7）所示：

計算分為q步.以某一步為例，首先Xi0在存儲帶有相同行的輸入的GPU 間被廣播；A0j在存儲有相同列的模型參數(shù)的GPU 間被廣播.然后每張GPU 上的模型和參數(shù)相乘.其他步驟的計算與之類似.最后每一步計算得到的子結(jié)果相加得到最終結(jié)果，如式（8）所示.

在2D 張量并行中，單個處理器內(nèi)存儲1/P份參數(shù)張量和 1/P份輸入張量.但在計算過程中需要傳輸q-1 份參數(shù)張量和輸入張量.

3）2.5D 張量并行

相比于1D 張量并行，2D 張量并行的通信開銷大，因此文獻(xiàn)[22]提出了2.5D 張量并行.2.5D 張量并行將輸入張量按行列切分后，再將切分后的分片按行分組.以GEMM 為例，在P=q×q×q張GPU 中，其中q=2，輸入張量首先被劃分為q×q×q份，然后得到的張量分片再被劃分成q組，劃分結(jié)果如式（9）所示.參數(shù)則仍按照式（7）劃分，并冗余q份存儲在與輸入q組對應(yīng)的GPU 中.

在計算過程中，每一組中按照2D 張量并行的方式進(jìn)行分布計算，由此得到的最后的結(jié)果如式（10）所示.

2.5 D 張量并行犧牲了少量存儲空間以減少了傳輸輸入張量帶來的通信開銷.在2.5D 張量并行中，單張GPU 需要存儲 1/q2份參數(shù)張量和 1/P份輸入張量；在計算過程中，2.5D 張量并行需要傳輸 (q-1)×q份的參數(shù)張量和q-1份的輸入張量.

4）3D 張量并行

3D 張量[23]并行在2.5D 張量并行的基礎(chǔ)上，將參數(shù)的分片按列進(jìn)行分組，并與輸入的分組對應(yīng).以GEMM 為例，P張GPU 按照q×q×q分組，其中q=2，則輸入的分組如式（11）所示，參數(shù)的分組如式（12）所示.其中Xijk和Akji被存儲在第 (i,j,k)張GPU 上.

在計算過程中，Xijk在j所在的GPU 組中被廣播，Akji在i所在 的GPU 組中被廣播.因此第 (i,j,k)張GPU中有Xik,Akj，并將這兩者相乘.最后的子結(jié)果被Reduce-Scatter 到所有GPU 中.

在3D 張量并行中，單個處理器內(nèi)存儲1/P份參數(shù)和1/P份輸入.3D 張量并行需要傳輸q-1 份的輸入張量和參數(shù)張量，以及q-1 份的中間結(jié)果.

Megatron[24]是由英偉達(dá)公司提出的一種基于Transformer 結(jié)構(gòu)的大模型訓(xùn)練框架，它使用了1D 張量并行策略.圖5 展示了Megatron 的模型切分方案.

Fig.5 ?llustration of 1D tensor parallelism in Megaton[24]圖5 Megatron 的1D 張量并行示意圖[24]

單個Transformer 塊中包含一個多層感知機(jī)（MLP）和一個注意力層.其中多層感知機(jī)分為2 個線性的GEMM 層和1 個非線性的GeLU函數(shù).由于GeLU需要按照列為粒度進(jìn)行計算，系統(tǒng)若對第1 個GEMM 層進(jìn)行行切分，則需要將多個分片計算的結(jié)果同步相加之后才能將結(jié)果輸入GeLU層進(jìn)行計算，增加了通信的時間.因此Megatron 使用列切分的方式對第1 個GEMM 層進(jìn)行切分.Megatron 使用行切分的方式對第2 個GEMM 進(jìn)行切分，可以直接使用本地GeLU計算的結(jié)果作為輸入，減少了等待所有GeLU分片結(jié)果并對其進(jìn)行拼接的開銷.

此外，Transformer 中的注意力層中包含了3 個分別表示Q，K，V的GEMM，此外還包含Softmax、Dropout和輸出前的GEMM.表示Q，K，V的GEMM 中包含了1 個或者多個頭（head）.Megatron 以頭為粒度對Q，K，V對應(yīng)的GEMM 按列進(jìn)行切分，以保證模型計算的正確性.對輸出前的GEMM 切分，Megatron 采用行切分的方式，以消除計算前輸入的通信同步.

2.2 大模型訓(xùn)練訪存感知的異構(gòu)存儲技術(shù)

若僅使用GPU 顯存存儲模型數(shù)據(jù)，則需要更大規(guī)模的GPU 集群以滿足更大規(guī)模模型訓(xùn)練時的存儲需求.

除GPU 顯存外，訓(xùn)練服務(wù)器中包含各種異構(gòu)的存儲資源，如DRAM 和SSD.模型參數(shù)、優(yōu)化器數(shù)據(jù)和訓(xùn)練產(chǎn)生的中間結(jié)果可以被存儲到這些異構(gòu)存儲資源中，以提高模型訓(xùn)練的規(guī)模.

但是，異構(gòu)存儲系統(tǒng)的數(shù)據(jù)移動開銷高.雖然GPU 內(nèi)的存儲資源的讀寫帶寬高，以匹配其計算資源的性能，但與外部存儲資源之間的互聯(lián)帶寬低.這導(dǎo)致數(shù)據(jù)在GPU 與外部設(shè)備間的搬移開銷大.如英偉達(dá)GPU 配備有帶寬高達(dá)數(shù)百GB/s 的DDR6 或HBM.但與外部的互聯(lián)往往使用PC?e，帶寬僅有數(shù)十個GBps，雖然有高速鏈路（如NVLink）被提出，但僅支持高端的GPU，硬件成本高昂，且擴(kuò)展性有限.如果在訓(xùn)練過程中僅簡單地使用GPU 外的存儲空間，其帶來的數(shù)據(jù)搬移開銷大，限制了GPU 內(nèi)的存算資源利用率，影響模型訓(xùn)練的效率.

現(xiàn)有工作利用大模型訓(xùn)練中訪存模式可預(yù)測的特性，以減少使用異構(gòu)存儲介質(zhì)帶來的數(shù)據(jù)移動開銷.在指定樣本批大小等其他模型訓(xùn)練的超參數(shù)后，大模型訓(xùn)練中的每一個操作的計算開銷和存儲開銷的變化相對固定.并且大模型的模型結(jié)構(gòu)固定，這使得每一個計算操作的順序也隨之固定.因此，根據(jù)訓(xùn)練第一個樣本批獲得的每個操作的訪存時刻和訪存大小信息，可以預(yù)測后續(xù)模型訓(xùn)練過程中的訪存模式[4,25].表3 對此類工作進(jìn)行了總結(jié).本節(jié)將按照異構(gòu)存儲技術(shù)中使用的存儲介質(zhì)類型介紹此類工作.

Table 3 Comparison of Different Large Model Training System Using Heterogeneous Storage Resources表3 不同使用異構(gòu)存儲資源的大模型訓(xùn)練系統(tǒng)對比

2.2.1 基于DRAM 的異構(gòu)存儲技術(shù)

相比于GPU 顯存，DRAM 的容量更大，因此，可以將模型數(shù)據(jù)或者訓(xùn)練的中間結(jié)果卸載到DRAM 中，并在計算前按需將卸載數(shù)據(jù)上傳到GPU 顯存之中，以支持更大規(guī)模的模型訓(xùn)練.

vDNN[26]使用DRAM 擴(kuò)充顯存容量以支持更大規(guī)模模型的訓(xùn)練.在模型訓(xùn)練過程中，前向傳播會產(chǎn)生大量的臨時變量，這些臨時變量需要被保存以用于反向傳播過程中的梯度計算，其帶來的顯存開銷的大小占整個訓(xùn)練過程中顯存開銷的40%～80%[26].為解決此問題，vDNN 將前向傳播計算得到的臨時變量卸載到DRAM 中，并在反向傳播開始前將這些臨時變量預(yù)取到顯存中.vDNN 提出了異步的顯存管理AP? 以降低臨時變量在顯存中釋放和分配的開銷；使用CUDA 的異步內(nèi)存拷貝，以利用計算掩蓋通信開銷；提出了顯存卸載策略以權(quán)衡顯存開銷和訓(xùn)練的性能.

1）異步顯存管理

臨時變量被卸載到DRAM 后，顯存空間需要被釋放，以用于模型后續(xù)的計算；臨時變量被上傳到顯存之前，系統(tǒng)需要預(yù)先分配一段顯存空間.而CUDA只提供同步的內(nèi)存釋放（cudaFree）與分配接口（cudaMalloc），這會在模型訓(xùn)練過程中增加額外的同步開銷.vDNN 則提出了異步的顯存管理接口.vDNN在系統(tǒng)啟動后會預(yù)分配一段顯存空間當(dāng)作顯存池.臨時變量需要分配顯存空間時，vDNN 會直接從顯存池中分配預(yù)先準(zhǔn)備好的顯存空間.被釋放的顯存空間會被重新放回顯存池中供之后的變量使用.這避免了使用CUDA 的同步顯存分配和釋放AP?.

2）異步數(shù)據(jù)傳輸

為降低數(shù)據(jù)傳輸帶來的額外開銷，在顯存卸載時，vDNN 首先會分配一段固定的DRAM 空間，避免了數(shù)據(jù)傳輸經(jīng)過回彈緩沖區(qū)（bounce buffer)，以最大化數(shù)據(jù)傳輸?shù)膸?并且，vDNN 使用異步顯存拷貝接口cudaMemcpyAsync 將顯存數(shù)據(jù)傳輸?shù)紻RAM上，以利用計算掩蓋一部分?jǐn)?shù)據(jù)通信所帶來的開銷.在卸載數(shù)據(jù)重新上傳回顯存時，除利用異步顯存拷貝接口外，vDNN 選取最近將被計算的層對應(yīng)的臨時變量進(jìn)行預(yù)取，以最小化預(yù)取帶來的顯存開銷的同時利用計算掩蓋通信帶來的開銷.

3）顯存卸載策略

為最大化降低顯存開銷，一種卸載策略是使用內(nèi)存友好的卷積算法并且將所有臨時變量卸載到DRAM 中（vDNNall）.但內(nèi)存友好的卷積算法可能會影響模型的收斂；并且卸載所有的臨時變量會增加卸載和預(yù)取帶來的通信開銷，影響訓(xùn)練的性能.為平衡顯存開銷和訓(xùn)練的性能，vDNN 提出了靜態(tài)和動態(tài)的卸載策略.

在靜態(tài)卸載策略中，vDNN 僅卸載計算開銷最大的卷積層（vDNNconv），以盡可能利用計算掩蓋通信開銷.在動態(tài)卸載策略中，vDNN 首先測試顯存開銷最小的方案是否可以滿足訓(xùn)練過程中的顯存限制，即使用vDNNall和內(nèi)存友好的卷積算法.若顯存開銷最小的方案不滿足訓(xùn)練過程中的顯存限制，則說明此模型不能在給定的硬件條件下進(jìn)行訓(xùn)練；若顯存開銷最小的方案滿足顯存限制，vDNN 再測試顯存開銷最大但訓(xùn)練性能最優(yōu)的方案是否滿足訓(xùn)練過程中的顯存限制，即不對顯存進(jìn)行卸載并使用性能最優(yōu)的卷積算法進(jìn)行訓(xùn)練.如果性能最優(yōu)的方案滿足顯存需求，則使用此方案；否則，vDNN 測試vDNNconv在使用算法性能友好而非內(nèi)存友好的卷積算法時是否滿足顯存限制.若不滿足，vDNN 使用貪心算法，將部分卷積層所使用的算法改為內(nèi)存友好的算法.當(dāng)所有vDNNconv中的卷積層使用顯存友好的算法后，訓(xùn)練過程仍不滿足顯存限制時，vDNN 開始使用vDNNall，并再次使用貪心算法，將部分卷積層所使用的算法改為內(nèi)存友好的算法，以搜索符合顯存限制的卸載策略.

SwapAdvisor[28]認(rèn)為vDNN 這類工作基于人工經(jīng)驗選擇卸載的數(shù)據(jù)，如vDNN 僅卸載了卷積層.這類工作一方面不具有通用性，另一方面不能達(dá)到最優(yōu)的性能.SwapAdvisor 從張量的角度對DNN 模型中的計算和存儲進(jìn)行建模，以支持模型訓(xùn)練中各種不同類型的模型數(shù)據(jù)的卸載；并采用遺傳算法搜索卸載和預(yù)取的近似最優(yōu)方案.

ZeRO-Offload[30]是利用異構(gòu)存儲方案的分布式訓(xùn)練框架.ZeRO-Offload 采用與ZeRO-DP 類似的訓(xùn)練流程，但是將模型數(shù)據(jù)和參數(shù)更新的過程卸載到DRAM 和CPU 中.如圖6 所示，ZeRO-Offload 在DRAM中保存精度為FP32 的參數(shù)與優(yōu)化器，在顯存中以ZeRO-DP 的方式存儲精度為FP16 的參數(shù).在完成一次訓(xùn)練之后，精度為FP16 的梯度信息會被傳輸?shù)紻RAM 中，并在CPU 中用作參數(shù)更新.更新得到的參數(shù)將被轉(zhuǎn)化成FP16，并傳輸?shù)斤@存中以更新舊版本的參數(shù).若CPU 的參數(shù)更新的速度慢于GPU 計算1步（step）的速度，ZeRO-Offload 將多步更新1 次，以降低CPU 參數(shù)更新帶來的開銷.

Fig.6 ?llustration of ZeRO-Offload[30]圖6 ZeRO-Offload 示意圖[30]

Mobius[32]是清華大學(xué)于2023 年提出的基于消費級GPU 服務(wù)器的大模型訓(xùn)練系統(tǒng).如圖7 所示，不同于數(shù)據(jù)中心GPU 服務(wù)器，消費級GPU 服務(wù)器的通信資源有限.一方面，消費級GPU 服務(wù)器中的GPU 不支持高帶寬的NVLink 鏈路，GPU 間的通信帶寬受限于PC?e 的鏈路帶寬；另一方面，消費級GPU 之間的數(shù)據(jù)通信不支持GPU P2P，這導(dǎo)致GPU 間的通信需要首先傳輸?shù)紻RAM，再被傳輸?shù)侥繕?biāo)GPU 中.若共享PC?e 鏈路的GPU 同時通信，會出現(xiàn)通信競爭.因此，在消費級GPU 中使用類似于ZeRO-Offload 此類GPU 間通信頻繁的系統(tǒng)，會帶來巨大的通信開銷，嚴(yán)重影響訓(xùn)練的性能.

Fig.7 Data center GPU server and commodity GPU server[32]圖7 數(shù)據(jù)中心GPU 服務(wù)器和消費級GPU 服務(wù)器[32]

Mobius[32]提出了基于流水線并行的異構(gòu)存儲訓(xùn)練策略.具體地，該策略將模型分成多份子模型存儲在DRAM 中，按照模型訓(xùn)練的順序依次將子模型上傳到負(fù)責(zé)訓(xùn)練的GPU 中，采用流水線并行的方式進(jìn)行訓(xùn)練.上傳后的子模型在完成所有子樣本的前向傳播后將從顯存中釋放.反向傳播的過程與前向傳播類似，得到的梯度數(shù)據(jù)會被傳輸?shù)紻RAM 中用于參數(shù)更新.為了降低數(shù)據(jù)傳輸帶來的通信開銷，Mobius 對上述訓(xùn)練過程進(jìn)行建模，并使用混合線性規(guī)劃的方式找出最優(yōu)的子模型劃分和預(yù)取方案.為避免通信競爭，Mobius 盡可能地將預(yù)取時間近似的子模型映射到不共享PC?e 鏈路的GPU 中.

2.2.2 基于SSD 的異構(gòu)存儲技術(shù)

隨著模型規(guī)模的增長，模型數(shù)據(jù)量和中間結(jié)果到達(dá)了TB 量級.現(xiàn)有工作嘗試將模型數(shù)據(jù)和中間變量卸載到SSD 以支持更大規(guī)模的模型訓(xùn)練.相較于DRAM，SSD 的容量更大、價格與能耗更低，可以支持更大規(guī)模的模型訓(xùn)練，但SSD 的讀寫帶寬低，給模型訓(xùn)練帶來了新的挑戰(zhàn).

FlashNeuron[4]將模型參數(shù)卸載到SSD 中.為最小化卸載帶來的開銷，F(xiàn)lashNeuron 設(shè)計了離線的卸載選擇策略，并提出了GPU 直訪SSD 的數(shù)據(jù)傳輸方案.

在制定卸載策略時，F(xiàn)lashNeuron 首先將每個張量卸載到SSD 中，以采樣每一個張量在GPU 中的顯存開銷、張量的計算開銷、卸載到SSD 的時間開銷和對張量進(jìn)行壓縮的時間開銷.根據(jù)上述采樣的結(jié)果，F(xiàn)lashNeuron 使用2 階段算法選擇需要被卸載的張量.在第1 階段中，F(xiàn)lashNeuron 將張量依次卸載到SSD 中，直到顯存空間滿足訓(xùn)練需求為止，如果卸載的時間小于訓(xùn)練需要的時間，則使用此卸載方案.否則，F(xiàn)lashNeuron 進(jìn)入第2 階段，使用壓縮張量比率最高的張量依次替換第1 階段得到的策略中不可被壓縮的張量，以降低卸載過程中數(shù)據(jù)傳輸?shù)拈_銷.為減少卸載與預(yù)取過程中CPU 的開銷，F(xiàn)lashNeuron 利用GDRCopy 和SPDK 實現(xiàn)了GPU 對SSD 的直接訪問.

圖8 中顏色越深的張量壓縮率越高，白色為不可壓縮的張量，箭頭長度表示張量卸載的時間開銷，斜花紋的方塊表示被卸載的張量.若對圖8 所示模型訓(xùn)練，在不卸載任何張量時，模型訓(xùn)練超出GPU 顯存16 MB.在第1 階段中，F(xiàn)lashNeuron 按照模型順序依次選擇需要被卸載的張量，直到第D個張量被選擇卸載之后，顯存可滿足訓(xùn)練所需的存儲空間需求.但此卸載策略下，模型訓(xùn)練計算的時間不能掩蓋卸載帶來的通信開銷，因此進(jìn)入第2 階段.FlashNeuron 使用壓縮率更高的第F個張量和第E個張量替換了第1 階段中壓縮不友好的第B個和第D個張量，此時卸載時間小于訓(xùn)練時間，因此此方案為最終卸載所使用的方案.

Fig.8 ?llustration of offloading strategy in FlashNeuron[4]圖8 FlashNeuron 卸載策略示意圖[4]

ZeRO-?nfinity[33]將模型數(shù)據(jù)存儲在SSD 中，它沿用了ZeRO 訓(xùn)練模式，同時也使用了ZeRO-Offload 中的CPU 更新參數(shù)的策略.在ZeRO-?nfinity 中，為了解決SSD 讀寫帶寬較低的問題，某一層的數(shù)據(jù)以圖9中的方式被拆分成更小的分片存儲到多個SSD 中，并由多張GPU 同時負(fù)責(zé)傳輸，以充分利用SSD 的聚合帶寬.不同于傳統(tǒng)的ZeRO 中使用廣播的方式，由于單層的數(shù)據(jù)分散在多張GPU 中，ZeRO-?nfinity 需要使用全局聚合的方式讓每一張GPU 獲得完整的用于訓(xùn)練的模型數(shù)據(jù).

2.3 大模型數(shù)據(jù)縮減技術(shù)

傳統(tǒng)的數(shù)據(jù)縮減技術(shù)往往通過壓縮的方式減少存儲開銷[34-36].但是大模型的數(shù)據(jù)稠密[37]難以被壓縮.現(xiàn)有工作根據(jù)大模型數(shù)據(jù)特征，從2 方面對大模型訓(xùn)練中的數(shù)據(jù)進(jìn)行縮減：增加計算量和犧牲模型精度.具體地，激活量檢查點與重算（checkpoint and recomputation）算法通過增加計算量減少了所需存儲的激活量數(shù)據(jù)；混合精度與量化通過犧牲模型精度以降低模型數(shù)據(jù)和中間變量的大小.本節(jié)將具體介紹這2 種不同的大模型數(shù)據(jù)縮減技術(shù).

2.3.1 激活量檢查點與重算

在模型訓(xùn)練過程中，前向傳播產(chǎn)生的激活量會被保存在顯存中，并在反向傳播過程中被用于計算梯度.為減少保存激活量帶來的存儲開銷，有研究提出在前向傳播時僅保存少量的激活量作為檢查點，并丟棄其他的激活量.在需要使用被丟棄的激活量進(jìn)行反向傳播時，系統(tǒng)利用激活量檢查點重新進(jìn)行前向傳播，以恢復(fù)之前所丟棄的激活量.

Chen 等人[38]在2016 年提出了激活量檢查點與重算的方法，并且提出了3 種選擇激活量檢查點的算法.第1 種算法是直接選擇計算開銷小的操作，如批歸一化和池化操作，這樣減少顯存開銷的同時帶來的額外計算開銷較小.第2 種算法是使用貪心算法找到平衡額外計算開銷和存儲開銷的最佳方案.第3種算法是采用遞歸的方式，通過不斷擴(kuò)大子模型的范圍以找到最優(yōu)的激活量檢查點選擇方案，但這種算法的搜索空間大、成本高.

Checkmate[39]是伯克利大學(xué)提出的復(fù)雜神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)的激活量檢查點選擇算法.隨著殘差神經(jīng)網(wǎng)絡(luò)的提出，基于線性網(wǎng)絡(luò)的激活量檢查點選擇方案不再適用.Checkmate 將模型中的數(shù)據(jù)依賴、存儲開銷、計算開銷與重算開銷進(jìn)行建模，使用線性規(guī)劃的方式找出最優(yōu)的激活量檢查點的選擇與重算的時刻.由于整數(shù)線性規(guī)劃是NP 難問題，Checkmate 提出了一種2 階段的近似求解方案.

2.3.2 混合精度訓(xùn)練與量化

模型訓(xùn)練過程可以使用半精度（FP16）以降低顯存開銷.但是FP16 的有效范圍和精度均比全精度小，如果直接使用FP16 會導(dǎo)致數(shù)據(jù)溢出和誤差的問題.一方面，隨著模型在訓(xùn)練過程中不斷收斂，梯度會隨之變小.此時，梯度和學(xué)習(xí)率的乘積可能過小導(dǎo)致數(shù)據(jù)下溢.另一方面，由于FP16 的間隔較小，若梯度較小，參數(shù)更新過程中會直接將結(jié)果舍入，進(jìn)而影響模型的收斂.百度公司和英偉達(dá)公司在2017 年提出了混合精度的訓(xùn)練方法[40]，其流程如圖10 所示.該方法通過權(quán)重備份、損失量（loss）擴(kuò)展和算數(shù)精度提高等方法解決了上述問題.

Fig.10 ?llustration of mixed precision training[40]圖10 混合精度訓(xùn)練示意圖[40]

1）權(quán)重備份

在存儲模型數(shù)據(jù)時，模型參數(shù)、激活量和梯度均以半精度的格式存儲.除此之外，系統(tǒng)中還保存一份全精度格式的模型參數(shù).模型訓(xùn)練過程使用半精度的模型參數(shù)進(jìn)行計算，并生成半精度的梯度數(shù)據(jù).在參數(shù)更新時，半精度的梯度數(shù)據(jù)與全精度版本的參數(shù)相加得到全精度的新模型參數(shù).新的全精度模型參數(shù)在下一輪訓(xùn)練前會先轉(zhuǎn)化成半精度版本.

2）損失量擴(kuò)展

為解決FP16 向下溢出的問題，反向傳播所產(chǎn)生的損失量會等比平移擴(kuò)展到FP16 的有效范圍中，擴(kuò)展的損失量會由反向傳播作用于梯度之中.在參數(shù)更新時，擴(kuò)展的梯度先被反向平移到全精度范圍中，再進(jìn)行參數(shù)更新.

3）算數(shù)精度提高

某些網(wǎng)絡(luò)結(jié)構(gòu)可以通過使用全精度矩陣?yán)奂影刖染仃嚦朔ǖ慕Y(jié)果，避免由2 個半精度矩陣相加導(dǎo)致的舍入誤差.

ActNN[41]將前向傳播過程中產(chǎn)生的激活量進(jìn)行壓縮以減少現(xiàn)存開銷.ActNN 使用2 比特壓縮量化算法，這是因為2 比特壓縮量化算法在GPU 上相較于其他壓縮算法開銷較小.ActNN 根據(jù)梯度信息，利用不同樣本、不同維度和不同層之間的異構(gòu)特性，對更重要的激活值分配更多比特，并且保證所有激活量中所有的浮點數(shù)平均大小為2 b，以此在降低存儲開銷的同時減少壓縮帶來的誤差.ActNN 提供5 個優(yōu)化等級，等級越高顯存開銷越小，但用于壓縮的計算開銷越大，運行速度越慢.

2.4 總結(jié)與比較

如表4 所示的各類存儲技術(shù)在單步訓(xùn)練中所需的顯存開銷、通信開銷以及引入的額外計算開銷各有差異.

Table 4 Comparison of Different Large Model Training Storage Technologies表4 各種不同大模型訓(xùn)練存儲技術(shù)的比較

基于大模型計算模式的分布式顯存管理技術(shù)僅適用于GPU 聚合顯存容量，可滿足大模型訓(xùn)練存儲需求的場景.其中基于模型層間并行的分布式顯存管理技術(shù)中的ZeRO-DP 訓(xùn)練模式需要在GPU 中廣播和聚合大量的模型梯度、優(yōu)化器和參數(shù)；而流水線并行訓(xùn)練模式僅需要在2 個GPU 間點對點地傳輸子模型間的少量激活量.因此，ZeRO-DP 訓(xùn)練模式更適用于高帶寬、全互聯(lián)的GPU 集群，而流水線并行訓(xùn)練模式更適合于非全互聯(lián)且卡間通信帶寬低的GPU 集群.與ZeRO-DP 類似，基于模型層內(nèi)并行訓(xùn)練分布式顯存管理技術(shù)需在GPU 集群中使用集合通信聚合和分散激活量和參數(shù)，但對數(shù)據(jù)存儲的粒度比ZeRO-DP更小，存儲的可擴(kuò)展性更高.因此，它更適用于高帶寬、全互聯(lián)的GPU 集群且模型中單層數(shù)據(jù)量龐大的場景.

大模型訓(xùn)練訪存感知的異構(gòu)存儲技術(shù)使用成本更低的DRAM 和SSD 擴(kuò)充高成本的GPU 顯存，以滿足大模型訓(xùn)練過程中的存儲需求.相比于SSD，DRAM的讀寫帶寬更高，單位容量的成本也更高.因此基于DRAM 的異構(gòu)存儲技術(shù)適用于對訓(xùn)練性能要求較高的場景，而基于SSD 的異構(gòu)存儲技術(shù)更適用于對成本控制更高的場景.大模型數(shù)據(jù)縮減技術(shù)中的激活量檢查點與重算技術(shù)通過增加計算開銷以降低存儲開銷；混合精度技術(shù)則通過犧牲模型訓(xùn)練的精度以減少模型數(shù)據(jù)大小.此類技術(shù)可以與大模型訓(xùn)練訪存感知的異構(gòu)存儲技術(shù)綜合使用以降低數(shù)據(jù)在GPU間和異構(gòu)存儲介質(zhì)間搬移的通信開銷.

圖11 展示了根據(jù)算法性能需求和所用硬件條件選擇最合適的存儲技術(shù)的流程.在實際訓(xùn)練過程中，可能會混合采用多種不同的存儲技術(shù).例如，在GPU集群中，通過高帶寬的NVLink 互聯(lián)的GPU 組使用基于層內(nèi)并行的分布式顯存管理技術(shù)；通過帶寬較低的PC?e 互聯(lián)的GPU 組使用流水線并行的方式.

Fig.11 Usage scenarios of various large model training storage technologies圖11 各類大模型訓(xùn)練存儲技術(shù)使用場景

3 面向大模型的存儲容錯技術(shù)

GPU 故障數(shù)量隨著GPU 集群規(guī)模的增大而提高.GPU 的頻繁故障一方面會導(dǎo)致訓(xùn)練得到的參數(shù)丟失；另一方面由于大模型訓(xùn)練中各GPU 間的數(shù)據(jù)存在依賴關(guān)系，單GPU 的故障會擴(kuò)散到整個GPU 集群中.有2 類主要的工作解決大模型訓(xùn)練故障的問題：參數(shù)檢查點和冗余計算.本節(jié)將具體介紹這2 類不同的大模型訓(xùn)練容錯技術(shù).

3.1 參數(shù)檢查點

參數(shù)檢查點技術(shù)以設(shè)定的頻率將訓(xùn)練得到的參數(shù)信息存儲到持久化的存儲介質(zhì)中，以對GPU 故障進(jìn)行容錯.在GPU 故障后，參數(shù)檢查點技術(shù)利用最新且完整的參數(shù)進(jìn)行恢復(fù).

參數(shù)檢查點技術(shù)存在2 方面的挑戰(zhàn)：1）參數(shù)持久化的開銷大，會阻塞計算任務(wù).在參數(shù)完全寫入到持久化介質(zhì)前，模型參數(shù)需要停止更新，以避免參數(shù)丟失.但是大模型參數(shù)量龐大，且持久化存儲介質(zhì)的讀寫帶寬遠(yuǎn)低于GPU 顯存的讀寫帶寬與GPU 的計算帶寬，這導(dǎo)致持久化模型參數(shù)開銷大，計算任務(wù)被長時間阻塞.例如，使用64 臺EC2 的搶占式實例訓(xùn)練GPT-2 模型，77%的訓(xùn)練時間將被用于參數(shù)的持久化操作[7].2）參數(shù)檢查點的頻率會影響訓(xùn)練與恢復(fù)的效率.若參數(shù)檢查點頻率較高，計算任務(wù)被參數(shù)持久化操作阻塞的頻率也會隨之提高，影響模型訓(xùn)練的效率；若參數(shù)檢查點的頻率較低，故障時模型只能使用較早的參數(shù)版本進(jìn)行恢復(fù)，需要重新對大量數(shù)據(jù)集進(jìn)行訓(xùn)練.

CheckFreq[42]針對上述2 方面挑戰(zhàn)分別提出了2階段參數(shù)檢查點技術(shù)和基于采樣的參數(shù)檢查點策略.

2 階段參數(shù)檢查點技術(shù)將檢查點的過程拆分為快照（snapshot）和持久化（persist）2 個階段.具體地，在快照階段，若GPU 顯存中有空余空間，系統(tǒng)將更新后的參數(shù)保存一份副本在顯存中；反之，系統(tǒng)則將更新后的參數(shù)傳輸?shù)紻RAM 中.由于下一次參數(shù)更新后模型參數(shù)才會被修改，快照可與下一次更新前的前向傳播和反向傳播并行進(jìn)行.在持久化階段，系統(tǒng)將快照的模型參數(shù)寫入持久化介質(zhì)中.由于快照的模型參數(shù)數(shù)據(jù)與訓(xùn)練過程的模型參數(shù)數(shù)據(jù)是相互獨立的，持久化操作可以在后臺進(jìn)行，不影響后續(xù)的計算任務(wù).如圖12 所示，第1 輪的參數(shù)更新后，數(shù)據(jù)先快照到顯存或者DRAM 中.快照過程可以與第2 輪的前向傳播與反向傳播并行進(jìn)行.快照結(jié)束之后，CheckFreq 再將快照的數(shù)據(jù)存儲至持久化介質(zhì)中.

Fig.12 ?llustration of CheckFreq[42]圖12 CheckFreq 示意圖[42]

ChekckFreq 依據(jù)訓(xùn)練過程中的系統(tǒng)采樣數(shù)據(jù)，動態(tài)指定檢查點的頻率.CheckFreq 指出不同的模型規(guī)模和硬件性能會導(dǎo)致不同的檢查點開銷與恢復(fù)時間.例如，模型規(guī)模越大，快照和持久化的數(shù)據(jù)量也會越多，這會增加參數(shù)檢查點的開銷.因此，參數(shù)檢查點的頻率需要依據(jù)模型和硬件情況設(shè)定.CheckFreq 首先動態(tài)采樣模型的訓(xùn)練時間、在顯存與內(nèi)存上快照的開銷、參數(shù)持久化的大小以及持久化介質(zhì)的帶寬等信息，然后根據(jù)所采樣的信息通過算法得出最優(yōu)的檢查點頻率.

3.2 冗余計算

參數(shù)檢查點技術(shù)需要大容量的持久化存儲設(shè)備以保存檢查點信息.并且在恢復(fù)階段，參數(shù)檢查點需要從持久化介質(zhì)中讀取之前版本的參數(shù)，這導(dǎo)致恢復(fù)開銷高.有工作利用冗余計算的方式，在多張GPU 中重復(fù)計算相同版本的參數(shù)，以對模型訓(xùn)練數(shù)據(jù)容錯.

Bamboo[7]在流水線并行訓(xùn)練模式中的每張GPU上冗余計算2 個相鄰的子模型，以達(dá)到容錯的目的.為了減少冗余計算帶來的額外開銷，Bamboo 利用流水線并行中天然存在的氣泡執(zhí)行前向傳播任務(wù)的冗余計算，并采用了惰性執(zhí)行反向傳播的冗余計算的策略.在模型訓(xùn)練中，反向傳播的計算開銷要遠(yuǎn)高于前向傳播的計算開銷，并且在PipeDream 中越后執(zhí)行前向傳播的GPU 計算開銷越大，這會導(dǎo)致流水線中存在大量的氣泡.如圖13，假設(shè)GPU1中前向傳播的時間是1 個時間單位，反向傳播的時間是2 個時間單位；GPU2中前向傳播的時間是1.2 個時間單位，反向傳播的時間是2.4 個時間單位.在這2 個GPU 中使用流水線并行會導(dǎo)致0.6 個時間單位的氣泡.Bamboo利用這些氣泡對相鄰的子模型進(jìn)行冗余的前向傳播.對于計算開銷較大的反向傳播，Bamboo 僅在出錯后利用冗余前向傳播得到的結(jié)果進(jìn)行反向傳播，以減少冗余計算對訓(xùn)練帶來的影響.

Fig.13 Bubbles in pipeline parallelism[7]圖13 流水線并行的氣泡[7]

在某GPU 故障后，其對應(yīng)的計算任務(wù)會被路由到負(fù)責(zé)冗余計算的GPU 上.若連續(xù)多個GPU 故障，導(dǎo)致某子模型數(shù)據(jù)從GPU 顯存中丟失，則需要從參數(shù)檢查點中重新加載之前版本的模型參數(shù).

4 總結(jié)與展望

本文詳細(xì)闡述了ChatGPT 時代下大模型訓(xùn)練過程中存在的存儲挑戰(zhàn)，并以此出發(fā)介紹了現(xiàn)有的大模型訓(xùn)練中所使用的存儲技術(shù).如今有大量大模型訓(xùn)練系統(tǒng)支持了本文所介紹的存儲技術(shù)[43-46]，并提供了簡潔高效的調(diào)用接口，降低了用戶進(jìn)行大模型訓(xùn)練的門檻.但隨著模型規(guī)模的急劇增加和硬件設(shè)備的發(fā)展，大模型訓(xùn)練中仍存在著一些存儲問題亟需解決：

1）大模型訓(xùn)練的存儲成本.一方面，為充分利用硬件的計算資源，訓(xùn)練需要使用讀寫帶寬更高的存儲介質(zhì)，如HBM.而這些存儲介質(zhì)的單位比特的價格高昂.另一方面，隨著模型規(guī)模的不斷攀升，訓(xùn)練所需的存儲空間也隨之增長，因此需要更大容量的存儲介質(zhì)，這增加了訓(xùn)練過程中用于存儲設(shè)備的成本.

2）大模型訓(xùn)練的綠色存儲.大模型訓(xùn)練能耗大，導(dǎo)致大量碳排放，從而影響全球環(huán)境.從存儲角度講，大模型訓(xùn)練需要使用數(shù)量龐大的高帶寬存儲介質(zhì)滿足大模型訓(xùn)練中的存儲需求，這些存儲介質(zhì)相較于讀寫帶寬較低的SSD 和HDD 能耗高.并且大模型訓(xùn)練中數(shù)據(jù)讀寫頻繁，進(jìn)一步提高了存儲能耗.