本文為英文版的機器翻譯版本,如內容有任何歧義或不一致之處,概以英文版為準。
模型平行化概念
模型平行化是一種分散式訓練方法,其中深度學習 (DL) 模型被分割到多個 GPU 和執行個體上。SageMaker 模型平行化程式庫 v2 (SMP v2) 與原生 PyTorch API 及功能相容。這可讓您輕鬆地將 PyTorch 全碎片資料平行化 (FSDP) 訓練指令碼調整為 SageMaker 訓練平台,並利用 SMP v2 提供的效能改善。此簡介頁面提供有關模型平行化的高階概觀,以及描述它如何幫助克服在訓練通常規模非常大的深度學習 (DL) 模型時出現的問題。它還提供 SageMaker 模型平行化程式庫供應的範例,以協助管理模型平行化策略和記憶體消耗。
什麼是模型平行化?
增加深度學習模型 (圖層和參數) 的大小可以為複雜的任務 (例如電腦視覺和自然語言處理) 提供更佳的準確性。不過,單一 GPU 記憶體可容納的最大模型大小有限制。訓練 DL 模型時,GPU 記憶體限制可能是瓶頸,狀況如下:
-
它們會限制您可以訓練的模型大小,因為模型的記憶體佔用量會與參數數量成比例擴展。
-
它們會在訓練期間限制每個 GPU 批次大小,進而降低 GPU 使用率和訓練效率。
為了克服在單一 GPU 上訓練模型的相關限制,SageMaker AI 提供了模型平行化程式庫,以協助在多個運算節點上有效地分散和訓練 DL 模型。此外,借助該程式庫,您可以使用支援 EFA 的裝置實現最佳化的分散式訓練,這可以透過低延遲、高輸送量和作業系統規避來增強節點間通訊的效能。
使用模型平行化之前預估記憶體
在您使用 SageMaker 模型平行程式庫之前,請考慮下列事項,以了解訓練大型 DL 模型的記憶體需求。
對於使用自動混合精確度 (例如 float16 (FP16) 或 bfloat16 (BF16) 和 Adam 最佳化工具) 的訓練任務,每個參數所需的 GPU 記憶體大約為 20 位元組,我們可以按下列方式細分:
-
一個 FP16 或 BF16 參數,約 2 位元組
-
一個 FP16 或 BF16 梯度,約 2 位元組
-
基於 Adam 最佳化工具的 FP32 最佳化工具狀態約 8 位元組
-
參數的 FP32 副本約 4 位元組 (
optimizer apply(OA) 操作需要) -
漸層的 FP32 副本約 4 位元組 (OA 操作需要)
即使對於具有 100 億個參數的相對較小之 DL 模型,它至少需要 200GB 的記憶體,這比單一 GPU 上的一般 GPU 記憶體大得多 (例如,具有 40GB/80GB 記憶體的 NVIDIA A100)。除了模型和最佳化工具狀態的記憶體需求之外,還有其他記憶體取用者,例如在轉送傳遞中產生的啟動。所需的記憶體可能大於 200GB。
對於分散式訓練,建議您分別使用具有 NVIDIA A100 和 H100 Tensor Core GPU 的 Amazon EC2 P4 和 P5 執行個體。如需 CPU 核心、RAM、已連接的儲存磁碟區和網路頻寬等規格的詳細資訊,請參閱 Amazon EC2 執行個體類型
即使使用加速運算執行個體,具有大約 100 億個參數 (例如 Megatron-LM 和 T5) 的模型,甚至是具有數千億個參數 (例如 GPT-3) 的大型模型,也無法在每個 GPU 裝置中使用模型複本。
程式庫如何運用模型平行化與記憶體節省技術
該程式庫包含各種類型的模型平行處理功能和記憶體節省功能,例如最佳化工具狀態碎片、啟動檢查點、啟動卸載。所有這些技術都可以結合起來,以有效率地訓練包含數千億個參數的大型模型。
碎片資料平行處理
碎片資料平行處理是一種節省記憶體的分散式訓練技術,可將模型狀態 (模型參數、漸層和最佳化工具狀態) 分割到資料平行群組中的 GPU。
SMP v2 透過 FSDP 實作碎片資料平行化,並延伸到實作部落格文章 AWS 上巨型模型訓練的近線性擴展
您可以將碎片資料平行化套用至模型,以作為獨立的策略。此外,如果您使用的是配有 NVIDIA A100 Tensor Core GPU 的最高效能 GPU 執行個體 (ml.p4d.24xlarge 和 ml.p4de.24xlarge),您可以利用 SageMaker 資料平行化 (SMDDP) 程式庫提供的 AllGather 操作來提高訓練速度。
若要深入了解碎片資料平行化,並學習如何設定資料,或結合使用碎片資料平行化與其他技術 (例如張量平行化和混合精確度訓練),請參閱混合碎片資料平行化。
專家平行化
SMP v2 與 NVIDIA Megatron
MoE 模型是一種轉換器模型,由多個專家組成,每個專家都由神經網路組成,通常是前饋網路 (FFN)。稱為路由器的閘道網路會決定要傳送哪些權杖給哪個專家。這些專家專門處理輸入資料的特定層面,使模型能夠更快地訓練、降低運算成本,同時實現與其對等密集模型相同的效能品質。專家平行化是一種平行化技術,可處理跨 GPU 裝置的 MoE 模型分割專家。
若要了解如何使用 SMP v2 訓練 MoE 模型,請參閱專家平行化。
張量平行處理
張量平行處理會跨裝置將各個層或 nn.Modules 分割,藉此平行執行。下圖顯示最簡單的範例,說明 SMP 程式庫如何將模型分割成四層,以達到雙向張量平行處理 ("tensor_parallel_degree": 2)。在下圖中,模型平行化群組、張量平行化群組和資料平行化群組的表示法分別為 MP_GROUP、TP_GROUP 和 DP_GROUP。每個模型複本的層會均分並分散為兩個 GPU。程式庫會管理張量分散式模型複本之間的通訊。
若要深入了解 PyTorch 的張量平行處理和其他可節省記憶體的功能,並學習如何設定核心功能的組合,請參閱張量平行化。
啟用檢查點和卸載
為了節省 GPU 記憶體,程式庫支援啟動檢查點,以避免在轉送傳遞期間將內部啟動儲存在使用者指定的模組之 GPU 記憶體中。程式庫會在向後傳遞期間重新運算這些啟動項目。此外,啟動卸載功能會將儲存的啟動卸載至 CPU 記憶體,並在向後傳遞期間擷取回 GPU,進一步減少啟用記憶體佔用量。如需如何使用這些功能的詳細資訊,請參閱啟用檢查點和啟用卸載。
為您的模型選擇正確的技術
如需選擇正確技術和組態的詳細資訊,請參閱SageMaker 分散式模型平行化最佳實務。