數(shù)據(jù)加載過程與性能優(yōu)化

數(shù)據(jù)加載的復雜性  

數(shù)據(jù)加載并不僅限于從存儲讀取數(shù)據(jù)，它涵蓋了解碼、格式轉(zhuǎn)換及數(shù)據(jù)增強等預處理環(huán)節(jié)。這些步驟通常在CPU上執(zhí)行，目的是將原始數(shù)據(jù)轉(zhuǎn)換為GPU可處理的張量格式。

數(shù)據(jù)加載的步驟  

數(shù)據(jù)加載流程可以概括為：

• ?從存儲系統(tǒng)讀取數(shù)據(jù)到系統(tǒng)內(nèi)存。
• 解碼原始數(shù)據(jù)。
• 應用數(shù)據(jù)變換和增強操作（如裁剪、尺寸調(diào)整）。
• 將處理后的數(shù)據(jù)傳輸至GPU。?

預取器的作用  

預取器顯著提升數(shù)據(jù)加載效率，它在模型訓練需要數(shù)據(jù)之前，預先加載數(shù)據(jù)到緩沖區(qū)，確保GPU始終有數(shù)據(jù)可用。然而，若訓練速度快于數(shù)據(jù)加載和預處理速度，預取緩沖區(qū)可能會被耗盡，導致GPU空閑等待。

性能影響因素  

多個因素會影響數(shù)據(jù)加載性能，包括樣本大小、批量大小、線程數(shù)量、I/O模式、使用的協(xié)議及緩存效果等。

I/O模式的特性  

I/O模式可分為順序和隨機，這取決于具體工作負載及訓練數(shù)據(jù)集的存儲方式。例如，3D U-Net模型的I/O軌跡顯示文件訪問為隨機，而文件內(nèi)部讀取為順序。

數(shù)據(jù)預處理瓶頸  

在處理非結(jié)構(gòu)化數(shù)據(jù)（如圖像、視頻）時，數(shù)據(jù)預處理環(huán)節(jié)可能成為性能瓶頸。研究表明，某些情況下數(shù)據(jù)預處理的功耗可與實際訓練功耗相當。

解決方案  

為應對數(shù)據(jù)預處理瓶頸，出現(xiàn)了多種解決方案，如將預處理任務卸載到GPU或在訓練前離線執(zhí)行部分預處理工作。

檢查點機制

檢查點的重要性  

檢查點是AI模型訓練中不可或缺的環(huán)節(jié)，定期保存模型狀態(tài)（包括模型權(quán)重、優(yōu)化器狀態(tài)等），以防數(shù)據(jù)丟失并確保訓練進度在意外情況下不會完全丟失。此外，檢查點有助于模型調(diào)試和評估，研究人員可通過分析檢查點數(shù)據(jù)監(jiān)控模型訓練過程。

檢查點大小與保存頻率  

檢查點大小主要由模型規(guī)模決定，通常每個參數(shù)需要約4個字節(jié)。檢查點保存頻率需要權(quán)衡數(shù)據(jù)安全性與訓練效率，保存過于頻繁可能增加存儲開銷與GPU等待時間，當前沒有統(tǒng)一的保存頻率標準，具體取決于模型規(guī)模和訓練時長等。

保存流程與性能優(yōu)化  

檢查點保存過程分為三個階段：

• ?訓練暫停，模型狀態(tài)從GPU轉(zhuǎn)移至系統(tǒng)內(nèi)存。
• 模型狀態(tài)序列化。
• 序列化數(shù)據(jù)寫入持久化存儲。?

性能瓶頸與優(yōu)化方法  

保存過程中的性能瓶頸通常在數(shù)據(jù)寫入速度上。優(yōu)化方法包括：

• ?并行化檢查點寫入：將檢查點數(shù)據(jù)分割，由多個GPU并行寫入。
• 內(nèi)存檢?查點保存：將數(shù)據(jù)復制至系統(tǒng)內(nèi)存后，由專用GPU異步寫入。?

存儲方式與壓縮  

AI框架支持文件存儲與對象存儲。文件存儲因高吞吐量和低延遲成為主流選擇，而對象存儲在可擴展性和易用性上表現(xiàn)優(yōu)異。壓縮方法（如量化與剪枝）可以有效減少存儲空間占用。

AI訓練中的存儲系統(tǒng)

文件存儲與對象存儲  

大多數(shù)AI框架支持文件存儲和對象存儲。文件存儲因高吞吐量、低延遲及對RDMA技術(shù)的支持，成為AI訓練的主要選擇。對象存儲在可擴展性和易用性方面具有優(yōu)勢，得益于新型存儲連接器（如Torch.data庫、WebDataset、PyTorch S3連接器等）的出現(xiàn)。

存儲系統(tǒng)性能優(yōu)化  

AI訓練對存儲系統(tǒng)性能要求極高，特別是在數(shù)據(jù)加載和檢查點保存階段。緩存機制、I/O訪問模式和RDMA技術(shù)等都是提升存儲系統(tǒng)性能的重要手段。

