本文為英文版的機器翻譯版本，如內容有任何歧義或不一致之處，概以英文版為準。
# (存檔) SageMaker 模型平行化程式庫 v1.x
**重要**
自 2023 年 12 月 19 日起，SageMaker 模型平行化 (SMP) 程式庫 v2 已發行。為了支援 SMP 程式庫 v2，未來版本不再支援 SMP v1 功能。以下章節和主題已存檔，並且特定於使用 SMP 程式庫 v1。如需使用 SMP 程式庫 v2 的詳細資訊，請參閱 [SageMaker 模型平行化程式庫 v2](model-parallel-v2.md)。
使用 Amazon SageMaker AI 的模型平行化程式庫來訓練因為 GPU 記憶體限制而難以訓練的大型深度學習 (DL) 模型。程式庫會自動且有效率地將模型分割至多個 GPU 和執行個體。使用程式庫，您可以透過有效率地訓練具有數十億或數兆個參數的較大型 DL 模型，更快速的達成目標預測準確度。
您可以使用這個程式庫，只需變更最少的程式碼，即可在多個 GPU 和多個節點之間自動分割您自己的 TensorFlow 和 PyTorch 模型。您可以透過 SageMaker Python SDK 存取程式庫的 API。
您可以使用以下章節進一步瞭解模型平行化和 SageMaker 模型平行化程式庫。此程式庫的 API 文件位於 *SageMaker Python SDK v2.199.0 文件*中的[分散式訓練 API](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html)。
**Topics**
+ [模型平行處理簡介](model-parallel-intro.md)
+ [支援的架構和 AWS 區域](distributed-model-parallel-support.md)
+ [SageMaker 模型平行化程式庫的核心功能](model-parallel-core-features.md)
+ [執行具有模型平行的 SageMaker 分散式訓練任務](model-parallel-use-api.md)
+ [對具有模型並行性的模型執行檢查點和微調](distributed-model-parallel-checkpointing-and-finetuning.md)
+ [Amazon SageMaker AI 模型平行化程式庫 v1 範例](distributed-model-parallel-examples.md)
+ [SageMaker 分散式模型平行處理最佳實務](model-parallel-best-practices.md)
+ [SageMaker 分散式模型平行化程式庫組態提示和缺陷](model-parallel-customize-tips-pitfalls.md)
+ [模型平行疑難排解](distributed-troubleshooting-model-parallel.md)
# 模型平行處理簡介
模型平行處理原則是一種分散式訓練方法，深度學習模型會在多個裝置間分割，不論是在執行個體內部或執行個體間皆如此。此簡介頁面提供有關模型平行處理原則的高階概觀、說明它如何協助克服訓練通常非常大的 DL 模型時所產生的問題，以及 SageMaker 模型平行程式庫所提供之協助管理模型平行策略以及記憶體使用的範例。
## 什麼是模型平行處理？
增加深度學習模型 (圖層和參數) 的大小可以為複雜的任務 (例如電腦視覺和自然語言處理) 提供更佳的準確性。不過，單一 GPU 記憶體可容納的最大模型大小有限制。訓練 DL 模型時，GPU 記憶體限制可能是瓶頸，狀況如下：
+ 它們會限制您可以訓練的模型大小，因為模型的記憶體佔用量會與參數數量成比例擴展。
+ 它們會在訓練期間限制每個 GPU 批次大小，進而降低 GPU 使用率和訓練效率。
為了克服在單一 GPU 上訓練模型的相關限制，SageMaker 提供了模型平行程式庫，以協助在多個運算節點上有效地分散和訓練 DL 模型。此外，借助該程式庫，您可以使用支援 EFA 的裝置實現最佳化的分散式訓練，這可以透過低延遲、高輸送量和作業系統規避來增強節點間通訊的效能。
## 使用模型平行處理之前預估記憶體
在您使用 SageMaker 模型平行程式庫之前，請考慮下列事項，以了解訓練大型 DL 模型的記憶體需求。
對於使用 AMP (FP16) 和 Adam 最佳化工具的訓練任務，每個參數所需的 GPU 記憶體大約為 20 位元組，我們可以按下列方式細分：
+ 一個 FP16 參數約 2 位元組
+ 一個 FP16 漸層約 2 位元組
+ 基於 Adam 最佳化工具的 FP32 最佳化工具狀態約 8 位元組
+ 參數的 FP32 副本約 4 位元組 (`optimizer apply` (OA) 操作需要)
+ 漸層的 FP32 副本約 4 位元組 (OA 操作需要)
即使對於具有 100 億個參數的相對較小之 DL 模型，它至少需要 200GB 的記憶體，這比單一 GPU 上的一般 GPU 記憶體大得多 (例如，具有 40GB/80GB 記憶體的 NVIDIA A100，以及具有 16/32 GB 的 V100)。請注意，除了模型和最佳化工具狀態的記憶體需求之外，還有其他記憶體取用者，例如在轉送傳遞中產生的啟動。所需的記憶體可能大於 200GB。
對於分散式訓練，建議您分別使用具有 NVIDIA V100 和 A100 張量核心 GPU 的 Amazon EC2 P3 和 P4 執行個體。如需 CPU 核心、RAM、已連接的儲存磁碟區和網路頻寬等規格的詳細資訊，請參閱 [Amazon EC2 執行個體類型](https://aws.amazon.com/ec2/instance-types/)頁面中的*加速運算*一節。
即使使用加速運算執行個體，明顯將了解具有大約 100 億個參數 (例如 Megatron-LM 和 T5) 的模型，甚至是具有數千億個參數 (例如 GPT-3) 的大型模型，也無法在每個 GPU 裝置中使用模型複本。
## 程式庫如何運用模型平行處理與記憶體節省技術
該程式庫包含各種類型的模型平行處理功能和記憶體節省功能，例如最佳化工具狀態碎片、啟動檢查點、啟動卸載。所有這些技術都可以結合起來，以有效率地訓練包含數千億個參數的大型模型。
**Topics**
+ [碎片資料平行處理 (適用於 PyTorch)](#model-parallel-intro-sdp)
+ [管道平行處理 (適用於 PyTorch 和 TensorFlow)](#model-parallel-intro-pp)
+ [張量平行處理 (適用於 PyTorch)](#model-parallel-intro-tp)
+ [最佳化工具狀態碎片 (適用於 PyTorch)](#model-parallel-intro-oss)
+ [啟動卸載和檢查點 (適用於 PyTorch)](#model-parallel-intro-activation-offloading-checkpointing)
+ [為您的模型選擇正確的技術](#model-parallel-intro-choosing-techniques)
### 碎片資料平行處理 (適用於 PyTorch)
*碎片資料平行處理*是一種節省記憶體的分散式訓練技術，可將模型狀態 (模型參數、漸層和最佳化工具狀態) 分割到資料平行群組中的 GPU。
SageMaker AI 透過實作 MiCS 來採用碎片資料平行化，MiCS 是一款程式庫，可讓您**最**小化**通**訊**規**模，並在關於 AWS巨型模型訓練的[近線性擴展部落格文章中討論那樣進行處理](https://www.amazon.science/blog/near-linear-scaling-of-gigantic-model-training-on-aws)。
您可以將碎片資料平行處理套用至模型，以作為獨立的策略。此外，如果您使用的是配有 NVIDIA A100 張量核心 GPU 的最高效能 GPU 執行個體 (`ml.p4d.24xlarge`)，您可以利用 SMDDP 系列產品提供的 `AllGather` 操作來提升訓練速度。
若要深入了解碎片資料平行處理，並學習如何設定資料，或結合使用碎片資料平行處理與其他技術 (例如張量平行處理和 FP16 訓練)，請參閱[碎片資料平行處理](model-parallel-extended-features-pytorch-sharded-data-parallelism.md)。
### 管道平行處理 (適用於 PyTorch 和 TensorFlow)
*管道平行處理*會跨裝置集合分割一組層或作業，讓每個作業完好無缺。當您指定模型分割區數量 (`pipeline_parallel_degree`) 的值時，GPU 的總數量 (`processes_per_host`) 必須可由模型分割區的數量整除。若要正確設定，您必須指定正確的 `pipeline_parallel_degree` 和 `processes_per_host` 參數值。簡單的數學原理如下：
```
(pipeline_parallel_degree) x (data_parallel_degree) = processes_per_host
```
在特定您提供的兩個輸入參數的情況下，該程式庫負責計算模型複本的數量 (亦稱為 `data_parallel_degree`)。
例如，如果您設定 `"pipeline_parallel_degree": 2` 和 `"processes_per_host": 8` 以使用具有八個 GPU 工作者的機器學習 (ML) 執行個體 (例如 `ml.p3.16xlarge`)，則程式庫會自動跨 GPU 和四向資料平行處理設定分散式模型。下列影像說明模型如何在八個 GPU 上分散，以達到四向資料平行處理和雙向管道平行處理。每個模型複本，我們將其定義為*管道平行群組*並將其標籤為 `PP_GROUP`，且跨兩個 GPU 進行分割。模型的每個分割區都會指派給四個 GPU，其中四個分割區複本位於*資料平行群組*中並標示為 `DP_GROUP`。如果沒有張量平行處理，管道平行群組本質上就是模型平行群組。
![\[模型如何在八個 GPU 上分散，以達到四向資料平行化和雙向管道平行化。\]](http://docs.aws.amazon.com/zh_tw/sagemaker/latest/dg/images/distributed/model-parallel/smdmp-pipeline-parallel-only.png)
若要深入了解管道平行處理，請參閱[SageMaker 模型平行化程式庫的核心功能](model-parallel-core-features.md)。
若要開始使用管道平行處理來執行模型，請參閱[使用 SageMaker 模型平行程式庫執行 SageMaker 分散式訓練任務](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-use-api.html)。
### 張量平行處理 (適用於 PyTorch)
*張量平行處理*會跨裝置將各個層或 `nn.Modules` 分割，藉此平行執行。下圖顯示最簡單的範例，說明程式庫如何將模型分割成四層，以達到雙向張量平行處理 (`"tensor_parallel_degree": 2`)。每個模型複本的層會均分並分散為兩個 GPU。在這個範例中，模型平行設定也包含 `"pipeline_parallel_degree": 1` 和 `"ddp": True` (在背景中使用 PyTorch DistributedDataParallel 套件)，因此資料平行處理的程度會變成 8。程式庫會管理張量分散式模型複本之間的通訊。
![\[說明程式庫如何將模型分割成四層，以達到雙向張量平行處理 ("tensor_parallel_degree": 2) 的最簡單範例。\]](http://docs.aws.amazon.com/zh_tw/sagemaker/latest/dg/images/distributed/model-parallel/smdmp-tensor-parallel-only.png)
此功能的實用性在於您可以選取特定的層或層子集，藉此套用張量平行處理。若要深入了解 PyTorch 的張量平行處理和其他可節省記憶體的功能，並探究如何設定管道和張量平行處理的組合，請參閱[張量平行處理](model-parallel-extended-features-pytorch-tensor-parallelism.md)。
### 最佳化工具狀態碎片 (適用於 PyTorch)
如要了解程式庫如何執行*最佳化工具狀態碎片*，不妨參考一個具有四層的簡單範例模型。最佳化狀態碎片的關鍵概念是，您不必在所有 GPU 中複製最佳化工具的狀態。相反地，最佳化工具狀態的單一複本會 data-parallel 等級進行資料分割，而不會跨裝置提供備援。例如，GPU 0 會保留第一層的最佳化工具狀態，下一個 GPU 1 會保留 L2 的最佳化工具狀態，依此類推。下列動畫圖顯示使用最佳化工具狀態碎片技術的向後傳播。在向後傳播結束時，`optimizer apply` (OA) 作業會更新最佳化工具狀態的運算和網路時間，而 `all-gather` (AG) 作業則會更新下一次反覆運算的模型參數。最重要的是，`reduce` 作業可能會與 GPU 0 上的運算重疊，因此可提高記憶體效率，並且加快向後傳播的速度。在目前的實作中，AG 和 OA 作業不會與 `compute` 重疊。這可能會導致在 AG 作業期間延伸運算，因此可能會有所取捨。
![\[使用最佳化工具狀態碎片技術進行向後傳播。\]](http://docs.aws.amazon.com/zh_tw/sagemaker/latest/dg/images/distributed/model-parallel/smdmp-optimizer-state-sharding.gif)
如需如何使用此功能的詳細資訊，請參閱[最佳化工具狀態碎片](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-optimizer-state-sharding.html)。
### 啟動卸載和檢查點 (適用於 PyTorch)
為了節省 GPU 記憶體，程式庫支援啟動檢查點，以避免在轉送傳遞期間將內部啟動儲存在使用者指定的模組之 GPU 記憶體中。程式庫會在向後傳遞期間重新運算這些啟動項目。此外，啟動卸載功能會將儲存的啟動卸載至 CPU 記憶體，並在向後傳遞期間擷取回 GPU，進一步減少啟用記憶體佔用量。如需如何使用這些功能的詳細資訊，請參閱[啟用檢查點](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-activation-checkpointing.html)和[啟用卸載](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-activation-offloading.html)。
### 為您的模型選擇正確的技術
如需有關選擇正確技術和組態的詳細資訊，請參閱 [SageMaker 分散式模型平行最佳實務](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-best-practices.html)和[組態提示與陷阱](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-customize-tips-pitfalls.html)。
# 支援的架構和 AWS 區域
使用 SageMaker 模型平行處理程式庫之前，請檢查支援的架構和執行個體類型，並判斷 AWS 您的帳戶和 中是否有足夠的配額 AWS 區域。
**注意**
若要查看程式庫的最新更新和發行說明，請參閱 *SageMaker Python SDK 文件*中的 [SageMaker 模型平行發行說明](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_release_notes/smd_model_parallel_change_log.html)。
## 支援的架構
SageMaker 模型平行處理程式庫支援下列深度學習架構，可在 AWS 深度學習容器 (DLC) 中取得或下載為二進位檔案。
SageMaker AI 和 SageMaker AI 模型平行化程式庫支援的 PyTorch 版本
| PyTorch 版本 | SageMaker 模型平行處理程式庫版本 | `smdistributed-modelparallel` 整合的 DLC 映像 URI | 二進位檔案的網址\$1\$1 |
| --- | --- | --- | --- |
| v2.0.0 | smdistributed-modelparallel==v1.15.0 | `763104351884.dkr.ecr..amazonaws.com/pytorch-training:2.0.0-gpu-py310-cu118-ubuntu20.04-sagemaker` | https://sagemaker-distributed-model-parallel.s3.us-west-2.amazonaws.com/pytorch-2.0.0/build-artifacts/2023-04-14-20-14/smdistributed\$1modelparallel-1.15.0-cp310-cp310-linux\$1x86\$164.whl |
| V1.13.1 | smdistributed-modelparallel==v1.15.0 | `763104351884.dkr.ecr..amazonaws.com/pytorch-training:1.13.1-gpu-py39-cu117-ubuntu20.04-sagemaker` | https://sagemaker-distributed-model-parallel.s3.us-west-2.amazonaws.com/pytorch-1.13.1/build-artifacts/2023-04-17-15-49/smdistributed\$1modelparallel-1.15.0-cp39-cp39-linux\$1x86\$164.whl |
| v1.12.1 | smdistributed-modelparallel==v1.13.0 | `763104351884.dkr.ecr..amazonaws.com/pytorch-training:1.12.1-gpu-py38-cu113-ubuntu20.04-sagemaker` | https://sagemaker-distributed-model-parallel.s3.us-west-2.amazonaws.com/pytorch-1.12.1/build-artifacts/2022-12-08-21-34/smdistributed\$1modelparallel-1.13.0-cp38-cp38-linux\$1x86\$164.whl |
| v1.12.0 | smdistributed-modelparallel==v1.11.0 | `763104351884.dkr.ecr..amazonaws.com/pytorch-training:1.12.0-gpu-py38-cu113-ubuntu20.04-sagemaker` | https://sagemaker-distributed-model-parallel.s3.us-west-2.amazonaws.com/pytorch-1.12.0/build-artifacts/2022-08-12-16-58/smdistributed\$1modelparallel-1.11.0-cp38-cp38-linux\$1x86\$164.whl |
| v1.11.0 | smdistributed-modelparallel==v1.10.0 | `763104351884.dkr.ecr..amazonaws.com/pytorch-training:1.11.0-gpu-py38-cu113-ubuntu20.04-sagemaker` | https://sagemaker-distributed-model-parallel.s3.us-west-2.amazonaws.com/pytorch-1.11.0/build-artifacts/2022-07-11-19-23/smdistributed\$1modelparallel-1.10.0-cp38-cp38-linux\$1x86\$164.whl |
| V1.10.2 | smdistributed-modelparallel==v1.7.0 | `763104351884.dkr.ecr..amazonaws.com/pytorch-training:1.10.2-gpu-py38-cu113-ubuntu20.04-sagemaker` | - |
| v1.10.0 | smdistributed-modelparallel==v1.5.0 | `763104351884.dkr.ecr..amazonaws.com/pytorch-training:1.10.0-gpu-py38-cu113-ubuntu20.04-sagemaker` | - |
| v1.9.1 | smdistributed-modelparallel==v1.4.0 | `763104351884.dkr.ecr..amazonaws.com/pytorch-training:1.9.1-gpu-py38-cu111-ubuntu20.04` | - |
| 1.8.1\$1 | smdistributed-modelparallel==v1.6.0 | `763104351884.dkr.ecr..amazonaws.com/pytorch-training:1.8.1-gpu-py36-cu111-ubuntu18.04` | - |
**注意**
SageMaker 模型平行處理程式庫 v1.6.0 及更新版本為 PyTorch 提供了擴充功能。如需詳細資訊，請參閱[SageMaker 模型平行化程式庫的核心功能](model-parallel-core-features.md)。
\$1\$1 二進位檔案的 URL 用於在自訂容器中安裝 SageMaker 模型平行處理程式庫。如需詳細資訊，請參閱[使用 SageMaker 分散式模型平行程式庫建立您自己的 Docker 容器](model-parallel-sm-sdk.md#model-parallel-bring-your-own-container)。
SageMaker AI 和 SageMaker 模型平行化程式庫支援的 TensorFlow 版本
| TensorFlow 版本 | SageMaker 模型平行處理程式庫版本 | `smdistributed-modelparallel` 整合的 DLC 映像 URI |
| --- | --- | --- |
| v2.6.0 | smdistributed-modelparallel==v1.4.0 | 763104351884.dkr.ecr..amazonaws.com/tensorflow-training:2.6.0-gpu-py38-cu112-ubuntu20.04 |
| v2.5.1 | smdistributed-modelparallel==v1.4.0 | 763104351884.dkr.ecr..amazonaws.com/tensorflow-training:2.5.1-gpu-py37-cu112-ubuntu18.04 |
**SageMaker AI 和 SageMaker 分散式資料庫平行程式庫支援的 Hugging Face 轉換器版本**
適用於 Hugging Face 的 AWS 深度學習容器使用適用於 PyTorch 和 TensorFlow 的 SageMaker 訓練容器作為其基礎映像。若要查詢 Hugging Face 轉換器程式庫版本及配對的 PyTorch 和 TensorFlow 版本，請參閱最新的 [Hugging Face 容器](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers)和[先前的 Hugging Face 容器版本](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#prior-hugging-face-container-versions)。
## AWS 區域
SageMaker 資料平行程式庫可在 SageMaker AWS 區域 深度學習容器提供服務的所有 中使用。 [AWS SageMaker](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#sagemaker-framework-containers-sm-support-only) 如需更多資訊，請參閱[可用的深度學習容器映像](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#available-deep-learning-containers-images)。
## 支援的執行個體類型
SageMaker 模型平行處理程式庫需要下列其中一種機器學習 (ML) 執行個體類型。
| 執行個體類型 |
| --- |
| ml.g4dn.12xlarge |
| ml.p3.16xlarge |
| ml.p3dn.24xlarge |
| ml.p4d.24xlarge |
| ml.p4de.24xlarge |
如需執行個體類型的空間，請參閱 [Amazon EC2 執行個體類型頁面](https://aws.amazon.com/ec2/instance-types/)中的**加速運算**區段。有關執行個體定價資訊，請參閱 [Amazon SageMaker AI 定價](https://aws.amazon.com/sagemaker/pricing/)。
如果您遇到類似下列內容的錯誤訊息，請按照[請求 SageMaker AI 資源服務配額增加](https://docs.aws.amazon.com/sagemaker/latest/dg/regions-quotas.html#service-limit-increase-request-procedure)的指示進行操作。
```
ResourceLimitExceeded: An error occurred (ResourceLimitExceeded) when calling
the CreateTrainingJob operation: The account-level service limit 'ml.p3dn.24xlarge
for training job usage' is 0 Instances, with current utilization of 0 Instances
and a request delta of 1 Instances.
Please contact AWS support to request an increase for this limit.
```
# SageMaker 模型平行化程式庫的核心功能
Amazon SageMaker AI 模型平行化程式庫提供發佈策略和節省記憶體技術，例如碎片資料平行化、張量平行化、管道排程模型依層分割，以及檢查點。模型平行化策略與技術有助於發佈大型模型到多個裝置，同時最佳化訓練速度及記憶體使用量。此程式庫也提供 Python 協助程式函式、內容管理員與包裝函式，以調整訓練指令碼來自動化或手動分割模型。
當您實作模型平行化至訓練任務時，您會保持相同的兩個步驟工作流程，如同[利用模型平行化執行 SageMaker 分散式訓練任務](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-use-api.html)區段所示。若要調整訓練指令碼，您需要新增零或幾行其他程式碼至訓練指令碼。若要啟動已調整訓練指令碼的訓練任務，您需要設定發佈設定參數，以便啟用節省記憶體功能，或傳遞平行處理程度的值。
若要開始使用範例，請參閱下列 Jupyter 筆記本，其中示範如何使用 SageMaker 模型平行化程式庫。
+ [PyTorch 範例筆記本](https://github.com/aws/amazon-sagemaker-examples/tree/main/training/distributed_training/pytorch/model_parallel)
+ [TensorFlow 範例筆記本](https://github.com/aws/amazon-sagemaker-examples/tree/main/training/distributed_training/tensorflow/model_parallel/mnist)
若要深入了解程式庫的核心功能，請參閱下列主題。
**注意**
SageMaker 分散式訓練程式庫可透過 SageMaker Training 平台中 PyTorch、Hugging Face 和 TensorFlow 的 AWS 深度學習容器取得。若要利用分散式訓練程式庫的功能，建議您使用 SageMaker Python SDK。如果您透過適用於 Python 的 SDK (Boto3) 或 AWS Command Line Interface使用 SageMaker API，您還可以手動設定 JSON 的請求語法。在整份文件，指示與範例著重於如何搭配 SageMaker Python SDK 使用分散式訓練程式庫。
**重要**
SageMaker 模型平行化程式庫支援 PyTorch 的所有核心功能，並支援 TensorFlow 的管道平行化。
**Topics**
+ [碎片資料平行處理](model-parallel-extended-features-pytorch-sharded-data-parallelism.md)
+ [管道模型](model-parallel-core-features-pipieline-parallelism.md)
+ [張量平行處理](model-parallel-extended-features-pytorch-tensor-parallelism.md)
+ [最佳化工具狀態碎片](model-parallel-extended-features-pytorch-optimizer-state-sharding.md)
+ [啟用檢查點](model-parallel-extended-features-pytorch-activation-checkpointing.md)
+ [啟用卸載](model-parallel-extended-features-pytorch-activation-offloading.md)
+ [使用模型平行化進行 FP16 訓練](model-parallel-extended-features-pytorch-fp16.md)
+ [支援 FlashAttention](model-parallel-attention-head-size-for-flash-attention.md)
# 碎片資料平行處理
*碎片資料平行處理*是節省記憶體的分散式訓練技術，可將模型狀態 (模型參數、漸層與最佳化工具狀態) 分割到資料平行群組的 GPU。
**注意**
碎片資料平行化適用 SageMaker 模型平行化程式庫 v1.11.0 及更高版本的 PyTorch。
當縱向擴展訓練任務到大型 GPU 叢集時，您可以透過碎片化模型的訓練狀態至多個 GPU 來減少模型的每個 GPU 記憶體的使用量。這會帶來兩個優點：您可以容納較大模型 (否則會耗盡標準資料平行處理的記憶體)，或者可以使用釋放的 GPU 記憶體來增加批次大小。
標準資料平行處理技術會複寫訓練狀態至資料平行群組的 GPU，並基於 `AllReduce` 作業執行漸層彙總。碎片資料平行處理會修改標準資料平行分散式訓練程序，以說明最佳化工具狀態的碎片性質。模型與最佳化工具狀態用以碎片化的一組等級稱為*碎片群組*。碎片資料平行處理技術會碎片化模型的可訓練參數、對應漸層以及最佳化工具狀態至*碎片群組*的 GPU。
SageMaker AI 透過實作 MiCS 來實現碎片資料平行處理，這在[巨型模型訓練的 AWS](https://www.amazon.science/blog/near-linear-scaling-of-gigantic-model-training-on-aws) AWS 部落格文章中進行了討論。在此實作，您可以設定碎片程度為設定參數，該參數必須小於資料平行處理程度。在每次向前與向後傳遞期間，MiCS 會透過 `AllGather` 作業暫時重新組合所有 GPU 的模型參數。在每一層的向前或向後傳遞之後，MiCS 會再次碎片化參數以節省 GPU 記憶體。在向後傳遞期間，MiCS 會縮減漸層，並透過 `ReduceScatter` 作業同時將其碎片化到各 GPU。最後，MiCS 會使用最佳化工具狀態的本機碎片，套用本機的縮減及碎片漸層至其對應本機參數碎片。為降低通訊額外負荷，SageMaker 模型平行化程式庫會在向前或向後傳遞階段預先擷取即將發生的圖層，並重疊網路通訊與運算。
模型的訓練狀態會跨碎片群組複寫。這表示在將漸層套用至參數之前，除在碎片群組進行的 `AllReduce` 作業之外，還必須在碎片群組進行 `ReduceScatter` 作業。
實際上，碎片資料平行處理必須在通訊額外負荷與 GPU 記憶體效率之間做取捨。使用碎片資料平行處理會增加通訊成本，然而每個 GPU 的記憶體使用量 (排除因啟用而導致的記憶體使用量) 會除以碎片資料平行處理程度，因此較大模型可納入 GPU 叢集。
**選取碎片資料平行處理程度**
當您針對碎片資料平行處理程度選取值時，該值必須能平均除碎片資料平行處理程度。例如，對於 8 向資料平行處理工作，請選擇 2、4 或 8 做為碎片資料平行處理程度。在選擇碎片資料平行處理程度時，建議您從小數字開始，然後逐漸增加，直到模型與所需的批次大小可一起放入記憶體。
**選取批次大小**
在設定碎片資料平行處理之後，請確定您找到可在 GPU 叢集成功執行的最佳訓練組態。若要訓練大型語言模型 (LLM)，請從批次大小 1 開始，然後逐漸增加，直到您收到記憶體不足 (OOM) 錯誤為止。如果即使最小批次大小也遇到 OOM 錯誤，請套用較高程度的碎片資料平行處理，或是組合碎片資料平行處理與張量平行處理。
**Topics**
+ [如何套用碎片資料平行處理至訓練任務](#model-parallel-extended-features-pytorch-sharded-data-parallelism-how-to-use)
+ [參考組態](#model-parallel-extended-features-pytorch-sharded-data-parallelism-how-to-use-config-sample)
+ [搭配 SMDDP 集體的碎片資料平行處理](#model-parallel-extended-features-pytorch-sharded-data-parallelism-smddp-collectives)
+ [搭配碎片資料平行處理的混合精確度訓練](#model-parallel-extended-features-pytorch-sharded-data-parallelism-16bits-training)
+ [搭配張量平行處理的碎片資料平行處理](#model-parallel-extended-features-pytorch-sharded-data-parallelism-with-tensor-parallelism)
+ [使用碎片資料平行處理的提示與考量事項](#model-parallel-extended-features-pytorch-sharded-data-parallelism-considerations)
## 如何套用碎片資料平行處理至訓練任務
若要開始使用碎片資料平行處理，請套用必要修改至訓練指令碼，並使用碎片資料平行處理特定參數來設定 SageMaker PyTorch 估算器。同時請考慮以參考值與範例筆記本為起點。
### 調整 PyTorch 訓練指令碼
遵循[步驟 1：修改 PyTorch 訓練指令碼](model-parallel-customize-training-script-pt.md)的指示，使用 `torch.nn.parallel` 與 `torch.distributed` 模組的 `smdistributed.modelparallel.torch` 包裝函式來包裝模型與最佳化工具物件。
**(選用) 註冊外部模型參數的其他修改**
如果您的模型是以 `torch.nn.Module` 構建，且使用未在模組類別內定義的參數，則應手動將其註冊到模組，以便 SMP 可收集完整參數。若要將參數註冊到模組，請使用 `smp.register_parameter(module, parameter)`。
```
class Module(torch.nn.Module):
def __init__(self, *args):
super().__init__(self, *args)
self.layer1 = Layer1()
self.layer2 = Layer2()
smp.register_parameter(self, self.layer1.weight)
def forward(self, input):
x = self.layer1(input)
# self.layer1.weight is required by self.layer2.forward
y = self.layer2(x, self.layer1.weight)
return y
```
### 設定 SageMaker PyTorch 估算器
在 [步驟 2：使用 SageMaker Python SDK 啟動訓練任務](model-parallel-sm-sdk.md) 設定 SageMaker PyTorch 估算器時，請新增碎片資料平行處理的參數。
若要開啟碎片資料平行處理，請新增 `sharded_data_parallel_degree` 參數至 SageMaker PyTorch 估算器。此參數指定訓練狀態要碎片化的 GPU 數量。`sharded_data_parallel_degree` 的值必須為整數並介於 1 與資料平行處理程度之間，且必須能平均除資料平行處理程度。請注意，程式庫會自動偵測 GPU 數量，因而也會自動偵測資料平行程度。下列其他參數可用於設定碎片資料平行處理。
+ `"sdp_reduce_bucket_size"`*(整數，預設值：5e8)* — 以預設 dtype 的元素數量指定 [PyTorch DDP 漸層儲存貯體](https://pytorch.org/docs/stable/notes/ddp.html#internal-design)的大小。
+ `"sdp_param_persistence_threshold"`*(整數，預設值：1e6)* — 針對可在每個 GPU 持續存在的元素數量指定參數張量大小。碎片資料平行處理會分割每個參數張量至資料平行群組的 GPU。如果參數張量的元素數量小於此閾值，則不會分割參數張量；這有助於減少通訊額外負荷，因為參數張量會跨資料平行 GPU 複寫。
+ `"sdp_max_live_parameters"`*(整數，預設值：1e9)* — 針對在向前與向後傳遞期間，可同時處於重新組合訓練狀態的參數指定其數量上限。當有效參數數量達到指定閾值時，使用 `AllGather` 作業擷取的參數會暫停。請注意，增加此參數會增加記憶體使用量。
+ `"sdp_hierarchical_allgather"`*(boo l，預設值：True)* — 如果設定為 `True`，則 `AllGather` 作業會以階層方式執行：會先在每個節點內執行，然後跨節點執行。對於多節點分散式訓練任務，會自動啟用階層式 `AllGather` 作業。
+ `"sdp_gradient_clipping"`*(浮動，預設值：1.0)* — 針對剪輯 L2 標準的漸層指定閾值，然後再透過模型參數向後傳播漸層。當啟用碎片資料平行處理時，也會啟用漸層剪輯。預設閾值為 `1.0`。如果您遇到漸層爆炸問題，請調整此參數。
下列程式碼範例顯示如何設定碎片資料平行處理。
```
import sagemaker
from sagemaker.pytorch import PyTorch
smp_options = {
"enabled": True,
"parameters": {
# "pipeline_parallel_degree": 1, # Optional, default is 1
# "tensor_parallel_degree": 1, # Optional, default is 1
"ddp": True,
# parameters for sharded data parallelism
"sharded_data_parallel_degree": 2, # Add this to activate sharded data parallelism
"sdp_reduce_bucket_size": int(5e8), # Optional
"sdp_param_persistence_threshold": int(1e6), # Optional
"sdp_max_live_parameters": int(1e9), # Optional
"sdp_hierarchical_allgather": True, # Optional
"sdp_gradient_clipping": 1.0 # Optional
}
}
mpi_options = {
"enabled" : True, # Required
"processes_per_host" : 8 # Required
}
smp_estimator = PyTorch(
entry_point="your_training_script.py", # Specify your train script
role=sagemaker.get_execution_role(),
instance_count=1,
instance_type='ml.p3.16xlarge',
framework_version='1.13.1',
py_version='py3',
distribution={
"smdistributed": {"modelparallel": smp_options},
"mpi": mpi_options
},
base_job_name="sharded-data-parallel-job"
)
smp_estimator.fit('s3://my_bucket/my_training_data/')
```
## 參考組態
SageMaker分散式訓練團隊提供下列參考組態，您可以使用這些組態做為起點。您可以從下列組態推斷，以便針對模型組態進行實驗並估算 GPU 記憶體使用量。
搭配 SMDDP 集體的碎片資料平行處理
| 模型/參數量 | 執行個體數 | 執行個體類型 | 序列長度 | 全域批次大小 | 最小批次大小 | 碎片資料平行程度 |
| --- | --- | --- | --- | --- | --- | --- |
| GPT-NEOX-20B | 2 | ml.p4d.24xlarge | 2048 | 64 | 4 | 16 |
| GPT-NEOX-20B | 8 | ml.p4d.24xlarge | 2048 | 768 | 12 | 32 |
例如，如果增加 200 億個參數模型的序列長度，或增加模型大小到 650 億個參數，則需要先嘗試縮小批次大小。如果模型仍無法容納最小批次大小 (批次大小為 1)，請嘗試增加模型平行化程度。
搭配張量平行處理與 NCCL 集體的碎片資料平行處理
| 模型/參數量 | 執行個體數 | 執行個體類型 | 序列長度 | 全域批次大小 | 最小批次大小 | 碎片資料平行程度 | 張量平行程度 | 啟用卸載 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| GPT-NEOX-65B | 64 | ml.p4d.24xlarge | 2048 | 512 | 8 | 16 | 8 | Y |
| GPT-NEOX-65B | 64 | ml.p4d.24xlarge | 4096 | 512 | 2 | 64 | 2 | Y |
當您希望將大型語言模型 (LLM) 納入大型擴展叢集，同時使用具較長序列長度的文字資料時，合併使用碎片資料平行處理與張量平行處理很有幫助，因為這會導致使用較小批次大小，因而處理 GPU 記憶體使用量來針對較長文字序列訓練 LLM。如需進一步了解，請參閱[搭配張量平行處理的碎片資料平行處理](#model-parallel-extended-features-pytorch-sharded-data-parallelism-with-tensor-parallelism)。
如需案例研究、基準測試與更多組態範例，請參閱部落格文章 [Amazon SageMaker AI 模型平行化程式庫的新效能改進](https://aws.amazon.com/blogs/machine-learning/new-performance-improvements-in-amazon-sagemaker-model-parallel-library/)。
## 搭配 SMDDP 集體的碎片資料平行處理
SageMaker 資料平行處理程式庫提供針對 AWS 基礎設施最佳化的集體通訊基本概念 (SMDDP 集體）。其透過使用 [Elastic Fabric Adapter (EFA)](https://aws.amazon.com/hpc/efa/) 採用全對所有類型的通訊模式來達成最佳化，進而產生高輸送量與較少延遲敏感的集體，卸載通訊相關處理至 CPU，並釋放 GPU 週期以便進行運算。在大型叢集，相較於 NCCL，SMDDP 集體可提升分散式訓練效能最多 40%。如需案例研究與基準測試結果，請參閱部落格 [Amazon SageMaker AI 模型平行化程式庫的新效能改進](https://aws.amazon.com/blogs/machine-learning/new-performance-improvements-in-amazon-sagemaker-model-parallel-library/)。
**注意**
搭配 SMDDP 集體的碎片資料平行化可在 SageMaker 模型平行化程式庫 v1.13.0 及更高版本取得，以及 SageMaker 資料平行化程式庫 1.6.0 及更高版本。另請參閲 [Supported configurations](#sharded-data-parallelism-smddp-collectives-supported-config)，以便搭配 SMDDP 集體使用碎片資料平行處理。
在碎片資料平行處理 (這是大規模分散式訓練常用的技術)，`AllGather` 集體用於重組碎片層參數，以便與 GPU 運算平行進行向前與向後傳遞運算。對於大型模型而言，重要的是應以有效率的方式執行 `AllGather` 作業，以利避免 GPU 瓶頸問題並降低訓練速度。當啟用碎片資料平行處理時，SMDDP 集體會進入這些關鍵效能 `AllGather` 集體，進而改善訓練輸送量。
**搭配 SMDDP 集體一起進行訓練**
當訓練任務已啟用碎片資料平行處理並符合時 [Supported configurations](#sharded-data-parallelism-smddp-collectives-supported-config)，SMDDP 集體就會自動啟用。在內部，SMDDP Collectives 會將`AllGather`集體最佳化為在 AWS 基礎設施上執行，並回到所有其他集體的 NCCL。此外，在不支援的組態，所有集體 (包括 `AllGather`) 都會自動使用 NCCL 後端。
自 SageMaker 模型平行化程式庫版本 1.13.0 起，已新增 `"ddp_dist_backend"` 參數至 `modelparallel` 選項。此設定參數的預設值為 `"auto"`，其會盡可能使用 SMDDP 集體，否則會回復為 NCCL。若要強制程式庫一律使用 NCCL，請指定`"nccl"` 至 `"ddp_dist_backend"` 設定參數。
下列程式碼範例示範如何使用碎片資料平行處理與 `"ddp_dist_backend"` 參數 (依預設設定為 `"auto"`，因此為選用新增) 來設定 PyTorch 估算器。
```
import sagemaker
from sagemaker.pytorch import PyTorch
smp_options = {
"enabled":True,
"parameters": {
"partitions": 1,
"ddp": True,
"sharded_data_parallel_degree": 64
"bf16": True,
"ddp_dist_backend": "auto" # Specify "nccl" to force to use NCCL.
}
}
mpi_options = {
"enabled" : True, # Required
"processes_per_host" : 8 # Required
}
smd_mp_estimator = PyTorch(
entry_point="your_training_script.py", # Specify your train script
source_dir="location_to_your_script",
role=sagemaker.get_execution_role(),
instance_count=8,
instance_type='ml.p4d.24xlarge',
framework_version='1.13.1',
py_version='py3',
distribution={
"smdistributed": {"modelparallel": smp_options},
"mpi": mpi_options
},
base_job_name="sharded-data-parallel-demo",
)
smd_mp_estimator.fit('s3://my_bucket/my_training_data/')
```
**支援的組態**
當符合下列所有組態需求時，會在訓練任務啟用 SMDDP 集體的 `AllGather` 作業。
+ 碎片資料平行處理程度大於 1
+ `Instance_count` 大於 1
+ `Instance_type` 等於 `ml.p4d.24xlarge`
+ 適用於 PyTorch v1.12.1 或更高版本的 SageMaker 訓練容器
+ SageMaker 資料平行處理程式庫 v1.6.0 或更高版本
+ SageMaker 模型平行化程式庫 v1.13.0 或更高版本
**效能與記憶體調整**
SMDDP 集體利用其他 GPU 記憶體。根據不同模型訓練使用案例，有兩個環境變數可設定 GPU 記憶體使用量。
+ `SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` – 在 SMDDP `AllGather` 作業期間，會複製 `AllGather` 輸入緩衝至暫時緩衝，以便進行節點間通訊。`SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` 變數控制此暫時緩衝的大小 (位元組)。如果暫時緩衝大小小於 `AllGather` 輸入緩衝大小，則 `AllGather` 集體會回復使用 NCCL。
+ 預設值：16 \$1 1024 \$1 1024 (16 MB)
+ 可接受值：8192 的任何倍數
+ `SMDDP_AG_SORT_BUFFER_SIZE_BYTES` – `SMDDP_AG_SORT_BUFFER_SIZE_BYTES` 變數用以調整暫時緩衝大小 (位元組)，以便保存從節點間通訊收集的資料。如果此暫時緩衝大小小於 `1/8 * sharded_data_parallel_degree * AllGather input size`，則 `AllGather` 集體會回復使用 NCCL。
+ 預設值：128 \$1 1024 \$1 1024 (128 MB)
+ 可接受值：8192 的任何倍數
**緩衝大小變數的調整指引**
環境變數的預設值應適用於多數使用案例。建議僅當訓練出現記憶體不足 (OOM) 錯誤時，才調整這些變數。
下列清單將討論部分調整提示，以減少 SMDDP 集體的 GPU 記憶體使用量，同時保留因此取得的效能提升。
+ 調校 `SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES`
+ 對於較小模型，`AllGather` 輸入緩衝大小較小。因此，對於具較少參數的模型，`SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` 的所需大小可更小。
+ `AllGather` 輸入緩衝大小會隨著 `sharded_data_parallel_degree` 增加而減少，因為模型會碎片化至更多 GPU。因此，對於具較大 `sharded_data_parallel_degree` 值的訓練任務，`SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` 的所需大小可更小。
+ 調校 `SMDDP_AG_SORT_BUFFER_SIZE_BYTES`
+ 對於參數較少的模型，從節點間通訊收集的資料量會較少。因此，對於具較少參數數量的此類模型，`SMDDP_AG_SORT_BUFFER_SIZE_BYTES` 的所需大小可更小。
部分集體可能會回復使用 NCCL; 因此，您可能無法從最佳化的 SMDDP 集體獲得效能提升。如果有其他 GPU 記憶體可用，您可以考慮增加 `SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` 與 `SMDDP_AG_SORT_BUFFER_SIZE_BYTES` 的值，以便從效能提升獲益。
下列程式碼示範如何針對 PyTorch 估算器附加環境變數至發佈參數的 `mpi_options`，藉此設定環境變數。
```
import sagemaker
from sagemaker.pytorch import PyTorch
smp_options = {
.... # All modelparallel configuration options go here
}
mpi_options = {
"enabled" : True, # Required
"processes_per_host" : 8 # Required
}
# Use the following two lines to tune values of the environment variables for buffer
mpioptions += " -x SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES=8192"
mpioptions += " -x SMDDP_AG_SORT_BUFFER_SIZE_BYTES=8192"
smd_mp_estimator = PyTorch(
entry_point="your_training_script.py", # Specify your train script
source_dir="location_to_your_script",
role=sagemaker.get_execution_role(),
instance_count=8,
instance_type='ml.p4d.24xlarge',
framework_version='1.13.1',
py_version='py3',
distribution={
"smdistributed": {"modelparallel": smp_options},
"mpi": mpi_options
},
base_job_name="sharded-data-parallel-demo-with-tuning",
)
smd_mp_estimator.fit('s3://my_bucket/my_training_data/')
```
## 搭配碎片資料平行處理的混合精確度訓練
若要利用半精確度浮點數與碎片資料平行處理來進一步節省 GPU 記憶體，您可以新增單一其他參數至分散式訓練組態來啟用 16 位元浮點格式 (FP16) 或 [Brain 浮點格式](https://en.wikichip.org/wiki/brain_floating-point_format) (BF16)。
**注意**
SageMaker 模型平行化程度程式庫 v1.11.0 及更高版本提供搭配碎片資料平行化的混合精確度訓練。
**對於搭配碎片資料平行處理的 FP16 訓練**
若要搭配碎片資料平行處理來執行 FP16 訓練，請新增 `"fp16": True"` 至 `smp_options` 組態字典。在訓練指令碼，您可以透過 `smp.DistributedOptimizer` 模組選擇靜態或動態損失縮放選項。如需更多資訊，請參閱[使用模型平行化進行 FP16 訓練](model-parallel-extended-features-pytorch-fp16.md)。
```
smp_options = {
"enabled": True,
"parameters": {
"ddp": True,
"sharded_data_parallel_degree": 2,
"fp16": True
}
}
```
**對於搭配碎片資料平行處理的 BF16 訓練**
SageMaker AI 的碎片資料平行化功能支援 BF16 資料類型的訓練。BF16 資料類型使用 8 位元來表示浮點數的指數，而 FP16 資料類型則使用 5 位元。保留指數的 8 位元可針對 32 位元單一精確度浮點 (FP32) 數的指數保持相同表示法。這可簡化 FP32 與 BF16 之間的轉換，且可大幅減少在 FP16 訓練過程經常出現的溢位及下溢問題，尤其是在訓練大型模型時。儘管這兩種資料類型總共使用 16 位元，但 BF16 格式的指數表示法範圍會增加，而精確度則會降低。對於訓練大型模型，這種降低的精確度通常被認為是為取得範圍與訓練穩定性的可接受折衷。
**注意**
目前，僅當啟用碎片資料平行處理時，BF16 訓練才有效。
若要搭配碎片資料平行處理執行 BF16 訓練，請新增 `"bf16": True` 至 `smp_options` 組態字典。
```
smp_options = {
"enabled": True,
"parameters": {
"ddp": True,
"sharded_data_parallel_degree": 2,
"bf16": True
}
}
```
## 搭配張量平行處理的碎片資料平行處理
如果您使用碎片資料平行處理，同時還需要減少全域批次大小，請考慮搭配碎片資料平行處理使用[張量平行處理](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-tensor-parallelism.html)。在非常大型的運算叢集 (通常為 128 個節點或以上) 使用碎片資料平行處理來訓練大型模型時，即使每個 GPU 批次大小很小仍會產生非常大的全域批次大小。這可能導致收斂問題或低運算效能問題。當單一批次已經很大且無法進一步縮減時，僅使用碎片資料平行處理，有時無法縮減每個 GPU 的批次大小。在這種情況，結合使用碎片資料平行處理與張量平行處理有助於減少全域批次大小。
選擇最佳碎片資料平行與張量平行程度取決於模型規模、執行個體類型，以及對於模型收斂合理的全域批次大小。建議您從較低張量平行程度開始，以便將全域批次大小納入運算叢集，解決 CUDA 記憶體不足錯誤並達到最佳效能。請參閱下列兩個範例案例，了解組合張量平行化與碎片資料平行化可如何協助您透過將 GPU 分組為模型平行化來調整全域批次大小，進而產生較少的模型複本數量及較小全域批次大小。
**注意**
您可從 SageMaker 模型平行化程式庫 1.15 取得此功能，並支援 PyTorch 第 1.13.1 版。
**注意**
此功能可透過程式庫的張量平行處理功能用於支援模型。若要尋找支援模型清單，請參閱對 [Hugging Face 轉換器模型支援](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-hugging-face.html)。另請注意，在修改訓練指令碼時，您需要傳遞 `tensor_parallelism=True` 至 `smp.model_creation` 引數。如需進一步了解，請參閱 *SageMaker AI 範例 GitHub 儲存庫*的訓練指令碼 [https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt2/train_gpt_simple.py#L793](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt2/train_gpt_simple.py#L793)。
### 範例 1
假設我們要在 1536 個 GPU 的叢集訓練模型 (192 個節點各有 8 個 GPU)，設定碎片資料平行處理程度為 32 (`sharded_data_parallel_degree=32`)，且每個 GPU 的批次大小為 1，其中每個批次的序列長度為 4096 個權杖。在這種情況，有 1536 個模型複本，全域批次大小變成 1536 個，每個全域批次包含約 600 萬個權杖。
```
(1536 GPUs) * (1 batch per GPU) = (1536 global batches)
(1536 batches) * (4096 tokens per batch) = (6,291,456 tokens)
```
新增張量平行處理至其中可降低全域批次大小。組態範例之一可以是設定張量平行程度為 8，而每個 GPU 的批次大小為 4。這會形成 192 個張量平行群組或 192 個模型複本，其中每個模型複本發佈至 8 個 GPU。批次大小 4 是每個反覆項目與每個張量平行群組的訓練資料量；也就是說，每個模型複本每次反覆會取用 4 個批次。在這種情況，全域批次大小變為 768，而每個全域批次包含約 300 萬個權杖。因此，相較於先前案例僅採用碎片資料平行處理，全域批次大小減少一半。
```
(1536 GPUs) / (8 tensor parallel degree) = (192 tensor parallelism groups)
(192 tensor parallelism groups) * (4 batches per tensor parallelism group) = (768 global batches)
(768 batches) * (4096 tokens per batch) = (3,145,728 tokens)
```
### 範例 2
當同時啟用碎片資料平行處理與張量平行處理時，程式庫會先套用張量平行處理並跨此維度來碎片化模型。對於每個張量平行等級，會按照每個 `sharded_data_parallel_degree` 來套用資料平行處理。
例如，假設我們想要設定 32 個 GPU，張量平行程度為 4 (形成 4 個 GPU 的群組)，碎片資料平行程度為 4，最後複寫度為 2。此指派會基於張量平行程度建立八個 GPU 群組，如下所示：`(0,1,2,3)`、`(4,5,6,7)`、`(8,9,10,11)`、`(12,13,14,15)`、`(16,17,18,19)`、`(20,21,22,23)`、`(24,25,26,27)`、`(28,29,30,31)`。也就是說，四個 GPU 會形成單一張量平行群組。在這種情況，張量平行群組第 0 級 GPU 所縮減的資料平行群組將是 `(0,4,8,12,16,20,24,28)`。縮減資料平行群組會基於碎片資料平行程度 4 進行碎片化，進而產生兩個資料平行處理的複寫群組。GPU `(0,4,8,12)` 形成單一碎片群組，該群組針對第 0 級張量平行集體保存所有參數的完整副本，且 GPU `(16,20,24,28)` 會形成另一此類群組。其他張量平行等級也有類似的碎片與複寫群組。
![\[圖 1：張量平行化群組。\]](http://docs.aws.amazon.com/zh_tw/sagemaker/latest/dg/images/distributed/model-parallel/sdp_tp_group_tp.jpg)
圖 1：張量平行化群組為 (節點、碎片資料平行度、張量平行度) = (4, 4, 4)，其中每個矩形代表一個 GPU，索引從 0 到 31。GPUs 會形成張量平行化群組，從 TPG0 到 TPG7。複寫群組是 (\$1TPG0, TPG4\$1, \$1TPG1, TPG5\$1, \$1TPG2, TPG6\$1 and \$1TPG3, TPG7\$1)；每個複寫群組對共用相同的顏色，但填充不同。
![\[圖 2：碎片資料平行化群組。\]](http://docs.aws.amazon.com/zh_tw/sagemaker/latest/dg/images/distributed/model-parallel/sdp_tp_group_sdp.jpg)
圖 2：碎片資料平行化群組為 (節點、碎片資料平行度、張量平行度) = (4, 4, 4)，其中每個矩形代表一個 GPU，索引從 0 到 31。GPUs 會形成從 SDPG0 到 SDPG 的碎片資料平行化群組7。複寫群組是 (\$1SDPG0, SDPG4\$1, \$1SDPG1, SDPG5\$1, \$1SDPG2, SDPG6\$1 and \$1SDPG3, SDPG7\$1)；每個複寫群組對共用相同的顏色，但填充不同。
### 如何利用張量平行處理啟用碎片資料平行處理
若要利用張量平行處理使用碎片資料平行處理，則當建立 SageMaker PyTorch 估算器類別的物件時，您需要在設定 `distribution` 的同時設定 `sharded_data_parallel_degree` 與 `tensor_parallel_degree`。
您還需要啟用 `prescaled_batch`。這表示，每個張量平行群組都會集體讀取所選批次大小的合併批次，而非每個 GPU 各自讀取其批次資料。實際上，其不會將資料集分成等於 GPU 數量 (或資料平行大小 `smp.dp_size()`) 的部分，而是分成等於 GPU 數量除以 `tensor_parallel_degree` 的部分 (也稱為縮減資料平行大小 `smp.rdp_size()`)。如需更多詳細資訊了解預先調整大小的批次，請參閱 *SageMaker Python SDK 文件*的[預先調整大小的批次](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#prescaled-batch)。另請參閲 *SageMaker AI 範例GitHub 儲存庫*的 GPT-2 範例訓練指令碼 [https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt2/train_gpt_simple.py#L164](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt2/train_gpt_simple.py#L164)。
下列程式碼片段為基於上述 [範例 2](#model-parallel-extended-features-pytorch-sharded-data-parallelism-with-tensor-parallelism-ex2) 的案例，此範例顥示如何建立 PyTorch 估算器物件。
```
mpi_options = "-verbose --mca orte_base_help_aggregate 0 "
smp_parameters = {
"ddp": True,
"fp16": True,
"prescaled_batch": True,
"sharded_data_parallel_degree": 4,
"tensor_parallel_degree": 4
}
pytorch_estimator = PyTorch(
entry_point="your_training_script.py",
role=role,
instance_type="ml.p4d.24xlarge",
volume_size=200,
instance_count=4,
sagemaker_session=sagemaker_session,
py_version="py3",
framework_version="1.13.1",
distribution={
"smdistributed": {
"modelparallel": {
"enabled": True,
"parameters": smp_parameters,
}
},
"mpi": {
"enabled": True,
"processes_per_host": 8,
"custom_mpi_options": mpi_options,
},
},
source_dir="source_directory_of_your_code",
output_path=s3_output_location
)
```
## 使用碎片資料平行處理的提示與考量事項
當使用 SageMaker 模型平行化程式庫的碎片資料平行化時，請考慮下列事項。
+ 碎片資料平行處理相容 FP16 訓練。若要執行 FP16 訓練，請參閱[使用模型平行化進行 FP16 訓練](model-parallel-extended-features-pytorch-fp16.md)區段。
+ 碎片資料平行處理相容張量平行處理。以下是搭配張量平行處理使用碎片資料平行處理時，可能需要考慮的項目。
+ 當搭配張量平行處理使用碎片資料平行處理時，內嵌層也會自動發佈至張量平行群組。換句話說，`distribute_embedding` 參數會自動設定為 `True`。如需張量平行處理的更多相關資訊，請參閱[張量平行處理](model-parallel-extended-features-pytorch-tensor-parallelism.md)。
+ 請注意，搭配張量平行處理的碎片資料平行處理目前會使用 NCCL 集體做為分散式訓練策略的後端。
如需進一步了解，請參閱[搭配張量平行處理的碎片資料平行處理](#model-parallel-extended-features-pytorch-sharded-data-parallelism-with-tensor-parallelism)區段。
+ 碎片資料平行處理目前不相容[管道平行處理](model-parallel-intro.md#model-parallel-intro-pp)或[最佳化工具狀態碎片](model-parallel-extended-features-pytorch-optimizer-state-sharding.md)。若要啟用碎片資料平行處理，請關閉最佳化工具狀態碎片，並設定管道平行程度為 1。
+ [啟用檢查點](model-parallel-extended-features-pytorch-activation-checkpointing.md)及[啟用卸載](model-parallel-extended-features-pytorch-activation-offloading.md)功能相容碎片資料平行處理。
+ 若要搭配漸層累積使用碎片資料平行處理，請在使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.DistributedModel](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.DistributedModel) 模組包裝模型時，設定 `backward_passes_per_step` 引數為累積步驟數。這可確保模型複寫群組 (碎片群組) 之間的漸層 `AllReduce` 作業發生在漸層累積的範圍。
+ 您可以使用程式庫的檢查點 API `smp.save_checkpoint` 與 `smp.resume_from_checkpoint` 來檢查使用碎片資料平行處理訓練的模型。如需更多資訊，請參閱[對分散式 PyTorch 模型執行檢查點 (適用於 SageMaker 模型平行處理程式庫 1.10.0 版及更新版本)](distributed-model-parallel-checkpointing-and-finetuning.md#model-parallel-extended-features-pytorch-checkpoint)。
+ [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.delay_param_initialization](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.delay_param_initialization) 設定參數的行為會在碎片資料平行處理下發生變更。當同時開啟這兩項功能時，參數會在建立模型時立即以碎片方式初始化，而不會延遲參數初始化，以便每個等級初始化並儲存各自的參數碎片。
+ 當啟用碎片資料平行處理時，在呼叫 `optimizer.step()` 時，程式庫會在內部執行漸層剪輯。您不需要使用公用程式 API 進行漸層剪輯，例如 [https://pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html](https://pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html)。若要調整漸層剪輯的閾值，您可以在建構 SageMaker PyTorch 估算器時透過發佈參數組態的 `sdp_gradient_clipping` 參數來加以設定，如 [如何套用碎片資料平行處理至訓練任務](#model-parallel-extended-features-pytorch-sharded-data-parallelism-how-to-use) 區段所示。
# 管道模型
SageMaker 模型平行化程式庫的其中一項核心功能是*管道平行化*，其可決定進行運算順序，以及在模型訓練期間跨裝置處理資料的順序。管道傳輸是在模型平行化達成真正平行化的技術，方法是讓 GPU 在不同資料範例同時運算，並克服因循序運算而導致的效能損失。當您使用管道平行處理時，訓練任務會以管道方式透過微批次執行，以利最大化 GPU 使用率。
**注意**
管道平行處理 (也稱為模型分割) 可用於 PyTorch 與 TensorFlow。如需架構的支援版本清單，請參閱[支援的架構和 AWS 區域](distributed-model-parallel-support.md)。
## 管道執行排程
管道傳輸是基於將迷你批次分割為微批次，將其逐一送入訓練管道，並遵循程式庫執行期所定義的執行排程。*微批次*是指定訓練迷你批次的較小子集。管道排程決定每個時段由哪個裝置執行哪個微批次。
例如，根據管道排程與模型分割區而定，GPU `i` 可能會在微批次 `b` 執行 (向前或向後) 運算，而 GPU `i+1` 則在微批次 `b+1` 執行運算，進而同時保持兩個 GPU 處於有效狀態。在單次向前或向後傳遞期間，單一微批次的執行流程可能多次造訪相同裝置，具體取決於分割決定。例如，模型開頭的作業可能與模型結束時間的作業位於同一裝置，而介於兩者之間的作業則位於不同裝置，這代表造訪該裝置兩次。
程式庫提供兩種不同管道排程，分別是*簡便式*與 *INTERLEAVED*，可使用 SageMaker Python SDK 的 `pipeline` 參數進行設定。在多數情況，INTERLEAVED 管道可更有效利用 GPU 來達到較佳效能。
### INTERLEAVED 管道
在 INTERLEAVED 管道，盡可能優先考量微批次的向後執行。這樣可更快釋放用於啟用的記憶體，以便更以有效的方式使用記憶體。這還允許更高度擴展微批次數量，進而減少 GPU 的閒置時間。在穩定狀態，每個裝置在向前與向後傳遞之間進行交替。這代表可能會先執行單一微批次的向後傳遞，即使另一微批次的向前傳遞尚未完成。
![\[2 個 GPU 上交錯管道的範例執行排程。\]](http://docs.aws.amazon.com/zh_tw/sagemaker/latest/dg/images/distributed/model-parallel/interleaved-pipeline-execution.png)
上圖為執行排程範例，說明 2 個 GPU 的 INTERLEAVED 管道。在圖中，F0 代表微批次 0 的向前傳遞，B1 代表微批次 1 的向後傳遞。**更新**代表參數的最佳化工具更新。GPU0 一律會盡可能優先處理向後傳遞 (例如，先執行 B0，然後再執行 F2)，這可允許清除之前用於啟用的記憶體。
### 簡便管道
相比之下，簡便管道會先完成執行每個微批次的向前傳遞，然後再開始向後傳遞。這代表其只在其管道內傳輸向前傳遞及向後傳遞階段。下圖範例說明這在 2 個 GPU 的運作方式。
![\[在開始後向傳遞之前，對每個微批次執行前向傳遞的管道範例。\]](http://docs.aws.amazon.com/zh_tw/sagemaker/latest/dg/images/distributed/model-parallel/simple-pipeline-execution.png)
### 在特定架構執行管道傳輸
請參閱下列各區段，了解 SageMaker 模型平行化程式庫針對TensorFlow 與 PyTorch 所做的特定架構管道排程決定。
#### 使用 TensorFlow 執行管道
下列映像範例說明採用自動化模型分割利用模型平行化程式庫分割TensorFlow 圖表。當分割圖表時，每個生成的子圖表會複寫 B 倍 (變數除外)，其中 B 是微批次數量。在此圖，每個子圖表均複寫 2 次 (B=2)。在子圖表的每個輸入處插入 `SMPInput` 作業，並在每個輸出處插入 `SMPOutput` 作業。這些作業會與程式庫後端進行通訊，以便彼此傳輸張量。
![\[使用自動模型分割，由模型平行化程式庫分割的 TensorFlow 圖表範例。\]](http://docs.aws.amazon.com/zh_tw/sagemaker/latest/dg/images/distributed/model-parallel/interleaved-pipeline-tf.png)
下列映像範例說明 2 個子圖表以 B=2 分割並新增漸層作業。`SMPInput` 作業的漸層是 `SMPOutput` 作業，反之亦然。這可讓漸層在反向傳播期間向後流動。
![\[以 B=2 分割的 2 個子圖的範例，並加入梯度操作。\]](http://docs.aws.amazon.com/zh_tw/sagemaker/latest/dg/images/distributed/model-parallel/interleaved-pipeline-tf.gif)
此 GIF 顯示範例 INTERLEAVED 管道執行排程，其中包含 B=2 微批次與 2 個子圖表。每個裝置都會循序執行其中一個子圖表複本，以便改善 GPU 使用率。隨著 B 變大，閒置時間時段的小部分會變為零。無論何時在特定子圖形複本執行 (向前或向後) 運算時，管道層都會向相應藍色 `SMPInput` 作業發出訊號以開始執行。
一旦運算單一最小批次所有微批次的漸層，程式庫會跨微批次結合漸層，然後可將其套用至參數。
#### 使用 PyTorch 執行管道
就概念而言，管道傳輸在 PyTorch 遵循類似想法。然而，由於 PyTorch 不涉及靜態圖表，因此模型平行化程式庫的 PyTorch 功能使用更加動態的管道傳輸範例。
與 TensorFlow 一樣，每個批次都會分成多個微批次，且每個裝置一次會執行一個微批次。不過，執行排程是透過啟動於每個裝置的執行伺服器來處理。每當目前裝置需要放置在另一裝置的子模組輸出時，執行請求將與子模組的輸入張量一起傳送至遠端裝置的執行伺服器。然後，伺服器使用指定輸入來執行此模組，並將回應傳回目前裝置。
由於目前裝置在遠端子模組執行期間處於閒置狀態，因此目前微批次的本機執行會暫停，且程式庫執行期將切換為執行目前裝置可有效處理的另一微批次。微批次的優先順序由所選擇的管道排程決定。對於 INTERLEAVED 管道排程，處於運算向後階段的微批次會盡可能優先考量。
# 張量平行處理
*張量平行化*是模型平行化類型，其中特定模型權重、漸層與最佳化工具狀態會跨裝置分割。有別於管道平行處理 (其可保持個別權重不變，但會分割權重*集*)，張量平行處理會分割個別權重。這通常涉及特定作業、模組或模型層的分散式運算。
如果單一參數使用多數 GPU 記憶體 (例如字彙量較大的大型嵌入資料表或具大量類別的大型 softmax 層)，則需要張量平行化。在這種情況，將此大型張量或作業視為原子單位不具效率，且會阻礙記憶體負載的平衡。
對於極大型模型而言，純管道傳輸完全不足以符合需求，此時，張量平行處理也很有幫助。例如，對於需要分割超過數十個執行個體的 GPT-3 規模模型，純微批次管道傳輸的效率不佳，因為管道深度過高，而且額外負荷變得過大。
**注意**
張量平行化適用 SageMaker 模型平行化程式庫 v1.6.0 及更高版本的 PyTorch。
**Topics**
+ [張量平行處理的運作方式](model-parallel-extended-features-pytorch-tensor-parallelism-how-it-works.md)
+ [使用張量平行化執行 SageMaker 分散式模型平行訓練任務](model-parallel-extended-features-pytorch-tensor-parallelism-examples.md)
+ [支援 Hugging Face 轉換器模型](model-parallel-extended-features-pytorch-hugging-face.md)
+ [在合併使用管道平行處理與張量平行處理組時的排名機制](model-parallel-extended-features-pytorch-ranking-mechanism.md)
# 張量平行處理的運作方式
張量平行處理可在 `nn.Modules` 層級運作；其可將模型的特定模組跨張量平行等級分割。這是管道平行處理所用的*模組集*現有分區之外的補充。
當模組透過張量平行處理進行分割時，會分散其向前與向後傳播。程式庫會處理所有裝置之間的必要通訊，以實作這些模組的分散式執行。模組會跨多個資料平行等級進行分割。與傳統工作負載發佈相反，當使用程式庫的張量平行處理時，每個資料平行等級均**無**完整模型複本。反之，除了整個未分散的模組，每個資料平行等級可能只有分散式模組的分區。
**範例：**考慮跨資料平行等級的張量平行處理，其中資料平行處理程度為 4，而張量平行處理程度為 2。假設您有資料平行群組，且在分割模組集之後，其包含下列模組樹狀目錄。
```
A
├── B
| ├── E
| ├── F
├── C
└── D
├── G
└── H
```
假設模組 B、G、H 支援張量平行處理，則此模型的張量平行分割可能結果之一如下：
```
dp_rank 0 (tensor parallel rank 0): A, B:0, C, D, G:0, H
dp_rank 1 (tensor parallel rank 1): A, B:1, C, D, G:1, H
dp_rank 2 (tensor parallel rank 0): A, B:0, C, D, G:0, H
dp_rank 3 (tensor parallel rank 1): A, B:1, C, D, G:1, H
```
每一行代表儲存在該 `dp_rank` 的模組集，表示法 `X:y` 代表模組 `X` 的第 `y` 個部分。注意下列事項：
1. 分割發生在跨資料平行級的子集 (稱為 `TP_GROUP`)，而非整個 `DP_GROUP`，以便跨 `dp_rank` 0 與 `dp_rank` 2 複寫確切的模型分割，並以類似方式跨 `dp_rank` 1 與 `dp_rank` 3 進行複寫。
1. 模組 `E` 與 `F` 不再是模型的一部分，這是因為其父模組 `B` 已分割，且在任何正常情況，屬於 `E` 與 `F` 一部分的任何執行會在 (已分割) `B` 模組進行。
1. 儘管張量平行處理支援 `H`，但在此範例並未進行分割，這強調顯示是否分割模組取決於使用者輸入。張量平行處理支援模組的事實未必代表會對其進行分割。
## 程式庫如何調整張量平行處理，使其適應 PyTorch `nn.Linear` 模組
當透過資料平行等級執行張量平行處理時，會*針對分割的模組*，在張量平行裝置之間分割參數、漸層與最佳化工具狀態的子集。對於模組的其餘部分，張量平行裝置會以一般資料平行方式操作。若要執行分割的模組，裝置會先在相同張量平行處理群組跨對等裝置收集*所有資料範例*的必要部分。然後，裝置會在所有這些資料範例執行模組的本機部分，接著再執行另一輪同步，這兩個階段都會合並每個資料範例的輸出部分，並將合併的資料範例傳回資料範例首次產生的 GPU。下圖顯示範例就已分割的 `nn.Linear` 模組說明此程序。
![\[兩個圖顯示兩個張量平行概念。\]](http://docs.aws.amazon.com/zh_tw/sagemaker/latest/dg/images/distributed/model-parallel/tensor-parallel-concept.png)
第一個圖顯示具大型 `nn.Linear` 模組的小型模型，在兩個張量平行處理等級進行資料平行處理。該 `nn.Linear` 模組複寫至兩個平行等級。
第二個圖顯示在分割 `nn.Linear` 模組時，套用張量平行處理至較大模型。每個 `tp_rank` 保留一半線性模組以及整個作業的其餘部分。當執行線性模組時，每個 `tp_rank` 都會針對所有資料範例收集相關的一半，並將其傳遞到其所屬 `nn.Linear` 模組的一半。結果需要減少散佈 (以求和作為減少作業)，以便每個等級都有其各自資料範例的最終線性輸出。模型的其餘部分以典型資料平行方式執行。
# 使用張量平行化執行 SageMaker 分散式模型平行訓練任務
在本區段，您會學習：
+ 如何設定 SageMaker PyTorch 預估器與 SageMaker 模型平行化選項，以便使用張量平行化。
+ 如何使用已延伸 `smdistributed.modelparallel` 模組來調整訓練指令碼，以達到張量平行處理。
若要進一步了解 `smdistributed.modelparallel` 模組，請參閱 [SageMaker Python SDK 文件](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html)的 *SageMaker 模型平行 API*。
**Topics**
+ [僅採用張量平行處理](#model-parallel-extended-features-pytorch-tensor-parallelism-alone)
+ [張量平行處理結合管道平行處理](#model-parallel-extended-features-pytorch-tensor-and-pipeline-parallelism)
## 僅採用張量平行處理
以下範例說明分散式訓練選項可單獨啟用張量平行處理，而無需管道平行處理。設定 `mpi_options` 與 `smp_options` 字典，指定分散式訓練選項至 SageMaker `PyTorch` 估算器。
**注意**
可透過適用 PyTorch 的 Deep Learning Containers 使用延伸節省記憶體功能，其會實作 SageMaker 模型平行化程式庫 v1.6.0 或更高版本。
**設定 SageMaker PyTorch 估算器**
```
mpi_options = {
"enabled" : True,
"processes_per_host" : 8, # 8 processes
"custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none "
}
smp_options = {
"enabled":True,
"parameters": {
"pipeline_parallel_degree": 1, # alias for "partitions"
"placement_strategy": "cluster",
"tensor_parallel_degree": 4, # tp over 4 devices
"ddp": True
}
}
smp_estimator = PyTorch(
entry_point='your_training_script.py', # Specify
role=role,
instance_type='ml.p3.16xlarge',
sagemaker_session=sagemaker_session,
framework_version='1.13.1',
py_version='py36',
instance_count=1,
distribution={
"smdistributed": {"modelparallel": smp_options},
"mpi": mpi_options
},
base_job_name="SMD-MP-demo",
)
smp_estimator.fit('s3://my_bucket/my_training_data/')
```
**提示**
若要尋找 `distribution` 的完整參數清單，請參閱 SageMaker Python SDK 文件的[模型平行化的設定參數](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html)。
**調整 PyTorch 訓練指令碼**
下列範例訓練指令碼顯示如何調整 SageMaker 模型平行化程式庫，以便適應訓練指令碼。在此範例，假設指令碼命名為 `your_training_script.py`。
```
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 Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 32, 3, 1)
self.conv2 = nn.Conv2d(32, 64, 3, 1)
self.fc1 = nn.Linear(9216, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = self.conv1(x)
x = F.relu(x)
x = self.conv2(x)
x = F.relu(x)
x = F.max_pool2d(x, 2)
x = torch.flatten(x, 1)
x = self.fc1(x)
x = F.relu(x)
x = self.fc2(x)
return F.log_softmax(x, 1)
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()
output = model(data)
loss = F.nll_loss(output, target, reduction="mean")
loss.backward()
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
if smp.local_rank() == 0:
dataset = datasets.MNIST("../data", train=True, download=False)
smp.barrier()
# 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()}")
train_loader = torch.utils.data.DataLoader(dataset, batch_size=64)
# smdistributed: Enable tensor parallelism for all supported modules in the model
# i.e., nn.Linear in this case. Alternatively, we can use
# smp.set_tensor_parallelism(model.fc1, True)
# to enable it only for model.fc1
with smp.tensor_parallelism():
model = Net()
# smdistributed: Use the DistributedModel wrapper to distribute the
# modules for which tensor parallelism is enabled
model = smp.DistributedModel(model)
optimizer = optim.AdaDelta(model.parameters(), lr=4.0)
optimizer = smp.DistributedOptimizer(optimizer)
train(model, device, train_loader, optimizer)
```
## 張量平行處理結合管道平行處理
以下是分散式訓練選項的範例，該選項可結合張量平行化與管道平行化。在設定 SageMaker `PyTorch` 估算器時，設定 `mpi_options` 與 `smp_options` 參數，以便指定具張量平行化的模型平行選項。
**注意**
可透過適用 PyTorch 的 Deep Learning Containers 使用延伸節省記憶體功能，其會實作 SageMaker 模型平行處理程式庫 v1.6.0 或更高版本。
**設定 SageMaker PyTorch 估算器**
```
mpi_options = {
"enabled" : True,
"processes_per_host" : 8, # 8 processes
"custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none "
}
smp_options = {
"enabled":True,
"parameters": {
"microbatches": 4,
"pipeline_parallel_degree": 2, # alias for "partitions"
"placement_strategy": "cluster",
"tensor_parallel_degree": 2, # tp over 2 devices
"ddp": True
}
}
smp_estimator = PyTorch(
entry_point='your_training_script.py', # Specify
role=role,
instance_type='ml.p3.16xlarge',
sagemaker_session=sagemaker_session,
framework_version='1.13.1',
py_version='py36',
instance_count=1,
distribution={
"smdistributed": {"modelparallel": smp_options},
"mpi": mpi_options
},
base_job_name="SMD-MP-demo",
)
smp_estimator.fit('s3://my_bucket/my_training_data/')
```
**調整 PyTorch 訓練指令碼**
下列範例訓練指令碼顯示如何調整 SageMaker 模型平行處理程式庫，以便適應訓練指令碼。請注意，訓練指令碼現在包含 `smp.step` 裝飾項目：
```
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 Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 32, 3, 1)
self.conv2 = nn.Conv2d(32, 64, 3, 1)
self.fc1 = nn.Linear(9216, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = self.conv1(x)
x = F.relu(x)
x = self.conv2(x)
x = F.relu(x)
x = F.max_pool2d(x, 2)
x = torch.flatten(x, 1)
x = self.fc1(x)
x = F.relu(x)
x = self.fc2(x)
return F.log_softmax(x, 1)
# 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
if smp.local_rank() == 0:
dataset = datasets.MNIST("../data", train=True, download=False)
smp.barrier()
# 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 = Net()
# smdistributed: enable tensor parallelism only for model.fc1
smp.set_tensor_parallelism(model.fc1, True)
# 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 = optim.AdaDelta(model.parameters(), lr=4.0)
optimizer = smp.DistributedOptimizer(optimizer)
train(model, device, train_loader, optimizer)
```
# 支援 Hugging Face 轉換器模型
SageMaker 模型平行化程式庫的張量平行化針對下列 Hugging Face 轉換器模型提供立即可用支援：
+ GPT-2、BERT 與 RoBERTa (適用 SageMaker 模型平行化程式庫 v1.7.0 及更高版本)
+ GPT-J (適用 SageMaker 模型平行化程度程式庫 v1.8.0 及更高版本)
+ GPT-NEO (適用 SageMaker 模型平行化程式庫 v1.10.0 及更高版本)
**注意**
對於任何其他轉換器模型，您需要使用 [smdistributed.modelparallel.torch.tp\$1register\$1with\$1module()](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch_tensor_parallel.html#smdistributed.modelparallel.torch.tp_register_with_module) API 來套用張量平行處理。
**注意**
若要使用張量平行化來訓練 Hugging Face 轉換器模型，請確定您使用 Hugging Face Deep Learning Containers (適用於具 SageMaker 模型平行化程式庫 v1.7.0 及更高版本的 PyTorch)。如需更多資訊，請參閱 [SageMaker 模型平行化程式庫版本備註](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_release_notes/smd_model_parallel_change_log.html)。
## 開箱即用的支援模型
對於開箱即用程式庫支援的 Hugging Face 轉換器模型，您不需要手動實作勾點，即可將轉換器 API 翻譯為 `smdistributed` 轉換器層。您可以透過使用內容管理器 [smdistributed.modelparallel.torch.tensor\$1parallelism()](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch_tensor_parallel.html#smdistributed.modelparallel.torch.tensor_parallelism) 來啟用張量平行處理，並利用 [smdistributed.modelparallel.torch.DistributedModel()](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.DistributedModel) 來包裝模型。您不需要使用 `smp.tp_register` API 手動註冊張量平行處理的勾點。
Hugging Face 轉換器之間與 `smdistributed.modelparallel` 之間的 `state_dict` 平移函式可按如下方式存取。
+ `smdistributed.modelparallel.torch.nn.huggingface.gpt2.translate_state_dict_to_hf_gpt2(state_dict, max_seq_len=None)`
+ `smdistributed.modelparallel.torch.nn.huggingface.gpt2.translate_hf_state_dict_to_smdistributed_gpt2(state_dict)`
+ `smdistributed.modelparallel.torch.nn.huggingface.bert.translate_state_dict_to_hf_bert(state_dict, max_seq_len=None)`
+ `smdistributed.modelparallel.torch.nn.huggingface.bert.translate_hf_state_dict_to_smdistributed_bert(state_dict)`
+ `smdistributed.modelparallel.torch.nn.huggingface.roberta.translate_state_dict_to_hf_roberta(state_dict, max_seq_len=None)`
+ `smdistributed.modelparallel.torch.nn.huggingface.roberta.translate_hf_state_dict_to_smdistributed_roberta(state_dict)`
+ `smdistributed.modelparallel.torch.nn.huggingface.gptj.translate_state_dict_to_hf_gptj(state_dict, max_seq_len=None)`(適用 SageMaker 模型平行化程式庫 v1.8.0 及更高版本)
+ `smdistributed.modelparallel.torch.nn.huggingface.gptj.translate_hf_gptj_state_dict_to_smdistributed_gptj`(適用 SageMaker 模型平行處理程式庫 v1.8.0 及更高版本)
+ `smdistributed.modelparallel.torch.nn.huggingface.gptneo.translate_state_dict_to_hf_gptneo(state_dict, max_seq_len=None)`(適用 SageMaker 模型平行化程式庫 v1.10.0 及更高版本)
+ `smdistributed.modelparallel.torch.nn.huggingface.gptneo.translate_hf_state_dict_to_smdistributed_gptneo(state_dict)`(適用 SageMaker 模型平行處理程式庫 v1.10.0 及更高版本)
**GPT-2 平移函式使用範例**
從包裝模型開始，如下列程式碼所示。
```
from transformers import AutoModelForCausalLM
with smp.tensor_parallelism():
model = AutoModelForCausalLM.from_config(hf_gpt2_config)
model = smp.DistributedModel(model)
```
指定來自 `DistributedModel` 物件的 `state_dict`，您可以使用下方所示程式碼的 `translate_state_dict_to_hf_gpt2` 功能載入權重至原始 Hugging Face GPT-2 模型。
```
from smdistributed.modelparallel.torch.nn.huggingface.gpt2 \
import translate_state_dict_to_hf_gpt2
max_seq_len = 1024
# [... code block for training ...]
if smp.rdp_rank() == 0:
state_dict = dist_model.state_dict()
hf_state_dict = translate_state_dict_to_hf_gpt2(state_dict, max_seq_len)
# can now call model.load_state_dict(hf_state_dict) to the original HF model
```
**RoBERTa 平移函式使用範例**
同樣地，指定支援的 HuggingFace 模型 `state_dict`，您可以透過 `translate_hf_state_dict_to_smdistributed` 函式將其轉換為 `smp.DistributedModel` 可讀取的格式。這在轉移學習使用案例很有用，其中將預先訓練模型載入 `smp.DistributedModel` 以便進行模型平行微調：
```
from smdistributed.modelparallel.torch.nn.huggingface.roberta \
import translate_state_dict_to_smdistributed
model = AutoModelForMaskedLM.from_config(roberta_config)
model = smp.DistributedModel(model)
pretrained_model = AutoModelForMaskedLM.from_pretrained("roberta-large")
translated_state_dict =
translate_state_dict_to_smdistributed(pretrained_model.state_dict())
# load the translated pretrained weights into the smp.DistributedModel
model.load_state_dict(translated_state_dict)
# start fine-tuning...
```
# 在合併使用管道平行處理與張量平行處理組時的排名機制
本區段說明模型平行化的排名機制如何搭配張量平行化。這是從 [SageMaker 模型平行化程式庫的核心功能](model-parallel-core-features.md) 的[排名基本資訊](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#ranking-basics)延伸而來。借助張量平行處理，程式庫引入三種類型的排名及處理程序群組 API：`smp.tp_rank()` 用於張量平行等級，`smp.pp_rank()` 用於管道並平行等級，以及 `smp.rdp_rank()` 用於縮減資料平行等級。對應的通訊處理程序群組是張量平行群組 (`TP_GROUP`)、管道平行群組 (`PP_GROUP`) 與縮減資料平行群組 (`RDP_GROUP`)。這些群組定義如下：
+ *張量平行群組* (`TP_GROUP`) 是資料並平行群組的可均分子集，模組在此進行張量平行發佈。當管道平行化程度為 1 時，`TP_GROUP` 與*模型平行群組* (`MP_GROUP`) 相同。
+ *管道平行群組* (`PP_GROUP`) 是進行管道平行處理的處理程序群組。當張量平行處理程度為 1 時，`PP_GROUP` 與 `MP_GROUP` 相同。
+ *縮減資料平行群組* (`RDP_GROUP`) 是一組處理程序，可同時保留相同管道平行處理分割區與相同張量平行分割區，並且彼此執行資料平行處理。這稱為縮減資料平行群組，因其為整個資料平行處理群組 `DP_GROUP` 的子集。對於分散在 `TP_GROUP` 內的模型參數，漸層 `allreduce` 作業僅針對縮減資料平行群組執行，而對於未分散的參數，漸層 `allreduce` 會在整個 `DP_GROUP` 進行。
+ 模型平行群組 (`MP_GROUP`) 是指集體儲存整個模型的處理程序群組。其包括在目前處理程序 `TP_GROUP` 內所有等級的 `PP_GROUP` 聯合。當張量平行處理程度為 1 時，`MP_GROUP` 等同於 `PP_GROUP`。它也與先前 `smdistributed` 發佈內容 `MP_GROUP` 的現有定義一致。請注意，目前 `TP_GROUP` 是目前 `DP_GROUP` 與目前 `MP_GROUP` 的子集。
若要進一步了解 SageMaker 模型平行化程式庫的通訊處理程序 API，請參閱 *SageMaker Python SDK 文件*的[常見 API](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_common_api.html#) 與 [Python 特定 API](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html)。
![\[張量平行化的排名機制、參數分佈和關聯的 AllReduce 操作。\]](http://docs.aws.amazon.com/zh_tw/sagemaker/latest/dg/images/distributed/model-parallel/tensor-parallel-ranking-mechanism.png)
例如，假設單一節點具 8 個 GPU 的處理程序群組，其中張量平行處理程度為 2，管道平行處理程度為 2，而資料平行處理程度為 4。上圖的上方中心部分顯示具 4 個圖層的模型範例。圖的左下與右下部分說明使用管道平行處理及張量平行處理分散在 4 個 GPU 的 4 層模型，其中張量平行處理用於中間兩層。這兩個下圖是簡便複本，用於說明不同群組範圍線。分割模型會跨 GPU 0-3 與 4-7 進行資料平行處理複寫。左下圖顯示 `MP_GROUP`、`PP_GROUP`、`TP_GROUP` 的定義。右下圖顯示同一組 GPU 的 `RDP_GROUP`、`DP_GROUP`、`WORLD`。具相同顏色圖層與圖層配量的漸層會一起 `allreduce`，以便進行資料平行處理。例如，第一層 (淺藍色) 跨 `DP_GROUP` 進行 `allreduce` 作業，而第二層深橘色配量則僅進行其處理程序 `RDP_GROUP` 內的 `allreduce` 作業。深紅色粗體箭頭代表具整個 `TP_GROUP` 批次的張量。
```
GPU0: pp_rank 0, tp_rank 0, rdp_rank 0, dp_rank 0, mp_rank 0
GPU1: pp_rank 1, tp_rank 0, rdp_rank 0, dp_rank 0, mp_rank 1
GPU2: pp_rank 0, tp_rank 1, rdp_rank 0, dp_rank 1, mp_rank 2
GPU3: pp_rank 1, tp_rank 1, rdp_rank 0, dp_rank 1, mp_rank 3
GPU4: pp_rank 0, tp_rank 0, rdp_rank 1, dp_rank 2, mp_rank 0
GPU5: pp_rank 1, tp_rank 0, rdp_rank 1, dp_rank 2, mp_rank 1
GPU6: pp_rank 0, tp_rank 1, rdp_rank 1, dp_rank 3, mp_rank 2
GPU7: pp_rank 1, tp_rank 1, rdp_rank 1, dp_rank 3, mp_rank 3
```
在此範例，管道平行處理是跨 GPU 配對 (0,1)；(2,3)；(4,5)；(6,7) 進行。此外，資料平行處理 (`allreduce`) 會跨 GPU 0、2、4、6 進行，並透過 GPU 1、3、5、7 獨立進行。張量平行處理發生於 `DP_GROUP` 子集及跨 GPU 配對 (0,2)；(1,3)；(4,6)；(5,7)。
對於這種混合管道與張量平行處理，`data_parallel_degree` 的數學運算仍是 `data_parallel_degree = number_of_GPUs / pipeline_parallel_degree`。程式庫會進一步從下列關聯 `reduced_data_parallel_degree * tensor_parallel_degree = data_parallel_degree` 計算縮減資料平行程度。
# 最佳化工具狀態碎片
*最佳化工具狀態碎片*是有用的節省記憶體技術，可跨資料平行裝置群組碎片化最佳化工具狀態 (描述最佳化工具狀態的權重集)。每當您使用具狀態的最佳化工具 (例如 Adam) 或 FP16 最佳化工具 (儲存參數的 FP16 與 FP32 副本) 時，都可使用最佳化工具狀態碎片。
**注意**
最佳化工具狀態碎片適用 SageMaker 模型平行化程式庫 v1.6.0 及更高版本的 PyTorch。
## 如何使用最佳化工具狀態碎片
您可以透過在 `modelparallel` 組態 設定 `"shard_optimizer_state": True` 來開啟*最佳化工具狀態碎片*。
當開啟此功能時，程式庫會基於資料平行處理程度分割模型參數集。對應於第 `i` 個分區的漸層僅會於第 `i` 個資料平行等級縮減。在第一次呼叫 `smp.step` 裝飾項目函式結束時，由 `smp.DistributedOptimizer` 包裝的最佳化工具會重新定義其參數，以便僅限對應目前資料平行等級分區的參數。重新定義的參數稱為*虛擬參數*，並與原始參數共用基礎儲存。在第一次呼叫 `optimizer.step` 期間，會基於這些重新定義的參數建立最佳化工具狀態，這些參數會因為原始分割而進行碎片處理。在更新最佳化工具之後，AllGather 作業 (做為 `optimizer.step` 呼叫的一部分) 會跨資料平行等級執行，以達到一致參數狀態。
**提示**
當資料平行處理的程度大於 1 且模型具十億個以上的參數時，最佳化工具狀態碎片很有幫助。
資料平行處理的程度由 `(processes_per_host * instance_count / pipeline_parallel_degree)` 計算，而 `smp.dp_size()` 函式會在背景處理大小。
**設定 SageMaker PyTorch 估算器**
```
mpi_options = {
"enabled" : True,
"processes_per_host" : 8, # 8 processes
"custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none "
}
smp_options = {
"enabled":True,
"parameters": {
"microbatches": 4,
"pipeline_parallel_degree": 2, # alias for "partitions"
"placement_strategy": "cluster",
"tensor_parallel_degree": 2, # tp over 2 devices
"ddp": True,
"shard_optimizer_state": True
}
}
```
**調整 PyTorch 訓練指令碼**
請參閱*結合管道平行處理的張量平行處理*區段的[調整 PyTorch 訓練指令碼](model-parallel-extended-features-pytorch-tensor-parallelism-examples.md#model-parallel-extended-features-pytorch-tensor-and-pipeline-parallelism-script)。此指令碼不需要其他修改。
# 啟用檢查點
*啟用檢查點* (或*漸層檢查點*) 是減少記憶體使用量的技術，方法是清除某些圖層的啟用，並在向後傳遞期間重新加以運算。實際上，這是以額外運算時間換取減少記憶體使用量。如果模組進行檢查點作業，則在向前傳遞結束時，模組的輸入與輸出都將保留在記憶體。在向前傳遞期間，原本會成為該模組內部運算一部分的任何中級張量都將被釋放。在檢查點模組的向後傳遞期間，會重新運算這些張量。此時，超出此檢查點模組的圖層已完成其向後傳遞，因此可降低運用檢查點的最高記憶體使用量。
**注意**
此功能適用 SageMaker 模型平行化程式庫 v1.6.0 及更高版本的 PyTorch。
## 如何使用啟用檢查點
當使用 `smdistributed.modelparallel` 時，您可以在模組的精細程度使用啟用檢查點。對於除 `torch.nn.Sequential` 外的所有 `torch.nn` 模組，僅當從管道平行處理的角度而言，模組樹狀目錄位於單一分割內時，您才能對其進行檢查點作業。對於 `torch.nn.Sequential` 模組，循序模組內部的每個模組樹狀目錄必須完全位於單一分割內，以便啟用檢查點作業。當您使用手動分割時，請注意這些限制。
當您使用[自動化模型分割](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-core-features.html#model-parallel-automated-model-splitting)時，您可以在訓練任務日誌找到開頭為 `Partition assignments:` 的分割指派日誌。如果跨多個等級分割模組 (例如，其中一個子代位於某一等級，另一子代位於不同等級)，程式庫會忽略而不嘗試對模組進行檢查點作業，並提出警告訊息，指出不會檢查該模組。
**注意**
SageMaker 模型平行化程式庫支援重疊與非重疊 `allreduce` 作業，並結合檢查點。
**注意**
PyTorch 的原生檢查點 API 不相容 `smdistributed.modelparallel`。
**範例 1：**下列範例程式碼示範當指令碼具模型定義時，如何使用啟用檢查點。
```
import torch.nn as nn
import torch.nn.functional as F
from smdistributed.modelparallel.torch.patches.checkpoint import checkpoint
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 32, 3, 1)
self.conv2 = nn.Conv2d(32, 64, 3, 1)
self.fc1 = nn.Linear(9216, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = F.max_pool2d(x, 2)
x = torch.flatten(x, 1)
# This call of fc1 will be checkpointed
x = checkpoint(self.fc1, x)
x = self.fc2(x)
return F.log_softmax(x, 1)
```
**範例 2：**下列範例程式碼示範當指令碼具循序模型時，如何使用啟用檢查點。
```
import torch.nn as nn
from smdistributed.modelparallel.torch.patches.checkpoint import checkpoint_sequential
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.seq = nn.Sequential(
nn.Conv2d(1,20,5),
nn.ReLU(),
nn.Conv2d(20,64,5),
nn.ReLU()
)
def forward(self, x):
# This call of self.seq will be checkpointed
x = checkpoint_sequential(self.seq, x)
return F.log_softmax(x, 1)
```
**範例 3：**下列範例程式碼示範當從程式庫 (例如 PyTorch 與 Hugging Face 轉換器) 匯入預先建置的模型時，如何使用啟用檢查點。無論您是否針對循序模組進行檢查點作業，請執行以下操作：
1. 以 `smp.DistributedModel()` 包裝模型。
1. 定義循序圖層物件。
1. 以 `smp.set_activation_checkpointig()` 包裝循序圖層物件。
```
import smdistributed.modelparallel.torch as smp
from transformers import AutoModelForCausalLM
smp.init()
model = AutoModelForCausalLM(*args, **kwargs)
model = smp.DistributedModel(model)
# Call set_activation_checkpointing API
transformer_layers = model.module.module.module.transformer.seq_layers
smp.set_activation_checkpointing(
transformer_layers, pack_args_as_tuple=True, strategy='each')
```
# 啟用卸載
當開啟啟用檢查點與管道平行處理，且微批次數量大於一時，*啟用卸載*是可進一步減少記憶體使用量的附加功能。啟用卸載會以非同步方式對應目前未在 CPU 執行的微批次來移動檢查點啟用。在 GPU 需要啟用以便進行微批次向後傳遞之前，此功能會從 CPU 預先取回已卸載的啟用。
**注意**
此功能適用 SageMaker 模型平行處理程式庫 v1.6.0 及更高版本的 PyTorch。
## 如何使用啟用卸載
當**微批次數量大於 1，且已開啟啟用檢查點**時，使用啟用卸載以減少記憶體使用量 (請參閱[啟用檢查點](model-parallel-extended-features-pytorch-activation-checkpointing.md))。當未使用啟用檢查點時，啟用卸載不具任何效果。當僅搭配單一微批次使用時，此作法不會節省記憶體。
若要使用啟用卸載，請在 `modelparallel` 配置設定 `"offload_activations": True`。
啟用卸載會以非同步方式移動 `nn.Sequential` 模組的檢查點啟用至 CPU。透過 PCIe 連結的資料傳輸會與 GPU 運算重疊。在運算特定檢查點層的向前傳遞之後，卸載會立即發生。在特定微批次的向後傳遞需要啟用之前不久，這些啟用就會載回 GPU。CPU-GPU 傳輸以相似方式與運算重疊。
若要調整啟用提早多久載回 GPU，您可以使用設定參數 `"activation_loading_horizon"` (預設值設定為 4，必須為大於 0 的 `int`)。較大的啟用載入時間將導致啟用更早載回 GPU。如果時間太大，可能減少啟用卸載所造成的節省記憶體影響。如果時間太小，則可能無法及時載入啟用，進而減少重疊並降低效能。
**提示**
對於具超過 1000 億個參數的大型模型，啟用卸載很有幫助。
**設定 SageMaker PyTorch 估算器**
```
mpi_options = {
"enabled" : True,
"processes_per_host" : 8, # 8 processes
"custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none "
}
smp_options = {
"enabled":True,
"parameters": {
"microbatches": 4,
"pipeline_parallel_degree": 2, # alias for "partitions"
"placement_strategy": "cluster",
"tensor_parallel_degree": 2, # tp over 2 devices
"ddp": True,
"offload_activations": True,
"activation_loading_horizon": 4 # optional. default is 4.
}
}
```
# 使用模型平行化進行 FP16 訓練
若要進行 FP16 訓練，請套用下列修改至訓練指令碼與估算器。
**注意**
此功能適用 SageMaker 模型平行化程式庫 v1.10.0 及更高版本的 PyTorch。
**調整 PyTorch 訓練指令碼**
1. 使用 [smdistributed.modelparallel.torch.model\$1creation()](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.model_creation) 內容管理器來包裝模型。
```
# fp16_training_script.py
import torch
import smdistributed.modelparallel.torch as smp
with smp.model_creation(
dtype=torch.float16 if args.fp16 else torch.get_default_dtype()
):
model = ...
```
**提示**
如果您使用張量平行處理，請新增 `tensor_parallelism=smp.tp_size() > 1` 至 `smp.model_creation` 內容管理器。新增此行也有助於自動偵測張量平行處理是否已啟動。
```
with smp.model_creation(
... ,
tensor_parallelism=smp.tp_size() > 1
):
model = ...
```
1. 當您使用 `smdistributed.modelparallel.torch.DistributedOptimizer` 包裝最佳化工具時，請設定 `static_loss_scaling` 或 `dynamic_loss_scaling` 引數。依預設，`static_loss_scaling` 設定為 `1.0`，`dynamic_loss_scaling` 設定為 `False`。如果設定 `dynamic_loss_scale=True`，您可以透過 `dynamic_loss_args` 引數饋送動態損失縮放選項作為字典。在多數情況，我們建議您使用動態損失縮放與預設選項。有關最佳化工具包裝函式的更多資訊、選項與範例，請參閱 [smdistributed.modelparallel.torch.DistributedOptimizer](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed-modelparallel-torch-distributedoptimizer) API。
下列程式碼範例示範使用 FP16 訓練的動態損失縮放來包裝 `Adadelta` 最佳化工具物件。
```
optimizer = torch.optim.Adadelta(...)
optimizer = smp.DistributedOptimizer(
optimizer,
static_loss_scale=None,
dynamic_loss_scale=True,
dynamic_loss_args={
"scale_window": 1000,
"min_scale": 1,
"delayed_shift": 2
}
)
```
**設定 SageMaker PyTorch 估算器**
在建立 SageMaker PyTorch 估算器物件時，新增 FP16 參數 (`"fp16"`) 至模型平行化的發佈組態。如需模型平行化設定參數的完整清單，請參閱 [`smdistributed` 的參數](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#parameters-for-smdistributed)。
```
from sagemaker.pytorch import PyTorch
smp_options = {
"enabled": True,
"parameters": {
"microbatches": 4,
"pipeline_parallel_degree": 2,
"tensor_parallel_degree": 2,
...,
"fp16": True
}
}
fp16_estimator = PyTorch(
entry_point="fp16_training_script.py", # Specify your train script
...,
distribution={
"smdistributed": {"modelparallel": smp_options},
"mpi": {...}
}
)
fp16_estimator.fit(...)
```
當 FP16 訓練開始時，模型與最佳化工具分別由 `FP16_Module` 與 `FP16_Optimizer` 包裝，這些都是 [Apex utils](https://nvidia.github.io/apex/fp16_utils.html#apex-fp16-utils) 的修改 `smdistributed` 版本。`FP16_Module` 將模型轉換為 FP16 dtype，並在 FP16 處理向前傳遞。
**提示**
您可以在 `optimizer.step` 之前透過呼叫 `clip_master_grads` 來套用漸層剪輯。
```
optimizer.clip_master_grads(max_norm) # max_norm(float or int): max norm of the gradients
```
**提示**
當使用 `torch.optim.lr_scheduler` 與 FP16 訓練時，您需要傳遞 `optimizer.optimizer` 至 LR 排程器，而非最佳化工具。請參閱以下範例程式碼。
```
from torch.optim.lr_scheduler import StepLR
scheduler = StepLR(
optimizer.optimizer if smp.state.cfg.fp16 else optimizer,
step_size=1,
gamma=args.gamma
)
```
# 支援 FlashAttention
支援 FlashAttention 是僅適用*分散式轉換器*模型程式庫的功能，該模型是以 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed-modelparallel-torch-distributedmodel](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed-modelparallel-torch-distributedmodel) 包裝的轉換器模型，用於模型平行訓練。此功能也相容 [張量平行處理](model-parallel-extended-features-pytorch-tensor-parallelism.md)。
僅當 `attention_head_size` 所設定的值為 8 的倍數且小於 128 時，[FlashAttention](https://github.com/HazyResearch/flash-attention) 程式庫才會支援模型。因此，當您訓練分散式轉換器並確保 FlashAttention 正常運作時，您應調整參數，以便注意頭大小符合要求。如需更多資訊，另請參閲 *FlashAttention GitHub 儲存庫*的[安裝與功能](https://github.com/HazyResearch/flash-attention#installation-and-features)。
例如，假設您使用 `hidden_width=864` 與 `num_heads=48` 設定轉換器模型。FlashAttention 的頭大小計算方式為 `attention_head_size = hidden_width / num_heads = 864 / 48 = 18`。若要啟用 FlashAttention，您需要調整 `num_heads` 參數為 `54`，以便 `attention_head_size = hidden_width / num_heads = 864 / 54 = 16` (這是 8 的倍數)。
# 執行具有模型平行的 SageMaker 分散式訓練任務
了解如何使用 SageMaker Python SDK 搭配 SageMaker 模型平行程式庫，針對您自己的訓練指令碼執行模型平行訓練任務。
執行 SageMaker 訓練任務有三種使用案例。
1. 您可以使用其中一個預先建置的 TensorFlow 和 PyTorch AWS 深度學習容器。如果您是首次使用模型平行程式庫，則建議您使用此選項。若要尋找有關如何執行 SageMaker AI 模型平行訓練任務的教學課程，請參閱[使用 Amazon SageMaker AI 模型平行化程式庫進行的 PyTorch 訓練](https://github.com/aws/amazon-sagemaker-examples/tree/main/training/distributed_training/pytorch/model_parallel)中的範例筆記本。
1. 您也可以擴充預先建置的容器，為您的演算法或預先建置 SageMaker Docker 映像不支援的模型處理任何額外的功能要求。若要尋找如何擴充預先建置容器的範例，請參閱[延伸預先建置的容器](prebuilt-containers-extend.md)。
1. 您可以使用 [SageMaker 訓練工具組](https://github.com/aws/sagemaker-training-toolkit)，調整您自己的 Docker 容器，以便與 SageMaker AI 搭配使用。有關範例，請參閱[調整您自己的訓練容器](https://docs.aws.amazon.com/sagemaker/latest/dg/adapt-training-container.html)。
如需上述清單中的選項 2 和 3 的說明，請參閱[擴充包含 SageMaker 分散式模型平行程式庫的預先建置 Docker 容器](model-parallel-sm-sdk.md#model-parallel-customize-container)，了解如何在擴充或自訂 Docker 容器中安裝模型平行程式庫。
在所有情況下，您都可以啟動訓練任務組態，設定 SageMaker `TensorFlow` 或 `PyTorch` 估算器以啟動程式庫。如需詳細資訊，請參閱下列主題。
**Topics**
+ [步驟 1：使用 SageMaker 的分散式模型平行程式庫修改訓練指令碼](model-parallel-customize-training-script.md)
+ [步驟 2：使用 SageMaker Python SDK 啟動訓練任務](model-parallel-sm-sdk.md)
# 步驟 1：使用 SageMaker 的分散式模型平行程式庫修改訓練指令碼
請使用此節來學習如何自訂訓練指令碼，以使用 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 模型平行處理程式庫分割訓練指令碼的模型
修改訓練指令碼以設定模型分割的方法有兩種：自動化分割或手動分割。
### 自動化模型分割
使用 SageMaker 模型平行處理程式庫時，您可以利用*自動化模型分割*，也稱為*自動化模型分區*。該程式庫採用的分割演算法能平衡記憶體的消耗，最小化裝置間的通訊並最佳化效能。您可以設定自動化磁碟分割演算法，以最佳化速度或記憶體用量。
或者，您可以手動進行模型分割。除非對模型架構極為熟悉，並清楚如何有效進行模型分割，否則，我們建議使用自動化模型分割。
#### 運作方式
在第一個訓練步驟期間，自動分割會在首次呼叫 `smp.step`-裝飾函式時進行。在呼叫期間，程式庫會先在 CPU RAM 上建構模型版本 (以避免 GPU 記憶體限制)，接著分析模型圖表，並做出分割決定。基於這個決定，每個模型分割區被載入至 GPU 上，接下來執行第一個步驟。由於這些分析和分割步驟，首次訓練步驟可能較為耗時。
在任一架構中，程式庫會透過針對 AWS 基礎設施最佳化的自有後端來管理裝置之間的通訊。
自動分割設計會因應架構特性調整，而程式庫在精細程度層進行分割，在各架構內更加自然。例如在 TensorFlow 中，每個特定操作都可以指派給不同裝置，而在 PyTorch 中，指派作業是在模組層完成的，其中每個模組由多個操作組成。下一節將檢閱每個架構中的設計細節。
##### 使用 PyTorch 自動化模型分割
在首次訓練步驟中，模型平行處理程式庫會在內部執行追蹤步驟，目的是建構模型圖表，並判定張量和參數形狀。完成追蹤步驟後，程式庫會建構一個樹狀圖，其中包含模型中的巢狀`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 自動化分割模型
模型平行處理程式庫會分析可訓練變數和圖表結構大小，並在內部使用圖表分割演算法。此演算法會為每項操作指派裝置，目標是在下列兩個限制條件下，將裝置間所需的通訊量降至最低：
+ 平衡各裝置中儲存的變數量
+ 平衡各裝置執行的操作數量
如果在 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)。
##### 比較不同架構之間的自動化模型分割
在 TensorFlow 中，運算的基本單位是 `tf.Operation`，TensorFlow 將模型表示為 `tf.Operation` 的有向無環圖 (DAG)，因此模型平行處理程式庫會分割此 DAG，讓每個節點各指派到一個裝置。關鍵是，`tf.Operation` 物件具備多樣化的可自訂屬性，並在某種程度上是通用的，因為每個模型都必定包含以這類物件組成的圖表。
而對 PyTorch 來說，則沒有多樣化且通用的對等操作概念。具有這些特性的 PyTorch 中最接近的運算單元是 `nn.Module`，它處於更高的粒度層，這就是為什麼程式庫在 PyTorch 中會在此層進行分割。
### 手動模型分割
若需手動指定模型在各裝置間的分割方式，請使用 `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 訓練指令碼
在本節中，您將學習如何修改 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 自動化分割
若要使用 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 進行混合模型和資料平行處理的自動化分割。
您可以透過 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 手動分割
使用 `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()
```
## 不支援架構功能
程式庫不支援下列 TensorFlow 功能：
+ 目前不支援 `tf.GradientTape()`。您可以使用 `Optimizer.get_gradients()` 或 `Optimizer.compute_gradients()` 來運算漸層。
+ 目前不支援 `tf.train.Checkpoint.restore()` API。如為檢查點，請改為使用 `smp.CheckpointManager` 提供相同的 API 和功能。請注意，`smp.CheckpointManager` 的檢查點還原應該在第一步驟後進行。
# 修改 PyTorch 訓練指令碼
在本節中，您將了解如何修改 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 自動化分割
若要使用 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 手動分割
使用 [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)
```
## 考量事項
使用 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
```
## 不支援架構功能
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 來使用這些程式庫的實作。
# 步驟 2：使用 SageMaker Python SDK 啟動訓練任務
SageMaker Python SDK 支援使用機器學習 (ML) 架構 (例如 TensorFlow 和 PyTorch) 模型進行受管訓練。若要使用其中一個架構啟動訓練任務，您可以定義 SageMaker [ TensorFlow 估算器](https://sagemaker.readthedocs.io/en/v2.199.0/frameworks/tensorflow/sagemaker.tensorflow.html#tensorflow-estimator)、SageMaker [PyTorch 估算器](https://sagemaker.readthedocs.io/en/v2.199.0/frameworks/pytorch/sagemaker.pytorch.html#pytorch-estimator)或 SageMaker 一般[估計器](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/estimators.html#sagemaker.estimator.Estimator)，以使用修改過的訓練指令碼和模型平行組態。
**Topics**
+ [使用 SageMaker TensorFlow 和 PyTorch 估算器](#model-parallel-using-sagemaker-pysdk)
+ [擴充包含 SageMaker 分散式模型平行程式庫的預先建置 Docker 容器](#model-parallel-customize-container)
+ [使用 SageMaker 分散式模型平行程式庫建立您自己的 Docker 容器](#model-parallel-bring-your-own-container)
## 使用 SageMaker TensorFlow 和 PyTorch 估算器
TensorFlow 和 PyTorch 估算器類別包含 `distribution` 參數，方便您用來指定使用分散式訓練架構的組態參數。SageMaker 模型平行程式庫的內部會使用 MPI 來處理混合資料和模型平行，因此您必須將 MPI 選項與程式庫搭配使用。
下列 TensorFlow 或 PyTorch 估算器範本會示範，如何設定 `distribution` 參數，藉此將 SageMaker 模型平行程式庫與 MPI 搭配使用。
------
#### [ Using the SageMaker TensorFlow estimator ]
```
import sagemaker
from sagemaker.tensorflow import TensorFlow
smp_options = {
"enabled":True, # Required
"parameters": {
"partitions": 2, # Required
"microbatches": 4,
"placement_strategy": "spread",
"pipeline": "interleaved",
"optimize": "speed",
"horovod": True, # Use this for hybrid model and data parallelism
}
}
mpi_options = {
"enabled" : True, # Required
"processes_per_host" : 8, # Required
# "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none"
}
smd_mp_estimator = TensorFlow(
entry_point="your_training_script.py", # Specify your train script
source_dir="location_to_your_script",
role=sagemaker.get_execution_role(),
instance_count=1,
instance_type='ml.p3.16xlarge',
framework_version='2.6.3',
py_version='py38',
distribution={
"smdistributed": {"modelparallel": smp_options},
"mpi": mpi_options
},
base_job_name="SMD-MP-demo",
)
smd_mp_estimator.fit('s3://my_bucket/my_training_data/')
```
------
#### [ Using the SageMaker PyTorch estimator ]
```
import sagemaker
from sagemaker.pytorch import PyTorch
smp_options = {
"enabled":True,
"parameters": { # Required
"pipeline_parallel_degree": 2, # Required
"microbatches": 4,
"placement_strategy": "spread",
"pipeline": "interleaved",
"optimize": "speed",
"ddp": True,
}
}
mpi_options = {
"enabled" : True, # Required
"processes_per_host" : 8, # Required
# "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none"
}
smd_mp_estimator = PyTorch(
entry_point="your_training_script.py", # Specify your train script
source_dir="location_to_your_script",
role=sagemaker.get_execution_role(),
instance_count=1,
instance_type='ml.p3.16xlarge',
framework_version='1.13.1',
py_version='py38',
distribution={
"smdistributed": {"modelparallel": smp_options},
"mpi": mpi_options
},
base_job_name="SMD-MP-demo",
)
smd_mp_estimator.fit('s3://my_bucket/my_training_data/')
```
------
若要啟用程式庫，您需要透過 SageMaker 估算器建構函式的 `distribution` 引數，將組態字典傳遞至 `"smdistributed"` 和 `"mpi"` 金鑰。
**SageMaker 模型平行的組態參數**
+ 若為 `"smdistributed"` 金鑰，利用 `"modelparallel"` 金鑰來傳遞字典及下列內部字典。
**注意**
系統不支援在單一訓練任務中使用 `"modelparallel"` 和 `"dataparallel"`。
+ `"enabled"` - 必要。若要啟用模型平行，請設定 `"enabled": True`。
+ `"parameters"` - 必要。為 SageMaker 模型平行指定一組參數。
+ 如需常用參數的完整清單，請參閱 *SageMaker Python SDK 文件中*的[針對 `smdistributed` 的參數](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#smdistributed-parameters)。
對於 TensorFlow，請參閱 [TensorFlow 特定的參數](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#tensorflow-specific-parameters)。
對於 PyTorch，請參閱 [PyTorch 特定的參數](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#pytorch-specific-parameters)。
+ `"pipeline_parallel_degree"` (或 `smdistributed-modelparallel
若要擴充預先建置的容器並使用 SageMaker 的模型平行處理程式庫，您必須使用其中一個可用的 AWS PyTorch 或 TensorFlow 深度學習容器 (DLC) 映像。SageMaker 模型平行程式庫隨附於 TensorFlow (2.3.0 及更新版本) 和具有 CUDA (`cuxyz`) 的 PyTorch (1.6.0 及更新版本) DLC 映像。如需 DLC 映像的完整清單，請參閱 *AWS 深度學習容器 GitHub 儲存庫*中的[可用深度學習容器映像](https://github.com/aws/deep-learning-containers/blob/master/available_images.md)。
**提示**
建議您使用包含 TensorFlow 或 PyTorch 最新版本的映像，來存取最新版本的 SageMaker 模型平行程式庫。
例如，您的 Dockerfile 應包含類似下列的 `FROM` 陳述式：
```
# Use the SageMaker DLC image URI for TensorFlow or PyTorch
FROM aws-dlc-account-id.dkr.ecr.aws-region.amazonaws.com/framework-training:{framework-version-tag}
# Add your dependencies here
RUN ...
ENV PATH="/opt/ml/code:${PATH}"
# this environment variable is used by the SageMaker AI container to determine our user code directory.
ENV SAGEMAKER_SUBMIT_DIRECTORY /opt/ml/code
```
此外，當您定義 PyTorch 或 TensorFlow 估算器時，您必須為您的訓練指令碼指定該 `entry_point` 項目。這應與 Dockerfile 中 `ENV SAGEMAKER_SUBMIT_DIRECTORY` 識別的路徑相同。
**提示**
您必須將此 Docker 容器推送到 Amazon Elastic Container Registry (Amazon ECR)，並使用映像 URI (`image_uri`) 來定義 SageMaker 估算器以進行訓練。如需詳細資訊，請參閱[延伸預先建置的容器](prebuilt-containers-extend.md)。
完成託管 Docker 容器並擷取容器的映象 URI 之後，請依照下列步驟建立 SageMaker `PyTorch` 估算器物件。此範例假設您已定義 `smp_options` 和 `mpi_options`。
```
smd_mp_estimator = Estimator(
entry_point="your_training_script.py",
role=sagemaker.get_execution_role(),
instance_type='ml.p3.16xlarge',
sagemaker_session=sagemaker_session,
image_uri='your_aws_account_id.dkr.ecr.region.amazonaws.com/name:tag'
instance_count=1,
distribution={
"smdistributed": smp_options,
"mpi": mpi_options
},
base_job_name="SMD-MP-demo",
)
smd_mp_estimator.fit('s3://my_bucket/my_training_data/')
```
## 使用 SageMaker 分散式模型平行程式庫建立您自己的 Docker 容器
若要建立自己的 Docker 容器以進行訓練並使用 SageMaker 模型平行程式庫，您必須在 Docerfile 中納入 SageMaker 分散式平行程式庫正確的相依性和二進位檔案。本節提供您必須包含的最少程式碼區塊組合，才能在您自己的 Docker 容器中妥善準備 SageMaker 訓練環境和模型平行程式庫。
**注意**
此自訂 Docker 選項將 SageMaker 模型平行程式庫做為二進位程式庫，僅適用於 PyTorch。
**使用 SageMaker 訓練工具組和模型平行程式庫建立 Dockerfile**
1. 從其中一個 [NVIDIA CUDA 基礎映像](https://hub.docker.com/r/nvidia/cuda)開始著手。
```
FROM
```
**提示**
官方 AWS 深度學習容器 (DLC) 映像是從 [NVIDIA CUDA 基礎映像](https://hub.docker.com/r/nvidia/cuda)建置而成。我們建議您查看[適用於 PyTorch 的 AWS 深度學習容器官方 Dockerfiles](https://github.com/aws/deep-learning-containers/tree/master/pytorch/training/docker)，以尋找您需要安裝的程式庫版本，以及如何設定它們。官方版 Dockerfiles 經過 SageMaker 和深度學習容器服務團隊完整的基準測試和管理。在提供的連結中，選擇您使用的 PyTorch 版本、選擇 CUDA (`cuxyz`) 資料夾，然後選擇以 `.gpu` 或 `.sagemaker.gpu` 結尾結尾的 Dockerfile。
1. 若要設定分散式訓練環境，您需要安裝適用於通訊和網路裝置的軟體，例如 [Elastic Fabric Adapter (EFA)](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa.html)、[NVIDIA 集合通訊程式庫 (NCCL)](https://developer.nvidia.com/nccl) 和 [Open MPI](https://www.open-mpi.org/)。視您選擇的 PyTorch 和 CUDA 版本而定，您必須安裝相容的程式庫版本。
**重要**
由於 SageMaker 模型平行程式庫需要在後續步驟中使用 SageMaker 資料平行程式庫，因此我們強烈建議您遵循[使用 SageMaker AI 分散式資料平行程式庫建立您自己的 Docker 容器](data-parallel-bring-your-own-container.md)的指示，為分散式訓練妥善設定 SageMaker 訓練環境。
如需有關使用 NCCL 和 Open MPI 設定 EFA 的詳細資訊，請參閱[開始使用 EFA 和 MPI](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa-start.html) 以及[開始使用 EFA 和 NCCL](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa-start-nccl.html)。
1. 新增下列引數，以指定適用於 PyTorch 的 SageMaker 分散式訓練套件的 URL。SageMaker 模型平行程式庫需要 SageMaker 資料平行程式庫，才能使用跨節點遠端直接記憶體存取 (RDMA)。
```
ARG SMD_MODEL_PARALLEL_URL=https://sagemaker-distributed-model-parallel.s3.us-west-2.amazonaws.com/pytorch-1.10.0/build-artifacts/2022-02-21-19-26/smdistributed_modelparallel-1.7.0-cp38-cp38-linux_x86_64.whl
ARG SMDATAPARALLEL_BINARY=https://smdataparallel.s3.amazonaws.com/binary/pytorch/1.10.2/cu113/2022-02-18/smdistributed_dataparallel-1.4.0-cp38-cp38-linux_x86_64.whl
```
1. 安裝 SageMaker 模型平行程式庫所需的相依性。
1. 安裝 [METIS](http://glaros.dtc.umn.edu/gkhome/metis/metis/overview) 程式庫。
```
ARG METIS=metis-5.1.0
RUN rm /etc/apt/sources.list.d/* \
&& wget -nv http://glaros.dtc.umn.edu/gkhome/fetch/sw/metis/${METIS}.tar.gz \
&& gunzip -f ${METIS}.tar.gz \
&& tar -xvf ${METIS}.tar \
&& cd ${METIS} \
&& apt-get update \
&& make config shared=1 \
&& make install \
&& cd .. \
&& rm -rf ${METIS}.tar* \
&& rm -rf ${METIS} \
&& rm -rf /var/lib/apt/lists/* \
&& apt-get clean
```
1. 安裝 [RAPIDS 記憶體管理員程式庫](https://github.com/rapidsai/rmm#rmm-rapids-memory-manager)。系統須使用 [CMake](https://cmake.org/) 3.14 或更高版本。
```
ARG RMM_VERSION=0.15.0
RUN wget -nv https://github.com/rapidsai/rmm/archive/v${RMM_VERSION}.tar.gz \
&& tar -xvf v${RMM_VERSION}.tar.gz \
&& cd rmm-${RMM_VERSION} \
&& INSTALL_PREFIX=/usr/local ./build.sh librmm \
&& cd .. \
&& rm -rf v${RMM_VERSION}.tar* \
&& rm -rf rmm-${RMM_VERSION}
```
1. 安裝 SageMaker 模型平行程式庫。
```
RUN pip install --no-cache-dir -U ${SMD_MODEL_PARALLEL_URL}
```
1. 安裝 SageMaker 資料平行程式庫。
```
RUN SMDATAPARALLEL_PT=1 pip install --no-cache-dir ${SMDATAPARALLEL_BINARY}
```
1. 安裝 [sagemaker-training 工具組](https://github.com/aws/sagemaker-training-toolkit)。此工具組包含建立與 SageMaker 訓練平台和 SageMaker Python SDK 相容的容器所需的常見功能。
```
RUN pip install sagemaker-training
```
1. 完成建立 Dockerfile 之後，請參閱[調整您自己的訓練容器](https://docs.aws.amazon.com/sagemaker/latest/dg/adapt-training-container.html)，了解如何建置 Docker 容器並將其託管在 Amazon ECR 中。
**提示**
如需有關在 SageMaker AI 中建立自訂 Dockerfile 以進行訓練的詳細資訊，請參閱[使用您自己的訓練演算法](https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-training-algo.html)。
# 對具有模型並行性的模型執行檢查點和微調
SageMaker 模型平行處理程式庫提供檢查點 API，可儲存模型狀態和依據各種模型平行處理策略分割的最佳化工具狀態，並從您要重新開始訓練和微調的位置載入繼續訓練的檢查點。這些 API 還支援部分或全部儲存模型和最佳化工具狀態的選項。
**Topics**
+ [對分散式模型進行檢查點](#distributed-model-parallel-checkpoint)
+ [微調分散式模型](#distributed-model-parallel-fine-tuning)
## 對分散式模型進行檢查點
根據 PyTorch 和 TensorFlow 之間的架構，以及您使用的 SageMaker 模型平行處理程式庫版本，選擇下列其中一個主題。
**Topics**
+ [對分散式 PyTorch 模型執行檢查點 (適用於 SageMaker 模型平行處理程式庫 1.10.0 版及更新版本)](#model-parallel-extended-features-pytorch-checkpoint)
+ [對分散式 PyTorch 模型執行檢查點 (適用於 1.6.0 版和 1.9.0 版之間的 SageMaker 模型平行處理程式庫)](#model-parallel-extended-features-pytorch-saving-loading-checkpoints)
+ [對分散式 TensorFlow 模型執行檢查點](#distributed-model-parallel-checkpoint-tensorflow)
### 對分散式 PyTorch 模型執行檢查點 (適用於 SageMaker 模型平行處理程式庫 1.10.0 版及更新版本)
SageMaker 模型平行處理程式庫提供檢查點 API，以儲存和載入分散式模型狀態及其最佳化工具狀態的完整或部分檢查點。
**注意**
如果您使用 PyTorch 和 SageMaker 模型平行處理程式庫 1.10.0 版或更新版本，則建議使用此檢查點方法。
**部分檢查點**
若要儲存使用模型平行處理訓練的模型檢查點，請使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.save_checkpoint](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.save_checkpoint)API，並將部分檢查點選項設定為 true (`partial=True`)。這會個別儲存每個模型分割區。除了模型和最佳化工具狀態之外，您還可以透過 `user_content` 引數儲存任何其他自訂資料。檢查點模型、最佳化工具和使用者內容都儲存為個別檔案。`save_checkpoint` API 呼叫會以下列結構建立檢查點資料夾。
```
- path
- ${tag}_partial (folder for partial checkpoints)
- model_rankinfo.pt
- optimizer_rankinfo.pt
- fp16_states_rankinfo.pt
- user_content.pt
- $tag (checkpoint file for full checkpoints)
- user_content_$tag (user_content file for full checkpoints)
- newest (a file that indicates the newest checkpoint)
```
若要從部分檢查點繼續訓練，請將 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.resume_from_checkpoint](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.resume_from_checkpoint) API 與 `partial=True` 搭配使用，並指定檢查點目錄和儲存部分檢查點時使用的標籤。請注意，模型權重的實際載入發生在模型分割之後，在 `smdistributed.modelparallel.torch.step` 裝飾訓練 Step Function 初次執行期間。
儲存部分檢查點時，程式庫也會將模型分割區決定儲存為具有 `.pt` 副檔名的檔案。相反地，從部分檢查點恢復時，程式庫會將分割決定檔案載入在一起。一旦已載入分割區決定，就無法變更分割區。
下列程式碼片段示範如何在 PyTorch 訓練指令碼中設定檢查點 API。
```
import smdistributed.modelparallel.torch as smp
model = ...
model = smp.DistributedModel(model)
optimizer = ...
optimizer = smp.DistributedOptimizer(optimizer)
user_content = ... # additional custom data
checkpoint_path = "/opt/ml/checkpoint/model_parallel"
# Save a checkpoint.
smp.save_checkpoint(
path=checkpoint_path,
tag=f"total_steps{total_steps}",
partial=True,
model=model,
optimizer=optimizer,
user_content=user_content
num_kept_partial_checkpoints=5
)
# Load a checkpoint.
# This automatically loads the most recently saved checkpoint.
smp_checkpoint = smp.resume_from_checkpoint(
path=checkpoint_path,
partial=True
)
```
**完整檢查點**
若要儲存最終模型成品以供推論，請使用 `smdistributed.modelparallel.torch.save_checkpoint` API 搭配 `partial=False`，該 API 結合模型分割區以建立單一模型成品。請注意，這不會結合最佳化工具狀態。
若要初始化具有特定權重的訓練，只要提供一個完整模型檢查點，您可以使用 `smdistributed.modelparallel.torch.resume_from_checkpoint` API 搭配 `partial=False`。請注意，這不會載入最佳化工具狀態。
**注意**
一般而言，使用張量平行處理，`state_dict` 必須在原始模型實作和 `DistributedModel` 實作之間進行轉譯。或者，您可以為 `smdistributed.modelparallel.torch.resume_from_checkpoint` 提供 `state_dict` 平移函式作為引數。但是，對於 [開箱即用的支援模型](model-parallel-extended-features-pytorch-hugging-face.md#model-parallel-extended-features-pytorch-hugging-face-out-of-the-box)，程式庫會自動處理此轉譯。
下列程式碼示範如何使用檢查點 API 對使用模型平行處理訓練的 PyTorch 模型執行完整檢查點。
```
import smdistributed.modelparallel.torch as smp
model = ...
model = smp.DistributedModel(model)
optimizer = ...
optimizer = smp.DistributedOptimizer(optimizer)
user_content = ... # additional custom data
checkpoint_path = "/opt/ml/checkpoint/model_parallel"
# Save a checkpoint.
smp.save_checkpoint(
path=checkpoint_path,
tag=f"total_steps{total_steps}",
partial=False,
model=model,
optimizer=optimizer,
user_content=user_content
num_kept_partial_checkpoints=5
)
# Load a checkpoint.
# This automatically loads the most recently saved checkpoint.
smp_checkpoint = smp.resume_from_checkpoint(
path=checkpoint_path,
partial=False
)
```
### 對分散式 PyTorch 模型執行檢查點 (適用於 1.6.0 版和 1.9.0 版之間的 SageMaker 模型平行處理程式庫)
SageMaker 模型平行處理程式庫提供 Python 函式，用於儲存具有張量平行處理的部分或完整訓練任務檢查點。下列程序顯示當您使用張量平行處理時，如何使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.save](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.save) 和 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.load](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.load) 儲存及載入檢查點。
**注意**
如果您使用 PyTorch、[張量平行處理](model-parallel-extended-features-pytorch-tensor-parallelism.md)，以及 1.6.0 版和 1.9.0 版之間的 SageMaker 模型平行處理程式庫，建議使用此檢查點方法。
1. 準備模型物件，並使用程式庫包裝函式 `smp.DistributedModel()` 包裝。
```
model = MyModel(...)
model = smp.DistributedModel(model)
```
1. 準備模型的最佳化工具。一組模型參數是最佳化工具函式所需的可迭代引數。若要準備一組模型參數，您必須處理 `model.parameters()` 以將不重複的 ID 指派給個別模型參數。
如果模型參數可迭代中存在具有重複 ID 的參數，則載入檢查點最佳化工具狀態將失敗。若要為最佳化工具建立具有不重複的 ID 的模型參數可迭代，請參閱以下內容：
```
unique_params = []
unique_params_set = set()
for p in model.parameters():
if p not in unique_params_set:
unique_params.append(p)
unique_params_set.add(p)
del unique_params_set
optimizer = MyOpt(unique_params, ...)
```
1. 使用程式庫的包裝函式 `smp.DistributedOptimizer()` 包裝最佳化工具。
```
optimizer = smp.DistributedOptimizer(optimizer)
```
1. 使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.save](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.save) 儲存模型和最佳化工具狀態。根據您想要儲存檢查點的方式，選擇下列兩個選項其中之一：
+ **選項 1：**為單一 `MP_GROUP` 在每個 `mp_rank` 上儲存部分模型。
```
model_dict = model.local_state_dict() # save a partial model
opt_dict = optimizer.local_state_dict() # save a partial optimizer state
# Save the dictionaries at rdp_rank 0 as a checkpoint
if smp.rdp_rank() == 0:
smp.save(
{"model_state_dict": model_dict, "optimizer_state_dict": opt_dict},
f"/checkpoint.pt",
partial=True,
)
```
透過張量平行處理，程式庫會儲存以下列格式命名的檢查點檔案：`checkpoint.pt_{pp_rank}_{tp_rank}`。
**注意**
使用張量平行處理，確保將 if 陳述式設定為 `if smp.rdp_rank() == 0` 而不是 `if smp.dp_rank() == 0`。當最佳化工具狀態以張量平行處理進行分割時，所有縮減的資料平行排名都必須儲存自己的最佳化工具狀態分割。使用錯誤的 *if* 陳述式執行檢查點，可能導致訓練任務停滯。如需在沒有張量平行處理的狀況下使用 `if smp.dp_rank() == 0` 的更多相關資訊，請參閱 *SageMaker Python SDK 文件*中的[儲存和載入的一般指示](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#general-instruction-for-saving-and-loading)。
+ **選項 2：**儲存完整模型。
```
if smp.rdp_rank() == 0:
model_dict = model.state_dict(gather_to_rank0=True) # save the full model
if smp.rank() == 0:
smp.save(
{"model_state_dict": model_dict},
"/checkpoint.pt",
partial=False,
)
```
**注意**
對於完整檢查點，請考量以下事項：
如果設定 `gather_to_rank0=True`，除了 `0` 以外的所有排名會傳回空白字典。
對於完整檢查點，您只能對模型執行檢查點。目前不支援最佳化工具狀態的完整檢查點。
完整模型只需要儲存於 `smp.rank() == 0`。
1. 使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.load](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.load) 載入檢查點。根據上一步中執行檢查點的方法，選擇下列兩個選項其中之一：
+ **選項 1：**載入部分檢查點。
```
checkpoint = smp.load("/checkpoint.pt", partial=True)
model.load_state_dict(checkpoint["model_state_dict"], same_partition_load=False)
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
```
如果您知道分割區不會變更，則可以設定 `model.load_state_dict()` 中的 `same_partition_load=True` 以取得更快的載入。
+ **選項 2：**載入完整檢查點。
```
if smp.rdp_rank() == 0:
checkpoint = smp.load("/checkpoint.pt", partial=False)
model.load_state_dict(checkpoint["model_state_dict"])
```
`if smp.rdp_rank() == 0` 條件為非必要，但它可以協助避免不同 `MP_GROUP` 之間的多餘載入。張量平行處理目前不支援完整檢查點最佳化工具狀態字典。
### 對分散式 TensorFlow 模型執行檢查點
若要在使用模型平行處理進行訓練時儲存 TensorFlow 模型，請使用 SageMaker 模型平行處理程式庫提供的下列功能。
+ [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_tensorflow.html#smp.DistributedModel.save_model](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_tensorflow.html#smp.DistributedModel.save_model)
+ [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_tensorflow.html#smp.CheckpointManager](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_tensorflow.html#smp.CheckpointManager)
## 微調分散式模型
微調需要在訓練指令碼中設定。下列程式碼片段顯示使用 [AutoModelForCausalLM](https://huggingface.co/docs/transformers/main/en/model_doc/auto#transformers.AutoModelForCausalLM) 類別 Hugging Face 轉換器的訓練指令碼的範例結構，其中包含註冊 `smdistributed.model.parallel.torch` 模組和設定以進行微調的修改。
**注意**
使用已啟動的 [smp.delayed\$1param\$1initialization](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.delay_param_initialization) 功能對分散式轉換器 (由 `smp.DistributedModel()` 包裝的轉換器模型) 進行微調，需要使用 FSx for Lustre 檔案系統來設定微調任務。如果您想要使用已延遲的參數初始化選項微調大規模模型，則應設定 FSx for Lustre 檔案系統。
```
import argparse
from transformers import AutoModelForCausalLM
import smdistributed.modelparallel
import smdistributed.modelparallel.torch as smp
def parse_args():
parser = argparse.ArgumentParser()
# set an arg group for model
model_grp = parser.add_argument_group(
title="model", description="arguments to describe model configuration"
)
... # set up numerous args to parse from the configuration dictionary to the script for training
# add arg for activating fine-tuning
model_grp.add_argument(
"--fine_tune",
type=int,
default=0,
help="Fine-tune model from checkpoint or pretrained model",
)
def main():
"""Main function to train GPT."""
args = parse_args()
... # parse numerous args
if args.fine_tune > 0 and args.delayed_param > 0 and smp.rank() == 0:
pretrained_model = AutoModelForCausalLM.from_pretrained(
args.model_name or args.model_dir
)
model_state_dict = pretrained_model.state_dict()
path = os.path.join(args.model_dir, "fullmodel.pt")
torch.save(model_state_dict, path)
# create a Transformer model and wrap by smp.model_creation()
# with options to configure model parallelism parameters offered by SageMaker AI
with smp.model_creation(
tensor_parallelism=smp.tp_size() > 1 or args.use_distributed_transformer > 0,
zero_init=args.use_distributed_transformer == 0,
dtype=dtype,
distribute_embedding=args.sharded_data_parallel_degree > 1 and smp.tp_size() > 1,
use_alibi=args.alibi > 0,
attention_in_fp32=args.attention_in_fp32 > 0,
fp32_residual_addition=args.residual_addition_in_fp32 > 0,
query_key_layer_scaling=args.query_key_layer_scaling > 0 and args.bf16 < 1,
fused_softmax=args.fused_softmax > 0,
fused_dropout=args.fused_dropout > 0,
fused_bias_gelu=args.fused_bias_gelu > 0,
flash_attention=args.flash_attention > 0,
):
if args.fine_tune > 0 and args.delayed_param == 0:
model = AutoModelForCausalLM.from_pretrained(
args.model_name or args.model_dir
)
else:
model = AutoModelForCausalLM.from_config(model_config)
# wrap the model by smp.DistributedModel() to apply SageMaker model parallelism
model = smp.DistributedModel(
model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation
)
# wrap the optimizer by smp.DistributedOptimizer() to apply SageMaker model parallelism
optimizer= ... # define an optimizer
optimizer = smp.DistributedOptimizer(
optimizer,
static_loss_scale=None,
dynamic_loss_scale=True,
dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2},
)
# for fine-tuning, use smp.resume_from_checkpoint() to load a pre-trained model
if args.fine_tune > 0 and args.delayed_param > 0:
smp.resume_from_checkpoint(args.model_dir, tag="fullmodel.pt", partial=False)
```
如需訓練指令碼和 Jupyter 筆記本的完整範例，請參閱 *SageMaker AI 範例 GitHub 儲存庫*中的 [PyTorch GPT-2 範例](https://github.com/aws/amazon-sagemaker-examples/tree/main/training/distributed_training/pytorch/model_parallel/gpt2)。
# Amazon SageMaker AI 模型平行化程式庫 v1 範例
此頁面提供部落格和 Jupyter 筆記本的清單，其中提供實作 SageMaker 模型平行化 (SMP) 程式庫 v1 以在 SageMaker AI 上執行分散式訓練任務的實際範例。
## 部落格與案例研究
下列部落格討論有關使用 SMP v1 的案例研究。
+ [Amazon SageMaker AI 模型平行化程式庫中的新效能改善](https://aws.amazon.com/blogs/machine-learning/new-performance-improvements-in-amazon-sagemaker-model-parallel-library/)，*AWS 機器學習部落格* (2022 年 12 月 16 日)
+ [在 Amazon SageMaker AI 上使用碎片資料平行化訓練具有近線性擴展的巨大模型](https://aws.amazon.com/blogs/machine-learning/train-gigantic-models-with-near-linear-scaling-using-sharded-data-parallelism-on-amazon-sagemaker/)，*AWS 機器學習部落格* (2022 年 10 月 31 日)
## 範例筆記本
範例筆記本在 [SageMaker AI 範例 GitHub 儲存庫](https://github.com/aws/amazon-sagemaker-examples/tree/master/training/distributed_training/)中提供。若要下載範例，請執行下列命令來複製儲存庫並前往 `training/distributed_training/pytorch/model_parallel`。
**注意**
在下列 SageMaker AI ML IDE 中複製並執行範例筆記本。
[SageMaker JupyterLab](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-updated-jl.html) (可在 2023 年 12 月之後建立的 [Studio](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-updated.html) 中使用)
[SageMaker 程式碼編輯器](https://docs.aws.amazon.com/sagemaker/latest/dg/code-editor.html) (可在 2023 年 12 月之後建立的 [Studio](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-updated.html) 中使用)
[Studio Classic](https://docs.aws.amazon.com/sagemaker/latest/dg/studio.html) (作為在 2023 年 12 月之後建立的 [Studio](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-updated.html) 中的應用程式提供)
[SageMaker 筆記本執行個體](https://docs.aws.amazon.com/sagemaker/latest/dg/nbi.html)
```
git clone https://github.com/aws/amazon-sagemaker-examples.git
cd amazon-sagemaker-examples/training/distributed_training/pytorch/model_parallel
```
**PyTorch 的 SMP v1 範例筆記本**
+ [使用 SageMaker 模型平行處理程式庫中的碎片資料平行處理技術，以近線性擴展訓練 GPT-2](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt2/smp-train-gpt-sharded-data-parallel.ipynb)
+ [使用 SageMaker 模型平行處理程式庫中的碎片資料平行處理技術，以近線性擴展微調 GPT-2](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt2/smp-fine-tune-gpt-sharded-data-parallel.ipynb)
+ [使用 SageMaker 模型平行處理程式庫中的碎片資料平行處理技術，以近線性擴展訓練 GPT-NeoX-20B](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt-neox/smp-train-gpt-neox-sharded-data-parallel.ipynb)
+ [使用 SageMaker 模型平行程式庫中的碎片資料平行處理和張量平行處理技術來訓練 GPT-J 6B](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt-j/smp-train-gptj-sharded-data-parallel-tp.ipynb)
+ [使用 SageMaker 模型平行處理程式庫中的碎片資料平行處理技術，以近線性擴展訓練 FLAN-T5](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/flan-t5/smp-train-t5-sharded-data-parallel.ipynb)
+ [使用 SageMaker 模型平行處理程式庫中的碎片資料平行處理技術，以近線性擴展訓練 Falcon](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/falcon/smp-train-falcon-sharded-data-parallel.ipynb)
**TensorFlow 的 SMP v1 範例筆記本**
+ [使用 TensorFlow 2.3.1 的 CNN 與 SageMaker 模型平行處理程式庫](https://sagemaker-examples.readthedocs.io/en/latest/training/distributed_training/tensorflow/model_parallel/mnist/tensorflow_smmodelparallel_mnist.html)
+ [在 SageMaker AI 上使用 TensorFlow 的 HuggingFace 分散式模型平行化程式庫訓練](https://github.com/huggingface/notebooks/blob/master/sagemaker/04_distributed_training_model_parallelism/sagemaker-notebook.ipynb)
# SageMaker 分散式模型平行處理最佳實務
當您使用 SageMaker 模型平行程式庫執行分散式訓練任務時，請遵循下列指南。
## 為指定的模型設定正確組態
擴展模型時，我們建議您依序瀏覽以下清單。每個清單項目都討論了使用程式庫技術以及可能出現的權衡的優勢。
**提示**
如果模型使用程式庫功能的子集運作良好，新增更多的模型平行處理或節省記憶體功能通常不會改善效能。
**使用大型 GPU 執行個體類型**
+ 在模型平行處理範圍中，最好使用具有大型 GPU 記憶體的強大執行個體來處理模型平行處理操作所產生的額外負荷，例如跨多個 GPU 分割模型。我們建議使用 `ml.p4d` 或 `ml.p3dn` 執行個體來訓練大型 DL 模型。這些執行個體也配備 Elastic Fabric Adapter (EFA)，提供較高的網路頻寬，並可進行透過模型平行處理的大規模訓練。
**分片最佳化工具狀態**
+ 分片最佳化工具狀態的影響取決於資料平行處理排名的數量。通常，較高的資料平行處理程度 (與運算節點的大小成正比) 可以改善記憶體使用量效率。
當您想要縮減叢集時，請務必檢查最佳化工具狀態分片組態。例如，具有最佳化工具狀態分片的大型 DL 模型適合具有 16 個 GPU 的運算叢集 (例如，兩個 p4d 或 p4de 執行個體)，可能不一定適合具有 8 個 GPU 的節點 (例如，單一 P4d 或 p4de 執行個體)。這是因為 8 個 GPU 的合併記憶體低於 16 個 GPU 的合併記憶體，而每個 GPU 分片 8 個 GPU 所需的記憶體也高於在 16 GPU 案例中每個 GPU 分片的記憶體。因此，增加的記憶體需求可能不適合較小的叢集。
如需詳細資訊，請參閱[最佳化工具狀態碎片](model-parallel-extended-features-pytorch-optimizer-state-sharding.md)。
**啟用檢查點**
+ 使用一組模組的啟用檢查點可以改善記憶體效率。您分組的模組越多，記憶體使用量就越有效率。當建立層的檢查點循序模組時，`smp.set_activation_checkpointing` 函式的 `strategy` 引數將分組層以建立檢查點。例如，將兩個或以上的層分組以進行檢查點，比一次一層檢查點的記憶體效率更好，而且這會犧牲額外運算時間以減少記憶體使用量。
如需詳細資訊，請參閱[啟用檢查點](model-parallel-extended-features-pytorch-activation-checkpointing.md)。
**張量平行**
+ 張量平行處理程度應為 2 次方 (2、4、8、...、2 n)，其中最大程度必須等於每個節點的 GPU 數量。例如，如果您使用具有 8 個 GPU 的節點，張量平行程度的可能數字為 2、4 和 8。對於張量平行程度，我們不建議任意數字 (例如 3、5、6 和 7)。當您使用多個節點時，錯誤設定張量平行處理程度可能導致跨節點執行張量平行處理；這會增加跨節點啟用通訊所產生的重大負荷，而且運算成本可能會變得非常昂貴。
如需詳細資訊，請參閱[張量平行處理](model-parallel-extended-features-pytorch-tensor-parallelism.md)。
**跨節點的管道平行處理**
+ 您可以在單一節點內和跨多個節點內執行管道平行處理。當您將管道平行處理與張量平行處理結合使用時，建議您跨多個節點執行管道平行處理，並在個別節點內保持張量平行處理。
+ 管道平行處理具有以下三個旋鈕：`microbatches`、`active_microbatches` 和 `prescaled_batch`。
+ 當您將張量平行處理與管道平行處理搭配使用時，建議您啟用 `prescaled_batch`，藉此增加每個模型平行群組的批次大小，以便有效率地管線傳輸。啟用 `prescaled_batch` 後，訓練指令碼中設定的批次大小，會成為每個沒有 `prescaled_batch` 的排名設定批次大小乘以 `tp_size`。
+ 增加的 `microbatches` 數量有助於達成有效的管線傳輸和更好的效能。請注意，有效的微批次大小是批次大小除以微批次數量。如果在保持批次大小常數不變的同時，增加微批次的數量，則每個微批次處理較少的範例。
+ `active_microbatches` 的數量是在管道傳輸過程中同步處理的微批次的最大數量。對於處理中的每個啟動中微批次，其啟用和漸層都會佔用 GPU 記憶體。因此，增加 `active_microbatches` 會佔用更多 GPU 記憶體。
+ 如果 GPU 和 GPU 記憶體使用量過低，請在管道傳輸期間增加 `active_microbatches` 以更好地平行化。
+ 如需如何將張量平行處理與管道平行處理搭配使用的更多相關資訊，請參閱[張量平行處理結合管道平行處理](model-parallel-extended-features-pytorch-tensor-parallelism-examples.md#model-parallel-extended-features-pytorch-tensor-and-pipeline-parallelism)。
+ 若要尋找上述參數的描述，請參閱 *SageMaker Python SDK 文件*中的 [`smdistributed` 參數](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#parameters-for-smdistributed)。
**卸載啟用至 CPU**
+ 請確認這用於與啟用檢查點和管線平行處理組合使用。若要確保卸載和預先載入在背景中進行，請為微批次參數指定大於 1 的值。
+ 當卸載啟用時, 你可能能夠增加 `active_microbatches`，而有時與微批次的總數相符。這取決於對哪些模組執行檢查點，以及模型的分割方式。
如需詳細資訊，請參閱[啟用卸載](model-parallel-extended-features-pytorch-activation-offloading.md)。
### 參考組態
SageMaker 模型平行訓練團隊根據使用 GPT-2 模型的實驗、序列長度 512 以及 50,000 個字彙大小，提供下列參考點。
| 模型參數數量 | 執行個體類型 | 管道平行處理 | 張量平行處理 | 最佳化工具狀態分片 | 啟用檢查點 | 預先擴充批次 | 批次大小 |
| --- | --- | --- | --- | --- | --- | --- | --- |
| 10 億 | 16 ml.p4d.24xlarge | 1 | 4 | True | 每個轉換器層 | True | batch\$1size=40 |
| 30 億 | 16 ml.p4d.24xlarge | 1 | 8 | True | 每個轉換器層 | True | batch\$1size=32 |
| 60 億 | 32 ml.p4d.24xlarge | 2 | 8 | True | 每個轉換器層 | True | batch\$1size=56, microbatches=4, active\$1microbatches=2 |
您可以從上述組態中推斷，以估算模型組態的 GPU 記憶體使用量。例如，如果您增加 100 億個參數模型的序列長度，或將模型的大小增加到 200 億，則可能需要先降低批次大小。如果模型仍然不符合，請嘗試增加張量平行處理程度。
## 修改訓練指令碼
+ 在訓練指令碼中使用 SageMaker 模型平行處理程式庫功能之前，請先檢閱 [SageMaker 分散式模型平行化程式庫組態提示和缺陷](model-parallel-customize-tips-pitfalls.md)。
+ 若要更快地啟動訓練任務，請使用 [SageMaker AI 本機模式](https://sagemaker.readthedocs.io/en/v2.199.0/overview.html?highlight=local%20mode#local-mode)。這可協助您在 SageMaker 筆記本執行個體上快速執行訓練任務。根據執行 SageMaker 筆記本執行個體的機器學習 (ML) 執行個體規模，您可能需要透過變更模型組態 (例如隱藏寬度、轉換器層數和注意力標題) 來調整模型大小。在使用大型叢集訓練完整模型之前，先驗證已縮減的模型是否在筆記本執行個體上執行良好。
## 使用 SageMaker AI 主控台和 Amazon CloudWatch 監控和記錄訓練任務
若要監控系統層級指標 (例如 CPU 記憶體使用率、GPU 記憶體使用率和 GPU 使用率)，請使用透過 [SageMaker AI 主控台](https://console.aws.amazon.com/sagemaker/)提供的視覺效果。
1. 在左側導覽窗格中，選擇**訓練**。
1. 選擇 **Training jobs (訓練任務)**。
1. 在主窗格中，選擇您要查看其更多詳細資訊的訓練任務名稱。
1. 瀏覽主窗格，並找到**監視器**區段以查看自動化視覺效果。
1. 若要查看訓練任務日誌，請選擇**監視器**區段中的**檢視日誌**。您可以在 CloudWatch 中存取訓練任務的分散式訓練任務日誌。如果您已啟動多節點分散式訓練，您應該會看到多個日誌串流，其標記格式為 **algo-n-1234567890**。**algo-1** 日誌串流會追蹤主節點 (第 0 個) 的訓練日誌。
如需詳細資訊，請參閱[用於監控和分析訓練任務的 Amazon CloudWatch 指標](training-metrics.md)。
## 許可
若要使用模型平行處理或 [SageMaker 分散式訓練範例筆記本](https://sagemaker-examples.readthedocs.io/en/latest/training/distributed_training/index.html)執行 SageMaker 訓練任務，請確認您有 IAM 角色的正確許可，例如：
+ 若要使用 [FSx for Lustre](https://aws.amazon.com/fsx/)，請新增 [https://console.aws.amazon.com/iam/home#/policies/arn%3Aaws%3Aiam%3A%3Aaws%3Apolicy%2FAmazonFSxFullAccess](https://console.aws.amazon.com/iam/home#/policies/arn%3Aaws%3Aiam%3A%3Aaws%3Apolicy%2FAmazonFSxFullAccess)。
+ 若要使用 Amazon S3 做為資料管道，請新增 [https://console.aws.amazon.com/iam/home#/policies/arn%3Aaws%3Aiam%3A%3Aaws%3Apolicy%2FAmazonS3FullAccess](https://console.aws.amazon.com/iam/home#/policies/arn%3Aaws%3Aiam%3A%3Aaws%3Apolicy%2FAmazonS3FullAccess)。
+ 若要使用 Docker，請建置您自己的容器，然後將其推送到 Amazon ECR，新增 [https://console.aws.amazon.com/iam/home#/policies/arn%3Aaws%3Aiam%3A%3Aaws%3Apolicy%2FAmazonEC2ContainerRegistryFullAccess](https://console.aws.amazon.com/iam/home#/policies/arn%3Aaws%3Aiam%3A%3Aaws%3Apolicy%2FAmazonEC2ContainerRegistryFullAccess)。
+ 若要取得完整存取權，以使用整個 SageMaker AI 功能套件，請新增 [https://console.aws.amazon.com/iam/home#/policies/iam/home#/policies/arn%3Aaws%3Aiam%3A%3Aaws%3Apolicy%2FAmazonSageMakerFullAccess](https://console.aws.amazon.com/iam/home#/policies/iam/home#/policies/arn%3Aaws%3Aiam%3A%3Aaws%3Apolicy%2FAmazonSageMakerFullAccess)。
# SageMaker 分散式模型平行化程式庫組態提示和缺陷
使用 Amazon SageMaker AI 模型平行化程式庫之前，請先檢閱下列提示和缺陷。此清單包含適用於跨架構的提示。如需 TensorFlow 和 PyTorch 的特定提示，請分別參閱 [修改 TensorFlow 訓練指令碼](model-parallel-customize-training-script-tf.md) 和 [修改 PyTorch 訓練指令碼](model-parallel-customize-training-script-pt.md)。
## 批次大小和微批次數量
+ 當批次大小增加時，程式庫效率最高。對於模型符合單一裝置但只能以小批次進行訓練的使用案例，在整合程式庫之後，批次大小可以且應該增加。模型平行化可節省大型模型的記憶體，讓您能夠使用先前不符合記憶體的批次大小進行訓練。
+ 選擇太小或太大的微批次數量可能會降低性能。該程式庫在每個裝置中循序執行每個微批次，因此微批大小 (批次大小除以微批次數量) 必須足夠以充分利用每個 GPU。同時，管道效率隨著微批次數量增加，因此保持正確的平衡非常重要。通常，好的起點是嘗試 2 或 4 微批次，將批次大小增加到記憶體限制，然後實驗使用更大的批次大小和微批次數量。隨著微批次數量增加，如果已使用交錯管道，則更大的批次大小可能會變得可行。
+ 您的批次大小必須始終可以被微批次數量整除。請注意，根據資料集的大小，有時每個週期的最後一批次可能比其餘更小，並且這個較小的批次也需要被微批次的數量整除。如果沒有，您可以在 `tf.Dataset.batch()` 呼叫中設定 `drop_remainder=True` (在 TensorFlow 中)，或在 `DataLoader` 中設定 `drop_last=True` (在 PyTorch 中)，以便不使用最後一個小批次。如果您為 Data Pipeline 使用不同的 API，則每當最後一個批次無法被微批次數量整除時，您可能需要手動略過最後一個批次。
## 手動分割
+ 如果您使用手動分割，請注意模型中多個作業和模組所使用的參數，例如轉換器架構中的內嵌資料表。共用相同參數的模組必須放置在同一裝置中，以確保正確性。使用自動分割時，程式庫會自動強制執行此限制。
## 資料準備
+ 如果模型需要多個輸入，請確保使用 `smp.dp_rank()` 在 Data Pipeline 中植入隨機操作 (例如，隨機顯示)。如果資料集在資料平行裝置之間進行確定性分割，請確保碎片已由 `smp.dp_rank()` 編製索引。這是為了確保在形成模型分區的所有排名上，看到的資料順序是一致的。
## 從 `smp.DistributedModel` 中傳回張量
+ 從 `smp.DistributedModel.call` (適用於 TensorFlow) 或 `smp.DistributedModel.forward` (適用於 PyTorch) 函式傳回的任何張量都會廣播至所有其他排名，包括運算特定張量的排名。因此，在呼叫和轉送方法 (例如中繼啟動) 之外不需要的任何張量都不應傳回，因為這會造成不必要的通訊和記憶體負荷，並且損害效能。
## `@smp.step` 裝飾項目
+ 如果 `smp.step` 裝飾函式具有沒有批次維度的張量引數，則在呼叫 `smp.step` 時必須在 `non_split_inputs` 清單中提供引數名稱。這樣可以防止程式庫嘗試將張量分成微批次。如需詳細資訊，請參閱 API 文件中的 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_common_api.html](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_common_api.html)。
## 延遲參數初始化
對於超過 100 億個參數的極大型模型，透過 CPU 記憶體進行加權初始化可能導致記憶體不足錯誤。若要因應此問題，程式庫提供 `smp.delay_param_initialization` 關聯內容管理員。這會延遲參數的實體配置，直到它們在 `smp.step` 裝飾函式第一次執行期間移動到 GPU。這樣可以避免訓練初始化期間 CPU 的不必要記憶體使用量。當您建立模型物件時，請使用關聯內容管理員，如下列程式碼所示。
```
with smp.delay_param_initialization(enabled=True):
model = MyModel()
```
## PyTorch 的張量平行處理
+ 如果您使用種子取得確定性結果，請根據 `smp.dp_rank()` (例如，`torch.manual_seed(42 + smp.dp_rank())`) 設定種子。如果您沒有這樣做，則會以相同的方式初始化 `nn.Parameter` 的不同分割區，從而影響收斂。
+ SageMaker 的模型平行化程式庫使用 NCCL 來實作模組發佈所需的集合。特別是較小型號，如果同時在 GPU 上排程過多的 NCCL 呼叫，則可能因為 NCCL 使用額外的空間，導致記憶體使用量增加。為了阻止這種情況發生，`smp` 會調節 NCCL 呼叫，以便在任何特定時間的進行中 NCCL 操作數量小於或等於特定限制。預設限制為 8，但可以使用環境變數 `SMP_NCCL_THROTTLE_LIMIT` 調整。如果您在使用張量平行處理時發現記憶體使用量超出預期，可以嘗試降低此限制。但是，選擇太小的限制可能導致輸送量損失。若要完全停用調節，您可以設定 `SMP_NCCL_THROTTLE_LIMIT=-1`。
+ 當張量平行處理程度為 1 時，以下身份成立；張量平行處理程度大於 1 時則不成立：`smp.mp_size() * smp.dp_size() == smp.size()`。這是因為張量平行群組是模型平行化群組和資料平行化群組的一部分。如果您的程式碼具有 `mp_rank`、`mp_size`、`MP_GROUP` 等的現有參考，並且如果您只想使用管道平行群組，則可能需要將參照取代為 `smp.pp_size()`。以下身分一律為 true：
+ `smp.mp_size() * smp.rdp_size() == smp.size()`
+ `smp.pp_size() * smp.dp_size() == smp.size()`
+ `smp.pp_size() * smp.tp_size() * smp.rdp_size() == smp.size()`
+ 由於 `smp.DistributedModel` 包裝函式在已啟用張量平行處理時修改模型參數，因此應使用分散式參數在呼叫 `smp.DistributedModel` 之後建立最佳化工具。例如，下列項目未運作：
```
## WRONG
model = MyModel()
optimizer = SomeOptimizer(model.parameters())
model = smp.DistributedModel(model) # optimizer now has outdated parameters!
```
相反地，最佳化工具應該使用下列 `smp.DistributedModel` 參數來建立：
```
## CORRECT
model = smp.DistributedModel(MyModel())
optimizer = SomeOptimizer(model.optimizers())
```
+ 當透過張量平行處理將模組取代為其分散式對應項目時，分散式模組不會從原始模組繼承其權重，並初始化新的權重。這表示，例如，如果權重需要在特定呼叫中初始化 (例如，透過 `load_state_dict` 呼叫)，一旦模組發佈發生，則需要在 `smp.DistributedModel` 呼叫之後進行。
+ 當直接存取分散式模組的參數時，請注意，權重不具有與原始模組的相同形狀。例如，
```
with smp.tensor_parallelism():
linear = nn.Linear(60, 60)
# will pass
assert tuple(linear.weight.shape) == (60, 60)
distributed_linear = smp.DistributedModel(linear)
# will fail. the number of input channels will have been divided by smp.tp_size()
assert tuple(distributed_linear.module.weight.shape) == (60, 60)
```
+ 強烈建議為張量平行處理使用 `torch.utils.data.distributed.DistributedSampler`。這可確保每個資料平行排名都會接收相同數量的資料範例，以防止因不同 `dp_rank` 採取不同步驟數而造成的懸置。
+ 如果您使用 PyTorch `DistributedDataParallel` 類別的 `join` API，處理不同資料平行排名具有不同批次數量的情況，您仍然需要確保相同 `TP_GROUP` 的排名具有相同批次數量；否則用於分散式模組執行的通訊集合可能會懸置。只要使用 `join` API，位於不同 `TP_GROUP` 的排名就可以有不同的批次數量。
+ 如果您想要對模型執行檢查點，並使用張量平行處理，請考慮下列事項：
+ 當您使用張量平行處理時，若要避免在儲存和載入模型時出現延遲和競爭狀態，請確認您從已降低的資料平行處理排名內的下列模型和最佳化工具狀態呼叫適當的函式。
+ 如果您正在傳輸現有管道平行指令碼，並為指令碼啟用張量平行，請務必修改用於使用 `if smp.rdp_rank() == 0` 儲存和載入的任何 `if smp.dp_rank() == 0` 區塊。否則，可能導致您的訓練任務停止。
如需對具有張量平行處理之模型執行檢查點的更多相關資訊，請參閱[對分散式模型進行檢查點](distributed-model-parallel-checkpointing-and-finetuning.md#distributed-model-parallel-checkpoint)。
# 模型平行疑難排解
如果遇到錯誤，可嘗試使用下列清單，對您的訓練任務進行疑難排解。如果問題持續，請聯絡 [AWS 支援](https://aws.amazon.com/premiumsupport)。
**Topics**
+ [將 SageMaker Debugger 與 SageMaker 模型平行化程式庫搭配使用的考量事項](#distributed-ts-model-parallel-debugger)
+ [儲存檢查點](#distributed-ts-model-parallel-checkpoints)
+ [使用模型平行和 TensorFlow 進行聚合](#distributed-ts-model-parallel-tf-convergence)
+ [分散式訓練任務停止或當機](#distributed-ts-model-parallel-training-issues)
+ [收到 PyTorch 訓練任務 NCCL 錯誤](#distributed-ts-model-parallel-nccl-error)
+ [收到 PyTorch 訓練任務的 `RecursionError`](#distributed-ts-model-parallel-super-forward-not-supported)
## 將 SageMaker Debugger 與 SageMaker 模型平行化程式庫搭配使用的考量事項
SageMaker Debugger 不適用於 SageMaker 模型平行化程式庫。偵錯工具對所有 SageMaker TensorFlow 和 PyTorch 訓練任務預設為啟用，您可能會看到類似以下的錯誤：
```
FileNotFoundError: [Errno 2] No such file or directory: '/opt/ml/checkpoints/metadata.json.sagemaker-uploading
```
若要解決此問題，在建立架構 `estimator` 時按 `debugger_hook_config=False` 停用偵錯工具，如以下範例所示。
```
bucket=sagemaker.Session().default_bucket()
base_job_name="sagemaker-checkpoint-test"
checkpoint_in_bucket="checkpoints"
# The S3 URI to store the checkpoints
checkpoint_s3_bucket="s3://{}/{}/{}".format(bucket, base_job_name, checkpoint_in_bucket)
estimator = TensorFlow(
...
distribution={"smdistributed": {"modelparallel": { "enabled": True }}},
checkpoint_s3_uri=checkpoint_s3_bucket,
checkpoint_local_path="/opt/ml/checkpoints",
debugger_hook_config=False
)
```
## 儲存檢查點
在 SageMaker AI 上儲存大型模型的檢查點時，可能會發生下列錯誤：
```
InternalServerError: We encountered an internal error. Please try again
```
這可能是由於在訓練期間將本機檢查點上傳到 Amazon S3 時，產生 SageMaker AI 限制所致。若要停用 SageMaker AI 中的檢查點，請使用下列範例明確上傳檢查點。
如果您遇到前面的錯誤，請勿與 SageMaker `estimator` 呼叫搭配使用 `checkpoint_s3_uri`。為較大的模型儲存檢查點時，我們建議將檢查點儲存到自訂目錄，並將其傳遞給輔助函式 (作為 `local_path` 引數)。
```
import os
def aws_s3_sync(source, destination):
"""aws s3 sync in quiet mode and time profile"""
import time, subprocess
cmd = ["aws", "s3", "sync", "--quiet", source, destination]
print(f"Syncing files from {source} to {destination}")
start_time = time.time()
p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
p.wait()
end_time = time.time()
print("Time Taken to Sync: ", (end_time-start_time))
return
def sync_local_checkpoints_to_s3(local_path="/opt/ml/checkpoints", s3_uri=os.path.dirname(os.path.dirname(os.getenv('SM_MODULE_DIR', '')))+'/checkpoints'):
""" sample function to sync checkpoints from local path to s3 """
import boto3
#check if local path exists
if not os.path.exists(local_path):
raise RuntimeError("Provided local path {local_path} does not exist. Please check")
#check if s3 bucket exists
s3 = boto3.resource('s3')
if not s3_uri.startswith("s3://"):
raise ValueError(f"Provided s3 uri {s3_uri} is not valid.")
s3_bucket = s3_uri.replace('s3://','').split('/')[0]
print(f"S3 Bucket: {s3_bucket}")
try:
s3.meta.client.head_bucket(Bucket=s3_bucket)
except Exception as e:
raise e
aws_s3_sync(local_path, s3_uri)
return
def sync_s3_checkpoints_to_local(local_path="/opt/ml/checkpoints", s3_uri=os.path.dirname(os.path.dirname(os.getenv('SM_MODULE_DIR', '')))+'/checkpoints'):
""" sample function to sync checkpoints from s3 to local path """
import boto3
#try to create local path if it does not exist
if not os.path.exists(local_path):
print(f"Provided local path {local_path} does not exist. Creating...")
try:
os.makedirs(local_path)
except Exception as e:
raise RuntimeError(f"Failed to create {local_path}")
#check if s3 bucket exists
s3 = boto3.resource('s3')
if not s3_uri.startswith("s3://"):
raise ValueError(f"Provided s3 uri {s3_uri} is not valid.")
s3_bucket = s3_uri.replace('s3://','').split('/')[0]
print(f"S3 Bucket: {s3_bucket}")
try:
s3.meta.client.head_bucket(Bucket=s3_bucket)
except Exception as e:
raise e
aws_s3_sync(s3_uri, local_path)
return
```
輔助函式的用法：
```
#base_s3_uri - user input s3 uri or save to model directory (default)
#curr_host - to save checkpoints of current host
#iteration - current step/epoch during which checkpoint is saved
# save checkpoints on every node using local_rank
if smp.local_rank() == 0:
base_s3_uri = os.path.dirname(os.path.dirname(os.getenv('SM_MODULE_DIR', '')))
curr_host = os.environ['SM_CURRENT_HOST']
full_s3_uri = f'{base_s3_uri}/checkpoints/{curr_host}/{iteration}'
sync_local_checkpoints_to_s3(local_path=checkpoint_dir, s3_uri=full_s3_uri)
```
## 使用模型平行和 TensorFlow 進行聚合
將 SageMaker AI 多節點訓練與 TensorFlow 和模型平行化程式庫搭配使用時，遺失可能無法如預期般聚合，因為每個節點上的訓練輸入檔案順序可能不同。這可能會導致同一模型平行群組中的不同排名處理不同的輸入檔案，進而導致不一致。為防止此情況發生，請確保輸入檔案轉換為 TensorFlow 資料集之前，在所有排名中都以相同的方式排序。達成此目標的一種方法是在訓練指令碼中對輸入檔案名稱進行排序。
## 分散式訓練任務停止或當機
如果您的訓練任務停止、當機或沒有回應，請閱讀下列疑難排解項目以找出問題的原因。如果您需要任何進一步支援，請透過 [AWS 支援](https://aws.amazon.com/premiumsupport)與 SageMaker 分散式訓練團隊聯絡。
+ 如果您**在 NCCL 初始化步驟中看到分散式訓練任務停止**，請考慮下列事項：
+ 如果您將其中一個已啟用 EFA 的執行個體 (`ml.p4d` 或 `ml.p3dn` 執行個體) 與自訂 VPC 及其子網路搭配使用，請確保所使用的安全群組具有往返相同 SG 的所有連接埠的傳入和傳出連線。通常，您還需要對任何 IP 的傳出連線做為單獨規則 (用於可存取的網際網路)。若要尋找如何針對 EFA 通訊新增傳入和傳出規則的指示，請參閱[初始化期間的 SageMaker AI 分散式訓練任務停滯](distributed-troubleshooting-data-parallel.md#distributed-ts-data-parallel-efa-sg)。
+ 如果您在**建立完整模型檢查點時看到分散式訓練任務停止**，這可能是因為模型或最佳化工具上的 `state_dict()` 呼叫並非以 `rdp_rank()==0` (使用張量平行處理時) 或 `dp_rank()==0` (僅使用管道平行處理時) 在所有排名上進行。這些排名需要溝通來建構要儲存的檢查點。如果啟用 `shard_optimizer_state`，在建立部分最佳化工具檢查點時，也可能發生類似的停止問題。
如需在以模型平行化的模型建立檢查點的詳細資訊，請參閱[儲存與載入的一般指示](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#general-instruction-for-saving-and-loading)和 [對分散式 PyTorch 模型執行檢查點 (適用於 1.6.0 版和 1.9.0 版之間的 SageMaker 模型平行處理程式庫)](distributed-model-parallel-checkpointing-and-finetuning.md#model-parallel-extended-features-pytorch-saving-loading-checkpoints)。
+ 如果訓練任務因 **CUDA 記憶體不足錯誤**而當機，這表示需要調整分散式訓練組態，以符合 GPU 叢集上的模型。如需更多資訊和最佳實務，請參閱[為指定的模型設定正確組態](model-parallel-best-practices.md#model-parallel-best-practices-configuration)。
+ 如果訓練工作因**無法修正的 [ECC 錯誤](https://docs.nvidia.com/deploy/a100-gpu-mem-error-mgmt/index.html)**而當機，這表示叢集中的其中一個 GPU 已經壞了。如果您需要技術支援，請與 AWS 團隊共享任務 ARN，並在可能的情況下從檢查點重新啟動訓練任務。
+ 在極少數情況下，先前運作但接近 GPU 記憶體限制的任務組態可能會因為 **CUDA 記憶體不足錯誤**而在不同的叢集中失敗。這可能是因為由於 ECC 錯誤，某些 GPU 的可用記憶體比平常低。
+ 執行未使用節點中所有 GPU 的多節點工作時，可能會發生**網路逾時當機**。若要解決此問題，請確定 `processes_per_host` 參數設定為每個執行個體中的 GPU 數量，以使用節點上的所有 GPU。例如，這是用於 `ml.p3.16xlarge`、`ml.p3dn.24xlarge` 和 `ml.p4d.24xlarge` 執行個體的 `processes_per_host=8`。
+ 如果您發現訓練任務在資料下載階段需要很長時間，請確保您為 SageMaker `Estimator` 類別的 `checkpoint_s3_uri` 提供的 Amazon S3 目前訓練任務的唯一路徑。如果在同時執行的多個訓練任務中重複使用此路徑，則所有這些檢查點都會上傳和下載到相同的 Amazon S3 路徑，並可能大幅增加檢查點載入時間。
+ 處理大型資料和模型時，請使用 FSx for Lustre。
+ 如果您的資料集很大且擷取需要很長時間，我們建議您將資料集保留在 [FSx for Lustre](https://aws.amazon.com/fsx/lustre/) 中。
+ 訓練模型超過 100 億個參數時，我們建議使用 FSx for Lustre 建立檢查點。
+ 建立檔案系統之後，請務必等待狀態變為**可用**，再開始使用該檔案系統的訓練任務。
## 收到 PyTorch 訓練任務 NCCL 錯誤
如果遇到下列錯誤，可能是由於處理耗盡 GPU 記憶體。
```
NCCL error in: ../torch/lib/c10d/ProcessGroupNCCL.cpp:825, unhandled system error, NCCL version 2.7.8
ncclSystemError: System call (socket, malloc, munmap, etc) failed.
```
您可以透過減少批次大小或 `active_microbatches` 來解決此問題。如果自動分割無法產生良好的平衡分割，您可能需要考慮手動分割。如需詳細資訊，請參閱[跨節點的管道平行處理](model-parallel-best-practices.md#model-parallel-best-practices-configuration-pipeline-across-nodes)。
## 收到 PyTorch 訓練任務的 `RecursionError`
該程式庫不支援在模組的正向呼叫內呼叫 `super.forward()`。如果您使用 `super.forward()`，可能會收到下列錯誤訊息。
```
RecursionError: maximum recursion depth exceeded
```
若要修復錯誤，您應呼叫 `super()._orig_forward()` 而不是呼叫 `super.forward()`。