

翻訳は機械翻訳により提供されています。提供された翻訳内容と英語版の間で齟齬、不一致または矛盾がある場合、英語版が優先します。

# SageMaker AI 分散データ並列処理ライブラリを使用した分散トレーニング
<a name="data-parallel-modify-sdp"></a>

SageMaker AI 分散データ並列処理 (SMDDP) ライブラリは、使いやすさと PyTorch とのシームレスな統合を目指して設計されています。

SageMaker AI の SMDDP ライブラリを使用して深層学習モデルをトレーニングする場合、トレーニングスクリプトの記述とモデルのトレーニングに専念できます。

まず、 AWS向けに最適化された集合演算を使用するために、SMDDP ライブラリをインポートしてください。以下のトピックでは、最適化したい集合演算に応じて、トレーニングスクリプトに何を追加すればよいかを説明します。

**Topics**
+ [SMDDP 集合演算を使用するようにトレーニングスクリプトを適応させる](data-parallel-modify-sdp-select-framework.md)
+ [SageMaker Python SDK を使用して SMDDP で分散トレーニングジョブを開始する](data-parallel-use-api.md)

# SMDDP 集合演算を使用するようにトレーニングスクリプトを適応させる
<a name="data-parallel-modify-sdp-select-framework"></a>

このセクションで紹介するトレーニングスクリプトのサンプルは簡略化されており、SageMaker AI 分散データ並列処理 (SMDDP) ライブラリをトレーニングスクリプトで有効にするために必要な変更点だけを取り上げています。SMDDP ライブラリを使用して分散トレーニングジョブを実行する方法を示すエンドツーエンドの Jupyter Notebook の例については、「[Amazon SageMaker AI データ並列処理ライブラリの例](distributed-data-parallel-v2-examples.md)」を参照してください。

**Topics**
+ [PyTorch トレーニングスクリプトで SMDDP ライブラリを使用する](data-parallel-modify-sdp-pt.md)
+ [PyTorch Lightning トレーニングスクリプトで SMDDP ライブラリを使用する](data-parallel-modify-sdp-pt-lightning.md)
+ [TensorFlow トレーニングスクリプトで SMDDP ライブラリを使用する (非推奨)](data-parallel-modify-sdp-tf2.md)

# PyTorch トレーニングスクリプトで SMDDP ライブラリを使用する
<a name="data-parallel-modify-sdp-pt"></a>