其他關(guān)注點  

• GPU直接存儲技術(shù)（GPUDirect Storage）：允許數(shù)據(jù)直接讀取到GPU內(nèi)存，提升加載效率，但尚未廣泛支持。
• 高速網(wǎng)絡：存儲網(wǎng)絡的速度需與存儲設備的寫入速度相匹配，以充分發(fā)揮存儲系統(tǒng)性能。

----------

在AI數(shù)據(jù)管道的各個階段，AI工作負載都需要與存儲系統(tǒng)進行交互。AI生命周期始于數(shù)據(jù)采集階段。如左側(cè)框所示，系統(tǒng)從多個來源收集數(shù)據(jù)并將其存儲在持久化存儲(Persistent Storage)中。隨后是數(shù)據(jù)準備階段，數(shù)據(jù)科學家將原始數(shù)據(jù)轉(zhuǎn)換為適合訓練和推理的格式。這個階段包括數(shù)據(jù)清洗(Data Cleaning)、特征提取(Feature Extraction)和數(shù)據(jù)標注(Data Labeling)等工作。由于每個項目需求不同，數(shù)據(jù)準備沒有統(tǒng)一的標準流程。

在第二個框中顯示的是訓練階段，這時需要訓練模型并進行驗證。此階段性能至關(guān)重要，存儲系統(tǒng)主要用于向GPU提供數(shù)據(jù)讀取服務或執(zhí)行檢查點保存。訓練完成后，模型即可部署用于推理。在推理階段仍需要為GPU提供數(shù)據(jù)，但對計算資源的需求比訓練階段低。此外，還需要快速讀取模型狀態(tài)，以便將其加載到執(zhí)行推理的機器上。

我們將重點討論模型訓練過程，特別是數(shù)據(jù)加載(Data Loading)和檢查點保存(Checkpoint Saving)這兩個環(huán)節(jié)。

首先，讓我們簡要回顧AI訓練工作負載的基本概念。AI訓練數(shù)據(jù)集由多個樣本組成，每個樣本包含模型輸入數(shù)據(jù)。這些數(shù)據(jù)集可以是文本、圖像、視頻或它們的組合形式，可以存儲在單個或多個文件中。訓練初始階段需要從存儲系統(tǒng)讀取數(shù)據(jù)到系統(tǒng)內(nèi)存，這個過程稱為數(shù)據(jù)加載。系統(tǒng)將訓練樣本組織成小批量(Mini-batch)輸入到模型中。這些批量數(shù)據(jù)經(jīng)過模型處理，即前向傳播(Forward Propagation)過程，并計算損失值(Loss Score)來評估模型性能。然后利用這個損失值更新模型權(quán)重。這個過程會不斷重復，直到模型收斂或達到預期精度。完整處理一遍數(shù)據(jù)集被稱為一個訓練周期(Epoch)。某些模型可能需要多個訓練周期才能收斂，這意味著需要多次讀取整個訓練數(shù)據(jù)集。

為防止數(shù)據(jù)丟失，系統(tǒng)會定期保存模型狀態(tài)，這就是檢查點保存過程。在此過程中，AI框架與存儲系統(tǒng)交互，寫入檢查點數(shù)據(jù)，或在發(fā)生故障時通過讀取存儲中的檢查點數(shù)據(jù)進行恢復。

這里展示了數(shù)據(jù)路徑和訓練工作負載的技術(shù)棧，包括數(shù)據(jù)加載和檢查點保存的細節(jié)。如圖所示，最上層是AI模型，最底層是存儲系統(tǒng)。影響訓練效率的因素遍布各個層次。從右下角的存儲系統(tǒng)開始，可以是本地存儲或分布式存儲系統(tǒng)，其上是操作系統(tǒng)層，再上是負責與存儲通信并連接AI框架的客戶端庫。

在數(shù)據(jù)加載方面，AI框架使用數(shù)據(jù)集(Dataset)的概念。訓練基于特定數(shù)據(jù)集進行，并提供相關(guān)API。內(nèi)部實現(xiàn)中，通過連接器（可以是工具、插件或庫）定義對底層存儲的訪問方式。數(shù)據(jù)集處理包括采樣(Sampling)、索引(Indexing)和數(shù)據(jù)轉(zhuǎn)換(Data Transformation)，這些操作在并行引擎上執(zhí)行。在檢查點保存階段，如右側(cè)所示，連接器通過客戶端庫與存儲系統(tǒng)通信。該階段涉及GPU與CPU之間的數(shù)據(jù)復制以及格式化等任務。

深入分析數(shù)據(jù)加載階段的具體流程：數(shù)據(jù)集包含一系列樣本（例如獨立的圖像文件）。首先，系統(tǒng)通過單線程或多線程方式從存儲讀取數(shù)據(jù)到系統(tǒng)內(nèi)存。原始數(shù)據(jù)需要轉(zhuǎn)換為張量(Tensor)格式以供GPU處理。第一步是數(shù)據(jù)解碼，然后應用各種變換和增強操作，如裁剪和尺寸調(diào)整。最后將數(shù)據(jù)傳輸?shù)紾PU。

