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

# 步驟 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` 上執行。這種情況下將會略過自動化分割邏輯，每項操作都是基於您的詳細設定進行配置。模型平行處理程式庫會基於產生的圖表結構，自動建立管道執行排程。

# 修改 TensorFlow 訓練指令碼
<a name="model-parallel-customize-training-script-tf"></a>

在本節中，您將學習如何修改 TensorFlow 訓練指令碼，以針對自動磁碟分割和手動磁碟分割設定 SageMaker 模型平行處理程式庫。這個範例選擇也包含與 Horovod 整合的範例，以用於混合模型和資料平行處理。

**注意**  
若需有關程式庫支援的 TensorFlow 版本資訊，請參閱[支援的架構和 AWS 區域](distributed-model-parallel-support.md)。

有關使用程式庫前必須對訓練指令碼進行的修改，已列入 [使用 TensorFlow 自動化分割](#model-parallel-customize-training-script-tf-23)。

想了解如何修改訓練指令碼，以便將混合模型和資料平行處理與 Horovod 搭配使用，請參閱[運用 TensorFlow 和 Horovod 進行混合模型和資料平行處理的自動化分割。](#model-parallel-customize-training-script-tf-2.3)。

若選擇手動磁碟分割，請參考[使用 TensorFlow 手動分割](#model-parallel-customize-training-script-tf-manual)。

下列主題提供訓練指令碼範例，您可以使用這些指令碼來設定 SageMaker 模型平行處理程式庫，以進行自動分割和手動分割 TensorFlow 模型。

**注意**  
自動分割預設為開啟。除非特別指定，否則範例指令碼都採用自動分割。

**Topics**
+ [使用 TensorFlow 自動化分割](#model-parallel-customize-training-script-tf-23)
+ [運用 TensorFlow 和 Horovod 進行混合模型和資料平行處理的自動化分割。](#model-parallel-customize-training-script-tf-2.3)
+ [使用 TensorFlow 手動分割](#model-parallel-customize-training-script-tf-manual)
+ [不支援架構功能](#model-parallel-tf-unsupported-features)

## 使用 TensorFlow 自動化分割
<a name="model-parallel-customize-training-script-tf-23"></a>

若要使用 SageMaker 模型平行處理程式庫執行 TensorFlow 模型，必須變更以下訓練指令碼：

1. 使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#smp.init](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#smp.init) 匯入和初始化程式庫。

1. 透過繼承自 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_tensorflow.html](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_tensorflow.html) 而不是 Keras 模型類別，來定義一個 Keras 模型。從 `smp.DistributedModel` 物件的呼叫方法傳回模型輸出。請注意，從呼叫方法傳回的任何張量都將跨模型平行裝置廣播，造成通訊開銷增加，因此不需傳回在呼叫方法之外的非必要張量 (例如中繼啟動)。

1. 在 `tf.Dataset.batch()` 方法中設定 `drop_remainder=True`。這是為了確保批次大小必然可以被微批次數量整除。

1. 使用 `smp.dp_rank()` (如 `shuffle(ds, seed=smp.dp_rank())`) 在 Data Pipeline 中植入隨機操作，確保存放不同模型分割的 GPU 上的資料範例保持一致。

1. 將轉送和向後邏輯放在 Step Function 中，並使用 `smp.step` 進行裝飾。

1. 使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#StepOutput](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#StepOutput) 方法 (如 `reduce_mean`) 對微批次的輸出上執行後處理。[https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#smp.init](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#smp.init) 函式的傳回值必須取決於 `smp.DistributedModel` 的輸出。

1. 如果有評估步驟，採取類似將轉送邏輯放在 `smp.step`-裝飾函式內的作法，並使用 [`StepOutput` API](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#StepOutput) 對輸出執行後處理。

若要進一步了解 SageMaker 模型平行處理程式庫 API，請參閱 [API 文件](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html)。

下列 Python 指令碼是進行變更之後的訓練指令碼範例。

```
import tensorflow as tf

# smdistributed: Import TF2.x API
import smdistributed.modelparallel.tensorflow as smp

# smdistributed: Initialize
smp.init()

# Download and load MNIST dataset.
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data(
    "MNIST-data-%d" % smp.rank()
)
x_train, x_test = x_train / 255.0, x_test / 255.0

# Add a channels dimension
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]

# smdistributed: If needed, seed the shuffle with smp.dp_rank(), and drop_remainder
# in batching to make sure batch size is always divisible by number of microbatches
train_ds = (
    tf.data.Dataset.from_tensor_slices((x_train, y_train))
    .shuffle(10000, seed=smp.dp_rank())
    .batch(256, drop_remainder=True)
)

# smdistributed: Define smp.DistributedModel the same way as Keras sub-classing API 
class MyModel(smp.DistributedModel):
    def __init__(self):
        super(MyModel, self).__init__()
        # define layers

    def call(self, x, training=None):
        # define forward pass and return the model output

model = MyModel()

loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
optimizer = tf.keras.optimizers.Adam()
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name="train_accuracy")

# smdistributed: Define smp.step. Return any tensors needed outside
@smp.step
def get_grads(images, labels):
    predictions = model(images, training=True)
    loss = loss_object(labels, predictions)

    grads = optimizer.get_gradients(loss, model.trainable_variables)
    return grads, loss, predictions


@tf.function
def train_step(images, labels):
    gradients, loss, predictions = get_grads(images, labels)

    # smdistributed: Accumulate the gradients across microbatches
    gradients = [g.accumulate() for g in gradients]
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # smdistributed: Merge predictions and average losses across microbatches
    train_accuracy(labels, predictions.merge())
    return loss.reduce_mean()


for epoch in range(5):
    # Reset the metrics at the start of the next epoch
    train_accuracy.reset_states()
    for images, labels in train_ds:
        loss = train_step(images, labels)
    accuracy = train_accuracy.result()
```

如果您已完成訓練指令碼的準備，請繼續執行 [步驟 2：使用 SageMaker Python SDK 啟動訓練任務](model-parallel-sm-sdk.md)。如果想要執行混合模型和資料平行訓練任務，請繼續至下一節。

## 運用 TensorFlow 和 Horovod 進行混合模型和資料平行處理的自動化分割。
<a name="model-parallel-customize-training-script-tf-2.3"></a>

您可以透過 SageMaker 模型平行處理程式庫搭配 Horovod 使用，以進行混合模型和資料平行處理。若要閱讀更多有關程式庫針對混合式平行處理分割模型方式的資訊，請參閱[管道平行處理 (適用於 PyTorch 和 TensorFlow)](model-parallel-intro.md#model-parallel-intro-pp)。

在此步驟中，我們將著重在介紹如何修改訓練指令碼，以適應 SageMaker 模型平行處理程式庫。

為正確設定訓練指令碼，以取得您在 [步驟 2：使用 SageMaker Python SDK 啟動訓練任務](model-parallel-sm-sdk.md) 要設定的混合式平行處理組態，請使用程式庫的輔助函式 `smp.dp_rank()` 和 `smp.mp_rank()`，即會分別自動偵測資料平行和模型平行的排名。

若要尋找程式庫支援的所有 MPI 基本，請參閱 SageMaker Python SDK 文件中的 [MPI 基本資訊](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#mpi-basics)。

指令碼中所需的必要變更如下：
+ 新增 `hvd.allreduce`
+ 根據 Horovod 的要求，在第一批次後廣播變數。
+ 在 `smp.dp_rank()` 的 Data Pipeline 中植入隨機顯示和/或碎片操作。

**注意**  
當您使用 Horovod 時，不得在訓練指令碼中直接呼叫 `hvd.init`。反之，您需要在 SageMaker Python SDK [步驟 2：使用 SageMaker Python SDK 啟動訓練任務](model-parallel-sm-sdk.md) 的 `modelparallel` 參數，將 `"horovod"` 設定為`True`。這讓程式庫可以基於模型分割的裝置指派狀況，在內部初始化 Horovod。直接在訓練指令碼中呼叫 `hvd.init()`，可能會造成問題。

**注意**  
直接在訓練指令碼中使用 `hvd.DistributedOptimizer` API，可能會導致訓練成效不佳與速度減緩，因為 API 背景作業下會將 `AllReduce` 操作置於 `smp.step` 中。我們建議您在取得 `smp.step` 傳回的漸層上呼叫 `accumulate()` 或 `reduce_mean()` 之後直接呼叫 `hvd.allreduce`，以搭配 Horovod 使用模型平行處理程式庫，如下列範例所示。

若要進一步了解 SageMaker 模型平行處理程式庫 API，請參閱 [API 文件](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html)。

```
import tensorflow as tf
import horovod.tensorflow as hvd

# smdistributed: Import TF2.x API 
import smdistributed.modelparallel.tensorflow as smp

# smdistributed: Initialize
smp.init()

# Download and load MNIST dataset.
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data(
    "MNIST-data-%d" % smp.rank()
)
x_train, x_test = x_train / 255.0, x_test / 255.0

# Add a channels dimension
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]

# smdistributed: Seed the shuffle with smp.dp_rank(), and drop_remainder
# in batching to make sure batch size is always divisible by number of microbatches
train_ds = (
    tf.data.Dataset.from_tensor_slices((x_train, y_train))
    .shuffle(10000, seed=smp.dp_rank())
    .batch(256, drop_remainder=True)
)

# smdistributed: Define smp.DistributedModel the same way as Keras sub-classing API 
class MyModel(smp.DistributedModel):
    def __init__(self):
        super(MyModel, self).__init__()
        # define layers

    def call(self, x, training=None):
        # define forward pass and return model outputs


model = MyModel()

loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
optimizer = tf.keras.optimizers.Adam()
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name="train_accuracy")

# smdistributed: Define smp.step. Return any tensors needed outside
@smp.step
def get_grads(images, labels):
    predictions = model(images, training=True)
    loss = loss_object(labels, predictions)

    grads = optimizer.get_gradients(loss, model.trainable_variables)
    return grads, loss, predictions


@tf.function
def train_step(images, labels, first_batch):
    gradients, loss, predictions = get_grads(images, labels)

    # smdistributed: Accumulate the gradients across microbatches
    # Horovod: AllReduce the accumulated gradients
    gradients = [hvd.allreduce(g.accumulate()) for g in gradients]
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # Horovod: Broadcast the variables after first batch 
    if first_batch:
        hvd.broadcast_variables(model.variables, root_rank=0)
        hvd.broadcast_variables(optimizer.variables(), root_rank=0)

    # smdistributed: Merge predictions across microbatches
    train_accuracy(labels, predictions.merge())
    return loss.reduce_mean()


for epoch in range(5):
    # Reset the metrics at the start of the next epoch
    train_accuracy.reset_states()

    for batch, (images, labels) in enumerate(train_ds):
        loss = train_step(images, labels, tf.constant(batch == 0))
```

## 使用 TensorFlow 手動分割
<a name="model-parallel-customize-training-script-tf-manual"></a>

使用 `smp.partition` 內容管理器，將操作指派至特定分割中。未放置在任何 `smp.partition` 內容中的任何操作都會放置在 `default_partition`。若要進一步了解 SageMaker 模型平行處理程式庫 API，請參閱 [API 文件](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html)。

```
import tensorflow as tf

# smdistributed: Import TF2.x API.
import smdistributed.modelparallel.tensorflow as smp

# smdistributed: Initialize
smp.init()

# Download and load MNIST dataset.
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data(
    "MNIST-data-%d" % smp.rank()
)
x_train, x_test = x_train / 255.0, x_test / 255.0

# Add a channels dimension
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]

# smdistributed: If needed, seed the shuffle with smp.dp_rank(), and drop_remainder
# in batching to make sure batch size is always divisible by number of microbatches.
train_ds = (
    tf.data.Dataset.from_tensor_slices((x_train, y_train))
    .shuffle(10000, seed=smp.dp_rank())
    .batch(256, drop_remainder=True)
)

# smdistributed: Define smp.DistributedModel the same way as Keras sub-classing API.
class MyModel(smp.DistributedModel):
    def __init__(self):
         # define layers

    def call(self, x):
        with smp.partition(0):
            x = self.layer0(x)
        with smp.partition(1):
            return self.layer1(x)


model = MyModel()

loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
optimizer = tf.keras.optimizers.Adam()
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name="train_accuracy")

# smdistributed: Define smp.step. Return any tensors needed outside
@smp.step
def get_grads(images, labels):
    predictions = model(images, training=True)
    loss = loss_object(labels, predictions)

    grads = optimizer.get_gradients(loss, model.trainable_variables)
    return grads, loss, predictions


@tf.function
def train_step(images, labels):
    gradients, loss, predictions = get_grads(images, labels)

    # smdistributed: Accumulate the gradients across microbatches
    gradients = [g.accumulate() for g in gradients]
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # smdistributed: Merge predictions and average losses across microbatches
    train_accuracy(labels, predictions.merge())
    return loss.reduce_mean()


for epoch in range(5):
    # Reset the metrics at the start of the next epoch
    train_accuracy.reset_states()
    for images, labels in train_ds:
        loss = train_step(images, labels)
    accuracy = train_accuracy.result()
```

## 不支援架構功能
<a name="model-parallel-tf-unsupported-features"></a>

程式庫不支援下列 TensorFlow 功能：
+ 目前不支援 `tf.GradientTape()`。您可以使用 `Optimizer.get_gradients()` 或 `Optimizer.compute_gradients()` 來運算漸層。
+ 目前不支援 `tf.train.Checkpoint.restore()` API。如為檢查點，請改為使用 `smp.CheckpointManager` 提供相同的 API 和功能。請注意，`smp.CheckpointManager` 的檢查點還原應該在第一步驟後進行。

# 修改 PyTorch 訓練指令碼
<a name="model-parallel-customize-training-script-pt"></a>

在本節中，您將了解如何修改 PyTorch 訓練指令碼，以針對自動磁碟分割和手動磁碟分割設定 SageMaker 模型平行處理程式庫。

**注意**  
若需有關程式庫支援的 PyTorch 版本資訊，請參閱[支援的架構和 AWS 區域](distributed-model-parallel-support.md)。

**提示**  
如需完整徹底的筆記本範例，示範如何將 PyTorch 訓練指令碼與 SageMaker 模型平行處理程式庫搭配使用，請參閱[Amazon SageMaker AI 模型平行化程式庫 v1 範例](distributed-model-parallel-examples.md)。

請注意，自動磁碟分割預設為啟用。除非特別指定，否則下列指令碼都採用自動分割。

**Topics**
+ [使用 PyTorch 自動化分割](#model-parallel-customize-training-script-pt-16)
+ [使用 PyTorch 手動分割](#model-parallel-customize-training-script-pt-16-hvd)
+ [考量事項](#model-parallel-pt-considerations)
+ [不支援架構功能](#model-parallel-pt-unsupported-features)

## 使用 PyTorch 自動化分割
<a name="model-parallel-customize-training-script-pt-16"></a>

若要使用 SageMaker 模型平行處理程式庫執行 PyTorch 訓練指令碼，必須變更以下訓練指令碼：

1. 使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#smp.init](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#smp.init) 匯入和初始化程式庫。

1. 以 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_pytorch.html#smp.DistributedModel](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_pytorch.html#smp.DistributedModel) 包裝模型。請注意，從基礎 `nn.Module` 物件的 `forward` 方法傳回的任何張量，都將跨模型平行裝置廣播，造成通訊開銷增加，因此不需傳回在呼叫方法之外的非必要張量 (例如中繼啟動)。
**注意**  
如為 FP16 訓練，您需要使用[smdistributed.modelparallel.torch.model\$1creation()](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html)內容管理器來包裝模型。如需詳細資訊，請參閱[使用模型平行化進行 FP16 訓練](model-parallel-extended-features-pytorch-fp16.md)。

1. 以 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_pytorch.html#smp.DistributedOptimizer](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_pytorch.html#smp.DistributedOptimizer) 包裝最佳化工具。
**注意**  
FP16 訓練時，您需要設定靜態或動態損失擴展功能。如需詳細資訊，請參閱[使用模型平行化進行 FP16 訓練](model-parallel-extended-features-pytorch-fp16.md)。

1. 使用傳回的 `DistributedModel` 物件，而非使用者模型。

1. 將轉送和向後邏輯放在 Step Function 中，並使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#smp.init](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#smp.init) 進行裝飾。

1. 透過 `torch.cuda.set_device(smp.local_rank())` 將每個程序限制在所屬裝置中。

1. 在 `smp.step` 呼叫之前，使用 `.to()` API 將輸入張量移動至 GPU (請參閱下列範例)。

1. 將 `torch.Tensor.backward` 和 `torch.autograd.backward` 取代為 `DistributedModel.backward`。

1. 使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#StepOutput](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#StepOutput) 方法 (如 `reduce_mean`) 對微批次的輸出上執行後處理。

1. 如果有評估步驟，採取類似將轉送邏輯放在 `smp.step`-裝飾函式內的作法，並使用 [`StepOutput` API](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#StepOutput) 對輸出執行後處理。

1. 在 `DataLoader` 中設定 `drop_last=True`。或者，如果批次大小不能被微批次數量整除，則手動略過訓練循環中的批次。

若要進一步了解 SageMaker 模型平行處理程式庫 API，請參閱 [API 文件](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html)。

```
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchnet.dataset import SplitDataset
from torchvision import datasets

import smdistributed.modelparallel.torch as smp

class GroupedNet(nn.Module):
    def __init__(self):
        super(GroupedNet, self).__init__()
        # define layers

    def forward(self, x):
        # define forward pass and return model outputs


# smdistributed: Define smp.step. Return any tensors needed outside.
@smp.step
def train_step(model, data, target):
    output = model(data)
    loss = F.nll_loss(output, target, reduction="mean")
    model.backward(loss)
    return output, loss


def train(model, device, train_loader, optimizer):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        # smdistributed: Move input tensors to the GPU ID used by the current process,
        # based on the set_device call.
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        # Return value, loss_mb is a StepOutput object
        _, loss_mb = train_step(model, data, target)

        # smdistributed: Average the loss across microbatches.
        loss = loss_mb.reduce_mean()

        optimizer.step()

# smdistributed: initialize the backend
smp.init()

# smdistributed: Set the device to the GPU ID used by the current process.
# Input tensors should be transferred to this device.
torch.cuda.set_device(smp.local_rank())
device = torch.device("cuda")

# smdistributed: Download only on a single process per instance.
# When this is not present, the file is corrupted by multiple processes trying
# to download and extract at the same time
dataset = datasets.MNIST("../data", train=True, download=False)

# smdistributed: Shard the dataset based on data-parallel ranks
if smp.dp_size() > 1:
    partitions_dict = {f"{i}": 1 / smp.dp_size() for i in range(smp.dp_size())}
    dataset = SplitDataset(dataset, partitions=partitions_dict)
    dataset.select(f"{smp.dp_rank()}")

# smdistributed: Set drop_last=True to ensure that batch size is always divisible
# by the number of microbatches
train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True)

model = GroupedNet()
optimizer = optim.Adadelta(model.parameters(), lr=4.0)

# smdistributed: Use the DistributedModel container to provide the model
# to be partitioned across different ranks. For the rest of the script,
# the returned DistributedModel object should be used in place of
# the model provided for DistributedModel class instantiation.
model = smp.DistributedModel(model)
optimizer = smp.DistributedOptimizer(optimizer)

train(model, device, train_loader, optimizer)
```

## 使用 PyTorch 手動分割
<a name="model-parallel-customize-training-script-pt-16-hvd"></a>

使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_pytorch.html#smp.DistributedOptimizer](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_pytorch.html#smp.DistributedOptimizer) 內容管理員，將模組放置在特定裝置中。未放置在任何 `smp.partition` 內容中的任何模組都會放置在 `default_partition`。如果 `auto_partition` 設定為 `False`，則需要提供 `default_partition`。在特定 `smp.partition` 內容中建立的模組，將被放在對應的分割上。

若要進一步了解 SageMaker 模型平行處理程式庫 API，請參閱 [API 文件](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html)。

```
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchnet.dataset import SplitDataset
from torchvision import datasets

import smdistributed.modelparallel.torch as smp

class GroupedNet(nn.Module):
    def __init__(self):
        super(GroupedNet, self).__init__()
        with smp.partition(0):
            # define child modules on device 0
        with smp.partition(1):
            # define child modules on device 1

    def forward(self, x):
        # define forward pass and return model outputs


# smdistributed: Define smp.step. Return any tensors needed outside.
@smp.step
def train_step(model, data, target):
    output = model(data)
    loss = F.nll_loss(output, target, reduction="mean")
    model.backward(loss)
    return output, loss


def train(model, device, train_loader, optimizer):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        # smdistributed: Move input tensors to the GPU ID used by the current process,
        # based on the set_device call.
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        # Return value, loss_mb is a StepOutput object
        _, loss_mb = train_step(model, data, target)

        # smdistributed: Average the loss across microbatches.
        loss = loss_mb.reduce_mean()

        optimizer.step()

# smdistributed: initialize the backend
smp.init()

# smdistributed: Set the device to the GPU ID used by the current process.
# Input tensors should be transferred to this device.
torch.cuda.set_device(smp.local_rank())
device = torch.device("cuda")

# smdistributed: Download only on a single process per instance.
# When this is not present, the file is corrupted by multiple processes trying
# to download and extract at the same time
dataset = datasets.MNIST("../data", train=True, download=False)

# smdistributed: Shard the dataset based on data-parallel ranks
if smp.dp_size() > 1:
    partitions_dict = {f"{i}": 1 / smp.dp_size() for i in range(smp.dp_size())}
    dataset = SplitDataset(dataset, partitions=partitions_dict)
    dataset.select(f"{smp.dp_rank()}")

# smdistributed: Set drop_last=True to ensure that batch size is always divisible
# by the number of microbatches
train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True)

model = GroupedNet()
optimizer = optim.Adadelta(model.parameters(), lr=4.0)

# smdistributed: Use the DistributedModel container to provide the model
# to be partitioned across different ranks. For the rest of the script,
# the returned DistributedModel object should be used in place of
# the model provided for DistributedModel class instantiation.
model = smp.DistributedModel(model)
optimizer = smp.DistributedOptimizer(optimizer)

train(model, device, train_loader, optimizer)
```

## 考量事項
<a name="model-parallel-pt-considerations"></a>

使用 SageMaker 模型平行處理程式庫設定 PyTorch 訓練指令碼時，有下列幾點需要注意：
+ 如果您使用根據全域漸層規範的最佳化技術，像是整個模型的漸層規範 (例如 LAMB 最佳化工具的某些變體或全域漸層剪輯)，您需要收集整個模型分割區中的所有規範以便進行勘誤。您可以透過程式庫的通訊基本資料類型來執行此操作。
+ 模型內 `nn.Modules` 的轉送方法中，所有 `torch.Tensor` 引數都必須用於模組輸出的運算中。換句話說，在有 `torch.Tensor` 引數的情況下，程式庫不支援模組輸出到未使用的模組。
+ `smp.DistributedModel.backward()` 呼叫的引數必須取決於所有模型輸出。換句話說，在轉送至 `smp.DistributedModel.backward` 呼叫的張量運算中未使用的 `smp.DistributedModel.forward` 呼叫輸出不應存在。
+ 如果程式碼中有`torch.cuda.synchronize()` 呼叫，您可能需要在同步呼叫之前立即呼叫 `torch.cuda.set_device(smp.local_rank())`。否則，可能會在裝置 0 中建立不必要的 CUDA 內容，這將額外消耗記憶體。
+ 由於程式庫將 `nn.Modules` 放置在不同的裝置上，因此模型中的模組不得依賴 `smp.step` 內部修改的任何全域狀態。在整個訓練期間保持固定的任何狀態，或在 `smp.step` 外以所有程序都可見的方式修改任何狀態，都是被允許的。
+ 使用程式庫時，您不需要將模型移動至 GPU (例如使用 `model.to(device)`)。如果您在分割模型前 (在第一次 `smp.step` 呼叫前)試圖將模型移至 GPU，則移動呼叫將被忽略。程式庫會自動將指派給某個階級的模型部分移動至其 GPU。一旦開始使用程式庫進行訓練，請不要將模型移動到 CPU 並使用，因為未指派給該程序持有分割的模組，將不會有正確的參數。如果您想在使用模型平行處理程式庫訓練模型之後，在沒有程式庫的前提下重新訓練模型或作為推論使用，建議使用檢查點 API 儲存完整模型，並將其載入至一般的 PyTorch 模組。
+ 如果您有一個模組清單，以便將一個輸出轉送到另一個模組中，則將該清單取代為 `nn.Sequential` 可以顯著改善效能。
+ 權重更新 (`optimizer.step()`) 必須在 `smp.step` 之外進行，因為屆時整個向後傳遞程序已完成，且漸層已準備就緒。當使用具有模型和資料平行處理的混合模型時，漸層的 AllReduce 也必然會完成。
+ 將程式庫與資料平行處理結合使用時，請確定所有資料平行排名上的批次數量都相同，如此一來，AllReduce 便不會懸置等待步驟內未參與的排名。
+ 如果您使用 ml.p4d 執行個體類型 (例如 ml.p4d.24xlarge) 啟動訓練任務，則必須設定資料載入器變數 `num_workers=0`。舉例來說，您可以將 `DataLoader` 定義如下：

  ```
  dataloader = torch.utils.data.DataLoader(
              data,
              batch_size=batch_size,
              num_workers=0,
              pin_memory=True,
              drop_last=True,
              shuffle=shuffle,
          )
  ```
+ `smp.step` 的輸入必須是由 `DataLoader` 產生的模型輸入。這是因為 `smp.step` 會沿批次維度內部分割輸入張量，然後將其管道化。這表示將 `DataLoader` 本身傳遞給 `smp.step` 函式以生成內部的模型輸入是不可行的。

  舉例來說，假如將 `DataLoader` 定義如下：

  ```
  train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True)
  ```

  您應該存取由 `train_loader` 產生的模型輸入，並將其傳遞給具備 `smp.step` 裝飾的函式。請勿直接將 `train_loader` 傳遞給 `smp.step` 函式。

  ```
  def train(model, device, train_loader, optimizer):
      model.train()
      for batch_idx, (data, target) in enumerate(train_loader):
          ...
          _, loss_mb = train_step(model, data, target)
          ...
  
  @smp.step
  def train_step(model, data, target):
      ...
      return output, loss
  ```
+ `smp.step` 的輸入張量，必須透過 `.to()` API 移動至目前裝置，這必須在 `torch.cuda.set_device(local_rank())` 呼叫後執行。

  舉例來說，可以將 `train` 函式定義如下。該函式在使用輸入張量呼叫 `train_step` 之前，使用 `.to()` API 將 `data` 和 `target` 增添至目前裝置。

  ```
  def train(model, device, train_loader, optimizer):
      model.train()
      for batch_idx, (data, target) in enumerate(train_loader):
          # smdistributed: Move input tensors to the GPU ID used by the current process,
          # based on the set_device call.
          data, target = data.to(device), target.to(device)
          optimizer.zero_grad()
          # Return value, loss_mb is a StepOutput object
          _, loss_mb = train_step(model, data, target)
  
          # smdistributed: Average the loss across microbatches.
          loss = loss_mb.reduce_mean()
  
          optimizer.step()
  ```

  此 `smp.set` 裝飾函式的輸入張量已移動至上述 `train` 函式中的目前裝置。此模型*不*需要移動到目前裝置。程式庫會自動將指派給某個階級的模型部分移動至其 GPU。

  ```
  @smp.step
  def train_step(model, data, target):
      output = model(data)
      loss = F.nll_loss(output, target, reduction="mean")
      model.backward(loss)
      return output, loss
  ```

## 不支援架構功能
<a name="model-parallel-pt-unsupported-features"></a>

SageMaker 模型平行處理程式庫不支援下列 PyTorch 功能：
+ 如果您使用資料平行處理搭配原生 [PyTorch DDP](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html)，則程式庫不支援 [https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) 包裝函式模組。程式庫內部管理整合 PyTorch DDP ，包含參數廣播和漸層 AllReduce。使用程式庫時，模組緩衝區僅在訓練開始時廣播一次。假如模型具備模組緩衝區，並需要在每個步驟中跨資料平行群組同步，您可以透過 `torch.distributed` API 執行此操作，使用透過 `smp.get_dp_process_group()` 取得的程序群組。
+ 如為混合精確度訓練，則不支援 `apex.amp` 模組。在自動混合精確度使用程式庫的情境中，建議使用 `torch.cuda.amp` 來操作，但在 Torch 中進行的實作則應使用 `smp.amp.GradScaler`。
+ `smp.DistributedModel` 不支援 `torch.jit.ScriptModules` 或 `ScriptFunctions`。
+ `apex`：`apex` 的 `FusedLayerNorm`、`FusedAdam`、`FusedLAMB` 和 `FusedNovoGrad` 不支援。您可以透過 `smp.optimizers` 和 `smp.nn` API 來使用這些程式庫的實作。