模型平行化概念 - Amazon SageMaker AI

本文為英文版的機器翻譯版本,如內容有任何歧義或不一致之處,概以英文版為準。

模型平行化概念

模型平行化是一種分散式訓練方法,其中深度學習 (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 執行個體類型頁面中的加速運算一節。如需 SMP v2 支援的執行個體類型,請參閱支援的執行個體類型

即使使用加速運算執行個體,具有大約 100 億個參數 (例如 Megatron-LM 和 T5) 的模型,甚至是具有數千億個參數 (例如 GPT-3) 的大型模型,也無法在每個 GPU 裝置中使用模型複本。

程式庫如何運用模型平行化與記憶體節省技術

該程式庫包含各種類型的模型平行處理功能和記憶體節省功能,例如最佳化工具狀態碎片、啟動檢查點、啟動卸載。所有這些技術都可以結合起來,以有效率地訓練包含數千億個參數的大型模型。

碎片資料平行處理

碎片資料平行處理是一種節省記憶體的分散式訓練技術,可將模型狀態 (模型參數、漸層和最佳化工具狀態) 分割到資料平行群組中的 GPU。

SMP v2 透過 FSDP 實作碎片資料平行化,並延伸到實作部落格文章 AWS 上巨型模型訓練的近線性擴展中討論的擴展感知混合碎片策略。

您可以將碎片資料平行化套用至模型,以作為獨立的策略。此外,如果您使用的是配有 NVIDIA A100 Tensor Core GPU 的最高效能 GPU 執行個體 (ml.p4d.24xlargeml.p4de.24xlarge),您可以利用 SageMaker 資料平行化 (SMDDP) 程式庫提供的 AllGather 操作來提高訓練速度。

若要深入了解碎片資料平行化,並學習如何設定資料,或結合使用碎片資料平行化與其他技術 (例如張量平行化和混合精確度訓練),請參閱混合碎片資料平行化

專家平行化

SMP v2 與 NVIDIA Megatron 整合,在支援原生 PyTorch FSDP API 的基礎上實現專家平行化。您可以保持 PyTorch FSDP 訓練程式碼不變,並在 SageMaker AI 中套用 SMP 專家平行化來訓練專家混合 (MoE) 模型。

MoE 模型是一種轉換器模型,由多個專家組成,每個專家都由神經網路組成,通常是前饋網路 (FFN)。稱為路由器的閘道網路會決定要傳送哪些權杖給哪個專家。這些專家專門處理輸入資料的特定層面,使模型能夠更快地訓練、降低運算成本,同時實現與其對等密集模型相同的效能品質。專家平行化是一種平行化技術,可處理跨 GPU 裝置的 MoE 模型分割專家。

若要了解如何使用 SMP v2 訓練 MoE 模型,請參閱專家平行化

張量平行處理

張量平行處理會跨裝置將各個層或 nn.Modules 分割,藉此平行執行。下圖顯示最簡單的範例,說明 SMP 程式庫如何將模型分割成四層,以達到雙向張量平行處理 ("tensor_parallel_degree": 2)。在下圖中,模型平行化群組、張量平行化群組和資料平行化群組的表示法分別為 MP_GROUPTP_GROUPDP_GROUP。每個模型複本的層會均分並分散為兩個 GPU。程式庫會管理張量分散式模型複本之間的通訊。

SMP 程式庫如何將模型分割為四層以達到雙向張量平行化 ("tensor_parallel_degree": 2) 的最簡單範例。

若要深入了解 PyTorch 的張量平行處理和其他可節省記憶體的功能,並學習如何設定核心功能的組合,請參閱張量平行化

啟用檢查點和卸載

為了節省 GPU 記憶體,程式庫支援啟動檢查點,以避免在轉送傳遞期間將內部啟動儲存在使用者指定的模組之 GPU 記憶體中。程式庫會在向後傳遞期間重新運算這些啟動項目。此外,啟動卸載功能會將儲存的啟動卸載至 CPU 記憶體,並在向後傳遞期間擷取回 GPU,進一步減少啟用記憶體佔用量。如需如何使用這些功能的詳細資訊,請參閱啟用檢查點啟用卸載

為您的模型選擇正確的技術

如需選擇正確技術和組態的詳細資訊,請參閱SageMaker 分散式模型平行化最佳實務