需要強調(diào)的是，數(shù)據(jù)加載不僅僅是簡單的存儲I/O操作，而是包含了存儲I/O和多個處理階段的復合過程，這些階段執(zhí)行實時數(shù)據(jù)預處理。目前，大多數(shù)預處理任務都在CPU上完成。

在數(shù)據(jù)加載階段，是什么原因?qū)е翯PU等待數(shù)據(jù)而處于空閑狀態(tài)？以PyTorch為例，其內(nèi)置了預取器(Prefetcher)組件。GPU與運行在CPU上的數(shù)據(jù)加載器協(xié)同工作。預取器在模型訓練需要數(shù)據(jù)之前，將數(shù)據(jù)預先加載到緩沖區(qū)。這種預加載機制確保GPU始終有可用數(shù)據(jù)，從而顯著減少空閑時間。用戶可以調(diào)整和監(jiān)控預取器參數(shù)以優(yōu)化系統(tǒng)性能。

訓練開始時，由于需要同時填充預取緩沖區(qū)和讀取第一批樣本，初始耗時較長。此時系統(tǒng)在等待存儲I/O完成。數(shù)據(jù)到達后開始預處理，在第一批樣本預處理完成之前，GPU處于空閑狀態(tài)。預處理完成后，數(shù)據(jù)被移至GPU，開始訓練這一批數(shù)據(jù)，同時數(shù)據(jù)加載器讀取下一批。這構(gòu)成了初始狀態(tài)下的工作流程。

達到穩(wěn)定狀態(tài)后，可能出現(xiàn)兩種情況。理想情況下，訓練時間大于或等于讀取時間（包括存儲I/O和處理時間）。這種情況下，預取器能夠保持數(shù)據(jù)持續(xù)供應。但如果訓練時間小于存儲和預處理的總時間，即使預取緩沖區(qū)很大，最終也會耗盡，導致數(shù)據(jù)饑餓(Data Starvation)，使GPU處于等待狀態(tài)。

存儲I/O性能受多個因素影響：樣本大小、批量大小、線程數(shù)量、I/O模式、使用的協(xié)議以及緩存效果。雖然通常認為讀取時間是主要瓶頸，但這種假設并不總是成立。例如，處理文本數(shù)據(jù)（如語言模型）時，由于只需進行分詞(Tokenization)，預處理時間較短。相比之下，處理圖像或視頻數(shù)據(jù)集時，預處理時間可能成為主要瓶頸。

這里是一些具體數(shù)據(jù)。左側(cè)圖表展示了ImageNet數(shù)據(jù)集解碼和轉(zhuǎn)換的成本分析，這是一篇來自MIT的研究論文。頂部柱狀圖中，深灰色區(qū)域代表總訓練耗時，包括訓練、讀取和預處理過程。中間柱狀圖顯示了數(shù)據(jù)讀取和預處理的時間開銷。數(shù)據(jù)表明，在這個工作負載中，GPU處理并非性能瓶頸，反而是存儲系統(tǒng)或數(shù)據(jù)預處理成為關(guān)鍵限制。然而，底部柱狀圖顯示的實際數(shù)據(jù)讀取時間非常短，這意味著對于該特定模型，數(shù)據(jù)處理環(huán)節(jié)是主要性能瓶頸。

Meta數(shù)據(jù)中心的研究也得出類似觀察結(jié)果。右側(cè)圖表展示了三種不同推薦模型的功耗分析。在底部柱狀圖中，推薦模型三的預處理功耗（藍色區(qū)域）與實際訓練功耗幾乎相當。他們還報告稱GPU無法快速處理數(shù)據(jù)以滿足計算需求。這充分表明，尤其在處理圖像和視頻數(shù)據(jù)的工作流中，預處理階段可能極其消耗計算資源。在這一研究領(lǐng)域，已經(jīng)存在專業(yè)的解決方案，例如專門的數(shù)據(jù)加載器可將預處理任務卸載到GPU，或在訓練前離線執(zhí)行部分預處理工作。

接下來分析I/O模式。訓練工作流在數(shù)據(jù)加載階段會產(chǎn)生順序和隨機I/O模式，本次我們重點關(guān)注讀取模式。每個訓練樣本都需要完整讀取，例如對于圖像數(shù)據(jù)，需讀取整個圖像文件。但樣本的檢索是隨機的，這導致樣本間形成隨機訪問模式。訓練通常涉及多個訓練周期(epoch)，每個周期數(shù)據(jù)加載器都會讀取整個數(shù)據(jù)集。為確保訓練過程的隨機性，數(shù)據(jù)總是被打亂并以隨機批次讀取。如果數(shù)據(jù)集可完全載入內(nèi)存，則可直接從內(nèi)存提供數(shù)據(jù)，避免冗余I/O操作。

