AI存儲:存儲系統在優化AI訓練中的關鍵作用 精華
數據加載過程與性能優化
數據加載的復雜性
數據加載并不僅限于從存儲讀取數據,它涵蓋了解碼、格式轉換及數據增強等預處理環節。這些步驟通常在CPU上執行,目的是將原始數據轉換為GPU可處理的張量格式。
數據加載的步驟
數據加載流程可以概括為:
- ?從存儲系統讀取數據到系統內存。
- 解碼原始數據。
- 應用數據變換和增強操作(如裁剪、尺寸調整)。
- 將處理后的數據傳輸至GPU。?
預取器的作用
預取器顯著提升數據加載效率,它在模型訓練需要數據之前,預先加載數據到緩沖區,確保GPU始終有數據可用。然而,若訓練速度快于數據加載和預處理速度,預取緩沖區可能會被耗盡,導致GPU空閑等待。
性能影響因素
多個因素會影響數據加載性能,包括樣本大小、批量大小、線程數量、I/O模式、使用的協議及緩存效果等。
I/O模式的特性
I/O模式可分為順序和隨機,這取決于具體工作負載及訓練數據集的存儲方式。例如,3D U-Net模型的I/O軌跡顯示文件訪問為隨機,而文件內部讀取為順序。
數據預處理瓶頸
在處理非結構化數據(如圖像、視頻)時,數據預處理環節可能成為性能瓶頸。研究表明,某些情況下數據預處理的功耗可與實際訓練功耗相當。
解決方案
為應對數據預處理瓶頸,出現了多種解決方案,如將預處理任務卸載到GPU或在訓練前離線執行部分預處理工作。
檢查點機制
檢查點的重要性
檢查點是AI模型訓練中不可或缺的環節,定期保存模型狀態(包括模型權重、優化器狀態等),以防數據丟失并確保訓練進度在意外情況下不會完全丟失。此外,檢查點有助于模型調試和評估,研究人員可通過分析檢查點數據監控模型訓練過程。
檢查點大小與保存頻率
檢查點大小主要由模型規模決定,通常每個參數需要約4個字節。檢查點保存頻率需要權衡數據安全性與訓練效率,保存過于頻繁可能增加存儲開銷與GPU等待時間,當前沒有統一的保存頻率標準,具體取決于模型規模和訓練時長等。
保存流程與性能優化
檢查點保存過程分為三個階段:
- ?訓練暫停,模型狀態從GPU轉移至系統內存。
- 模型狀態序列化。
- 序列化數據寫入持久化存儲。?
性能瓶頸與優化方法
保存過程中的性能瓶頸通常在數據寫入速度上。優化方法包括:
- ?并行化檢查點寫入:將檢查點數據分割,由多個GPU并行寫入。
- 內存檢?查點保存:將數據復制至系統內存后,由專用GPU異步寫入。?
存儲方式與壓縮
AI框架支持文件存儲與對象存儲。文件存儲因高吞吐量和低延遲成為主流選擇,而對象存儲在可擴展性和易用性上表現優異。壓縮方法(如量化與剪枝)可以有效減少存儲空間占用。
AI訓練中的存儲系統
文件存儲與對象存儲
大多數AI框架支持文件存儲和對象存儲。文件存儲因高吞吐量、低延遲及對RDMA技術的支持,成為AI訓練的主要選擇。對象存儲在可擴展性和易用性方面具有優勢,得益于新型存儲連接器(如Torch.data庫、WebDataset、PyTorch S3連接器等)的出現。
存儲系統性能優化
AI訓練對存儲系統性能要求極高,特別是在數據加載和檢查點保存階段。緩存機制、I/O訪問模式和RDMA技術等都是提升存儲系統性能的重要手段。
其他關注點
- GPU直接存儲技術(GPUDirect Storage):允許數據直接讀取到GPU內存,提升加載效率,但尚未廣泛支持。
- 高速網絡:存儲網絡的速度需與存儲設備的寫入速度相匹配,以充分發揮存儲系統性能。
----------
在AI數據管道的各個階段,AI工作負載都需要與存儲系統進行交互。AI生命周期始于數據采集階段。如左側框所示,系統從多個來源收集數據并將其存儲在持久化存儲(Persistent Storage)中。隨后是數據準備階段,數據科學家將原始數據轉換為適合訓練和推理的格式。這個階段包括數據清洗(Data Cleaning)、特征提取(Feature Extraction)和數據標注(Data Labeling)等工作。由于每個項目需求不同,數據準備沒有統一的標準流程。
在第二個框中顯示的是訓練階段,這時需要訓練模型并進行驗證。此階段性能至關重要,存儲系統主要用于向GPU提供數據讀取服務或執行檢查點保存。訓練完成后,模型即可部署用于推理。在推理階段仍需要為GPU提供數據,但對計算資源的需求比訓練階段低。此外,還需要快速讀取模型狀態,以便將其加載到執行推理的機器上。
我們將重點討論模型訓練過程,特別是數據加載(Data Loading)和檢查點保存(Checkpoint Saving)這兩個環節。
首先,讓我們簡要回顧AI訓練工作負載的基本概念。AI訓練數據集由多個樣本組成,每個樣本包含模型輸入數據。這些數據集可以是文本、圖像、視頻或它們的組合形式,可以存儲在單個或多個文件中。訓練初始階段需要從存儲系統讀取數據到系統內存,這個過程稱為數據加載。系統將訓練樣本組織成小批量(Mini-batch)輸入到模型中。這些批量數據經過模型處理,即前向傳播(Forward Propagation)過程,并計算損失值(Loss Score)來評估模型性能。然后利用這個損失值更新模型權重。這個過程會不斷重復,直到模型收斂或達到預期精度。完整處理一遍數據集被稱為一個訓練周期(Epoch)。某些模型可能需要多個訓練周期才能收斂,這意味著需要多次讀取整個訓練數據集。
為防止數據丟失,系統會定期保存模型狀態,這就是檢查點保存過程。在此過程中,AI框架與存儲系統交互,寫入檢查點數據,或在發生故障時通過讀取存儲中的檢查點數據進行恢復。
這里展示了數據路徑和訓練工作負載的技術棧,包括數據加載和檢查點保存的細節。如圖所示,最上層是AI模型,最底層是存儲系統。影響訓練效率的因素遍布各個層次。從右下角的存儲系統開始,可以是本地存儲或分布式存儲系統,其上是操作系統層,再上是負責與存儲通信并連接AI框架的客戶端庫。
在數據加載方面,AI框架使用數據集(Dataset)的概念。訓練基于特定數據集進行,并提供相關API。內部實現中,通過連接器(可以是工具、插件或庫)定義對底層存儲的訪問方式。數據集處理包括采樣(Sampling)、索引(Indexing)和數據轉換(Data Transformation),這些操作在并行引擎上執行。在檢查點保存階段,如右側所示,連接器通過客戶端庫與存儲系統通信。該階段涉及GPU與CPU之間的數據復制以及格式化等任務。
深入分析數據加載階段的具體流程:數據集包含一系列樣本(例如獨立的圖像文件)。首先,系統通過單線程或多線程方式從存儲讀取數據到系統內存。原始數據需要轉換為張量(Tensor)格式以供GPU處理。第一步是數據解碼,然后應用各種變換和增強操作,如裁剪和尺寸調整。最后將數據傳輸到GPU。
需要強調的是,數據加載不僅僅是簡單的存儲I/O操作,而是包含了存儲I/O和多個處理階段的復合過程,這些階段執行實時數據預處理。目前,大多數預處理任務都在CPU上完成。
在數據加載階段,是什么原因導致GPU等待數據而處于空閑狀態?以PyTorch為例,其內置了預取器(Prefetcher)組件。GPU與運行在CPU上的數據加載器協同工作。預取器在模型訓練需要數據之前,將數據預先加載到緩沖區。這種預加載機制確保GPU始終有可用數據,從而顯著減少空閑時間。用戶可以調整和監控預取器參數以優化系統性能。
訓練開始時,由于需要同時填充預取緩沖區和讀取第一批樣本,初始耗時較長。此時系統在等待存儲I/O完成。數據到達后開始預處理,在第一批樣本預處理完成之前,GPU處于空閑狀態。預處理完成后,數據被移至GPU,開始訓練這一批數據,同時數據加載器讀取下一批。這構成了初始狀態下的工作流程。
達到穩定狀態后,可能出現兩種情況。理想情況下,訓練時間大于或等于讀取時間(包括存儲I/O和處理時間)。這種情況下,預取器能夠保持數據持續供應。但如果訓練時間小于存儲和預處理的總時間,即使預取緩沖區很大,最終也會耗盡,導致數據饑餓(Data Starvation),使GPU處于等待狀態。
存儲I/O性能受多個因素影響:樣本大小、批量大小、線程數量、I/O模式、使用的協議以及緩存效果。雖然通常認為讀取時間是主要瓶頸,但這種假設并不總是成立。例如,處理文本數據(如語言模型)時,由于只需進行分詞(Tokenization),預處理時間較短。相比之下,處理圖像或視頻數據集時,預處理時間可能成為主要瓶頸。
這里是一些具體數據。左側圖表展示了ImageNet數據集解碼和轉換的成本分析,這是一篇來自MIT的研究論文。頂部柱狀圖中,深灰色區域代表總訓練耗時,包括訓練、讀取和預處理過程。中間柱狀圖顯示了數據讀取和預處理的時間開銷。數據表明,在這個工作負載中,GPU處理并非性能瓶頸,反而是存儲系統或數據預處理成為關鍵限制。然而,底部柱狀圖顯示的實際數據讀取時間非常短,這意味著對于該特定模型,數據處理環節是主要性能瓶頸。
Meta數據中心的研究也得出類似觀察結果。右側圖表展示了三種不同推薦模型的功耗分析。在底部柱狀圖中,推薦模型三的預處理功耗(藍色區域)與實際訓練功耗幾乎相當。他們還報告稱GPU無法快速處理數據以滿足計算需求。這充分表明,尤其在處理圖像和視頻數據的工作流中,預處理階段可能極其消耗計算資源。在這一研究領域,已經存在專業的解決方案,例如專門的數據加載器可將預處理任務卸載到GPU,或在訓練前離線執行部分預處理工作。
接下來分析I/O模式。訓練工作流在數據加載階段會產生順序和隨機I/O模式,本次我們重點關注讀取模式。每個訓練樣本都需要完整讀取,例如對于圖像數據,需讀取整個圖像文件。但樣本的檢索是隨機的,這導致樣本間形成隨機訪問模式。訓練通常涉及多個訓練周期(epoch),每個周期數據加載器都會讀取整個數據集。為確保訓練過程的隨機性,數據總是被打亂并以隨機批次讀取。如果數據集可完全載入內存,則可直接從內存提供數據,避免冗余I/O操作。
底部圖表展示了3D U-Net模型的I/O軌跡,這是醫學圖像分割領域的標準基準測試。每個文件存儲一個醫學圖像樣本,平均文件大小約150MB。Y軸顯示唯一文件標識,X軸代表時間。不同顏色表示不同文件,說明文件訪問是隨機進行的。但若細究單個文件的訪問模式,會發現文件內部讀取是順序的。
情況并非總是如此,因為有時訓練數據會將多個樣本聚合存儲在單一數據集中。這里我展示了來自訓練基準的推薦模型I/O訪問模式。這是一個包含大量樣本的單一訓練文件,文件規模達3.8TB。訓練過程中,Y軸代表文件偏移量,X軸表示時間。
工作負載讀取訓練數據、執行訓練,隨后暫停以評估模型精度,如此往復。從藍色條可見,訓練文件訪問看似順序。然而放大觀察,會發現工作負載實際上通過滑動窗口(sliding window)進行隨機數據訪問。這個基準會在存儲系統中產生大量小規模I/O。因此,訪問模式可以是順序的,也可以是隨機的,這取決于具體工作負載和訓練數據集的存儲方式。
檢查點(Checkpoint)保存對于大型模型訓練至關重要,訓練過程可能持續數天甚至數周。檢查點保存確保即使出現電力故障或系統崩潰等意外情況,也不會導致全部訓練進展丟失。這是一個周期性過程,用于保存模型當前狀態,包括模型權重(已學習的參數)和優化器狀態。檢查點還可能包含隨機狀態、調度器狀態及其他相關信息。訓練時間越長,模型權重的價值就越高,因此避免丟失這些權重、重新開始訓練的需求變得尤為緊迫。
檢查點保存有多重目的,容錯是最主要原因,但同時也可用于模型調試和評估。通過監控訓練過程,研究者可以判斷模型是否正在收斂或發散。若發現模型性能未如預期改善,可回退到表現更好的先前檢查點。通常,檢查點在訓練過程中被保留,后續模型可以還原到任何先前版本。根據故障具體情況,可從任意檢查點恢復模型。
檢查點過程的一個關鍵方面是檢查點大小,這取決于模型規模,而非訓練數據集、GPU內存或GPU數量。一般經驗法則是每個參數約需四個字節。假設每個模型參數使用兩個字節表示低精度值,另有十二個字節用于優化器和其他狀態信息。以GPT-3為例,其1750億參數模型生成的檢查點數據約為2.4TB;Megatron-Turing NLG模型擁有5300億參數,其檢查點存儲需求約為7.4TB。檢查點大小具有可預測性,在固定硬件帶寬下,檢查點保存吞吐量可精確估算。檢查點過程代價高昂,因為在完成保存之前會暫停訓練。
檢查點保存過程包括三個階段:首先,模型在GPU上運行。訓練被暫停,模型狀態從GPU轉移到系統內存,進行序列化,隨后保存到持久存儲。發生故障或需要恢復時,將從存儲中讀取檢查點數據到系統內存,反序列化,然后將模型狀態移回GPU內存。
檢查點可以保存在單個或多個文件中,這取決于模型并行策略。AI領域存在兩種并行性類型:數據并行(Data Parallelism),即每個GPU持有完整模型副本并處理不同數據批次。通常由一個GPU(稱為"rank zero")負責寫入檢查點。在這種情況下,無需寫入所有GPU的全部內存內容,這也是當前的默認行為。
對于超大模型,如果無法裝載到單個GPU,可將模型分割到多個GPU,每個GPU寫入其模型參數部分的檢查點。這需要精確協調以確保模型各部分正確保存。目前也出現許多混合方法,結合數據和模型并行性。這種情況下,檢查點機制可能更為復雜,涉及多個GPU寫入模型不同部分,具體取決于并行策略。在檢查點場景中,多個GPU可參與寫入,但檢查點文件始終由單個線程順序寫入。恢復檢查點時,所有GPU從存儲讀取檢查點數據。
檢查點保存極其重要,正如前文提到的TB級檢查點數據。存儲性能在此至關重要。觀察當前趨勢,特別是大型語言模型,模型規模持續增長,檢查點大小也隨之增長。更大模型需要更多GPU,故障概率也隨之提高。Meta最新研究表明,檢查點保存可能使訓練速度降低43%,這一數字凸顯了高效檢查點對訓練的重要性。存儲系統應優化以高效寫入檢查點數據,減少檢查點過程中的GPU等待時間。
Meta同一研究還報告,完整恢復開銷占總訓練時間的12%。Llama 3模型訓練論文也強調了類似挑戰,指出高突發檢查點寫入可能導致整個存儲系統飽和。阿里巴巴報告稱,在大規模語言模型訓練作業中,43%的故障發生在檢查點階段。這些故障可能源于軟件問題(如bug或數據錯誤)或硬件問題(如網絡故障)。研究者希望增加檢查點頻率,但同時也要權衡檢查點保存的高額成本。
檢查點保存性能直接影響檢查點頻率。低頻檢查點策略缺乏吸引力。因此,高效檢查點保存對大規模訓練至關重要。當前檢查點方法往往導致GPU停頓。以數據并行檢查點策略為例,當使用torch.save作為默認檢查點機制時,訓練會被暫停。檢查點數據被移動到內存,序列化,寫入硬盤。這一過程產生大量順序寫入,使GPU處于空閑等待狀態。學術界和產業集群已提出多種優化方法,旨在減少檢查點保存期間的GPU等待時間。雖然無法詳盡列舉所有方法,但我將重點闡述幾種常見技術。
第一種技術是并行化檢查點寫入(Parallel Checkpoint Writing)。這種方法利用數據并行性,因為每個GPU持有相同的檢查點數據。它將檢查點創建分配給多個并行數據并行的GPU,使得每個GPU僅寫入檢查點文件的一部分,而非由單個GPU完成。因此,復制到存儲的數據規模和寫入量減少,您將從更多并行存儲I/O中受益,而不是生成單一的存儲I/O。基于這些變化,檢查點保存的性能和效率得到提升。然而,需要注意的是,這種方法會對存儲產生大量并行I/O,因此您的存儲系統必須具備可擴展性,并能高效處理這些并行I/O。此外,每個單獨文件生成的檢查點數據量也變小。
第二種優化是內存檢查點保存(In Memory Checkpointing)。這種技術解決了檢查點保存過程中硬盤或網絡I/O成為瓶頸的情況。在這種情況下,您暫停檢查點,將數據復制到系統內存中,然后繼續訓練,而不等待檢查點數據寫入持久存儲。之后,專用GPU異步地將此檢查點數據寫入持久存儲。采用這種方法存在丟失檢查點數據的風險,但文獻提供了確保持久性的方法。通常,至少確保先前的檢查點已寫入存儲。此外,還需權衡:與等待硬盤相比,丟失一個檢查點是否值得?
我們已討論數據加載和檢查點保存,現在聚焦存儲部分。大多數現有AI框架支持文件存儲(File Storage)和對象存儲(Object Storage)訪問,用于數據加載和檢查點保存。為提供高吞吐量和低延遲,存儲集群通常配備高性能閃存存儲。大多數AI訓練工作負載使用基于文件的存儲系統來存儲和訪問訓練數據,以確保昂貴的GPU得以高效利用。AI框架設計為訪問文件系統,期待使用文件樣式的API來加載和恢復檢查點。其主要原因在于性能。對象存儲由于其高且不可預測的延遲,阻礙了其在AI訓練工作流中的直接使用。對象存儲的典型部署模型涉及一個更快的緩存層,可以是本地驅動器或緩存,將數據從對象層移動到為GPU提供支持的更快層。
盡管存在更快速的對象存儲解決方案,但基于文件的解決方案在為GPU提供數據時仍占主導地位。然而,對象存儲具有顯著優勢,包括可擴展性和容量。隨著數據集規模的持續增長,許多AI應用正在云中開發,導致基于云的AI應用增加。此外,對象存儲解決方案用戶友好,運行在用戶空間,使數據科學家能夠避免處理遠程存儲的復雜性。近期,我們看到許多新的存儲連接器,使得高效訪問對象存儲成為可能,表明對象存儲在AI訓練工作負載中的使用案例正在增加。
如果查看PyTorch如何使用不同連接器訪問基于文件和對象的存儲解決方案,可以概括當前的技術格局。所謂存儲連接器,是指能夠使AI框架訪問存儲系統讀寫數據的特殊工具或庫。這些連接器旨在架起AI框架與各種存儲解決方案之間的橋梁,確保以優化方式訪問存儲。
可以看到,頂部是PyTorch框架,下方有多種連接器選項。選擇合適的工具可能令人困惑,因此我將重點介紹最新和最流行的幾種。首先,從左側的三個黃色框來看對象存儲連接器。
PyTorch提供Torch.data庫,包含數據加載工具。可以配置其直接從對象存儲讀取數據并進行流式傳輸,支持S3及其他對象存儲協議。Web數據集是一個示例,允許直接從云存儲訪問TAR文件,無需解壓。另一個選項是PyTorch的S3連接器,支持檢查點保存和數據加載。該連接器允許用戶以順序和隨機方式訪問存儲,使用高性能庫,顯著提升標準GET操作效率。它包含各種優化,如高級網絡管理和有效處理重試與超時,從而加快檢查點計算。
第三個選項是基于FUSE的文件系統客戶端,允許用戶將S3存儲桶掛載為本地文件系統。這使應用程序能夠使用標準文件操作(如打開和讀取)訪問S3桶中的對象存儲。可用選項眾多,Mountpoint-S3是最新的之一。這些基于FUSE的客戶端在性能和POSIX兼容性間進行平衡。值得注意的是,它們提供有限的POSIX兼容性,僅支持少數易于映射到對象協議API的文件操作。例如,大多數不支持重命名或修改現有文件,因為S3和對象協議不允許。然而,其POSIX兼容性足以滿足AI框架在數據加載和檢查點保存方面的需求。
第二個連接器是基于文件的存儲,這是默認方法。可執行標準的POSIX文件API調用,如讀取和寫入。幾乎所有AI框架都偏好這種方法。例如,如果您使用Torch.save、Torch.load或PyTorch的分布式檢查點模塊,它們都與基于文件的存儲一起使用。此外,PyTorch提供了一整套優化庫,用于使用這些文件API訪問不同類型的數據集,如圖像數據集和文本數據集。基于文件的存儲系統的主要優點之一是能夠使用RDMA(遠程直接內存訪問)技術,這使得框架在訪問文件系統時能夠透明地受益于RDMA。
RDMA到GPU內存是另一回事。在左側,我展示了哪些連接器支持RDMA到GPU內存。這項技術稱為GPU直接存儲(GPUDirect Storage),目前尚未得到像PyTorch這樣的框架廣泛支持。要啟用此功能,需要額外的配置和集成。這項技術僅在使用基于文件的存儲解決方案時存在。例如,HSDLY是一個數據加載庫,支持GPU直接存儲,允許將數據集直接讀取到GPU內存中。另一個支持GPU直接存儲的庫是Kyo。這兩個連接器都使用K文件API進行GPU直接存儲。
這里展示了基于文件的連接器和基于對象的連接器之間的高層次架構差異。以NFS作為文件表示的示例,比較對象連接器與基于文件的連接器,其中Torch Data 3連接器和Mountpoint作為對象連接器的代表。盡管存在諸多差異,但可以將它們分為三大類。
第一個是緩存。緩存是主要區別之一。基于文件的存儲得益于操作系統頁面緩存,這避免了在數據集適合緩存時對同一數據的重復訪問。這也改善了順序訪問與讀取頭的關系。相反,對象連接器在用戶空間運行,并繞過虛擬文件系統層,因此它們無法利用操作系統頁面緩存。因此,這些對象連接器必須部署自己的獨立緩存機制。例如,Mountpoint S3提供讀取緩存。
第二個區別在于I/O訪問模式,這非常重要,并可能影響存儲吞吐量。例如,通過S3連接器或Mountpoint,我們看到這兩種解決方案都得益于異步S3客戶端,生成許多小的并行I/O到存儲。它們不是發送單個GET請求,而是創建多個并行范圍GET請求。底層庫使用多部分上傳,而不是執行單個PUT操作。我們觀察到數據以大約8 MB的部分大小被讀取或寫入,并且進行了大量的異步I/O。這是這兩個連接器的一項優化。
第三個區別是協議。這些連接器使用不同的協議。對象協議的延遲較高,認證和授權過程在對象存儲中非常昂貴。對于每個請求,您必須執行IAM認證和授權,這可能會產生高昂的成本。此外,正如前面的幻燈片提到的,基于文件的存儲解決方案受益于RDMA技術,而對象存儲解決方案則不然。
總結我們今天討論的內容:我們談到了數據加載,這不僅包括存儲I/O,還包括數據轉換。I/O訪問模式依賴于模型和數據集。存儲系統必須提供高吞吐量和低延遲,以確保數據盡可能快速地送入GPU。
檢查點保存至關重要;大型模型需要高讀寫帶寬,以高效保存和恢復檢查點。檢查點文件可以保存在一個或多個文件中,但每個檢查點文件由單個寫入者寫入。同時,必須注意,累積的檢查點可能會產生顯著的存儲需求,尤其是在頻繁檢查點保存的情況下。因此,存儲解決方案必須高效處理這些需求。
我們還討論了文件存儲和對象存儲;目前,AI框架期望使用類似文件的接口。然而,近年來,對訪問對象存儲解決方案的支持明顯增加。
----------
問題:訓練結束后,檢查點(Checkpoint)是否會保留?
回答:這取決于具體情況,不能一概而論。通常,在訓練過程中會保留檢查點。由于模型調試、評估等原因,您可能希望在訓練結束后繼續保留這些檢查點。如果檢查點僅用于容錯,則可以不必長期保留。總的來說,這取決于具體用例。在模型調試、驗證等工作中,檢查點往往會保留更長時間。
問題:塊存儲(Block Storage)真的有適用場景嗎?
回答:在人工智能訓練方面,目前觀察較少。您可以使用塊存儲,但訓練框架通常期望類似文件系統的訪問方式。我尚未看到對象存儲(Object Storage)在訓練過程中的廣泛使用。通常,您需要使用文件系統(File System)或對象存儲。這是當前的技術觀察。
問題:檢查點保存的頻率是多少?
回答:這是一個常見的問題。根據在線資源和文獻,檢查點保存通常發生在每個訓練周期(Epoch)結束后。在完成對訓練數據集的完整遍歷時進行檢查點保存。對于大型語言模型(LLMs),通常在某些訓練步驟后進行檢查點保存。例如,Meta在其研究論文中提到每30分鐘進行一次檢查點保存。然而,這難以作為通用標準。檢查點保存頻率可能是每30分鐘一次,或每幾個小時一次。這很大程度上取決于寫入吞吐量以及根據模型規模進行檢查點保存的需求。遺憾的是,很難給出一個確切的數字。
問題:檢查點機制在SSD和HDD之間有何差異?
回答:在人工智能場景中,檢查點保存需要盡快完成。檢查點保存通常產生順序寫入,而非隨機寫入。盡管HDD可以進行順序I/O,但AI框架通常與SSD或高速NVMe驅動器交互。HDD通常用于備份或存檔。對于即時存儲需求,推薦使用快速存儲層,如NVMe SSD。如果需要長期保存,可以將檢查點后續移動到HDD,但目前觀察顯示,訓練階段主要使用閃存存儲。
問題:訓練中I/O性能的典型要求是什么,具體包括IOPS和每秒字節數?
回答:對于數據加載或文本類工作負載的LLMs,I/O性能要求遠低于圖像或視頻數據集。根據MLPerf最近的基準測試,訓練用于醫學圖像分割的3D U-Net模型時,每個GPU的平均讀取吞吐量約為2.7至3 GB/s。相比之下,對于LLMs或文本類數據,通常約為每個GPU每秒1 MB,因為這類模型通常是計算密集型而非存儲密集型。對于圖像或視頻工作負載,存儲要求的每GPU帶寬可能達到數GB/s。
問題:為什么檢查點保存需要序列化?
回答:檢查點保存需要序列化,是因為需要將GPU上的張量(Tensor)數據提取出來,通常需要將這些數據序列化為數據流,然后以單個文件形式寫入存儲。這是基于當前的技術理解,但可能需要進一步確認。
問題:您認為S3對象存儲會采用類似RDMA的技術,以便更快地直接讀寫數據到GPU內存嗎?
回答:目前,對象存儲尚未支持RDMA技術。我尚未在相關文獻中看到相關實現。然而,隨著人工智能框架中對象存儲使用的增加,人們可能會開始優化存儲和網絡性能。據目前所知,現有的對象存儲解決方案中尚未支持RDMA。
問題:每個GPU的存儲吞吐量(Storage Throughput)的經驗法則是什么?
回答:在數據加載階段,文本數據的存儲吞吐量通常在幾兆字節到更少的范圍內,具體取決于數據量。對于圖像或視頻類工作負載,吞吐量可達到幾GB。以醫學圖像工作負載為例,每個GPU的讀取吞吐量約為3 GB/s,這能使GPU利用率達到90%。醫學圖像通常對存儲吞吐量要求最高。在檢查點(Checkpoint)保存方面,存儲吞吐量取決于寫入存儲的速度需求。顯然,更高的吞吐量意味著更好的性能,但難以給出具體的吞吐量標準。
問題:由于檢查點文件較大,許多環境是否會集中使用磁帶存儲(Tape Storage)以保留多個檢查點副本?
回答:檢查點的初始寫入必須在高速存儲介質上進行,因為整個檢查點數據需要寫入遠程存儲并進行冗余復制。因此,第一個檢查點需要存儲在高速存儲層。幾乎所有檢查點都會寫入這一快速存儲層。如遇存儲容量限制,有幾種處理策略:出于容錯目的,可以刪除早期檢查點;若用于模型調試或評估且存儲空間受限,可將早期檢查點遷移至低速存儲。磁帶存儲通常不用于訓練階段,但訓練完成后,可用于長期備份。
問題:存儲如何處理序列化(Serialization)的檢查點保存?是單線程還是將作業分配給單個存儲設備?
回答:默認情況下為單線程模式。序列化檢查點數據后,單線程將其寫入存儲。對于內存檢查點保存,數據移動到系統內存后,持久存儲的刷新機制可能會有所不同。如果采用異步寫入機制,處理方式可能有所變化。通常,這些框架的默認行為是單線程寫入,每個檢查點文件由單獨的線程寫入。
問題:高速網絡是否顯著增強了I/O和檢查點保存過程?
回答:檢查點保存發生在GPU權重同步之后。此處討論的網絡是指GPU與服務器和存儲間的網絡連接。性能可能受網絡速度和存儲設備驅動器速度的共同影響。兩者需要平衡:不能有超高速存儲驅動器卻配備低速網絡,反之亦然。因此,存儲網絡的高速確實至關重要。
問題:檢查點保存的頻率如何?是否可以進行檢查點緩存以便快速恢復?
回答:檢查點緩存確實具有價值,尤其是最新的檢查點。通常不會頻繁讀取檢查點,僅在發生故障時使用。除調試或驗證模型時,如發現模型性能不佳或未收斂,才會讀取檢查點。根據具體情況,檢查點緩存很有意義,特別是在故障恢復場景中。可將最新檢查點保存在緩存中,需要時可快速恢復。但在選擇緩存哪個檢查點時,需謹慎考慮模型準確性,可能需要比較前一個、前兩個或前三個檢查點,這取決于具體用例。
問題:檢查點保存中壓縮的好處是什么?
回答:傳統無損壓縮技術(如Gzip)在人工智能訓練工作負載中壓縮效果有限,因為檢查點數據高度隨機,缺乏重復模式。但存在其他有效技術:
- 量化(Quantization):將浮點數轉換為低精度,減少位使用和存儲需求。例如,從單精度轉換為半精度。
- 剪枝(Pruning):去除模型中不重要的部分,減小檢查點體積。
這兩種技術是常見的檢查點壓縮方法,能有效減少存儲開銷。
參考資料: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.
本文轉載自 ??Andy730??,作者: 常華Andy