SageMaker AI 分散データ並列処理 (SMDDP) ライブラリ v1.4.0 以降は、[PyTorch の分散パッケージ](https://pytorch.org/tutorials/beginner/dist_overview.html)用のバックエンドオプションとして使用できます。SMDDP の `AllReduce` および `AllGather` 集合演算を使用するには、トレーニングスクリプトの冒頭で SMDDP ライブラリをインポートし、プロセスグループの初期化中に SMDDP を PyTorch 分散モジュールのバックエンドとして設定するだけで済みます。1 行でバックエンドを指定するだけで、ネイティブの PyTorch 分散モジュールはすべてそのまま残すことができ、トレーニングスクリプト全体を変更する必要はありません。次のコードスニペットは、PyTorch ベースの分散トレーニングパッケージ ([PyTorch Distributed Data Parallel (DDP)](https://pytorch.org/docs/stable/notes/ddp.html)、[PyTorch Fully Sharded Data Parallelism (FSDP)](https://pytorch.org/docs/stable/fsdp.html)、[DeepSpeed](https://github.com/microsoft/DeepSpeed)、[Megatron-DeepSpeed](https://github.com/microsoft/Megatron-DeepSpeed)) のバックエンドとして SMDDP ライブラリを使用する方法を示しています。

## PyTorch DDP または FSDP の場合
<a name="data-parallel-enable-for-ptddp-ptfsdp"></a>

プロセスグループを次のように初期化します。

```
import torch.distributed as dist
import smdistributed.dataparallel.torch.torch_smddp

dist.init_process_group(backend="smddp")
```

**注記**  
(PyTorch DDP ジョブのみ) `smddp` バックエンドは、現時点では、`torch.distributed.new_group()` API によるサブプロセスグループの作成には対応していません。`smddp` バックエンドは、`NCCL` や `Gloo` などの他のプロセスグループバックエンドと同時に使用することはできません。

## DeepSpeed または Megatron-DeepSpeed の場合
<a name="data-parallel-enable-for-deepspeed"></a>

プロセスグループを次のように初期化します。

```
import deepspeed
import smdistributed.dataparallel.torch.torch_smddp

deepspeed.init_distributed(dist_backend="smddp")
```

**注記**  
SMDDP の `AllGather` を [SageMaker Python SDK を使用して SMDDP で分散トレーニングジョブを開始する](data-parallel-use-api.md) の `mpirun` ベースのランチャー (`smdistributed` と `pytorchddp`) で使用するには、トレーニングスクリプトで次の環境変数を設定する必要があります。  

```
export SMDATAPARALLEL_OPTIMIZE_SDP=true
```

PyTorch FSDP トレーニングスクリプトの記述に関する一般的なガイダンスについては、PyTorch ドキュメントの「[Advanced Model Training with Fully Sharded Data Parallel (FSDP)](https://pytorch.org/tutorials/intermediate/FSDP_adavnced_tutorial.html)」を参照してください。

PyTorch DDP トレーニングスクリプトの記述に関する一般的なガイダンスについては、PyTorch ドキュメントの「[Getting started with distributed data parallel](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html)」を参照してください。

トレーニングスクリプトの調整が完了したら、[SageMaker Python SDK を使用して SMDDP で分散トレーニングジョブを開始する](data-parallel-use-api.md) に進みます。

# PyTorch Lightning トレーニングスクリプトで SMDDP ライブラリを使用する
<a name="data-parallel-modify-sdp-pt-lightning"></a>

[PyTorch Lightning](https://pytorch-lightning.readthedocs.io/en/latest/starter/introduction.html) トレーニングスクリプトを持ち込んで SageMaker AI で分散データ並列トレーニングジョブを実行したい場合は、トレーニングスクリプトに最小限の変更を加えてトレーニングジョブを実行できます。必要な変更として、例えば、`smdistributed.dataparallel` ライブラリの PyTorch モジュールをインポートし、SageMaker トレーニングツールキットによって事前設定された SageMaker AI 環境変数を受け入れるように PyTorch Lightning の環境変数を設定して、プロセスグループのバックエンドを `"smddp"` に設定して SMDDP をライブラリを有効にします。詳細については、コード例を使って各ステップごとに説明します。

**注記**  
PyTorch Lightning サポートは SageMaker AI データ並列ライブラリ v1.5.0 以降で利用できます。

## PyTorch Lightning == v2.1.0 および PyTorch == 2.0.1
<a name="smddp-pt-201-lightning-210"></a>

1. `pytorch_lightning` ライブラリと `smdistributed.dataparallel.torch` モジュールをインポートします。

   ```
   import lightning as pl
   import smdistributed.dataparallel.torch.torch_smddp
   ```

1. [LightningEnvironment](https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.plugins.environments.LightningEnvironment.html) をインスタンス化します。

   ```
   from lightning.fabric.plugins.environments.lightning import LightningEnvironment
   
   env = LightningEnvironment()
   env.world_size = lambda: int(os.environ["WORLD_SIZE"])
   env.global_rank = lambda: int(os.environ["RANK"])
   ```

1. **PyTorch DDP の場合** – [DDPStrategy](https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.strategies.DDPStrategy.html) クラスのオブジェクトを作成し、`process_group_backend` には `"smddp"`、`accelerator` には `"gpu"` を指定します。このオブジェクトを [Trainer](https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html) クラスに渡します。

   ```
   import lightning as pl
   from lightning.pytorch.strategies import DDPStrategy
   
   ddp = DDPStrategy(
       cluster_environment=env, 
       process_group_backend="smddp", 
       accelerator="gpu"
   )
   
   trainer = pl.Trainer(
       max_epochs=200, 
       strategy=ddp, 
       devices=num_gpus, 
       num_nodes=num_nodes
   )
   ```

   **PyTorch FSDP の場合** – [FSDPStrategy](https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.strategies.FSDPStrategy.html) クラスのオブジェクトを (選択した[ラッピングポリシー](https://pytorch.org/docs/stable/fsdp.html)で) 作成し、`process_group_backend` には `"smddp"`、`accelerator` には `"gpu"` を指定します。このオブジェクトを [Trainer](https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html) クラスに渡します。

   ```
   import lightning as pl
   from lightning.pytorch.strategies import FSDPStrategy
   
   from functools import partial
   from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
   
   policy = partial(
       size_based_auto_wrap_policy, 
       min_num_params=10000
   )
   
   fsdp = FSDPStrategy(
       auto_wrap_policy=policy,
       process_group_backend="smddp", 
       cluster_environment=env
   )
   
   trainer = pl.Trainer(
       max_epochs=200, 
       strategy=fsdp, 
       devices=num_gpus, 
       num_nodes=num_nodes
   )
   ```

トレーニングスクリプトの調整が完了したら、[SageMaker Python SDK を使用して SMDDP で分散トレーニングジョブを開始する](data-parallel-use-api.md) に進みます。

**注記**  
[SageMaker Python SDK を使用して SMDDP で分散トレーニングジョブを開始する](data-parallel-use-api.md) で SageMaker AI PyTorch 推定器を作成してトレーニングジョブリクエストを送信する場合、SageMaker AI PyTorch トレーニングコンテナに `pytorch-lightning` および `lightning-bolts` をインストールするために `requirements.txt` を指定する必要があります。  

```
# requirements.txt
pytorch-lightning
lightning-bolts
```
トレーニングスクリプトやジョブ送信と一緒に `requirements.txt` ファイルを配置するソースディレクトリを指定する方法の詳細については、Amazon SageMaker AI Python SDK ドキュメントの「[Using third-party libraries](https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#id12)」を参照してください。

# TensorFlow トレーニングスクリプトで SMDDP ライブラリを使用する (非推奨)
<a name="data-parallel-modify-sdp-tf2"></a>

**重要**  
SMDDP ライブラリは TensorFlow のサポートを終了し、v2.11.0 より新しい TensorFlow の DLC では使用できなくなりました。SMDDP ライブラリがインストールされている以前の TensorFlow DLC を調べるには、「[サポートされるフレームワーク](distributed-data-parallel-support.md#distributed-data-parallel-supported-frameworks)」を参照してください。

次のステップは、TensorFlow トレーニングスクリプトを変更して SageMaker AI の分散データ並列ライブラリを利用する方法を示しています。  

ライブラリの API は Horovod の API と同様になるように設計されています。ライブラリが TensorFlow に提供する各 API の詳細については、[SageMaker AI 分散データ並列 TensorFlow API ドキュメント](https://sagemaker.readthedocs.io/en/stable/api/training/smd_data_parallel.html#api-documentation)を参照してください。

**注記**  
SageMaker AI 分散データ並列は、`tf.keras` モジュールを除く `tf` コアモジュールで構成される TensorFlow トレーニングスクリプトに適応可能です。SageMaker AI 分散データ並列は、Keras 実装を使った TensorFlow はサポートしていません。

**注記**  
SageMaker AI の分散データ並列処理ライブラリは、Automatic Mixed Precision (AMP) に標準対応しています。AMP を有効にするには、トレーニングスクリプトにフレームワークレベルの変更を加える以外、特別なアクションは必要ありません。勾配が FP16 の場合、SageMaker AI データ並列処理ライブラリは、FP16 で `AllReduce` オペレーションを実行します。トレーニングスクリプトへの AMP API の実装に関する詳細は、次のリソースを参照してください。  
「[Frameworks - TensorFlow](https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html#tensorflow)」(「NVIDIA Deep Learning Performance ドキュメント」)
[深層学習の自動混合精度](https://developer.nvidia.com/automatic-mixed-precision) (NVIDIA デベロッパードキュメント)**
[TensorFlow 混合精度 API](https://www.tensorflow.org/guide/mixed_precision) (TensorFlow ドキュメント)**

1. ライブラリの TensorFlow クライアントをインポートし、初期化します。

   ```
   import smdistributed.dataparallel.tensorflow as sdp 
   sdp.init()
   ```

1. 各 GPU を 1 つの `smdistributed.dataparallel` プロセスに `local_rank` を付けて固定します。これは、特定のノード内のプロセスの相対ランクを参照します。`sdp.tensorflow.local_rank()` API により、デバイスのローカルランクを取得できます。リーダーノードはランク 0 で、ワーカーノードはランク 1、2、3 などとなります。これは、次のコードブロックで `sdp.local_rank()` として呼び出されます。`set_memory_growth` は SageMaker AI 分散とは直接関係ありませんが、TensorFlow を使用した分散トレーニングのために設定する必要があります。

   ```
   gpus = tf.config.experimental.list_physical_devices('GPU')
   for gpu in gpus:
       tf.config.experimental.set_memory_growth(gpu, True)
   if gpus:
       tf.config.experimental.set_visible_devices(gpus[sdp.local_rank()], 'GPU')
   ```

1. 学習レートをワーカー数でスケールします。`sdp.tensorflow.size()` API により、クラスター内のワーカー数を取得できます。これは、次のコードブロックで `sdp.size()` として呼び出されます。

   ```
   learning_rate = learning_rate * sdp.size()
   ```

1. ライブラリの `DistributedGradientTape` を使用して、トレーニング中の `AllReduce` オペレーションを最適化します。これは `tf.GradientTape` をラップします。  

   ```
   with tf.GradientTape() as tape:
         output = model(input)
         loss_value = loss(label, output)
       
   # SageMaker AI data parallel: Wrap tf.GradientTape with the library's DistributedGradientTape
   tape = sdp.DistributedGradientTape(tape)
   ```

1. 初期モデル変数をリーダーノード (ランク 0) からすべてのワーカーノード (ランク 1～n) にブロードキャストします。これは、すべてのワーカーランクにわたる一貫した初期化を保証するために必要です。モデル変数とオプティマイザ変数が初期化された後、`sdp.tensorflow.broadcast_variables` API を使用します。これは、次のコードブロックで `sdp.broadcast_variables()` として呼び出されます。

   ```
   sdp.broadcast_variables(model.variables, root_rank=0)
   sdp.broadcast_variables(opt.variables(), root_rank=0)
   ```

1. 最後に、チェックポイントをリーダーノードのみに保存するようにスクリプトを変更します。リーダーノードには同期されたモデルがあります。これにより、ワーカーノードがチェックポイントを上書きしてチェックポイントを破損する可能性を回避できます。

   ```
   if sdp.rank() == 0:
       checkpoint.save(checkpoint_dir)
   ```

次に、ライブラリを使用した分散トレーニングの TensorFlow トレーニングスクリプトの例を示します。

```
import tensorflow as tf

# SageMaker AI data parallel: Import the library TF API
import smdistributed.dataparallel.tensorflow as sdp

# SageMaker AI data parallel: Initialize the library
sdp.init()

gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)
if gpus:
    # SageMaker AI data parallel: Pin GPUs to a single library process
    tf.config.experimental.set_visible_devices(gpus[sdp.local_rank()], 'GPU')

# Prepare Dataset
dataset = tf.data.Dataset.from_tensor_slices(...)

# Define Model
mnist_model = tf.keras.Sequential(...)
loss = tf.losses.SparseCategoricalCrossentropy()

# SageMaker AI data parallel: Scale Learning Rate
# LR for 8 node run : 0.000125
# LR for single node run : 0.001
opt = tf.optimizers.Adam(0.000125 * sdp.size())

@tf.function
def training_step(images, labels, first_batch):
    with tf.GradientTape() as tape:
        probs = mnist_model(images, training=True)
        loss_value = loss(labels, probs)

    # SageMaker AI data parallel: Wrap tf.GradientTape with the library's DistributedGradientTape
    tape = sdp.DistributedGradientTape(tape)

    grads = tape.gradient(loss_value, mnist_model.trainable_variables)
    opt.apply_gradients(zip(grads, mnist_model.trainable_variables))

    if first_batch:
       # SageMaker AI data parallel: Broadcast model and optimizer variables
       sdp.broadcast_variables(mnist_model.variables, root_rank=0)
       sdp.broadcast_variables(opt.variables(), root_rank=0)

    return loss_value

...

# SageMaker AI data parallel: Save checkpoints only from master node.
if sdp.rank() == 0:
    checkpoint.save(checkpoint_dir)
```

トレーニングスクリプトの調整が完了したら、「[SageMaker Python SDK を使用して SMDDP で分散トレーニングジョブを開始する](data-parallel-use-api.md)」に進みます。

# SageMaker Python SDK を使用して SMDDP で分散トレーニングジョブを開始する
<a name="data-parallel-use-api"></a>

「[SMDDP 集合演算を使用するようにトレーニングスクリプトを適応させる](data-parallel-modify-sdp-select-framework.md)」で適応させたスクリプトを使用して分散トレーニングジョブを実行するには、SageMaker Python SDK のフレームワークや汎用推定器を使用し、準備したトレーニングスクリプトをエントリポイントスクリプトおよび分散トレーニング設定として指定します。

このページでは、[SageMaker AI Python SDK](https://sagemaker.readthedocs.io/en/stable/api/training/index.html)の 2 とおりの使い方を説明します。
+ SageMaker AI で分散型トレーニングジョブを迅速に導入したい場合は、SageMaker AI の [PyTorch](https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/sagemaker.pytorch.html#sagemaker.pytorch.estimator.PyTorch) または [TensorFlow](https://sagemaker.readthedocs.io/en/stable/frameworks/tensorflow/sagemaker.tensorflow.html#tensorflow-estimator) フレームワーク推定クラスを設定します。フレームワーク推定器はトレーニングスクリプトを取得し、`framework_version` パラメータに指定された値を指定し、[pre-built PyTorch or TensorFlow Deep Learning Containers (DLC)](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#sagemaker-framework-containers-sm-support-only) の適切な画像 URI を自動的に照合します。
+ 事前構築済みコンテナのいずれかの拡張や、カスタムコンテナを構築し SageMaker AI で独自の ML 環境を作成する場合は、SageMaker AI 汎用 `Estimator` クラスを使用して、Amazon Elastic Container Registry (Amazon ECR) でホストされているカスタム Docker コンテナのイメージ URI を指定します。

トレーニングデータセットは、トレーニングジョブを起動する の Amazon S3  [または Amazon FSx for Lustre](https://docs.aws.amazon.com/fsx/latest/LustreGuide/what-is.html) AWS リージョン に保存する必要があります。Jupyter Notebook を使用する場合は、SageMaker ノートブックインスタンスまたは SageMaker Studio Classic アプリケーションが同じ AWS リージョンで実行されている必要があります。トレーニングデータの保存の詳細については、「[SageMaker Python SDK data inputs](https://sagemaker.readthedocs.io/en/stable/overview.html#use-file-systems-as-training-input)」ドキュメントを参照してください。

**ヒント**  
トレーニングのパフォーマンスを向上させるために、Amazon S3 の代わりに Amazon FSx for Lustre を使用することをお勧めします。Amazon FSx は Amazon S3 よりも高スループットで、低レイテンシーです。

**ヒント**  
EFA 対応のインスタンスタイプで分散トレーニングを適切に実行するには、VPC のセキュリティグループとの間のインバウンドトラフィックとアウトバウンドトラフィックをすべて許可するように設定し、インスタンス間のトラフィックを有効にする必要があります。セキュリティグループのルールを設定する方法については、「*Amazon EC2 ユーザーガイド*」の「[ステップ 1: EFA 対応のセキュリティグループを準備する](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa-start.html#efa-start-security)」を参照してください。

トレーニングスクリプトの分散トレーニングジョブを実行する方法については、以下のトピックの中から選択して手順を確認してください。トレーニングジョブを起動した後、[Amazon SageMaker デバッガー](train-debugger.md) または Amazon CloudWatch を使ってシステム使用率とモデルパフォーマンスをモニタリングできます。

技術的な詳細については、次のトピックの手順に従ってください。また、開始するには「[Amazon SageMaker AI データ並列処理ライブラリの例](distributed-data-parallel-v2-examples.md)」を試してみることをお勧めします。

**Topics**
+ [SageMaker Python SDK で PyTorch フレームワーク推定器を使用する](data-parallel-framework-estimator.md)
+ [SageMaker AI 汎用推定器を使用して構築済みの DLC コンテナを拡張する](data-parallel-use-python-skd-api.md)
+ [SageMaker AI 分散データ並列ライブラリを使用して独自の Docker コンテナを作成する](data-parallel-bring-your-own-container.md)

# SageMaker Python SDK で PyTorch フレームワーク推定器を使用する
<a name="data-parallel-framework-estimator"></a>

分散トレーニングを開始するには、SageMaker AI フレームワーク推定器 ([https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/sagemaker.pytorch.html#sagemaker.pytorch.estimator.PyTorch](https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/sagemaker.pytorch.html#sagemaker.pytorch.estimator.PyTorch) または [https://sagemaker.readthedocs.io/en/stable/frameworks/tensorflow/sagemaker.tensorflow.html#tensorflow-estimator](https://sagemaker.readthedocs.io/en/stable/frameworks/tensorflow/sagemaker.tensorflow.html#tensorflow-estimator)) に `distribution` 引数を追加します。詳細については、次の選択肢から、SageMaker AI 分散データ並列処理 (SMDDP) ライブラリでサポートされているフレームワークのいずれかを選択してください。

------
#### [ PyTorch ]

PyTorch 分散トレーニングを開始するには、次のランチャーオプションを使用できます。
+ `pytorchddp` – このオプションは、`mpirun` を実行し、SageMaker AI で PyTorch 分散トレーニングを実行するために必要な環境変数を設定します。このオプションを使用するには、次のディクショナリを `distribution` パラメータに渡してください。

  ```
  { "pytorchddp": { "enabled": True } }
  ```
+ `torch_distributed` – このオプションは、`torchrun` を実行し、SageMaker AI で PyTorch 分散トレーニングを実行するために必要な環境変数を設定します。このオプションを使用するには、次のディクショナリを `distribution` パラメータに渡してください。

  ```
  { "torch_distributed": { "enabled": True } }
  ```
+ `smdistributed` – このオプションも `mpirun` を実行しますが、`smddprun` を使用して、SageMaker AI で PyTorch 分散トレーニングを実行するために必要な環境変数を設定します。

  ```
  { "smdistributed": { "dataparallel": { "enabled": True } } }
  ```

NCCL の `AllGather` を SMDDP の `AllGather` に置き換えることにした場合は、3 つのオプションすべてを使用できます。ユースケースに合ったオプションを選択してください。

NCCL の `AllReduce` を SMDDP の `AllReduce` に置き換えることにした場合は、`mpirun` ベースのオプションのいずれか (`smdistributed` または `pytorchddp`)を選択する必要があります。次のように MPI オプションを追加することもできます。

```
{ 
    "pytorchddp": {
        "enabled": True, 
        "custom_mpi_options": "-verbose -x NCCL_DEBUG=VERSION"
    }
}
```

```
{ 
    "smdistributed": { 
        "dataparallel": {
            "enabled": True, 
            "custom_mpi_options": "-verbose -x NCCL_DEBUG=VERSION"
        }
    }
}
```

次のコード例は、分散トレーニングのオプションを指定した PyTorch 推定器の基本構造を示しています。

```
from sagemaker.pytorch import PyTorch

pt_estimator = PyTorch(
    base_job_name="training_job_name_prefix",
    source_dir="subdirectory-to-your-code",
    entry_point="adapted-training-script.py",
    role="SageMakerRole",
    py_version="py310",
    framework_version="2.0.1",

    # For running a multi-node distributed training job, specify a value greater than 1
    # Example: 2,3,4,..8
    instance_count=2,

    # Instance types supported by the SageMaker AI data parallel library: 
    # ml.p4d.24xlarge, ml.p4de.24xlarge
    instance_type="ml.p4d.24xlarge",

    # Activate distributed training with SMDDP
    distribution={ "pytorchddp": { "enabled": True } }  # mpirun, activates SMDDP AllReduce OR AllGather
    # distribution={ "torch_distributed": { "enabled": True } }  # torchrun, activates SMDDP AllGather
    # distribution={ "smdistributed": { "dataparallel": { "enabled": True } } }  # mpirun, activates SMDDP AllReduce OR AllGather
)

pt_estimator.fit("s3://bucket/path/to/training/data")
```

**注記**  
PyTorch Lightning と Lightning Bolts などのユーティリティライブラリは、SageMaker AI PyTorch DLC にはプリインストールされていません。次の `requirements.txt` ファイルを作成し、トレーニングスクリプトを保存するソースディレクトリに保存します。  

```
# requirements.txt
pytorch-lightning
lightning-bolts
```
例えば、ツリー構造のディレクトリは次のようになります。  

```
├── pytorch_training_launcher_jupyter_notebook.ipynb
└── sub-folder-for-your-code
    ├──  adapted-training-script.py
    └──  requirements.txt
```
トレーニングスクリプトやジョブ送信と一緒に `requirements.txt` ファイルを配置するソースディレクトリを指定する方法の詳細については、Amazon SageMaker AI Python SDK ドキュメントの「[Using third-party libraries](https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#id12)」を参照してください。

**SMDDP 集合演算を有効にし、適切な分散トレーニングランチャーオプションを使用するための考慮事項**
+ SMDDP の `AllReduce` と SMDDP の `AllGather` は、現時点では相互互換性がありません。
+ SMDDP の `AllReduce` は、`mpirun` ベースのランチャーである `smdistributed` または `pytorchddp` を使用する場合はデフォルトで有効になり、NCCL の `AllGather` が使用されます。
+ SMDDP の `AllGather` は `torch_distributed` ランチャーの使用時にデフォルトで有効になり、`AllReduce` は NCCL にフォールバックされます。
+ SMDDP の `AllGather` は、`mpirun` ベースのランチャーを使用する場合も、次のように追加の環境変数を設定することで有効化できます。

  ```
  export SMDATAPARALLEL_OPTIMIZE_SDP=true
  ```

------
#### [ TensorFlow ]

**重要**  
SMDDP ライブラリは TensorFlow のサポートを終了し、v2.11.0 より新しい TensorFlow の DLC では使用できなくなりました。SMDDP ライブラリがインストールされている以前の TensorFlow DLC を調べるには、「[TensorFlow (非推奨)](distributed-data-parallel-support.md#distributed-data-parallel-supported-frameworks-tensorflow)」を参照してください。

```
from sagemaker.tensorflow import TensorFlow

tf_estimator = TensorFlow(
    base_job_name = "training_job_name_prefix",
    entry_point="adapted-training-script.py",
    role="SageMakerRole",
    framework_version="2.11.0",
    py_version="py38",

    # For running a multi-node distributed training job, specify a value greater than 1
    # Example: 2,3,4,..8
    instance_count=2,

    # Instance types supported by the SageMaker AI data parallel library: 
    # ml.p4d.24xlarge, ml.p3dn.24xlarge, and ml.p3.16xlarge
    instance_type="ml.p3.16xlarge",

    # Training using the SageMaker AI data parallel distributed training strategy
    distribution={ "smdistributed": { "dataparallel": { "enabled": True } } }
)

tf_estimator.fit("s3://bucket/path/to/training/data")
```

------

# SageMaker AI 汎用推定器を使用して構築済みの DLC コンテナを拡張する
<a name="data-parallel-use-python-skd-api"></a>

SageMaker AI の構築済みコンテナはカスタマイズまたは拡張が可能で、構築済み SageMaker AI Docker イメージではサポートされていないアルゴリズムやモデルに対するどのような追加の機能要件でも扱うことができます。構築済みコンテナを拡張する方法の例については、「[構築済みコンテナを拡張する](https://docs.aws.amazon.com/sagemaker/latest/dg/prebuilt-containers-extend.html)」を参照してください。

ビルド済みコンテナの拡張や、独自のコンテナをライブラリを使用するよう適応させるには、「[サポートされるフレームワーク](distributed-data-parallel-support.md#distributed-data-parallel-supported-frameworks)」に記載されているイメージのいずれかを使用する必要があります。

**注記**  
TensorFlow 2.4.1 および PyTorch 1.8.1 から、SageMaker AI フレームワーク DLC は EFA 対応のインスタンスタイプをサポートしています。TensorFlow 2.4.1 以降および PyTorch 1.8.1 以降を含む DLC イメージを使用することをお勧めします。

例えば、PyTorch を使用している場合、Dockerfile には次のような `FROM` ステートメントを記述する必要があります。

```
# SageMaker AI PyTorch image
FROM 763104351884.dkr.ecr.<aws-region>.amazonaws.com/pytorch-training:<image-tag>

ENV PATH="/opt/ml/code:${PATH}"

# this environment variable is used by the SageMaker AI PyTorch container to determine our user code directory.
ENV SAGEMAKER_SUBMIT_DIRECTORY /opt/ml/code

# /opt/ml and all subdirectories are utilized by SageMaker AI, use the /code subdirectory to store your user code.
COPY train.py /opt/ml/code/train.py

# Defines cifar10.py as script entrypoint
ENV SAGEMAKER_PROGRAM train.py
```

[SageMaker Training toolkit](https://github.com/aws/sagemaker-training-toolkit) と SageMaker AI 分散データ並列ライブラリのバイナリファイルを使用して、SageMaker AI で動作するように独自の Docker コンテナをさらにカスタマイズできます。詳細については、次のセクションを参照してください。

# SageMaker AI 分散データ並列ライブラリを使用して独自の Docker コンテナを作成する
<a name="data-parallel-bring-your-own-container"></a>

トレーニング用に独自の Docker コンテナを構築して SageMaker AI データ並列ライブラリを使用するには、SageMaker AI 分散並列ライブラリの正しい依存関係とバイナリファイルを Dockerfile に含める必要があります。このセクションでは、データ並列ライブラリを使用して SageMaker AI の分散トレーニングを行うための最小限に依存関係を含む完全な Dockerfile を作成する方法について説明します。

**注記**  
SageMaker AI データ並列ライブラリをバイナリとして使用するこのカスタム Docker オプションは PyTorch でのみ使用できます。

**SageMaker トレーニングツールキットとデータ並列ライブラリを使用して Dockerfile を作成する**

1. [NVIDIA CUDA](https://hub.docker.com/r/nvidia/cuda) の Docker イメージから始めます。CUDA ランタイムと開発ツール (ヘッダーとライブラリ) を含む cuDNN 開発者バージョンを使用し、[PyTorch source code](https://github.com/pytorch/pytorch#from-source) からビルドします。

   ```
   FROM nvidia/cuda:11.3.1-cudnn8-devel-ubuntu20.04
   ```
**ヒント**  
公式の AWS Deep Learning Container (DLC) イメージは、[NVIDIA CUDA ベースイメージ](https://hub.docker.com/r/nvidia/cuda)から構築されています。残りの手順を実行しながら、ビルド済みの DLC イメージをリファレンスとして使用する場合は、「[AWS Deep Learning Containers for PyTorch Dockerfiles](https://github.com/aws/deep-learning-containers/tree/master/pytorch)」をご覧ください。

1. 以下の引数を追加して PyTorch と他のパッケージのバージョンを指定します。SageMaker AI のデータ並列ライブラリや Amazon S3 プラグインなどの AWS リソースを使用するソフトウェアの Amazon S3 バケットパスも指定します。

   次のコード例で提供されているバージョン以外のサードパーティーライブラリを使用するには、[PyTorch 用の AWS Deep Learning Container の公式 Dockerfiles](https://github.com/aws/deep-learning-containers/tree/master/pytorch/training/docker) を調べて、テスト済み、互換性があり、アプリケーションに適したバージョンを見つけることをお勧めします。

   `SMDATAPARALLEL_BINARY` 引数に指定する URL を調べるには、「[サポートされるフレームワーク](distributed-data-parallel-support.md#distributed-data-parallel-supported-frameworks)」のルックアップテーブルを参照してください。

   ```
   ARG PYTORCH_VERSION=1.10.2
   ARG PYTHON_SHORT_VERSION=3.8
   ARG EFA_VERSION=1.14.1
   ARG SMDATAPARALLEL_BINARY=https://smdataparallel.s3.amazonaws.com/binary/pytorch/${PYTORCH_VERSION}/cu113/2022-02-18/smdistributed_dataparallel-1.4.0-cp38-cp38-linux_x86_64.whl
   ARG PT_S3_WHL_GPU=https://aws-s3-plugin.s3.us-west-2.amazonaws.com/binaries/0.0.1/1c3e69e/awsio-0.0.1-cp38-cp38-manylinux1_x86_64.whl
   ARG CONDA_PREFIX="/opt/conda"
   ARG BRANCH_OFI=1.1.3-aws
   ```

1. SageMaker トレーニングコンポーネントを適切にビルドしてデータ並列ライブラリを実行するには、以下の環境変数を設定します。以降のステップでは、これらの変数をコンポーネントに使用します。

   ```
   # Set ENV variables required to build PyTorch
   ENV TORCH_CUDA_ARCH_LIST="7.0+PTX 8.0"
   ENV TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
   ENV NCCL_VERSION=2.10.3
   
   # Add OpenMPI to the path.
   ENV PATH /opt/amazon/openmpi/bin:$PATH
   
   # Add Conda to path
   ENV PATH $CONDA_PREFIX/bin:$PATH
   
   # Set this enviroment variable for SageMaker AI to launch SMDDP correctly.
   ENV SAGEMAKER_TRAINING_MODULE=sagemaker_pytorch_container.training:main
   
   # Add enviroment variable for processes to be able to call fork()
   ENV RDMAV_FORK_SAFE=1
   
   # Indicate the container type
   ENV DLC_CONTAINER_TYPE=training
   
   # Add EFA and SMDDP to LD library path
   ENV LD_LIBRARY_PATH="/opt/conda/lib/python${PYTHON_SHORT_VERSION}/site-packages/smdistributed/dataparallel/lib:$LD_LIBRARY_PATH"
   ENV LD_LIBRARY_PATH=/opt/amazon/efa/lib/:$LD_LIBRARY_PATH
   ```

1. `curl`、`wget`、`git` をインストールまたは更新し、以降のステップでパッケージをダウンロードしてビルドします。

   ```
   RUN --mount=type=cache,id=apt-final,target=/var/cache/apt \
       apt-get update && apt-get install -y  --no-install-recommends \
           curl \
           wget \
           git \
       && rm -rf /var/lib/apt/lists/*
   ```

1. Amazon EC2 ネットワーク通信用の [Elastic Fabric Adapter (EFA)](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa.html) ソフトウェアをインストールします。

   ```
   RUN DEBIAN_FRONTEND=noninteractive apt-get update
   RUN mkdir /tmp/efa \
       && cd /tmp/efa \
       && curl --silent -O https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_VERSION}.tar.gz \
       && tar -xf aws-efa-installer-${EFA_VERSION}.tar.gz \
       && cd aws-efa-installer \
       && ./efa_installer.sh -y --skip-kmod -g \
       && rm -rf /tmp/efa
   ```

1. [Conda](https://docs.conda.io/en/latest/) をインストールしてパッケージ管理を行います。

   ```
   RUN curl -fsSL -v -o ~/miniconda.sh -O  https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh  && \
       chmod +x ~/miniconda.sh && \
       ~/miniconda.sh -b -p $CONDA_PREFIX && \
       rm ~/miniconda.sh && \
       $CONDA_PREFIX/bin/conda install -y python=${PYTHON_SHORT_VERSION} conda-build pyyaml numpy ipython && \
       $CONDA_PREFIX/bin/conda clean -ya
   ```

1. PyTorch とその依存関係を取得し、ビルドおよびインストールします。NCCL のバージョンを管理して [AWS OFI NCCL plug-in](https://github.com/aws/aws-ofi-nccl) との互換性を保証するために、[PyTorch from the source code](https://github.com/pytorch/pytorch#from-source) をビルドします。

   1. [PyTorch official dockerfile](https://github.com/pytorch/pytorch/blob/master/Dockerfile) の手順に従って、ビルドの依存関係をインストールし、再コンパイルを高速化するために [ccache](https://ccache.dev/) を設定します。

      ```
      RUN DEBIAN_FRONTEND=noninteractive \
          apt-get install -y --no-install-recommends \
              build-essential \
              ca-certificates \
              ccache \
              cmake \
              git \
              libjpeg-dev \
              libpng-dev \
          && rm -rf /var/lib/apt/lists/*
        
      # Setup ccache
      RUN /usr/sbin/update-ccache-symlinks
      RUN mkdir /opt/ccache && ccache --set-config=cache_dir=/opt/ccache
      ```

   1. [PyTorch’s common and Linux dependencies](https://github.com/pytorch/pytorch#install-dependencies) をインストールします。

      ```
      # Common dependencies for PyTorch
      RUN conda install astunparse numpy ninja pyyaml mkl mkl-include setuptools cmake cffi typing_extensions future six requests dataclasses
      
      # Linux specific dependency for PyTorch
      RUN conda install -c pytorch magma-cuda113
      ```

   1. [PyTorch GitHub repository](https://github.com/pytorch/pytorch) のクローンを作成します。

      ```
      RUN --mount=type=cache,target=/opt/ccache \
          cd / \
          && git clone --recursive https://github.com/pytorch/pytorch -b v${PYTORCH_VERSION}
      ```

   1. 特定の [NCCL](https://developer.nvidia.com/nccl) バージョンをインストールしてビルドします。これを行うには、PyTorch のデフォルトの NCCL フォルダ (`/pytorch/third_party/nccl`) の内容を、NVIDIA リポジトリの特定の NCCL バージョンに置き換えます。NCCL のバージョンは、このガイドのステップ 3 で設定しました。

      ```
      RUN cd /pytorch/third_party/nccl \
          && rm -rf nccl \
          && git clone https://github.com/NVIDIA/nccl.git -b v${NCCL_VERSION}-1 \
          && cd nccl \
          && make -j64 src.build CUDA_HOME=/usr/local/cuda NVCC_GENCODE="-gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_80,code=sm_80" \
          && make pkg.txz.build \
          && tar -xvf build/pkg/txz/nccl_*.txz -C $CONDA_PREFIX --strip-components=1
      ```

   1. PyTorch をビルドしてインストールします。このプロセスが完了するまで、通常は 1 時間強かかります。前のステップでダウンロードした NCCL バージョンを使用してビルドされます。

      ```
      RUN cd /pytorch \
          && CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" \
          python setup.py install \
          && rm -rf /pytorch
      ```

1. [AWS OFI NCCL プラグイン](https://github.com/aws/aws-ofi-nccl)をビルドしてインストールします。これにより、SageMaker AI 分散データ並列ライブラリの [libfabric](https://github.com/ofiwg/libfabric) サポートが有効になります。

   ```
   RUN DEBIAN_FRONTEND=noninteractive apt-get update \
       && apt-get install -y --no-install-recommends \
           autoconf \
           automake \
           libtool
   RUN mkdir /tmp/efa-ofi-nccl \
       && cd /tmp/efa-ofi-nccl \
       && git clone https://github.com/aws/aws-ofi-nccl.git -b v${BRANCH_OFI} \
       && cd aws-ofi-nccl \
       && ./autogen.sh \
       && ./configure --with-libfabric=/opt/amazon/efa \
       --with-mpi=/opt/amazon/openmpi \
       --with-cuda=/usr/local/cuda \
       --with-nccl=$CONDA_PREFIX \
       && make \
       && make install \
       && rm -rf /tmp/efa-ofi-nccl
   ```

1. [TorchVision](https://github.com/pytorch/vision.git) をビルドしてインストールします。

   ```
   RUN pip install --no-cache-dir -U \
       packaging \
       mpi4py==3.0.3
   RUN cd /tmp \
       && git clone https://github.com/pytorch/vision.git -b v0.9.1 \
       && cd vision \
       && BUILD_VERSION="0.9.1+cu111" python setup.py install \
       && cd /tmp \
       && rm -rf vision
   ```

1. OpenSSH をインストールおよび設定します。MPI がコンテナ間で通信するには OpenSSH が必要です。確認なしに OpenSSH がコンテナと通信できるようになります。

   ```
   RUN apt-get update \
       && apt-get install -y  --allow-downgrades --allow-change-held-packages --no-install-recommends \
       && apt-get install -y --no-install-recommends openssh-client openssh-server \
       && mkdir -p /var/run/sshd \
       && cat /etc/ssh/ssh_config | grep -v StrictHostKeyChecking > /etc/ssh/ssh_config.new \
       && echo "    StrictHostKeyChecking no" >> /etc/ssh/ssh_config.new \
       && mv /etc/ssh/ssh_config.new /etc/ssh/ssh_config \
       && rm -rf /var/lib/apt/lists/*
   
   # Configure OpenSSH so that nodes can communicate with each other
   RUN mkdir -p /var/run/sshd && \
    sed 's@session\s*required\s*pam_loginuid.so@session optional pam_loginuid.so@g' -i /etc/pam.d/sshd
   RUN rm -rf /root/.ssh/ && \
    mkdir -p /root/.ssh/ && \
    ssh-keygen -q -t rsa -N '' -f /root/.ssh/id_rsa && \
    cp /root/.ssh/id_rsa.pub /root/.ssh/authorized_keys \
    && printf "Host *\n StrictHostKeyChecking no\n" >> /root/.ssh/config
   ```

1. PT S3 プラグインをインストールして、Amazon S3 のデータセットに効率的にアクセスします。

   ```
   RUN pip install --no-cache-dir -U ${PT_S3_WHL_GPU}
   RUN mkdir -p /etc/pki/tls/certs && cp /etc/ssl/certs/ca-certificates.crt /etc/pki/tls/certs/ca-bundle.crt
   ```

1. [libboost](https://www.boost.org/) ライブラリをインストールします。本パッケージは SageMaker AI データ並列ライブラリの非同期 IO 機能をネットワーク化するために必要です。

   ```
   WORKDIR /
   RUN wget https://sourceforge.net/projects/boost/files/boost/1.73.0/boost_1_73_0.tar.gz/download -O boost_1_73_0.tar.gz \
       && tar -xzf boost_1_73_0.tar.gz \
       && cd boost_1_73_0 \
       && ./bootstrap.sh \
       && ./b2 threading=multi --prefix=${CONDA_PREFIX} -j 64 cxxflags=-fPIC cflags=-fPIC install || true \
       && cd .. \
       && rm -rf boost_1_73_0.tar.gz \
       && rm -rf boost_1_73_0 \
       && cd ${CONDA_PREFIX}/include/boost
   ```

1. PyTorch トレーニング用に次の SageMaker AI ツールをインストールしてください。

   ```
   WORKDIR /root
   RUN pip install --no-cache-dir -U \
       smclarify \
       "sagemaker>=2,<3" \
       sagemaker-experiments==0.* \
       sagemaker-pytorch-training
   ```

1. 最後に、SageMaker AI データ並列ライブラリと残りの依存関係をインストールします。

   ```
   RUN --mount=type=cache,id=apt-final,target=/var/cache/apt \
     apt-get update && apt-get install -y  --no-install-recommends \
     jq \
     libhwloc-dev \
     libnuma1 \
     libnuma-dev \
     libssl1.1 \
     libtool \
     hwloc \
     && rm -rf /var/lib/apt/lists/*
   
   RUN SMDATAPARALLEL_PT=1 pip install --no-cache-dir ${SMDATAPARALLEL_BINARY}
   ```

1. Dockerfile の作成が完了したら、「[Adapting Your Own Training Container](https://docs.aws.amazon.com/sagemaker/latest/dg/adapt-training-container.html)」を参照し、Docker コンテナを構築する方法、Amazon ECR でホストする方法、および SageMaker Python SDK を使用してトレーニングジョブを実行する方法を確認してください。

以下のサンプルコードは、前述のコードブロックをすべて組み合わせた後の Dockerfile 全体を示しています。

```
# This file creates a docker image with minimum dependencies to run SageMaker AI data parallel training
FROM nvidia/cuda:11.3.1-cudnn8-devel-ubuntu20.04

# Set appropiate versions and location for components
ARG PYTORCH_VERSION=1.10.2
ARG PYTHON_SHORT_VERSION=3.8
ARG EFA_VERSION=1.14.1
ARG SMDATAPARALLEL_BINARY=https://smdataparallel.s3.amazonaws.com/binary/pytorch/${PYTORCH_VERSION}/cu113/2022-02-18/smdistributed_dataparallel-1.4.0-cp38-cp38-linux_x86_64.whl
ARG PT_S3_WHL_GPU=https://aws-s3-plugin.s3.us-west-2.amazonaws.com/binaries/0.0.1/1c3e69e/awsio-0.0.1-cp38-cp38-manylinux1_x86_64.whl
ARG CONDA_PREFIX="/opt/conda"
ARG BRANCH_OFI=1.1.3-aws

# Set ENV variables required to build PyTorch
ENV TORCH_CUDA_ARCH_LIST="3.7 5.0 7.0+PTX 8.0"
ENV TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
ENV NCCL_VERSION=2.10.3

# Add OpenMPI to the path.
ENV PATH /opt/amazon/openmpi/bin:$PATH

# Add Conda to path
ENV PATH $CONDA_PREFIX/bin:$PATH

# Set this enviroment variable for SageMaker AI to launch SMDDP correctly.
ENV SAGEMAKER_TRAINING_MODULE=sagemaker_pytorch_container.training:main

# Add enviroment variable for processes to be able to call fork()
ENV RDMAV_FORK_SAFE=1

# Indicate the container type
ENV DLC_CONTAINER_TYPE=training

# Add EFA and SMDDP to LD library path
ENV LD_LIBRARY_PATH="/opt/conda/lib/python${PYTHON_SHORT_VERSION}/site-packages/smdistributed/dataparallel/lib:$LD_LIBRARY_PATH"
ENV LD_LIBRARY_PATH=/opt/amazon/efa/lib/:$LD_LIBRARY_PATH

# Install basic dependencies to download and build other dependencies
RUN --mount=type=cache,id=apt-final,target=/var/cache/apt \
  apt-get update && apt-get install -y  --no-install-recommends \
  curl \
  wget \
  git \
  && rm -rf /var/lib/apt/lists/*

# Install EFA.
# This is required for SMDDP backend communication
RUN DEBIAN_FRONTEND=noninteractive apt-get update
RUN mkdir /tmp/efa \
    && cd /tmp/efa \
    && curl --silent -O https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_VERSION}.tar.gz \
    && tar -xf aws-efa-installer-${EFA_VERSION}.tar.gz \
    && cd aws-efa-installer \
    && ./efa_installer.sh -y --skip-kmod -g \
    && rm -rf /tmp/efa

# Install Conda
RUN curl -fsSL -v -o ~/miniconda.sh -O  https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh  && \
    chmod +x ~/miniconda.sh && \
    ~/miniconda.sh -b -p $CONDA_PREFIX && \
    rm ~/miniconda.sh && \
    $CONDA_PREFIX/bin/conda install -y python=${PYTHON_SHORT_VERSION} conda-build pyyaml numpy ipython && \
    $CONDA_PREFIX/bin/conda clean -ya

# Install PyTorch.
# Start with dependencies listed in official PyTorch dockerfile
# https://github.com/pytorch/pytorch/blob/master/Dockerfile
RUN DEBIAN_FRONTEND=noninteractive \
    apt-get install -y --no-install-recommends \
        build-essential \
        ca-certificates \
        ccache \
        cmake \
        git \
        libjpeg-dev \
        libpng-dev && \
    rm -rf /var/lib/apt/lists/*

# Setup ccache
RUN /usr/sbin/update-ccache-symlinks
RUN mkdir /opt/ccache && ccache --set-config=cache_dir=/opt/ccache

# Common dependencies for PyTorch
RUN conda install astunparse numpy ninja pyyaml mkl mkl-include setuptools cmake cffi typing_extensions future six requests dataclasses

# Linux specific dependency for PyTorch
RUN conda install -c pytorch magma-cuda113

# Clone PyTorch
RUN --mount=type=cache,target=/opt/ccache \
    cd / \
    && git clone --recursive https://github.com/pytorch/pytorch -b v${PYTORCH_VERSION}
# Note that we need to use the same NCCL version for PyTorch and OFI plugin.
# To enforce that, install NCCL from source before building PT and OFI plugin.

# Install NCCL.
# Required for building OFI plugin (OFI requires NCCL's header files and library)
RUN cd /pytorch/third_party/nccl \
    && rm -rf nccl \
    && git clone https://github.com/NVIDIA/nccl.git -b v${NCCL_VERSION}-1 \
    && cd nccl \
    && make -j64 src.build CUDA_HOME=/usr/local/cuda NVCC_GENCODE="-gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_80,code=sm_80" \
    && make pkg.txz.build \
    && tar -xvf build/pkg/txz/nccl_*.txz -C $CONDA_PREFIX --strip-components=1

# Build and install PyTorch.
RUN cd /pytorch \
    && CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" \
    python setup.py install \
    && rm -rf /pytorch

RUN ccache -C

# Build and install OFI plugin. \
# It is required to use libfabric.
RUN DEBIAN_FRONTEND=noninteractive apt-get update \
    && apt-get install -y --no-install-recommends \
        autoconf \
        automake \
        libtool
RUN mkdir /tmp/efa-ofi-nccl \
    && cd /tmp/efa-ofi-nccl \
    && git clone https://github.com/aws/aws-ofi-nccl.git -b v${BRANCH_OFI} \
    && cd aws-ofi-nccl \
    && ./autogen.sh \
    && ./configure --with-libfabric=/opt/amazon/efa \
        --with-mpi=/opt/amazon/openmpi \
        --with-cuda=/usr/local/cuda \
        --with-nccl=$CONDA_PREFIX \
    && make \
    && make install \
    && rm -rf /tmp/efa-ofi-nccl

# Build and install Torchvision
RUN pip install --no-cache-dir -U \
    packaging \
    mpi4py==3.0.3
RUN cd /tmp \
    && git clone https://github.com/pytorch/vision.git -b v0.9.1 \
    && cd vision \
    && BUILD_VERSION="0.9.1+cu111" python setup.py install \
    && cd /tmp \
    && rm -rf vision

# Install OpenSSH.
# Required for MPI to communicate between containers, allow OpenSSH to talk to containers without asking for confirmation
RUN apt-get update \
    && apt-get install -y  --allow-downgrades --allow-change-held-packages --no-install-recommends \
    && apt-get install -y --no-install-recommends openssh-client openssh-server \
    && mkdir -p /var/run/sshd \
    && cat /etc/ssh/ssh_config | grep -v StrictHostKeyChecking > /etc/ssh/ssh_config.new \
    && echo "    StrictHostKeyChecking no" >> /etc/ssh/ssh_config.new \
    && mv /etc/ssh/ssh_config.new /etc/ssh/ssh_config \
    && rm -rf /var/lib/apt/lists/*
# Configure OpenSSH so that nodes can communicate with each other
RUN mkdir -p /var/run/sshd && \
    sed 's@session\s*required\s*pam_loginuid.so@session optional pam_loginuid.so@g' -i /etc/pam.d/sshd
RUN rm -rf /root/.ssh/ && \
    mkdir -p /root/.ssh/ && \
    ssh-keygen -q -t rsa -N '' -f /root/.ssh/id_rsa && \
    cp /root/.ssh/id_rsa.pub /root/.ssh/authorized_keys \
    && printf "Host *\n StrictHostKeyChecking no\n" >> /root/.ssh/config

# Install PT S3 plugin.
# Required to efficiently access datasets in Amazon S3
RUN pip install --no-cache-dir -U ${PT_S3_WHL_GPU}
RUN mkdir -p /etc/pki/tls/certs && cp /etc/ssl/certs/ca-certificates.crt /etc/pki/tls/certs/ca-bundle.crt

# Install libboost from source.
# This package is needed for smdataparallel functionality (for networking asynchronous IO).
WORKDIR /
RUN wget https://sourceforge.net/projects/boost/files/boost/1.73.0/boost_1_73_0.tar.gz/download -O boost_1_73_0.tar.gz \
    && tar -xzf boost_1_73_0.tar.gz \
    && cd boost_1_73_0 \
    && ./bootstrap.sh \
    && ./b2 threading=multi --prefix=${CONDA_PREFIX} -j 64 cxxflags=-fPIC cflags=-fPIC install || true \
    && cd .. \
    && rm -rf boost_1_73_0.tar.gz \
    && rm -rf boost_1_73_0 \
    && cd ${CONDA_PREFIX}/include/boost

# Install SageMaker AI PyTorch training.
WORKDIR /root
RUN pip install --no-cache-dir -U \
    smclarify \
    "sagemaker>=2,<3" \
    sagemaker-experiments==0.* \
    sagemaker-pytorch-training

# Install SageMaker AI data parallel binary (SMDDP)
# Start with dependencies
RUN --mount=type=cache,id=apt-final,target=/var/cache/apt \
    apt-get update && apt-get install -y  --no-install-recommends \
        jq \
        libhwloc-dev \
        libnuma1 \
        libnuma-dev \
        libssl1.1 \
        libtool \
        hwloc \
    && rm -rf /var/lib/apt/lists/*

# Install SMDDP
RUN SMDATAPARALLEL_PT=1 pip install --no-cache-dir ${SMDATAPARALLEL_BINARY}
```

**ヒント**  
SageMaker AI でのトレーニング用のカスタム Dockerfile の作成に関する一般的な情報については、「[Use Your Own Training Algorithms](https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-training-algo.html)」を参照してください。

**ヒント**  
カスタム Dockerfile を拡張して SageMaker AI モデル並列ライブラリを組み込む場合は、「[SageMaker 分散モデル並列ライブラリを使用した独自の Docker コンテナの作成](model-parallel-sm-sdk.md#model-parallel-bring-your-own-container)」を参照してください。