底部圖表展示了3D U-Net模型的I/O軌跡，這是醫(yī)學圖像分割領(lǐng)域的標準基準測試。每個文件存儲一個醫(yī)學圖像樣本，平均文件大小約150MB。Y軸顯示唯一文件標識，X軸代表時間。不同顏色表示不同文件，說明文件訪問是隨機進行的。但若細究單個文件的訪問模式，會發(fā)現(xiàn)文件內(nèi)部讀取是順序的。

情況并非總是如此，因為有時訓練數(shù)據(jù)會將多個樣本聚合存儲在單一數(shù)據(jù)集中。這里我展示了來自訓練基準的推薦模型I/O訪問模式。這是一個包含大量樣本的單一訓練文件，文件規(guī)模達3.8TB。訓練過程中，Y軸代表文件偏移量，X軸表示時間。

工作負載讀取訓練數(shù)據(jù)、執(zhí)行訓練，隨后暫停以評估模型精度，如此往復。從藍色條可見，訓練文件訪問看似順序。然而放大觀察，會發(fā)現(xiàn)工作負載實際上通過滑動窗口(sliding window)進行隨機數(shù)據(jù)訪問。這個基準會在存儲系統(tǒng)中產(chǎn)生大量小規(guī)模I/O。因此，訪問模式可以是順序的，也可以是隨機的，這取決于具體工作負載和訓練數(shù)據(jù)集的存儲方式。

檢查點(Checkpoint)保存對于大型模型訓練至關(guān)重要，訓練過程可能持續(xù)數(shù)天甚至數(shù)周。檢查點保存確保即使出現(xiàn)電力故障或系統(tǒng)崩潰等意外情況，也不會導致全部訓練進展丟失。這是一個周期性過程，用于保存模型當前狀態(tài)，包括模型權(quán)重（已學習的參數(shù)）和優(yōu)化器狀態(tài)。檢查點還可能包含隨機狀態(tài)、調(diào)度器狀態(tài)及其他相關(guān)信息。訓練時間越長，模型權(quán)重的價值就越高，因此避免丟失這些權(quán)重、重新開始訓練的需求變得尤為緊迫。

檢查點保存有多重目的，容錯是最主要原因，但同時也可用于模型調(diào)試和評估。通過監(jiān)控訓練過程，研究者可以判斷模型是否正在收斂或發(fā)散。若發(fā)現(xiàn)模型性能未如預期改善，可回退到表現(xiàn)更好的先前檢查點。通常，檢查點在訓練過程中被保留，后續(xù)模型可以還原到任何先前版本。根據(jù)故障具體情況，可從任意檢查點恢復模型。

檢查點過程的一個關(guān)鍵方面是檢查點大小，這取決于模型規(guī)模，而非訓練數(shù)據(jù)集、GPU內(nèi)存或GPU數(shù)量。一般經(jīng)驗法則是每個參數(shù)約需四個字節(jié)。假設每個模型參數(shù)使用兩個字節(jié)表示低精度值，另有十二個字節(jié)用于優(yōu)化器和其他狀態(tài)信息。以GPT-3為例，其1750億參數(shù)模型生成的檢查點數(shù)據(jù)約為2.4TB；Megatron-Turing NLG模型擁有5300億參數(shù)，其檢查點存儲需求約為7.4TB。檢查點大小具有可預測性，在固定硬件帶寬下，檢查點保存吞吐量可精確估算。檢查點過程代價高昂，因為在完成保存之前會暫停訓練。

檢查點保存過程包括三個階段：首先，模型在GPU上運行。訓練被暫停，模型狀態(tài)從GPU轉(zhuǎn)移到系統(tǒng)內(nèi)存，進行序列化，隨后保存到持久存儲。發(fā)生故障或需要恢復時，將從存儲中讀取檢查點數(shù)據(jù)到系統(tǒng)內(nèi)存，反序列化，然后將模型狀態(tài)移回GPU內(nèi)存。

檢查點可以保存在單個或多個文件中，這取決于模型并行策略。AI領(lǐng)域存在兩種并行性類型：數(shù)據(jù)并行(Data Parallelism)，即每個GPU持有完整模型副本并處理不同數(shù)據(jù)批次。通常由一個GPU（稱為"rank zero"）負責寫入檢查點。在這種情況下，無需寫入所有GPU的全部內(nèi)存內(nèi)容，這也是當前的默認行為。

對于超大模型，如果無法裝載到單個GPU，可將模型分割到多個GPU，每個GPU寫入其模型參數(shù)部分的檢查點。這需要精確協(xié)調(diào)以確保模型各部分正確保存。目前也出現(xiàn)許多混合方法，結(jié)合數(shù)據(jù)和模型并行性。這種情況下，檢查點機制可能更為復雜，涉及多個GPU寫入模型不同部分，具體取決于并行策略。在檢查點場景中，多個GPU可參與寫入，但檢查點文件始終由單個線程順序?qū)懭?。恢復檢查點時，所有GPU從存儲讀取檢查點數(shù)據(jù)。

