步驟 1:使用 SageMaker的分散式模型平行程式庫修改您自己的訓練指令碼 - Amazon SageMaker

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

步驟 1:使用 SageMaker的分散式模型平行程式庫修改您自己的訓練指令碼

使用本節瞭解如何自訂訓練指令碼,以使用 Amazon SageMaker 模型平行程式庫的核心功能。若要使用程式庫特定的 API 函數和參數,我們建議您將此文件與 SageMaker Python SDK 文件中的SageMaker 模型 parallel 程式庫 API 一起使用。

這些章節中所提供的訓練指令碼範例經過簡化,僅點出您使用程式庫時的必要變更。如需示範如何搭配 SageMaker 模型平行程式庫使用 TensorFlow 或 PyTorch 訓練指令碼的可執行筆記本範例 end-to-end,請參閱。Amazon SageMaker 模型並行程式庫 v2 範例

使用模型平行程式庫分割訓練指令碼的 SageMaker 模型

修改訓練指令碼以設定模型分割的方法有兩種:自動化分割或手動分割。

自動化模型分割

使用 SageMaker的模型平行程度程式庫時,您可以利用自動化模型分割,也稱為自動化模型分割。該程式庫採用的分割演算法能平衡記憶體的消耗,最小化裝置間的通訊並最佳化效能。您可以設定自動化磁碟分割演算法,以最佳化速度或記憶體用量。

或者,您可以手動進行模型分割。除非對模型架構極為熟悉,並清楚如何有效進行模型分割,否則,我們建議使用自動化模型分割。

運作方式

在第一個訓練步驟期間,自動分割會在首次呼叫 smp.step-裝飾函式時進行。在呼叫期間,程式庫會先在 CPU RAM 上建構模型版本 (以避免 GPU 記憶體限制),接著分析模型圖表,並做出分割決定。基於這個決定,每個模型分割區被載入至 GPU 上,接下來執行第一個步驟。由於這些分析和分割步驟,首次訓練步驟可能較為耗時。

在任何一個框架中,庫都通過其自己的後端管理設備之間的通信,該後端針對 AWS 基礎結構進行了優化。

自動分割設計會因應架構特性調整,而程式庫在精細程度層進行分割,在各架構內更加自然。例如,在中 TensorFlow,每個特定操作都可以分配給不同的設備,而在中 PyTorch,分配是在模塊級別完成,其中每個模塊由多個操作組成。下一節將檢閱每個架構中的設計細節。

在首次訓練步驟中,模型平行處理程式庫會在內部執行追蹤步驟,目的是建構模型圖表,並判定張量和參數形狀。完成追蹤步驟後,程式庫會建構一個樹狀圖,其中包含模型中的巢狀nn.Module物件,以及從追蹤收集到的其他資料,像是nn.Parameters的儲存數量,以及每個nn.Module物件的執行時間。

接著,程式庫由下往上走完樹狀圖並執行分割演算法,為每個 nn.Module 指派一個裝置以平衡運算負載(以模組執行時間計算)和記憶體使用量(由總儲存 nn.Parameter 大小和啟動次數計算)。如果多個nn.Modules共用相同的nn.Parameter,這些模組將被放在同一裝置內,以避免維護相同參數的多個版本。決定分割內容後,指派的模組和加權將載入至對應裝置上。

如需有關如何將smp.step裝飾器註冊到 PyTorch 訓練指令碼的指示,請參閱自動拆分 PyTorch

模型平行處理程式庫會分析可訓練變數和圖表結構大小,並在內部使用圖表分割演算法。此演算法會為每項操作指派裝置,目標是在下列兩個限制條件下,將裝置間所需的通訊量降至最低:

  • 平衡各裝置中儲存的變數量

  • 平衡各裝置執行的操作數量

如果在 Python SDK 的模型平行參數中指定 speed 作為 optimize,程式庫會試著平衡每個裝置中的操作數量及 tf.Variable 物件的數量。否則,它會試著平衡 tf.Variables 的大小總計。

決定分割內容後,程式庫會為每個裝置需要執行的子圖建立一個序列化表示法,並匯入至各裝置。磁碟分割時,程式庫會將消耗相同 tf.Variable 的操作,以及屬於同一 Keras 層的操作指派到同一裝置上。它還遵守由 TensorFlow強加的託管約束。舉例來說,這表示如果有兩個 Keras 圖層共用一個 tf.Variable,那麼這些圖層中的所有操作都會放置在單一裝置上。

如需有關如何將smp.step裝飾器註冊到 PyTorch 訓練指令碼的指示,請參閱自動拆分 TensorFlow

比較不同架構之間的自動化模型分割

在中 TensorFlow,計算的基本單位是 atf.Operation,並將模型 TensorFlow 表示為 tf.Operation s 的有向無環圖 (DAG),因此模型平行程式庫會分割此 DAG,以便每個節點移至一個裝置。關鍵是,tf.Operation 物件具備多樣化的可自訂屬性,並在某種程度上是通用的,因為每個模型都必定包含以這類物件組成的圖表。

PyTorch 另一方面,沒有足夠豐富和普遍的操作的等效概念。具有這些特性的最接近的計算單元是一個nn.Module,它處於更高的粒度級別,這就是為什麼庫在 PyTorch. PyTorch

手動模型分割

若需手動指定模型在各裝置間的分割方式,請使用 smp.partition 內容管理器。如需有關如何設定手動分割內容管理器的指示,請參閱下列頁面。

若要在修改之後使用此選項,在步驟 2 中,您需要在 SageMaker Python SDK 的架構估算器類別default_partition中設auto_partition定並定義。False未透過 smp.partition 內容管理器明確指派至磁碟分割上的任何操作,都會在 default_partition 上執行。這種情況下將會略過自動化分割邏輯,每項操作都是基於您的詳細設定進行配置。模型平行處理程式庫會基於產生的圖表結構,自動建立管道執行排程。