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

# 步驟 1：使用 SageMaker 的分散式模型平行程式庫修改訓練指令碼
<a name="model-parallel-customize-training-script"></a>

請使用此節來學習如何自訂訓練指令碼，以使用 Amazon SageMaker AI 模型平行化程式庫的核心功能。若要使用程式庫專用的 API 函式和參數，我們建議您將本文件搭配 *SageMaker Python SDK 文件*的 [SageMaker 模型平行程式庫 API](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html) 合併使用。

這些章節中所提供的訓練指令碼範例經過簡化，僅點出您使用程式庫時的必要變更。如需完整徹底可執行的筆記本範例，示範如何將 TensorFlow 或 PyTorch 訓練指令碼與 SageMaker 模型平行處理程式庫搭配使用，請參閱[Amazon SageMaker AI 模型平行化程式庫 v2 範例](distributed-model-parallel-v2-examples.md)。

**Topics**
+ [使用 SageMaker 模型平行處理程式庫分割訓練指令碼的模型](#model-parallel-model-splitting-using-smp-lib)
+ [修改 TensorFlow 訓練指令碼](model-parallel-customize-training-script-tf.md)
+ [修改 PyTorch 訓練指令碼](model-parallel-customize-training-script-pt.md)

## 使用 SageMaker 模型平行處理程式庫分割訓練指令碼的模型
<a name="model-parallel-model-splitting-using-smp-lib"></a>

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

### 自動化模型分割
<a name="model-parallel-automated-model-splitting"></a>

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

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

#### 運作方式
<a name="model-parallel-automated-model-splitting-how-it-works"></a>

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

在任一架構中，程式庫會透過針對 AWS 基礎設施最佳化的自有後端來管理裝置之間的通訊。

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

##### 使用 PyTorch 自動化模型分割
<a name="model-parallel-auto-model-split-pt"></a>

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

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

如需如何將`smp.step`裝飾項目註冊到 PyTorch 訓練指令碼的指示，請參閱[使用 PyTorch 自動化分割](model-parallel-customize-training-script-pt.md#model-parallel-customize-training-script-pt-16)。

##### 使用 TensorFlow 自動化分割模型
<a name="model-parallel-auto-model-split-tf"></a>

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

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

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

如需有關如何將 `smp.step` 裝飾項目註冊到 PyTorch 訓練指令碼的指示，請參閱[使用 TensorFlow 自動化分割](model-parallel-customize-training-script-tf.md#model-parallel-customize-training-script-tf-23)。

##### 比較不同架構之間的自動化模型分割
<a name="model-parallel-auto-model-split-comparison"></a>

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

而對 PyTorch 來說，則沒有多樣化且通用的對等操作概念。具有這些特性的 PyTorch 中最接近的運算單元是 `nn.Module`，它處於更高的粒度層，這就是為什麼程式庫在 PyTorch 中會在此層進行分割。

### 手動模型分割
<a name="model-parallel-manual-model-splitting"></a>

若需手動指定模型在各裝置間的分割方式，請使用 `smp.partition` 內容管理器。如需有關如何設定手動分割內容管理器的指示，請參閱下列頁面。
+ [使用 TensorFlow 手動分割](model-parallel-customize-training-script-tf.md#model-parallel-customize-training-script-tf-manual)
+ [使用 PyTorch 手動分割](model-parallel-customize-training-script-pt.md#model-parallel-customize-training-script-pt-16-hvd)

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