檢查點保存極其重要，正如前文提到的TB級檢查點數(shù)據(jù)。存儲性能在此至關(guān)重要。觀察當前趨勢，特別是大型語言模型，模型規(guī)模持續(xù)增長，檢查點大小也隨之增長。更大模型需要更多GPU，故障概率也隨之提高。Meta最新研究表明，檢查點保存可能使訓練速度降低43%，這一數(shù)字凸顯了高效檢查點對訓練的重要性。存儲系統(tǒng)應優(yōu)化以高效寫入檢查點數(shù)據(jù)，減少檢查點過程中的GPU等待時間。

Meta同一研究還報告，完整恢復開銷占總訓練時間的12%。Llama 3模型訓練論文也強調(diào)了類似挑戰(zhàn)，指出高突發(fā)檢查點寫入可能導致整個存儲系統(tǒng)飽和。阿里巴巴報告稱，在大規(guī)模語言模型訓練作業(yè)中，43%的故障發(fā)生在檢查點階段。這些故障可能源于軟件問題（如bug或數(shù)據(jù)錯誤）或硬件問題（如網(wǎng)絡故障）。研究者希望增加檢查點頻率，但同時也要權(quán)衡檢查點保存的高額成本。

檢查點保存性能直接影響檢查點頻率。低頻檢查點策略缺乏吸引力。因此，高效檢查點保存對大規(guī)模訓練至關(guān)重要。當前檢查點方法往往導致GPU停頓。以數(shù)據(jù)并行檢查點策略為例，當使用torch.save作為默認檢查點機制時，訓練會被暫停。檢查點數(shù)據(jù)被移動到內(nèi)存，序列化，寫入硬盤。這一過程產(chǎn)生大量順序?qū)懭?，使GPU處于空閑等待狀態(tài)。學術(shù)界和產(chǎn)業(yè)集群已提出多種優(yōu)化方法，旨在減少檢查點保存期間的GPU等待時間。雖然無法詳盡列舉所有方法，但我將重點闡述幾種常見技術(shù)。

第一種技術(shù)是并行化檢查點寫入（Parallel Checkpoint Writing）。這種方法利用數(shù)據(jù)并行性，因為每個GPU持有相同的檢查點數(shù)據(jù)。它將檢查點創(chuàng)建分配給多個并行數(shù)據(jù)并行的GPU，使得每個GPU僅寫入檢查點文件的一部分，而非由單個GPU完成。因此，復制到存儲的數(shù)據(jù)規(guī)模和寫入量減少，您將從更多并行存儲I/O中受益，而不是生成單一的存儲I/O。基于這些變化，檢查點保存的性能和效率得到提升。然而，需要注意的是，這種方法會對存儲產(chǎn)生大量并行I/O，因此您的存儲系統(tǒng)必須具備可擴展性，并能高效處理這些并行I/O。此外，每個單獨文件生成的檢查點數(shù)據(jù)量也變小。

第二種優(yōu)化是內(nèi)存檢查點保存（In Memory Checkpointing）。這種技術(shù)解決了檢查點保存過程中硬盤或網(wǎng)絡I/O成為瓶頸的情況。在這種情況下，您暫停檢查點，將數(shù)據(jù)復制到系統(tǒng)內(nèi)存中，然后繼續(xù)訓練，而不等待檢查點數(shù)據(jù)寫入持久存儲。之后，專用GPU異步地將此檢查點數(shù)據(jù)寫入持久存儲。采用這種方法存在丟失檢查點數(shù)據(jù)的風險，但文獻提供了確保持久性的方法。通常，至少確保先前的檢查點已寫入存儲。此外，還需權(quán)衡：與等待硬盤相比，丟失一個檢查點是否值得？

我們已討論數(shù)據(jù)加載和檢查點保存，現(xiàn)在聚焦存儲部分。大多數(shù)現(xiàn)有AI框架支持文件存儲（File Storage）和對象存儲（Object Storage）訪問，用于數(shù)據(jù)加載和檢查點保存。為提供高吞吐量和低延遲，存儲集群通常配備高性能閃存存儲。大多數(shù)AI訓練工作負載使用基于文件的存儲系統(tǒng)來存儲和訪問訓練數(shù)據(jù)，以確保昂貴的GPU得以高效利用。AI框架設計為訪問文件系統(tǒng)，期待使用文件樣式的API來加載和恢復檢查點。其主要原因在于性能。對象存儲由于其高且不可預測的延遲，阻礙了其在AI訓練工作流中的直接使用。對象存儲的典型部署模型涉及一個更快的緩存層，可以是本地驅(qū)動器或緩存，將數(shù)據(jù)從對象層移動到為GPU提供支持的更快層。

盡管存在更快速的對象存儲解決方案，但基于文件的解決方案在為GPU提供數(shù)據(jù)時仍占主導地位。然而，對象存儲具有顯著優(yōu)勢，包括可擴展性和容量。隨著數(shù)據(jù)集規(guī)模的持續(xù)增長，許多AI應用正在云中開發(fā)，導致基于云的AI應用增加。此外，對象存儲解決方案用戶友好，運行在用戶空間，使數(shù)據(jù)科學家能夠避免處理遠程存儲的復雜性。近期，我們看到許多新的存儲連接器，使得高效訪問對象存儲成為可能，表明對象存儲在AI訓練工作負載中的使用案例正在增加。

