從BERT 到ChatGPT：大模型訓練中的存儲系統挑戰與技術發展

馮楊洋 汪 慶 謝旻暉 舒繼武

（清華大學計算機科學與技術系 北京 100084）

（fyy21@mails.tsinghua.edu.cn）

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

Fig.1 ?ncrease of model size圖1 模型規模增長

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

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

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

1 大模型訓練的存儲挑戰

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

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

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

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

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

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

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

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

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

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

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

2 面向大模型的存儲加速技術

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

大模型訓練中的計算模式可以分為層間并行和層內并行.這2 種計算模式對應不同的任務劃分方式和任務依賴關系.其中，層間并行以張量為粒度將模型劃分到多張GPU 上，單個計算任務僅依賴于負責計算相鄰張量的任務；層內并行則是以張量的某一維度為粒度，將單個張量的計算拆分到多張GPU 上，單個計算任務依賴于多個負責計算相鄰張量分片的任務.

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

因此，現有工作依據大模型訓練的2 種計算模式的任務劃分方式和任務間的依賴關系，對模型數據進行劃分.其中：模型層間分片根據層間并行的計算模式，以張量為粒度對模型數據進行拆分；模型層內分片根據層內并行的計算模式，以張量中的某一維度為粒度對模型數據進行拆分.

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

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

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

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

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

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

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

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

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

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

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

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

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

2.1.2 基于模型層內并行的分布式顯存管理技術

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

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

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

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

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

1）1D 張量并行

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

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

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

2）2D 張量并行

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

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

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

3）2.5D 張量并行

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

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

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

4）3D 張量并行

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

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

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

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

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

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

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

2.2 大模型訓練訪存感知的異構存儲技術

若僅使用GPU 顯存存儲模型數據，則需要更大規模的GPU 集群以滿足更大規模模型訓練時的存儲需求.

除GPU 顯存外，訓練服務器中包含各種異構的存儲資源，如DRAM 和SSD.模型參數、優化器數據和訓練產生的中間結果可以被存儲到這些異構存儲資源中，以提高模型訓練的規模.

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

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

Table 3 Comparison of Different Large Model Training System Using Heterogeneous Storage Resources表3 不同使用異構存儲資源的大模型訓練系統對比

2.2.1 基于DRAM 的異構存儲技術

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

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

1）異步顯存管理

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

2）異步數據傳輸

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

3）顯存卸載策略

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

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

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

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

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

Mobius[32]是清華大學于2023 年提出的基于消費級GPU 服務器的大模型訓練系統.如圖7 所示，不同于數據中心GPU 服務器，消費級GPU 服務器的通信資源有限.一方面，消費級GPU 服務器中的GPU 不支持高帶寬的NVLink 鏈路，GPU 間的通信帶寬受限于PC?e 的鏈路帶寬；另一方面，消費級GPU 之間的數據通信不支持GPU P2P，這導致GPU 間的通信需要首先傳輸到DRAM，再被傳輸到目標GPU 中.若共享PC?e 鏈路的GPU 同時通信，會出現通信競爭.因此，在消費級GPU 中使用類似于ZeRO-Offload 此類GPU 間通信頻繁的系統，會帶來巨大的通信開銷，嚴重影響訓練的性能.

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

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

2.2.2 基于SSD 的異構存儲技術

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

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

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

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

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

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

2.3 大模型數據縮減技術

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

2.3.1 激活量檢查點與重算

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

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

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

2.3.2 混合精度訓練與量化

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

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

1）權重備份

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

2）損失量擴展

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

3）算數精度提高

某些網絡結構可以通過使用全精度矩陣累加半精度矩陣乘法的結果，避免由2 個半精度矩陣相加導致的舍入誤差.

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

2.4 總結與比較

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

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

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

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

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

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

3 面向大模型的存儲容錯技術

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

3.1 參數檢查點

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

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

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

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

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

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

3.2 冗余計算

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

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

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

在某GPU 故障后，其對應的計算任務會被路由到負責冗余計算的GPU 上.若連續多個GPU 故障，導致某子模型數據從GPU 顯存中丟失，則需要從參數檢查點中重新加載之前版本的模型參數.

4 總結與展望

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

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

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