如果查看PyTorch如何使用不同連接器訪問基于文件和對象的存儲解決方案，可以概括當前的技術(shù)格局。所謂存儲連接器，是指能夠使AI框架訪問存儲系統(tǒng)讀寫數(shù)據(jù)的特殊工具或庫。這些連接器旨在架起AI框架與各種存儲解決方案之間的橋梁，確保以優(yōu)化方式訪問存儲。

可以看到，頂部是PyTorch框架，下方有多種連接器選項。選擇合適的工具可能令人困惑，因此我將重點介紹最新和最流行的幾種。首先，從左側(cè)的三個黃色框來看對象存儲連接器。

PyTorch提供Torch.data庫，包含數(shù)據(jù)加載工具?？梢耘渲闷渲苯訌膶ο蟠鎯ψx取數(shù)據(jù)并進行流式傳輸，支持S3及其他對象存儲協(xié)議。Web數(shù)據(jù)集是一個示例，允許直接從云存儲訪問TAR文件，無需解壓。另一個選項是PyTorch的S3連接器，支持檢查點保存和數(shù)據(jù)加載。該連接器允許用戶以順序和隨機方式訪問存儲，使用高性能庫，顯著提升標準GET操作效率。它包含各種優(yōu)化，如高級網(wǎng)絡管理和有效處理重試與超時，從而加快檢查點計算。

第三個選項是基于FUSE的文件系統(tǒng)客戶端，允許用戶將S3存儲桶掛載為本地文件系統(tǒng)。這使應用程序能夠使用標準文件操作（如打開和讀?。┰L問S3桶中的對象存儲?？捎眠x項眾多，Mountpoint-S3是最新的之一。這些基于FUSE的客戶端在性能和POSIX兼容性間進行平衡。值得注意的是，它們提供有限的POSIX兼容性，僅支持少數(shù)易于映射到對象協(xié)議API的文件操作。例如，大多數(shù)不支持重命名或修改現(xiàn)有文件，因為S3和對象協(xié)議不允許。然而，其POSIX兼容性足以滿足AI框架在數(shù)據(jù)加載和檢查點保存方面的需求。

第二個連接器是基于文件的存儲，這是默認方法?？蓤?zhí)行標準的POSIX文件API調(diào)用，如讀取和寫入。幾乎所有AI框架都偏好這種方法。例如，如果您使用Torch.save、Torch.load或PyTorch的分布式檢查點模塊，它們都與基于文件的存儲一起使用。此外，PyTorch提供了一整套優(yōu)化庫，用于使用這些文件API訪問不同類型的數(shù)據(jù)集，如圖像數(shù)據(jù)集和文本數(shù)據(jù)集?；谖募拇鎯ο到y(tǒng)的主要優(yōu)點之一是能夠使用RDMA（遠程直接內(nèi)存訪問）技術(shù)，這使得框架在訪問文件系統(tǒng)時能夠透明地受益于RDMA。

RDMA到GPU內(nèi)存是另一回事。在左側(cè)，我展示了哪些連接器支持RDMA到GPU內(nèi)存。這項技術(shù)稱為GPU直接存儲（GPUDirect Storage），目前尚未得到像PyTorch這樣的框架廣泛支持。要啟用此功能，需要額外的配置和集成。這項技術(shù)僅在使用基于文件的存儲解決方案時存在。例如，HSDLY是一個數(shù)據(jù)加載庫，支持GPU直接存儲，允許將數(shù)據(jù)集直接讀取到GPU內(nèi)存中。另一個支持GPU直接存儲的庫是Kyo。這兩個連接器都使用K文件API進行GPU直接存儲。

這里展示了基于文件的連接器和基于對象的連接器之間的高層次架構(gòu)差異。以NFS作為文件表示的示例，比較對象連接器與基于文件的連接器，其中Torch Data 3連接器和Mountpoint作為對象連接器的代表。盡管存在諸多差異，但可以將它們分為三大類。

第一個是緩存。緩存是主要區(qū)別之一?；谖募拇鎯Φ靡嬗诓僮飨到y(tǒng)頁面緩存，這避免了在數(shù)據(jù)集適合緩存時對同一數(shù)據(jù)的重復訪問。這也改善了順序訪問與讀取頭的關(guān)系。相反，對象連接器在用戶空間運行，并繞過虛擬文件系統(tǒng)層，因此它們無法利用操作系統(tǒng)頁面緩存。因此，這些對象連接器必須部署自己的獨立緩存機制。例如，Mountpoint S3提供讀取緩存。

第二個區(qū)別在于I/O訪問模式，這非常重要，并可能影響存儲吞吐量。例如，通過S3連接器或Mountpoint，我們看到這兩種解決方案都得益于異步S3客戶端，生成許多小的并行I/O到存儲。它們不是發(fā)送單個GET請求，而是創(chuàng)建多個并行范圍GET請求。底層庫使用多部分上傳，而不是執(zhí)行單個PUT操作。我們觀察到數(shù)據(jù)以大約8 MB的部分大小被讀取或?qū)懭?，并且進行了大量的異步I/O。這是這兩個連接器的一項優(yōu)化。

第三個區(qū)別是協(xié)議。這些連接器使用不同的協(xié)議。對象協(xié)議的延遲較高，認證和授權(quán)過程在對象存儲中非常昂貴。對于每個請求，您必須執(zhí)行IAM認證和授權(quán)，這可能會產(chǎn)生高昂的成本。此外，正如前面的幻燈片提到的，基于文件的存儲解決方案受益于RDMA技術(shù)，而對象存儲解決方案則不然。

總結(jié)我們今天討論的內(nèi)容：我們談到了數(shù)據(jù)加載，這不僅包括存儲I/O，還包括數(shù)據(jù)轉(zhuǎn)換。I/O訪問模式依賴于模型和數(shù)據(jù)集。存儲系統(tǒng)必須提供高吞吐量和低延遲，以確保數(shù)據(jù)盡可能快速地送入GPU。

檢查點保存至關(guān)重要；大型模型需要高讀寫帶寬，以高效保存和恢復檢查點。檢查點文件可以保存在一個或多個文件中，但每個檢查點文件由單個寫入者寫入。同時，必須注意，累積的檢查點可能會產(chǎn)生顯著的存儲需求，尤其是在頻繁檢查點保存的情況下。因此，存儲解決方案必須高效處理這些需求。

我們還討論了文件存儲和對象存儲；目前，AI框架期望使用類似文件的接口。然而，近年來，對訪問對象存儲解決方案的支持明顯增加。

----------

問題：訓練結(jié)束后，檢查點（Checkpoint）是否會保留？

回答：這取決于具體情況，不能一概而論。通常，在訓練過程中會保留檢查點。由于模型調(diào)試、評估等原因，您可能希望在訓練結(jié)束后繼續(xù)保留這些檢查點。如果檢查點僅用于容錯，則可以不必長期保留?？偟膩碚f，這取決于具體用例。在模型調(diào)試、驗證等工作中，檢查點往往會保留更長時間。

問題：塊存儲（Block Storage）真的有適用場景嗎？

回答：在人工智能訓練方面，目前觀察較少。您可以使用塊存儲，但訓練框架通常期望類似文件系統(tǒng)的訪問方式。我尚未看到對象存儲（Object Storage）在訓練過程中的廣泛使用。通常，您需要使用文件系統(tǒng)（File System）或?qū)ο蟠鎯?。這是當前的技術(shù)觀察。

問題：檢查點保存的頻率是多少？

回答：這是一個常見的問題。根據(jù)在線資源和文獻，檢查點保存通常發(fā)生在每個訓練周期（Epoch）結(jié)束后。在完成對訓練數(shù)據(jù)集的完整遍歷時進行檢查點保存。對于大型語言模型（LLMs），通常在某些訓練步驟后進行檢查點保存。例如，Meta在其研究論文中提到每30分鐘進行一次檢查點保存。然而，這難以作為通用標準。檢查點保存頻率可能是每30分鐘一次，或每幾個小時一次。這很大程度上取決于寫入吞吐量以及根據(jù)模型規(guī)模進行檢查點保存的需求。遺憾的是，很難給出一個確切的數(shù)字。

問題：檢查點機制在SSD和HDD之間有何差異？

回答：在人工智能場景中，檢查點保存需要盡快完成。檢查點保存通常產(chǎn)生順序?qū)懭?，而非隨機寫入。盡管HDD可以進行順序I/O，但AI框架通常與SSD或高速NVMe驅(qū)動器交互。HDD通常用于備份或存檔。對于即時存儲需求，推薦使用快速存儲層，如NVMe SSD。如果需要長期保存，可以將檢查點后續(xù)移動到HDD，但目前觀察顯示，訓練階段主要使用閃存存儲。

問題：訓練中I/O性能的典型要求是什么，具體包括IOPS和每秒字節(jié)數(shù)？

回答：對于數(shù)據(jù)加載或文本類工作負載的LLMs，I/O性能要求遠低于圖像或視頻數(shù)據(jù)集。根據(jù)MLPerf最近的基準測試，訓練用于醫(yī)學圖像分割的3D U-Net模型時，每個GPU的平均讀取吞吐量約為2.7至3 GB/s。相比之下，對于LLMs或文本類數(shù)據(jù)，通常約為每個GPU每秒1 MB，因為這類模型通常是計算密集型而非存儲密集型。對于圖像或視頻工作負載，存儲要求的每GPU帶寬可能達到數(shù)GB/s。

問題：為什么檢查點保存需要序列化？

回答：檢查點保存需要序列化，是因為需要將GPU上的張量（Tensor）數(shù)據(jù)提取出來，通常需要將這些數(shù)據(jù)序列化為數(shù)據(jù)流，然后以單個文件形式寫入存儲。這是基于當前的技術(shù)理解，但可能需要進一步確認。

問題：您認為S3對象存儲會采用類似RDMA的技術(shù)，以便更快地直接讀寫數(shù)據(jù)到GPU內(nèi)存嗎？

回答：目前，對象存儲尚未支持RDMA技術(shù)。我尚未在相關(guān)文獻中看到相關(guān)實現(xiàn)。然而，隨著人工智能框架中對象存儲使用的增加，人們可能會開始優(yōu)化存儲和網(wǎng)絡性能。據(jù)目前所知，現(xiàn)有的對象存儲解決方案中尚未支持RDMA。

問題：每個GPU的存儲吞吐量（Storage Throughput）的經(jīng)驗法則是什么？

回答：在數(shù)據(jù)加載階段，文本數(shù)據(jù)的存儲吞吐量通常在幾兆字節(jié)到更少的范圍內(nèi)，具體取決于數(shù)據(jù)量。對于圖像或視頻類工作負載，吞吐量可達到幾GB。以醫(yī)學圖像工作負載為例，每個GPU的讀取吞吐量約為3 GB/s，這能使GPU利用率達到90%。醫(yī)學圖像通常對存儲吞吐量要求最高。在檢查點（Checkpoint）保存方面，存儲吞吐量取決于寫入存儲的速度需求。顯然，更高的吞吐量意味著更好的性能，但難以給出具體的吞吐量標準。

問題：由于檢查點文件較大，許多環(huán)境是否會集中使用磁帶存儲（Tape Storage）以保留多個檢查點副本？

回答：檢查點的初始寫入必須在高速存儲介質(zhì)上進行，因為整個檢查點數(shù)據(jù)需要寫入遠程存儲并進行冗余復制。因此，第一個檢查點需要存儲在高速存儲層。幾乎所有檢查點都會寫入這一快速存儲層。如遇存儲容量限制，有幾種處理策略：出于容錯目的，可以刪除早期檢查點；若用于模型調(diào)試或評估且存儲空間受限，可將早期檢查點遷移至低速存儲。磁帶存儲通常不用于訓練階段，但訓練完成后，可用于長期備份。

問題：存儲如何處理序列化（Serialization）的檢查點保存？是單線程還是將作業(yè)分配給單個存儲設備？

回答：默認情況下為單線程模式。序列化檢查點數(shù)據(jù)后，單線程將其寫入存儲。對于內(nèi)存檢查點保存，數(shù)據(jù)移動到系統(tǒng)內(nèi)存后，持久存儲的刷新機制可能會有所不同。如果采用異步寫入機制，處理方式可能有所變化。通常，這些框架的默認行為是單線程寫入，每個檢查點文件由單獨的線程寫入。

問題：高速網(wǎng)絡是否顯著增強了I/O和檢查點保存過程？

回答：檢查點保存發(fā)生在GPU權(quán)重同步之后。此處討論的網(wǎng)絡是指GPU與服務器和存儲間的網(wǎng)絡連接。性能可能受網(wǎng)絡速度和存儲設備驅(qū)動器速度的共同影響。兩者需要平衡：不能有超高速存儲驅(qū)動器卻配備低速網(wǎng)絡，反之亦然。因此，存儲網(wǎng)絡的高速確實至關(guān)重要。

問題：檢查點保存的頻率如何？是否可以進行檢查點緩存以便快速恢復？

回答：檢查點緩存確實具有價值，尤其是最新的檢查點。通常不會頻繁讀取檢查點，僅在發(fā)生故障時使用。除調(diào)試或驗證模型時，如發(fā)現(xiàn)模型性能不佳或未收斂，才會讀取檢查點。根據(jù)具體情況，檢查點緩存很有意義，特別是在故障恢復場景中?？蓪⒆钚聶z查點保存在緩存中，需要時可快速恢復。但在選擇緩存哪個檢查點時，需謹慎考慮模型準確性，可能需要比較前一個、前兩個或前三個檢查點，這取決于具體用例。

問題：檢查點保存中壓縮的好處是什么？

回答：傳統(tǒng)無損壓縮技術(shù)（如Gzip）在人工智能訓練工作負載中壓縮效果有限，因為檢查點數(shù)據(jù)高度隨機，缺乏重復模式。但存在其他有效技術(shù)：

• 量化（Quantization）：將浮點數(shù)轉(zhuǎn)換為低精度，減少位使用和存儲需求。例如，從單精度轉(zhuǎn)換為半精度。
• 剪枝（Pruning）：去除模型中不重要的部分，減小檢查點體積。

這兩種技術(shù)是常見的檢查點壓縮方法，能有效減少存儲開銷。

參考資料：Kaynar, Ugur. "AI Storage: The Critical Role of Storage in Optimizing AI Training Workloads." YouTube, October 30, 2024. https://www.youtube.com/watch?v=vycUUhlRmVs.

本文轉(zhuǎn)載自 ??Andy730??，作者： 常華Andy