翻訳は機械翻訳により提供されています。提供された翻訳内容と英語版の間で齟齬、不一致または矛盾がある場合、英語版が優先します。 # (アーカイブ済み) SageMaker モデル並列処理ライブラリ v1.x **重要** 2023 年 12 月 19 日に、SageMaker モデル並列処理 (SMP) ライブラリ v2 がリリースされました。SMP v1 の機能は今後のリリースではサポートされなくなります。代わりに SMP ライブラリ v2 をご利用ください。以下のセクションとトピックはアーカイブ済みであり、SMP ライブラリ v1 の使用に特化した内容です。SMP ライブラリ v2 の使用方法については、「[SageMaker モデル並列処理ライブラリ v2](model-parallel-v2.md)」を参照してください。 Amazon SageMaker AI のモデル並列ライブラリを使用して、GPU メモリの制限によりトレーニングが困難な大規模な深層学習 (DL) モデルをトレーニングします。ライブラリは、モデルを複数の GPU およびインスタンス間で自動的かつ効率的に分割します。ライブラリを使用すると、数十億または数兆のパラメータを使って大規模な DL モデルを効率的にトレーニングすることで、目標予測精度をより速く達成できます。 ライブラリを使用して、最小限のコード変更で、複数の GPU および複数のノード間で、独自の TensorFlow および PyTorch モデルを自動的にパーティション化できます。ライブラリの API には、SageMaker Python SDK からアクセスできます。 以降のセクションでは、モデル並列処理と SageMaker モデル並列ライブラリについてさらに学習します。このライブラリの API ドキュメントは、*SageMaker Python SDK v2.199.0 ドキュメント*の「[Distributed Training APIs](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 の使用率とトレーニング効率が低下します。 1 つの GPU でのモデルのトレーニングに関連する制限を克服するために、SageMaker はモデル並列ライブラリを提供し、DL モデルを複数のコンピューティングノードに分散して効率的にトレーニングすることを支援します。さらに、このライブラリでは、EFA 対応デバイスを使用して最適化された分散トレーニングを実現できます。これにより、低遅延、高スループット、OS バイパスによるノード間通信のパフォーマンスが向上します。 ## モデル並列処理を使用する前にメモリ要件を見積もる SageMaker モデル並列ライブラリを使用する前に、以下の点を考慮して、大きな DL モデルの学習に必要なメモリ量を把握してください。 AMP (FP16) と Adam オプティマイザを使用するトレーニングジョブでは、パラメータ 1 つあたりに必要な GPU メモリは約 20 バイトで、以下のように分類できます。 + FP16 パラメータ (最大 2 バイト) + FP16 勾配 (最大 2 バイト) + FP32 オプティマイザステート (Adam オプティマイザに基づく最大 8 バイト) + パラメータの FP32 コピー (最大 4 バイト) (`optimizer apply` (OA) オペレーションに必要) + 勾配の FP32 コピー (最大 4 バイト) (OA オペレーションに必要) 100 億個のパラメータを持つ比較的小規模な DL モデルでも、少なくとも 200GB のメモリが必要になることがあります。これは、単一の GPU で利用できる一般的な GPU メモリ (例えば、40GB/80GB のメモリを搭載した NVIDIA A100 と 16/32GB のメモリを搭載した V100) よりもはるかに大きいメモリです。モデルとオプティマイザの状態に必要なメモリ量に加えて、フォワードパスで生成されるアクティベーションなど、他のメモリ消費量もあることに注意してください。必要なメモリは 200GB をはるかに超える場合があります。 分散型トレーニングでは、それぞれ NVIDIA V100 と A100 Tensor コア GPU を搭載した Amazon EC2 P3 インスタンスと P4 インスタンスを使用することをお勧めします。CPU コア、RAM、アタッチされたストレージボリューム、ネットワーク帯域幅などの仕様の詳細については、[Amazon EC2 インスタンスタイプ](https://aws.amazon.com/ec2/instance-types/)ページの「*Accelerated Computing*」セクションを参照してください。 高速コンピューティングインスタンスを使用しても、Megatron-LM や T5 のように約 100 億個のパラメータを持つモデルや、GPT-3 のように数千億個のパラメータを持つさらに大規模なモデルでは、各 GPU デバイスのモデルレプリカに適合しないことは明らかです。 ## ライブラリがモデルの並列処理とメモリ節約の手法を採用している方法 このライブラリは、さまざまなタイプのモデル並列処理機能と、オプティマイザ状態のシャーディング、アクティベーションチェックポイント、アクティベーションオフロードなどのメモリ節約機能で構成されています。これらの手法をすべて組み合わせることで、数千億のパラメータで構成される大規模なモデルを効率的にトレーニングできます。 **Topics** + [シャーディングデータ並列処理 (PyTorch で利用可能)](#model-parallel-intro-sdp) + [パイプライン並列処理 (PyTorch と TensorFlow で利用可能)](#model-parallel-intro-pp) + [Tensor 並列処理 (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 を実装することでシャーディングデータ並列処理を実現しています。これについては、ブログ記事「[Near-linear scaling of gigantic-model training on AWS](https://www.amazon.science/blog/near-linear-scaling-of-gigantic-model-training-on-aws)」で説明しています。 シャーディングデータ並列処理は、スタンドアロン戦略としてモデルに適用できます。さらに、NVIDIA A100 Tensor Core GPU、`ml.p4d.24xlarge` を搭載した最もパフォーマンスの高い GPU インスタンスを使用している場合は、SMDDP Collectives が提供する `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 ``` 指定した 2 つの入力パラメータを指定すると、ライブラリがモデルレプリカ (別名 `data_parallel_degree`) の数を計算します。 例えば、`ml.p3.16xlarge` のような 8 つの GPU ワーカーを持つ ML インスタンスを `"pipeline_parallel_degree": 2` および `"processes_per_host": 8` に設定して使用すると、ライブラリは GPU 全体に分散モデルと 4 方向データ並列処理を自動的にセットアップします。次の図は、モデルを 8 つの GPU に分散して 4 方向データ並列処理と双方向パイプライン並列処理を実現する方法を示しています。パイプライン並列グループとして定義し、`PP_GROUP` というラベルを付けた各モデルレプリカは、2 つの GPU に分割されます。モデルの各パーティションは 4 つの GPU に割り当てられ、4 つのパーティションレプリカはデータ並列グループに属し、`DP_GROUP` というラベルが付けられます。テンソル並列処理がなければ、パイプライン並列グループは基本的にモデル並列グループになります。 ![\[モデルを 8 つの GPU に分散させて 4 方向のデータ並列処理と双方向パイプライン並列処理を実現する方法。\]](http://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/images/distributed/model-parallel/smdmp-pipeline-parallel-only.png) パイプラインの並列処理について詳しくは、「[SageMaker モデル並列処理ライブラリの主要機能](model-parallel-core-features.md)」を参照してください。 パイプライン並列処理によるモデルの実行を開始するには、「[Run a SageMaker Distributed Training Job with the SageMaker Model Parallel Library](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-use-api.html)」を参照してください。 ### Tensor 並列処理 (PyTorch で利用可能) *Tensor 並列処理*は、個々のレイヤー (`nn.Modules`) をデバイス間で分割し、並列に実行します。次の図は、ライブラリがモデルを 4 つのレイヤーに分割して双方向のテンソル並列処理 (`"tensor_parallel_degree": 2`) を実現する方法の最も単純な例を示しています。各モデルレプリカのレイヤーは 2 分割され、2 つの GPU に分散されます。この例では、モデルの並列構成には `"pipeline_parallel_degree": 1` と `"ddp": True` (バックグラウンドでは PyTorch DistributedDataParallel パッケージを使用) も含まれているため、データの並列度は 8 になります。このライブラリは、テンソル分散モデルのレプリカ間の通信を管理します。 ![\[ライブラリがモデルを 4 つのレイヤーに分割して双方向のテンソル並列処理 ("tensor_parallel_degree": 2) を実現する方法の最も単純な例。\]](http://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/images/distributed/model-parallel/smdmp-tensor-parallel-only.png) この機能の便利さは、特定のレイヤーまたはレイヤーのサブセットを選択してテンソル並列処理を適用できる点にあります。PyTorch のテンソル並列処理やその他のメモリ節約機能の詳細と、パイプラインとテンソルの並列処理の組み合わせを設定する方法については、「[テンソル並列処理](model-parallel-extended-features-pytorch-tensor-parallelism.md)」を参照してください。 ### オプティマイザステートのシャーディング (PyTorch で利用可能) ライブラリが*オプティマイザステートシャーディング*をどのように実行するのかを理解するために、4 つのレイヤーからなる単純なモデル例を考えてみましょう。ステートシャーディングを最適化するうえで重要なのは、オプティマイザの状態をすべての GPU に複製する必要がないということです。代わりに、オプティマイザの状態の単一のレプリカが、デバイス間の冗長性なくデータ並列ランクにわたってシャーディングされます。例えば、GPU 0 はレイヤー 1 のオプティマイザ状態を保持し、次の GPU 1 は L2 のオプティマイザ状態を保持する、という具合です。次のアニメーション図は、オプティマイザ状態のシャーディング手法による後方伝播を示しています。後方伝播の最後には、`optimizer apply` (OA) オペレーションでオプティマイザの状態を更新し、`all-gather` (AG) オペレーションで次の反復に備えてモデルパラメータを更新するための計算時間とネットワーク時間が残ります。最も重要なのは、この `reduce` 処理が GPU 0 での計算と重複する可能性があるため、メモリ効率が高まり、後方伝播が速くなることです。現在の実装では、AG と OA のオペレーションは `compute` と重複しません。その結果、AG オペレーション中の計算時間が長くなる可能性があるため、トレードオフが発生する可能性があります。 ![\[オプティマイザ状態のシャーディング手法による後方伝播。\]](http://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/images/distributed/model-parallel/smdmp-optimizer-state-sharding.gif) この機能の使用方法の詳細については、「[Optimizer State Sharding](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-optimizer-state-sharding.html)」を参照してください。 ### アクティベーションオフロードとチェックポイント (PyTorch で利用可能) GPU メモリを節約するため、ライブラリはアクティベーションチェックポイントをサポートしています。これにより、フォワードパス中にユーザー指定モジュールの内部アクティベーションが GPU メモリに保存されるのを防ぐことができます。ライブラリはバックワードパス中にこれらのアクティベーションを再計算します。さらに、アクティベーションオフロード機能は、保存されているアクティベーションを CPU メモリにオフロードし、バックワードパス中に GPU にフェッチバックすることで、アクティベーションメモリのフットプリントをさらに削減します。これらの機能の使用方法について詳しくは、「[Activation Checkpointing](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-activation-checkpointing.html)」と「[Activation Offloading](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-activation-offloading.html)」を参照してください。 ### モデルに適した手法の選択 適切な手法と構成の選択について詳しくは、「[SageMaker Distributed Model Parallel Best Practices](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-best-practices.html)」と「[Configuration Tips and Pitfalls](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-customize-tips-pitfalls.html)」を参照してください。 # サポートされているフレームワークと AWS リージョン SageMaker モデル並列処理ライブラリを使用する前に、サポートされているフレームワークとインスタンスタイプを確認し、 AWS アカウントと に十分なクォータがあるかどうかを確認します AWS リージョン。 **注記** ライブラリの最新の更新を確認するには、「SageMaker Python SDK ドキュメント」の「[SageMaker Model Parallel Release Notes](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 モデル並列処理ライブラリでサポートされている PyTorch バージョン | PyTorch バージョン | SageMaker モデル並列処理ライブラリバージョン | `smdistributed-modelparallel` 統合 DLC イメージ URI | バイナリファイルの URL\$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` | - | | v1.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 Transformer のバージョン** Hugging Face の AWS 深層学習コンテナは、PyTorch と TensorFlow の SageMaker トレーニングコンテナをベースイメージとして使用します。Hugging Face Transformers ライブラリのバージョンとペアになる 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 データ並列ライブラリは、[AWS SageMaker の Deep Learning Containers ](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#sagemaker-framework-containers-sm-support-only) AWS リージョン が稼働しているすべての で使用できます。詳細については、「[Available Deep Learning Containers Images](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 Pricing](https://aws.amazon.com/sagemaker/pricing/)」を参照してください。 次のようなエラーメッセージが表示された場合は、「[Request a service quota increase for SageMaker AI resources](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)」セクションに示されているのと同じ 2 段階のワークフローを維持します。トレーニングスクリプトを適合させるために、トレーニングスクリプトにゼロまたは数行のコード行を追加します。適合させたトレーニングスクリプトのトレーニングジョブを開始するには、メモリ節約機能を有効にしたり、並列処理の制限値を渡すため、ディストリビューション設定パラメータを設定する必要があります。 サンプルを使って開始するには、SageMaker モデル並列処理ライブラリの使用方法を示す次の Jupyter Notebook を参照してください。 + [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 トレーニングプラットフォーム内の PyTorch、Hugging Face、TensorFlow の AWS 深層学習コンテナを通じて利用できます。分散トレーニングライブラリの機能を利用するには、SageMaker Python SDK の使用をお勧めします。また、SDK for Python (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 ごとのメモリフットプリントを削減できます。これには 2 つのメリットがあります。1 つは、標準のデータ並列処理ではメモリが不足する大規模なモデルに対応できること、もう 1 つは、解放された GPU メモリを使用してバッチサイズを増やせることです。 標準のデータ並列処理手法は、データ並列グループ内の GPU 全体でトレーニング状態を複製し、`AllReduce` オペレーションに基づいて勾配集計を実行します。シャーディングデータ並列処理は、オプティマイザ状態のシャーディングの性質を考慮して、標準のデータ並列分散トレーニング手順を変更します。モデル状態とオプティマイザ状態がシャーディングされるランクのグループは、*シャーディンググループ*と呼ばれます。シャーディングデータ並列処理手法では、モデルのトレーニング可能なパラメータと、それに対応する勾配とオプティマイザ状態を*シャーディンググループ*内の GPU 全体でシャーディングします。 SageMaker AI は、 AWS ブログ記事[「巨大モデルトレーニングのほぼ線形スケーリング AWS](https://www.amazon.science/blog/near-linear-scaling-of-gigantic-model-training-on-aws)」で説明されている MiCS の実装を通じてシャーディングデータ並列処理を実現します。この実装では、シャーディング度を設定可能なパラメータとして設定できます。このパラメータは、データ並列処理度よりも小さくなければなりません。MiCS は、フォワードパスとバックワードパスのたびに、`AllGather` オペレーションを通じてすべての GPU のモデルパラメータを一時的に再結合します。各レイヤーのフォワードパスまたはバックワードパスの後、MiCS はパラメータを再度シャーディングして GPU メモリを節約します。バックワードパス中、MiCS は勾配を縮小し、同時に `ReduceScatter` オペレーションを通じて GPU 間で勾配をシャーディングします。最後に、MiCS は、オプティマイザ状態のローカルシャードを使用して、ローカルの縮小勾配とシャーディング勾配を対応するローカルパラメータシャードに適用します。通信オーバーヘッドを減らすために、SageMaker モデル並列処理ライブラリは、フォワードパスまたはバックワードパスで次のレイヤーをプリフェッチし、ネットワーク通信と計算を重複させます。 モデルのトレーニング状態はシャーディンググループ全体で複製されます。つまり、勾配をパラメータに適用する前に、シャーディンググループ内で行われる `ReduceScatter` オペレーションに加えて、シャーディンググループ全体で `AllReduce` オペレーションを実行する必要があります。 事実上、シャーディングデータ並列処理では、通信オーバーヘッドと 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 Collectives によるシャーディングデータ並列処理](#model-parallel-extended-features-pytorch-sharded-data-parallelism-smddp-collectives) + [シャーディングデータ並列処理による混合精度トレーニング](#model-parallel-extended-features-pytorch-sharded-data-parallelism-16bits-training) + [テンソル並列処理と NCCL Collectives によるシャーディングデータ並列処理](#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 推定器を設定するときは、シャーティングデータ並列処理用のパラメータを追加します。 シャードディングデータ並列処理を有効にするには、SageMaker PyTorch 推定器に `sharded_data_parallel_degree` パラメータを追加します。このパラメータは、トレーニング状態をシャーディングする GPU の数を指定します。`sharded_data_parallel_degree` の値は、1 とデータ並列処理度の間の整数で、データ並列処理度を均等に割る必要があります。ライブラリは GPU の数を自動的に検出するため、データ並列度も検出されることに注意してください。シャーディングデータ並列処理には、次の追加パラメータが利用できます。 + `"sdp_reduce_bucket_size"` *(int、デフォルト: 5e8)* — [PyTorch DDP 勾配バケット](https://pytorch.org/docs/stable/notes/ddp.html#internal-design)のサイズをデフォルトの dtype の要素の数で指定します。 + `"sdp_param_persistence_threshold"` *(int、デフォルト: 1e6)* — 各 GPU で保持できる要素の数でパラメータテンソルのサイズを指定します。シャーディングデータ並列処理は、各パラメータテンソルをデータ並列グループの GPU 間で分割します。パラメータテンソルの要素の数がこのしきい値よりも小さい場合、パラメータテンソルは分割されません。パラメータテンソルは、データ並列 GPU 間で複製されるため、通信オーバーヘッドを削減できます。 + `"sdp_max_live_parameters"` *(int、デフォルト: 1e9)* — フォワードパスとバックワードパス中に再結合されたトレーニング状態に同時に配置できるパラメータの最大数を指定します。アクティブなパラメータの数が所定のしきい値に達すると、`AllGather` オペレーションによるパラメータの取得は一時停止します。このパラメータをの数を増加すると、メモリ使用量が増えることに注意してください。 + `"sdp_hierarchical_allgather"` *(bool、デフォルト: True)* — `True` に設定すると、`AllGather` オペレーションは階層的に実行されます。最初に各ノード内で実行され、次にノード間で実行されます。複数ノードの分散型トレーニングジョブでは、階層型 `AllGather` オペレーションが自動的に有効になります。 + `"sdp_gradient_clipping"` *(float、デフォルト: 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 Collectives によるシャーディングデータ並列処理 | モデル/パラメータの数 | Num インスタンス | インスタンスタイプ | シーケンスの長さ | グローバルバッチサイズ | ミニバッチサイズ | シャーディングデータ並列度 | | --- | --- | --- | --- | --- | --- | --- | | 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 Collectives によるシャーディングデータ並列処理 | モデル/パラメータの数 | Num インスタンス | インスタンスタイプ | シーケンスの長さ | グローバルバッチサイズ | ミニバッチサイズ | シャーディングデータ並列度 | テンソル並列度 | アクティベーションオフロード | | --- | --- | --- | --- | --- | --- | --- | --- | --- | | GPT-NEOX-65B | 64 | ml.p4d.24xlarge | 2048 | 512 | 8 | 16 | 8 | はい | | GPT-NEOX-65B | 64 | ml.p4d.24xlarge | 4096 | 512 | 2 | 64 | 2 | はい | シャーディングデータ並列処理とテンソル並列処理を組み合わせて使用すると、大規模言語モデル (LLM) を大規模なクラスターに適合させ、一方でシーケンス長の長いテキストデータを使用し (この場合、バッチサイズが小さくなります)、結果として長いテキストシーケンスに対して LLM をトレーニングするために GPU メモリ使用量を処理する場合に役立ちます。詳細については[テンソル並列処理と NCCL Collectives によるシャーディングデータ並列処理](#model-parallel-extended-features-pytorch-sharded-data-parallelism-with-tensor-parallelism)を参照してください。 導入事例、ベンチマーク、その他の設定例については、ブログ記事「[New performance improvements in Amazon SageMaker AI model parallel library](https://aws.amazon.com/blogs/machine-learning/new-performance-improvements-in-amazon-sagemaker-model-parallel-library/)」を参照してください。 ## SMDDP Collectives によるシャーディングデータ並列処理 SageMaker データ並列処理ライブラリは、 AWS インフラストラクチャに最適化された集合通信プリミティブ (SMDDP 集合) を提供します。[Elastic Fabric Adapter (EFA)](https://aws.amazon.com/hpc/efa/) を使用してオールツーオールタイプの通信パターンを採用することで最適化を実現しています。これにより、スループットが高くレイテンシーの影響を受けにくい集合体が得られ、通信関連の処理が CPU にオフロードされ、GPU サイクルがコンピューティングに解放されます。大規模なクラスターでは、SMDDP Collectives は分散トレーニングのパフォーマンスを NCCL と比較して最大 40% 向上させることができます。導入事例およびベンチマークの結果については、ブログ「[New performance improvements in the Amazon SageMaker AI model parallelism library](https://aws.amazon.com/blogs/machine-learning/new-performance-improvements-in-amazon-sagemaker-model-parallel-library/)」を参照してください。 **注記** SMDDP Collectives によるシャーディングデータ並列処理は、SageMaker モデル並列処理ライブラリ v1.13.0 以降、および SageMaker データ並列処理ライブラリ v1.6.0 以降で使用できます。SMDDP Collectives でシャーディングデータ並列処理を使用するには、「[Supported configurations](#sharded-data-parallelism-smddp-collectives-supported-config)」も参照してください。 大規模な分散トレーニングで一般的に使用される手法であるシャーディングデータ並列処理では、`AllGather` 集合体を使用して、GPU 計算と並列して、フォワードパス計算とバックワードパス計算のシャーディングレイヤーパラメータを再構成します。大規模なモデルでは、GPU のボトルネックの問題やトレーニング速度の低下を回避するために、`AllGather` オペレーションを効率的に実行することが重要です。シャーディングデータ並列処理を有効にすると、SMDDP Collectives はパフォーマンスが重要な `AllGather` 集合体にまとめられ、トレーニングのスループットが向上します。 **SMDDP Collectives によるトレーニング** トレーニングジョブでシャーディングデータ並列処理が有効になっていて、[Supported configurations](#sharded-data-parallelism-smddp-collectives-supported-config) を満たすと、SMDDP Collectives が自動的に有効になります。内部的には、SMDDP Collectives は AWS 、インフラストラクチャでパフォーマンスを発揮するように`AllGather`集合体を最適化し、他のすべての集合体の NCCL にフォールバックします。さらに、サポートされていない設定では、`AllGather` を含むすべての集合体が自動的に NCCL バックエンドを使用します。 SageMaker モデル並列処理ライブラリバージョン 1.13.0 以降、`"ddp_dist_backend"` パラメータは `modelparallel` オプションに追加されています。この設定パラメータのデフォルト値は `"auto"` で、可能な限り SMDDP Collectives を使用し、それ以外の場合は NCCL にフォールバックします。ライブラリが常に NCCL を使用するように強制するには、`"ddp_dist_backend"` 設定パラメータに `"nccl"` を指定します。 次のコード例は、シャーディングデータ並列処理と `"ddp_dist_backend"` パラメータを使用して PyTorch 推定器を設定する方法を示しています。このパラメータはデフォルトで `"auto"` に設定されているため、オプションで追加できます。 ``` 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 Collectives の `AllGather` オペレーションは、以下の設定要件がすべて満たされるとトレーニングジョブで有効になります。 + シャーディングデータ並列処理度が 1 を超えている。 + `Instance_count` が 1 より大きい + `Instance_type` が `ml.p4d.24xlarge` に等しい + PyTorch v1.12.1 以降の SageMaker トレーニングコンテナ + SageMaker データ並列処理ライブラリ v1.6.0 またはそれ以降 + SageMaker モデルの並列処理ライブラリ v1.13.0 またはそれ以降 **パフォーマンスとメモリのチューニング** SMDDP Collectives は追加の GPU メモリを使用します。モデルトレーニングのさまざまなユースケースに応じて、GPU メモリ使用量を設定する環境変数が 2 つあります。 + `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 Collectives によるパフォーマンスの向上を維持しながら、その GPU メモリ使用量を削減するためのチューニングのヒントをいくつか紹介します。 + `SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` のチューニング + `AllGather` 入力バッファのサイズは、モデルが小さいほど小さくなります。したがって、パラメータの数が少ないモデルでは、`SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` に必要なサイズを小さくすることができます。 + モデルは多くの GPU にシャーディングされるため、`AllGather` 入力バッファサイズは `sharded_data_parallel_degree` が大きくなるにつれて小さくなります。そのため、`sharded_data_parallel_degree` の値が大きいトレーニングジョブでは、`SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` に必要なサイズを小さくすることができます。 + `SMDDP_AG_SORT_BUFFER_SIZE_BYTES` のチューニング + ノード間通信から収集されるデータ量は、パラメータの数が少ないモデルほど少なくなります。したがって、パラメータの数が少ないこうしたモデルでは、`SMDDP_AG_SORT_BUFFER_SIZE_BYTES` に必要なサイズを小さくすることができます。 一部の集合体は NCCL を使用するようフォールバックする場合があるため、最適化された SMDDP Collectives ではパフォーマンスが向上しない場合があります。追加の 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 メモリをさらに節約するには、分散トレーニング設定にパラメータを 1 つ追加して 16 ビット浮動小数点形式 (FP16) または [Brain 浮動小数点形式](https://en.wikichip.org/wiki/brain_floating-point_format) (BF16) を有効にします。 **注記** シャーディングデータ並列処理による混合精度トレーニングは、SageMaker モデル並列処理ライブラリ v1.11.0 以降で使用できます。 **シャーディングデータ並列処理による FP16 トレーニング** シャーディングデータ並列処理で FP16 トレーニングを実行するには、`smp_options` 設定ディクショナリに `"fp16": True"` を追加します。トレーニングスクリプトでは、`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 トレーニングを実行するには、`smp_options` 設定ディクショナリに `"bf16": True` を追加します。 ``` smp_options = { "enabled": True, "parameters": { "ddp": True, "sharded_data_parallel_degree": 2, "bf16": True } } ``` ## テンソル並列処理と NCCL Collectives によるシャーディングデータ並列処理 シャーディングデータ並列処理を使用し、グローバルバッチサイズも小さくする必要がある場合は、[テンソル並列処理](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-tensor-parallelism.html)とシャーディングデータ並列処理の使用を検討してください。非常に大規模な計算クラスター (通常は 128 ノード以上) でシャーディングデータ並列処理を使用する大規模なモデルをトレーニングする場合、GPU あたりのバッチサイズが小さくても、グローバルバッチサイズが非常に大きくなります。これにより、収束の問題や計算パフォーマンスの低下の問題が発生する可能性があります。1 つのバッチがすでに大きく、それ以上削減できない場合、シャーディングデータ並列処理だけでは GPU あたりのバッチサイズを小さくできないことがあります。このような場合、シャーディングデータ並列処理とテンソル並列処理を組み合わせて使用すると、グローバルバッチサイズを小さくすることができます。 最適なシャーディングデータ並列度とテンソル並列度の選択は、モデルの規模、インスタンスタイプ、およびモデルが収束するのに適したグローバルバッチサイズによって異なります。CUDA のメモリ不足エラーを解決して最高のパフォーマンスを実現するには、低いテンソル並列度から始めてグローバルバッチサイズを計算クラスターに合わせることをお勧めします。次の 2 つのケース例を参照し、テンソル並列処理とシャーディングデータ並列処理の組み合わせにより、GPU をモデル並列処理用にグループ化することでグローバルバッチサイズを調整し、モデルレプリカの数を減らしてグローバルバッチサイズを小さくする方法を確認してください。 **注記** この機能は SageMaker モデル並列処理ライブラリ v1.15 から提供されており、PyTorch v1.13.1 をサポートしています。 **注記** この機能は、ライブラリのテンソル並列処理機能によりサポート対象のモデルで使用できます。サポートされているモデルのリストについては、「[Hugging Face Transformer モデルのサポート](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-hugging-face.html)」を参照してください。また、トレーニングスクリプトを変更する際は、`smp.model_creation` 引数に `tensor_parallelism=True` を渡す必要があることにも注意してください。詳細については、「*SageMaker AI Examples 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) ``` テンソル並列処理を追加すると、グローバルバッチサイズを小さくすることができます。設定例の 1 つは、テンソル並列度を 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 に設定するとします。この割り当てにより、テンソル並列度に基づいて、`(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)` という形式で 8 個の GPU が作成されます。つまり、4 つの GPU が 1 つのテンソル並列グループを形成します。この場合、テンソル並列グループの 0 ランク GPU の縮小データ並列グループは `(0,4,8,12,16,20,24,28)` となります。縮小データ並列グループは、シャーディングデータ並列度 4 に基づいてシャーディングされ、データ並列処理用の 2 つのレプリケーショングループとなります。GPU `(0,4,8,12)` は 1 つのシャーディンググループを形成し、0 番目のテンソル並列ランクの全パラメータの完全なコピーをそれにまとめて保持します。一方、GPU `(16,20,24,28)` は別のそうしたグループを形成します。他のテンソル並列ランクにも同様のシャーディンググループとレプリケーショングループがあります。 ![\[図 1: テンソル並列処理グループ。\]](http://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/images/distributed/model-parallel/sdp_tp_group_tp.jpg) 図 1: (ノード、シャーディングデータ並列度、テンソル並列度) = (4、4、4) のテンソル並列処理グループ。各四角形は、0 ～ 31 のインデックスを持つ GPU を表しています。これらの GPU が、TPG7 ～ TPG0 のテンソル並列処理グループを形成しています。レプリケーショングループは (\$1TPG0、TPG4\$1、\$1TPG1、TPG5\$1、\$1TPG2、TPG6\$1、\$1TPG3、TPG7\$1) です。各レプリケーショングループペアは同じ色ですが、塗りつぶし方が異なります。 ![\[図 2: シャーディングデータ並列処理グループ。\]](http://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/images/distributed/model-parallel/sdp_tp_group_sdp.jpg) 図 2: (ノード、シャーディングデータ並列度、テンソル並列度) = (4、4、4) のシャーディングデータ並列処理グループ。各四角形は、0 ～ 31 のインデックスを持つ GPU を表しています。これらの GPU が、SDPG0 ～ SDPG7 のシャーディングデータ並列処理グループを形成しています。レプリケーショングループは (\$1SDPG0、SDPG4\$1、\$1SDPG1、SDPG5\$1、\$1SDPG2、SDPG6\$1、\$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 Examples 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 Collectives が使用されていることに注意してください。 詳細については、「[テンソル並列処理と NCCL Collectives によるシャーディングデータ並列処理](#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 モデルのチェックポイント機能 (v1.10.0 以降の SageMaker モデル並列処理ライブラリ用)](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) 設定パラメータの挙動が変わります。これら 2 つの機能を同時にオンにすると、パラメータの初期化を遅らせるのではなく、モデル作成時にパラメータがすぐにシャーディング方式で初期化されるため、ランクごとに独自のパラメータのシャードが初期化されて保存されます。 + シャーディングデータ並列処理を有効にすると、`optimizer.step()` の呼び出しの実行時にライブラリが内部で勾配クリッピングを実行します。勾配クリッピングには、[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) などのユーティリティ API を使用する必要はありません。勾配クリッピングのしきい値を調整するには、「[シャーディングデータ並列処理をトレーニングジョブに適用する方法](#model-parallel-extended-features-pytorch-sharded-data-parallelism-how-to-use)」セクションに示すように、SageMaker PyTorch 推定器を作成するときに、分散パラメータ設定の `sdp_gradient_clipping` パラメータを使用してしきい値を設定できます。 # モデルのパイプライン化 SageMaker のモデル並列ライブラリの主要機能の 1 つは、モデルトレーニング中にデバイス間で行われる計算やデータ処理の順序を決定する*パイプライン並列処理*です。パイプライン化は、複数の GPU にそれぞれ異なるデータサンプルで同時に計算させることで、モデル並列処理における真の並列化を実現し、逐次計算によるパフォーマンスの低下を克服する手法です。パイプライン並列処理を使うと、GPU 使用率を最大化するように、トレーニングジョブはすべてのマイクロバッチにわたってパイプライン方式で実行されます。 **注記** パイプライン並列処理は、モデルパーティショニングとも呼ばれ、PyTorch と TensorFlow の両方で使用できます。サポートされているフレームワークのバージョンについては、「[サポートされているフレームワークと AWS リージョン](distributed-model-parallel-support.md)」を参照してください。 ## パイプライン実行スケジュール パイプライン化は、ミニバッチのマイクロバッチへの分割に基づいており、マイクロバッチは、トレーニングパイプラインに 1 つずつ投入され、ライブラリランタイムで定義された実行スケジュールに従います。マイクロバッチは、トレーニングミニバッチのより小さなサブセットです。**パイプラインのスケジュールは、タイムスロットごとにどのマイクロバッチがどのデバイスで実行されるかを決定します。 例えば、パイプラインのスケジュールとモデルパーティションに応じて、GPU `i` はマイクロバッチ `b` で (フォワードまたはバックワードの) 計算を実行し、一方、GPU `i+1` はマイクロバッチ `b+1` で計算を実行するため、両方の GPU が同時にアクティブのままになります。1 回のフォワードパスまたはバックワードパスの間に、パーティショニングの決定に応じて、1 つのマイクロバッチの実行フローが同じデバイスに複数回アクセスする場合があります。例えば、モデルの先頭にあるオペレーションは、モデルの末尾にあるオペレーションと同じデバイス上に配置され、一方、その間のオペレーションは異なるデバイス上に配置されるため、このデバイスは 2 回アクセスされることになります。 ライブラリには、シンプルとインターリーブという 2 種類のパイプラインスケジュールがあり、SageMaker Python SDK の `pipeline` パラメータを使って設定できます。****ほとんどの場合、インターリーブパイプラインの方が GPU をより効率的に利用することで、より良いパフォーマンスを実現できます。 ### インターリーブパイプライン インターリーブパイプラインでは、可能な限りマイクロバッチのバックワード実行が優先されます。これにより、アクティベーションに使用したメモリを迅速に解放し、メモリをより効率的に使用できるようになります。また、マイクロバッチ数をより多くスケールでき、GPU のアイドル時間を短縮できます。定常状態では、各デバイスはフォワードパスとバックワードパスを交互に実行します。これは、あるマイクロバッチのバックワードパスが、別のマイクロバッチのフォワードパスが終了する前に実行される可能性があることを意味します。 ![\[2 つの GPU にわたるインターリーブパイプラインの実行スケジュールの例。\]](http://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/images/distributed/model-parallel/interleaved-pipeline-execution.png) 上の図は、2 つの GPU にわたるインターリーブパイプラインの実行スケジュールの例を示しています。図では、F0 はマイクロバッチ 0 のフォワードパス、B1 はマイクロバッチ 1 のバックワードパスを表しています。**Update** は、オプティマイザによるパラメータの更新を表しています。GPU0 は、可能な限り常にバックワードパスを優先 (例えば、F2 の前に B0 を実行)します。これにより、アクティベーションに使用したメモリをより早くクリアできます。 ### 単純なパイプライン 対照的に、シンプルパイプラインは、各マイクロバッチのフォワードパスの実行を終了してからバックワードパスを開始します。つまり、フォワードパスとバックワードパスのステージをそれぞれのステージ内でパイプライン化します。次の図は、2 つの GPU にわたる動作の例を示しています。 ![\[各マイクロバッチのフォワードパスの実行後にバックワードパスを開始するパイプラインの例。\]](http://docs.aws.amazon.com/ja_jp/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/ja_jp/sagemaker/latest/dg/images/distributed/model-parallel/interleaved-pipeline-tf.png) 次のイメージは、B=2 で分割された 2 つのサブグラフに勾配演算を追加した例です。`SMPInput` オペレーションの勾配は `SMPOutput` オペレーションであり、その逆も同様です。これにより、バックプロパゲーションの間に勾配を逆方向に流すことができます。 ![\[B=2 で分割された 2 つのサブグラフに勾配演算を追加した例。\]](http://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/images/distributed/model-parallel/interleaved-pipeline-tf.gif) この GIF は、B=2 マイクロバッチと 2 つのサブグラフを含むインターリーブパイプライン実行スケジュールの例を示しています。各デバイスは、サブグラフのレプリカを 1 つずつ順番に実行し、GPU 使用率を向上させます。B が大きくなると、アイドル時間のスロットの割合はゼロになります。特定のサブグラフレプリカで (フォワードまたはバックワードの) 計算を実行するときはいつでも、パイプラインレイヤーは対応する青の `SMPInput` オペレーションに信号を送って実行を開始します。 1 つのミニバッチ内のすべてのマイクロバッチからの勾配が計算されると、ライブラリはマイクロバッチ全体の勾配を結合し、それをパラメータに適用できます。 #### PyTorch におけるパイプライン実行 概念的には、パイプラインは PyTorch でも同様の考えに従います。ただし、PyTorch は静的グラフを含まないため、モデル並列処理ライブラリの PyTorch 機能はより動的なパイプラインパラダイムを使います。 TensorFlow と同様に、各バッチは複数のマイクロバッチに分割され、各デバイスで一度に 1 つずつ実行されます。ただし、実行スケジュールは、各デバイスで起動された実行サーバーを介して処理されます。別のデバイスに配置されたサブモジュールの出力が現在のデバイスで必要になるたびに、実行リクエストがサブモジュールへの入力テンソルと共にリモートデバイスの実行サーバーに送信されます。次に、サーバーは指定された入力でこのモジュールを実行し、現在のデバイスに応答を返します。 リモートサブモジュールの実行中、現在のデバイスはアイドルであるため、現在のマイクロバッチのローカル実行は一時停止し、ライブラリランタイムは実行を現在のデバイスがアクティブに動作できる別のマイクロバッチに切り替えます。マイクロバッチの優先順位は、選択されたパイプラインスケジュールによって決定されます。インターリーブパイプラインスケジュールの場合は、計算のバックワードステージに存在するマイクロバッチが可能な限り優先されます。 # テンソル並列処理 *テンソル並列処理*は、特定のモデルの重み、勾配、オプティマイザ状態がデバイス間で分割されるモデル並列処理の一種です。パイプライン並列処理とは対照的に、個々の重みはそのまま維持するが、*セットする*重みのうち、テンソル並列処理は個々の重みを分割します。これには、通常、モデルの特定のオペレーション、モジュール、またはレイヤーの分散計算が含まれます。 テンソル並列処理は、単一のパラメータが GPU メモリの大部分を消費する場合 (ボキャブラリーサイズが大きい大きな埋め込みテーブルや、クラス数が多い大きなソフトマックス層など) に必要です。この場合、この大きなテンソルまたはオペレーションを原子単位として扱うことは非効率的であり、メモリ負荷のバランスを妨げます。 テンソル並列処理は、純粋なパイプラインだけでは十分ではない非常に大きなモデルにも役立ちます。例えば、数十インスタンスのパーティション化を必要とする 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 Transformer モデルのサポート](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 でテンソル並列処理がサポートされていると仮定します。このモデルのテンソル並列パーティションの結果として考えられる 1 つは、次のようになります。 ``` 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 を表します。次の点に注意してください。 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` モジュールにおけるこのプロセスの例を示しています。 ![\[2 つのテンソル並列概念を示す 2 つの図。\]](http://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/images/distributed/model-parallel/tensor-parallel-concept.png) 最初の図は、大きな `nn.Linear` モジュールのある小さなモデルで、2 つのテンソル並列処理ランクに対してデータ並列処理を行います。`nn.Linear` モジュールは 2 つの並列ランクにレプリケートされます。 2 番目の図は、より大きなモデルに適用されたテンソル並列処理を示しており、`nn.Linear` モジュールが分割されています。各 `tp_rank` はリニアモジュールの半分と、残りのオペレーション全体を保持します。リニアモジュールの実行中、各 `tp_rank` は、すべてのデータサンプルの関連する半分を収集し、それを `nn.Linear` モジュールの半分に渡します。各ランクが独自のデータサンプルの最終的な線形出力を持つように、結果をリダクション散乱 (総和を縮小オペレーションとして) する必要があります。残りのモデルは、一般的なデータ並列方式で実行されます。 # テンソル並列処理を使用して SageMaker 分散モデル並列トレーニングジョブを実行する このセクションでは、以下を行います。 + テンソル並列処理を使用するように SageMaker PyTorch 推定器と SageMaker モデルの並列処理オプションを構成する方法。 + 拡張 `smdistributed.modelparallel` モジュールを使用してトレーニングスクリプトをテンソル並列処理に適合させる方法。 `smdistributed.modelparallel` モジュールに関する詳細ついては、*SageMaker Python SDK ドキュメント*の「[SageMaker 分散モデル並列 API](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html)」を参照してください。 **Topics** + [テンソル並列処理のみ](#model-parallel-extended-features-pytorch-tensor-parallelism-alone) + [テンソル並列処理とパイプライン並列処理の組み合わせ](#model-parallel-extended-features-pytorch-tensor-and-pipeline-parallelism) ## テンソル並列処理のみ 以下は、パイプライン並列処理を使用せずに、テンソル並列処理を単独でアクティブにする分散トレーニングオプションの例です。`mpi_options` および `smp_options` ディクショナリを設定して、SageMaker `PyTorch` 推定器に分散トレーニングオプションを指定します。 **注記** 拡張メモリ保存機能は、SageMaker モデル並列処理ライブラリ v1.6.0 以降を実装する PyTorch 用の Deep Learning Containers を通じて利用できます。 **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` パラメータを設定してテンソル並列処理によるモデル並列オプションを指定します。 **注記** 拡張メモリ保存機能は、SageMaker モデル並列処理ライブラリ v1.6.0 以降を実装する PyTorch 用の Deep Learning Containers を通じて利用できます。 **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 Transformer モデルのサポート SageMaker モデル並列処理ライブラリのテンソル並列処理は、次の Hugging Face Transformer モデルに対するすぐに利用可能なサポートを提供します。 + GPT-2、BERT、および RoBERTa (SageMaker モデル並列処理ライブラリ v1.7.0 以降で利用可能) + GPT-J (SageMaker モデル並列処理ライブラリ v1.8.0 以降で利用可能) + GPT-Neo (SageMaker モデル並列処理ライブラリ v1.10.0 以降で利用可能) **注記** その他の Transformer モデルでは、[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 Transformer モデルのトレーニングにテンソル並列処理を使用するには、必ず SageMaker モデル並列処理ライブラリ v1.7.0 以降を含む PyTorch 用 Hugging Face Deep Learning Containers を使用してください。詳細については、[SageMaker モデル並列処理ライブラリのリリースノート](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_release_notes/smd_model_parallel_change_log.html)を参照してください。 ## すぐに利用可能なサポート対象モデル ライブラリですぐに利用可能なサポート対象の Hugging Face Transformer モデルの場合、Transformer 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 Transformer と `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)の延長です。テンソル並列処理では、ライブラリは 3 種類のランキングとプロセスグループ 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#)」および「[PyTorch 固有 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/ja_jp/sagemaker/latest/dg/images/distributed/model-parallel/tensor-parallel-ranking-mechanism.png) 例えば、8 GPU を持つ単一ノードのプロセスグループについて考慮してみましょう。テンソル並列度 2、パイプライン並列度は 2、データ並列度の度合いは 4 とします。前の図の中央上部には、4 レイヤーモデルの例が示されています。図の左下と右下部分には、4 つの GPU に分散された 4 レイヤーモデルが描かれており、パイプライン並列処理とテンソル並列処理の両方が使用され、テンソル並列処理が中間の 2 レイヤーに使用されています。これらの 2 つの図は、異なるグループ境界線を示す単純なコピーです。パーティション化されたモデルは、GPU 0-3 および 4-7 のデータ並列処理のためにレプリケートされます。左下の図は、`MP_GROUP`、`PP_GROUP`、および `TP_GROUP` の定義を示しています。右下の図は、同じ GPU のセット上にある `RDP_GROUP`、`DP_GROUP`、および `WORLD` をしめしています。同じ色を持つレイヤーとレイヤースライスの勾配は、データ並列処理に対して一緒に `allreduce` されます。例えば、最初のレイヤー (水色) では、`DP_GROUP` 全体にわたって `allreduce` オペレーションが実行されます。一方、2 つ目のレイヤーの濃いオレンジのスライスでは、そのプロセスの `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 では個別に実行されます。テンソル並列処理は、GPU ペア (0,2)、(1,3)、(4,6)、(5,7) に全体で `DP_GROUP` のサブセットに対して行われます。 このようなハイブリッドパイプラインとテンソル並列処理では、`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 で使用できます。 ## オプティマイザ状態シャーディングの使用方法 *オプティマイザ状態のシャーディング*は、`"shard_optimizer_state": True` を `modelparallel` 設定で設定することで有効にできます。 この機能をオンにすると、ライブラリはデータ並列度に基づいてモデルパラメータのセットをパーティション化します。第 `i` パーティションに対応する勾配は、第 `i` データ並列ランクでのみ縮小します。`smp.step` デコレータ関数の最初の呼び出しの最後に、`smp.DistributedOptimizer` でラップされたオプティマイザは、現在のデータ並列ランクのパーティションに対応するパラメータのみに制限されるようにパラメータを再定義します。再定義されたパラメータは*仮想パラメータ*と呼ばれ、基礎となるストレージを元のパラメータと共有します。`optimizer.step` を初めて呼び出すとき、これらの再定義されたパラメータに基づいてオプティマイザ状態が作成されます。これらのパラメータは、元のパーティションのためにシャードされます。オプティマイザの更新後、(`optimizer.step` 呼び出しの一部として) AllGather オペレーションがデータ並列ランク全体で実行され、一貫したパラメータ状態を実現します。 **ヒント** オプティマイザ状態のシャーディングは、データ並列度が 1 より大きく、モデルに 10 億を超えるパラメータがある場合に役立ちます。 データ並列処理度は `(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 トレーニングスクリプトを適合させる** 「*Tensor 並列処理とパイプライン並列処理の組み合わせ*」セクションの「[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` モジュールでは、パイプライン並列処理の観点から、モジュールツリーが 1 つのパーティション内にある場合にのみモジュールツリーをチェックポイントできます。`torch.nn.Sequential` モジュールの場合、アクティベーションチェックポイントを機能させるためには、シーケンシャルモジュール内の各モジュールツリーが完全に 1 つのパーティション内にある必要があります。手動パーティショニングを使用するときは、次の制限に注意してください。 [自動モデルパーティショニング](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 Transformer などのライブラリからビルド済みのモデルをインポートするときに、アクティベーションチェックポイントを使用する方法を示しています。シーケンシャルモジュールをチェックポイントするかどうかにかかわらず、次の操作を行います。 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') ``` # アクティベーションオフロード アクティベーションチェックポイントとパイプライン並列処理がオンで、マイクロバッチの数が 1 より大きい場合、*アクティベーションオフロード*は、メモリ使用量をさらに削減できる追加機能です。アクティベーションオフロードは、CPU で現在実行されていないマイクロバッチに対応するチェックポイントされたアクティベーションを非同期的に移動します。GPU がマイクロバッチのバックワードパスのアクティベーションを必要とする直前に、この機能はオフロードされたアクティベーションを CPU からプリフェッチして戻します。 **注記** この機能は、SageMaker モデル並列処理ライブラリ v1.6.0 以降の PyTorch で利用できます。 ## アクティベーションオフロードの使用方法 **マイクロバッチの数が 1 を超え、アクティベーションチェックポイントがオンになっている場合は**、アクティベーションオフロードを使用してメモリ使用量を減らします ([アクティベーションチェックポイント](model-parallel-extended-features-pytorch-activation-checkpointing.md) を参照)。アクティベーションチェックポイントを使用しない場合、アクティベーションオフロードは効果がありません。マイクロバッチを 1 つだけ使用する場合、メモリの節約にはなりません。 アクティベーションオフロードを使用するには、`modelparallel` 設定で `"offload_activations": True` を設定してください。 アクティベーションオフロードは、`nn.Sequential` モジュール内のチェックポイントしたアクティベーションを非同期的に CPU に移動します。PCIe リンクを介したデータ転送は GPU 計算と重複します。オフロードは、特定のチェックポイントしたレイヤーのフォワードパスが計算されるとすぐに行われます。アクティベーションは、特定のマイクロバッチのバックワードパスに必要になる直前に GPU に読み込まれます。CPU - GPU 間の転送も同様に計算と重複します。 アクティベーションが GPU に読み込まれるまでの時間を調整するには、設定パラメータ `"activation_loading_horizon"` を使用できます (デフォルトは 4 に設定されていますが、0 より `int` 大きい必要があります)。アクティベーションのロード範囲が大きいと、アクティベーションは GPU に早期にロードバックされます。範囲が大きすぎると、アクティベーションオフロードによるメモリ節約効果が低下する可能性があります。範囲が小さすぎると、アクティベーションが時間内にロードされず、重複の量が減少し、パフォーマンスが低下することがあります。 **ヒント** アクティベーションオフロードは、パラメータが 1,000 億を超える大規模なモデルに役立ちます。 **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. Wrap your model using the [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 トレーニングを使用する場合は、オプティマイザではなく LR スケジューラーに `optimizer.optimizer` を渡す必要があります。次のコード例を参照してください。 ``` 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) によってラップされた Transformer モデルである*分散トランスフォーマー*モデルにのみ適用されます。この機能は [テンソル並列処理](model-parallel-extended-features-pytorch-tensor-parallelism.md) とも互換性があります。 [FlashAttention](https://github.com/HazyResearch/flash-attention) ライブラリは、`attention_head_size` が 8 の倍数で 128 未満の値に設定されているモデルのみをサポートしています。したがって、分散トランスフォーマーをトレーニングして FlashAttention が正しく動作することを確認するときは、アテンションヘッドサイズが要件を満たすようにパラメータを調整する必要があります。詳細については、「*FlashAttention GitHub リポジトリ*」の「[インストールと機能](https://github.com/HazyResearch/flash-attention#installation-and-features)」も参照してください。 例えば、`hidden_width=864` と `num_heads=48` を使用して Transformer モデルを設定すると仮定します。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 トレーニングジョブを実行するには、次の 3 つのユースケースシナリオがあります。 1. TensorFlow および PyTorch 用に構築済みの AWS 深層学習コンテナのいずれかを使用できます。このオプションは、モデル並列ライブラリを初めて使用する場合に推奨されます。SageMaker モデルの並列トレーニングジョブの実行方法に関するチュートリアルについては、[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)」と共に使用することをお勧めします。 これらのセクションで提供されるトレーニングスクリプトの例は簡略化されており、ライブラリを使うために必要な変更を強調するように考えられています。SageMaker モデル並列処理ライブラリと共に TensorFlow または PyTorch トレーニングスクリプトを使用する方法を示す、エンドツーエンドで実行可能なノートブックサンプルについては、「[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 モデル並列処理ライブラリを使用してトレーニングスクリプトのモデルを分割する トレーニングスクリプトを変更してモデル分割を設定するには、自動分割と手動分割の 2 つの方法があります。 ### 自動モデル分割 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 による自動モデル分割 モデル並列処理ライブラリは、トレーニング可能な変数のサイズとグラフ構造を分析し、内部でグラフパーティショニングアルゴリズムを使用します。このアルゴリズムは、デバイス間で必要な通信量を最小化することを目的に、次の 2 つの制約条件の下での、各オペレーションに対するデバイスの割り当てを導き出します。 + 各デバイスに保存されている変数の数のバランスを取る + 各デバイスで実行されるオペレーションの数のバランスを取る `optimize` に `speed` を指定した場合 (Python SDK のモデル並列処理パラメータで)、ライブラリは各デバイスにおけるオペレーションと `tf.Variable` オブジェクトの数のバランスを取ろうとします。それ以外の場合は、`tf.Variables` の合計サイズのバランスを取ろうとします。 パーティショニングの決定が行われると、ライブラリは各デバイスの実行に必要なサブグラフのシリアル化された表現を作成し、各デバイスにインポートします。パーティショニング中、ライブラリは同じ `tf.Variable` を使用するオペレーションと、同じ Keras レイヤーの一部であるオペレーションを同じデバイス上に配置します。また、TensorFlow によって課されるコロケーション制約にも対応します。これは、例えば、`tf.Variable` を共有する 2 つの Keras レイヤーがある場合、これらのレイヤーに含まれるすべてのオペレーションは、1 つのデバイスに配置されることを意味します。 `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 を分割して、各ノードが 1 つのデバイスに向かうようにします。非常に重要な点は、`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)」も確認してください。 次のトピックでは、TensorFlow モデルの自動パーティショニングおよび手動パーティショニング用の SageMaker のモデル並列処理ライブラリを設定するために使用できるトレーニングスクリプトの例を示します。 **注記** 自動パーティショニングはデフォルトで有効になっています。特に指定がない限り、サンプルスクリプトでは自動パーティショニングを使用します。 **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. Keras Model クラスの代わりに [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 モデルを定義します。`smp.DistributedModel` オブジェクトの call メソッドからのモデル出力を返します。call メソッドから返されたテンソルはどれもモデル並列デバイス間でブロードキャストされるため、通信オーバーヘッドが発生します。したがって、call メソッド範囲外の不要なテンソル (中間アクティベーションなど) は返さないように注意してください。 1. `tf.Dataset.batch()` メソッドで `drop_remainder=True` を設定します。これは、バッチサイズが常にマイクロバッチ数で割り切れるようにするためです。 1. データパイプラインのランダムオペレーションを `smp.dp_rank()` を使ってシードします (例: `shuffle(ds, seed=smp.dp_rank())`)。これにより、異なるモデルパーティションを保持する GPU 間でデータサンプルの一貫性が確保されます。 1. フォワードおよびバックワードロジックをステップ関数に入れ、`smp.step` で修飾します。 1. `reduce_mean` などの [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) メソッドを使って、マイクロバッチ全体の出力に対して後処理を実行します。[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()` を使ってシードする。 **注記** Horovod を使用するときは、トレーニングスクリプトで `hvd.init` を直接呼び出すことはできません。代わりに、[ステップ 2: SageMaker Python SDK を使用してトレーニングジョブを起動する](model-parallel-sm-sdk.md) の SageMaker Python SDK の `modelparallel` パラメータで、`"horovod"` を `True` に設定する必要があります。これにより、ライブラリはモデルパーティションのデバイス割り当てに基づいて Horovod を内部的に初期化できます。`hvd.init()` をトレーニングスクリプト内で直接呼び出すと、問題が発生する可能性があります。 **注記** トレーニングスクリプト内で `hvd.DistributedOptimizer` API を直接使用すると、API が黙示的に `AllReduce` オペレーションを `smp.step` 内に配置するため、トレーニングのパフォーマンスと速度が低下する可能性があります。Horovod でモデル並列処理ライブラリを使用する場合は、次の例で示すように、`smp.step` から返された勾配に対して `accumulate()` または `reduce_mean()` を呼び出した後に `hvd.allreduce` を直接呼び出すことをお勧めします。 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 は現在サポートされていません。チェックポイントには、代わりに同じ API と機能を備えた `smp.CheckpointManager` を使用してください。`smp.CheckpointManager` によるチェックポイントの復元は、最初のステップの後で実行することに注意してください。 # PyTorch トレーニングスクリプトを変更する このセクションでは、PyTorch トレーニングスクリプトを変更して、自動パーティショニングおよび手動パーティショニングに使用する SageMaker モデル並列処理ライブラリを設定する方法について学習します。 **注記** ライブラリでサポートされている PyTorch のバージョンを確認するには、「[サポートされているフレームワークと AWS リージョン](distributed-model-parallel-support.md)」を参照してください。 **ヒント** SageMaker モデル並列処理ライブラリで PyTorch トレーニングスクリプトを使用する方法を示すエンドツーエンドのノートブックサンプルについては、「[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` メソッドから返されたテンソルはモデル並列デバイス間でブロードキャストされるため、通信オーバーヘッドが発生します。したがって、call メソッド範囲外の不要なテンソル (中間アクティベーションなど) は返さないように注意してください。 **注記** 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. フォワードおよびバックワードロジックをステップ関数に入れ、[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. `reduce_mean` などの [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) メソッドを使って、マイクロバッチ全体の出力に対して後処理を実行します。 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` の forward メソッドに渡すすべての `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 に移動しようとすると、move 呼び出しは無視されます。ライブラリは、ランクに割り当てられたモデルの一部を 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` 関数を次のように定義します。この関数は、`.to()` API を使って現在のデバイスに `data` と `target` を追加してから、これらの入力テンソルを使って `train_step` を呼び出します。 ``` 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 ``` ## サポートされていないフレームワーク機能 次の PyTorch 機能は、SageMaker のモデル並列処理ライブラリではサポートされていません。 + ネイティブの [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) ラッパーモジュールは、ライブラリでサポートされていません。ライブラリは、パラメータブロードキャストや勾配 AllReduce など、PyTorch DDP との統合を内部的に管理しています。ライブラリを使う場合、モジュールバッファはトレーニングの開始時に 1 回だけブロードキャストされます。モデルにモジュールバッファがあり、各ステップでデータ並列グループ間で同期させる必要がある場合は、`smp.get_dp_process_group()` で取得できるプロセスグループを使って、`torch.distributed` API を使ってそれを行えます。 + 混合精度トレーニングでは、`apex.amp` モジュールはサポートされていません。自動混合精度でライブラリを使う推奨方法は、`torch.cuda.amp` を使用することです。ただし、Torch の実装の代わりに `smp.amp.GradScaler` を使用することを除きます。 + `torch.jit.ScriptModules` または `ScriptFunctions` は、`smp.DistributedModel` ではサポートされていません。 + `apex` : `apex` からの `FusedLayerNorm`、`FusedAdam`、`FusedLAMB`、`FusedNovoGrad` はサポートされていません。代わりに、`smp.optimizers` および `smp.nn` API を介してこれらのライブラリ実装を使用できます。 # ステップ 2: SageMaker Python SDK を使用してトレーニングジョブを起動する SageMaker Python SDK は、TensorFlow および PyTorch などの ML フレームワークによるモデルのマネージドトレーニングをサポートしています。これらのフレームワークのいずれかを使用してトレーニングジョブを開始するには、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 推定器のテンプレートは、SageMaker モデル並列ライブラリを MPI で使用するための `distribution` パラメータを構成する方法を示しています。 ------ #### [ 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/') ``` ------ ライブラリを有効にするには、`"smdistributed"` キーと `"mpi"` キーの設定ディクショナリを SageMaker Python SDK の推定器コンストラクタの `distribution` 引数に渡す必要があります。 **SageMaker モデル並列処理の設定パラメータ** + `"smdistributed"` キーについては、`"modelparallel"` キーと次の内部ディクショナリを含むディクショナリを渡します。 **注記** 1 つのトレーニングジョブでの `"modelparallel"` と `"dataparallel"` の使用はサポートされていません。 + `"enabled"` – 必須。モデル並列処理を有効にするには、 `"enabled": True` を設定します。 + `"parameters"` – 必須。SageMaker モデル並列処理の一連のパラメータを指定します。 + 一般的なパラメータの完全なリストについては、「SageMaker Python SDK ドキュメント」の「[Parameters for `smdistributed`](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#smdistributed-parameters)」を参照してください。 TensorFlow については、「[TensorFlow-specific Parameters](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#tensorflow-specific-parameters)」を参照してください。 PyTorch については、「[PyTorch-specific Parameters](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#pytorch-specific-parameters)」を参照してください。 + `"pipeline_parallel_degree"` (または `smdistributed-modelparallel 構築済みのコンテナを拡張し、SageMaker のモデル並列処理ライブラリを使用するには、PyTorch または TensorFlow に使用可能な AWS 深層学習コンテナ (DLC) イメージのいずれかを使用する必要があります。SageMaker モデル並列処理ライブラリは、CUDA (`cuxyz`) を使う TensorFlow (2.3.0 以降) と PyTorch (1.6.0 以降) の DLC イメージに含まれています。DLC イメージの完全なリストについては、「AWS Deep Learning Containers GitHub リポジトリ」の「[Available Deep Learning Containers Images](https://github.com/aws/deep-learning-containers/blob/master/available_images.md)」を参照してください。 **ヒント** 最新バージョンの SageMaker モデル並列処理ライブラリにアクセスするには、最新バージョンの TensorFlow または PyTorch を含むイメージを使うことをお勧めします。 例えば、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 モデル並列ライブラリを使用するには、SageMaker 分散並列ライブラリの正しい依存関係とバイナリファイルを Dockerfile に含める必要があります。このセクションでは、SageMaker トレーニング環境とモデル並列ライブラリを独自の Docker コンテナに適切に準備するために含める必要のある最小限のコードブロックセットについて説明します。 **注記** SageMaker モデル並列ライブラリをバイナリとして使用するこのカスタム Docker オプションは PyTorch でのみ使用できます。 **SageMaker トレーニングツールキットとモデル並列ライブラリを使用して Dockerfile を作成する方法** 1. [NVIDIA CUDA ベースイメージ](https://hub.docker.com/r/nvidia/cuda)の 1 つから始めます。 ``` FROM ``` **ヒント** 公式の AWS Deep Learning Container (DLC) イメージは、[NVIDIA CUDA ベースイメージ](https://hub.docker.com/r/nvidia/cuda)から構築されています。Deep [Learning Container for PyTorch の公式 Dockerfiles AWS](https://github.com/aws/deep-learning-containers/tree/master/pytorch/training/docker)を調べて、インストールする必要があるライブラリのバージョンと設定方法を確認することをお勧めします。公式の Dockerfile は、SageMaker と Deep Learning Containers のサービスチームによって完成され、ベンチマークテストされ、管理されています。表示されたリンクで、使用する PyTorch バージョンを選択し、CUDA (`cuxyz`) フォルダーを選択し、`.gpu` または `.sagemaker.gpu` で終わる Dockerfile を選択します。 1. 分散型トレーニング環境を設定するには、[Elastic Fabric Adapter (EFA)](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa.html)、[NVIDIA Collective Communications Library (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-toolkit](https://github.com/aws/sagemaker-training-toolkit) をインストールします。ツールキットには、SageMaker トレーニングプラットフォームおよび SageMaker Python SDK と互換性のあるコンテナを作成するために必要な共通の機能が含まれています。 ``` RUN pip install sagemaker-training ``` 1. Dockerfile の作成が完了したら、「[Adapting Your Own Training Container](https://docs.aws.amazon.com/sagemaker/latest/dg/adapt-training-container.html)」を参照して、Docker コンテナを構築し、Amazon ECR でホストする方法について確認してください。 **ヒント** SageMaker AI でのトレーニング用のカスタム Dockerfile の作成に関する一般的な情報については、「[Use Your Own Training Algorithms](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 モデルのチェックポイント機能 (v1.10.0 以降の SageMaker モデル並列処理ライブラリ用)](#model-parallel-extended-features-pytorch-checkpoint) + [分散 PyTorch モデルのチェックポイント機能 (v1.6.0 から v1.9.0 の SageMaker モデル並列処理ライブラリ用)](#model-parallel-extended-features-pytorch-saving-loading-checkpoints) + [分散 TensorFlow モデルのチェックポイント](#distributed-model-parallel-checkpoint-tensorflow) ### 分散 PyTorch モデルのチェックポイント機能 (v1.10.0 以降の SageMaker モデル並列処理ライブラリ用) SageMaker モデル並列処理ライブラリは、分散モデルの状態およびそのオプティマイザの状態のチェックポイントを完全にまたは部分的に保存および読み込むためのチェックポイント API を提供します。 **注記** PyTorch と v1.10.0 以降の SageMaker モデル並列処理ライブラリを使用している場合は、このチェックポイント手法をお勧めします。 **部分チェックポイント** モデル並列処理を使用してトレーニングされたモデルのチェックポイントを保存するには、部分チェックポイントオプションを true (`partial=True`) に設定して、[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 を使用します。これにより、各モデルのパーティションが個別に保存されます。モデルとオプティマイザの状態に加えて、`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) ``` 部分チェックポイントからトレーニングを再開するには、`partial=True` として [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 を使用し、チェックポイントディレクトリと部分チェックポイントの保存時に使用したタグを指定します。モデルの重みが実際に読み込まれるのは、モデルパーティショニングの後、`smdistributed.modelparallel.torch.step` - 修飾トレーニングステップ関数の初回実行時に行われることにご注意ください。 部分チェックポイントを保存するとき、ライブラリは `.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 ) ``` **完全チェックポイント** 推論の目的で最終モデルアーティファクトを保存するには、`partial=False` として `smdistributed.modelparallel.torch.save_checkpoint` API を使用します。これにより、モデルパーティションが組み合わされて単一のモデルアーティファクトが作成されます。これはオプティマイザの状態は組み合わせないことに注意してください。 完全なモデルチェックポイントを指定して、特定の重みでトレーニングを初期化するには、`partial=False` として `smdistributed.modelparallel.torch.resume_from_checkpoint` API を使用できます。これはオプティマイザの状態は読み込まないことに注意してください。 **注記** テンソル並列処理では、一般的に元のモデル実装と `DistributedModel` 実装との間で `state_dict` を変換する必要があります。オプションで、`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 モデルのチェックポイント機能 (v1.6.0 から v1.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)、および v1.6.0 から v1.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) を使用して、モデルとオプティマイザの状態を保存します。チェックポイントの保存方法に応じて、次の 2 つのオプションのいずれかを選択します: + **オプション 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.dp_rank() == 0` ではなく `if smp.rdp_rank() == 0` のように設定していることを確認してください。オプティマイザの状態をテンソル並列処理でシャーディングする場合、すべてのリダクションデータの並列ランクは、オプティマイザの状態の独自のパーティションを保存する必要があります。チェックポイントに間違った if ステートメントを使用すると、トレーニングジョブが停止する場合があります。テンソル並列処理なしで `if smp.dp_rank() == 0` を使用する方法の詳細については、「*SageMaker Python SDK ドキュメント*」の「[General Instruction for Saving and Loading](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) を使用してチェックポイントを読み込みます。前のステップでチェックポイントした方法に応じて、次の 2 つのオプションのいずれかを選択します: + **オプション 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` 間での冗長な読み込みを回避できます。現在、完全チェックポイントのオプティマイザの状態 dict は、テンソル並列処理ではサポートされていません。 ### 分散 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) ## 分散モデルの微調整 微調整はトレーニングスクリプトで設定する必要があります。次のコードスニペットは、`smdistributed.model.parallel.torch` モジュールと微調整用の設定を登録するよう修正した Hugging Face Transformers の [AutoModelForCausalLM](https://huggingface.co/docs/transformers/main/en/model_doc/auto#transformers.AutoModelForCausalLM) クラスを使用するトレーニングスクリプトの構造体の例を示しています。 **注記** [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 Notebook の完全な例については、「*SageMaker AI Examples GitHub リポジトリ*」の「[GPT-2 examples for PyTorch](https://github.com/aws/amazon-sagemaker-examples/tree/main/training/distributed_training/pytorch/model_parallel/gpt2)」を参照してください。 # Amazon SageMaker AI モデル並列処理ライブラリ v1 の例 このページには、SageMaker モデル並列処理 (SMP) ライブラリ v1 を実装して SageMaker AI で分散トレーニングジョブを実行する方法について、実践的な例を紹介したブログ記事や Jupyter Notebook の一覧を掲載しています。 ## ブログと導入事例 以下のブログでは、SMP v1 の導入事例について説明しています。 + [New performance improvements in the Amazon SageMaker AI model parallelism library](https://aws.amazon.com/blogs/machine-learning/new-performance-improvements-in-amazon-sagemaker-model-parallel-library/)、*AWS Machine Learning ブログ* (2022 年 12 月 16 日) + [Train gigantic models with near-linear scaling using sharded data parallelism on 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 Machine Learning ブログ* (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 Code Editor](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 サンプルノートブック** + [Train GPT-2 with near-linear scaling using the sharded data parallelism technique in the SageMaker model parallelism library](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt2/smp-train-gpt-sharded-data-parallel.ipynb) + [Fine-tune GPT-2 with near-linear scaling using sharded data parallelism technique in the SageMaker model parallelism library](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt2/smp-fine-tune-gpt-sharded-data-parallel.ipynb) + [Train GPT-NeoX-20B with near-linear scaling using the sharded data parallelism technique in the SageMaker model parallelism library](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) + [Train GPT-J 6B using the sharded data parallelism and tensor parallelism techniques in the SageMaker model parallelism library](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) + [Train FLAN-T5 with near-linear scaling using sharded data parallelism technique in the SageMaker model parallelism library](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/flan-t5/smp-train-t5-sharded-data-parallel.ipynb) + [Train Falcon with near-linear scaling using sharded data parallelism technique in the SageMaker model parallelism library](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 と SageMaker モデル並列処理ライブラリを使用した CNN](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 メモリを持つ強力なインスタンスを使用することをお勧めします。大規模な DL モデルをトレーニングするには、`ml.p4d` または `ml.p3dn` インスタンスを使用することをお勧めします。これらのインスタンスには Elastic Fabric Adapter (EFA) も装備されており、ネットワーク帯域幅が向上し、モデルの並列処理による大規模なトレーニングが可能になります。 **シャーディングオプティマイザの状態** + シャーディングオプティマイザの状態の影響は、データの並列ランク数によって異なります。通常、データの並列性 (コンピュートノードのサイズに比例する) が高いほど、メモリ使用量の効率が向上します。 クラスターをダウンサイズする場合は、オプティマイザ状態のシャーディングの構成を確認してください。例えば、オプティマイザー状態のシャーディングを使用した大規模な DL モデルが、16 GPU (2 つの P4d または P4de インスタンスなど) のコンピューティングクラスターに適合する場合でも、8 GPU (単一の P4d または P4de インスタンスなど) を備えたノードに適合するとは限りません。これは、8 GPU の合計メモリが 16 GPU の合計メモリよりも小さく、8 GPU でのシャーディングに必要な GPU ごとのメモリも、16 GPU のシナリオの場合より多く必要なためです。その結果、必要なメモリが増加し、より小さなクラスターには収まらない可能性があります。 詳細については、「[オプティマイザ状態シャーディング](model-parallel-extended-features-pytorch-optimizer-state-sharding.md)」を参照してください。 **アクティベーションチェックポイント** + モジュールのグループに対してアクティベーションチェックポイントを使用することにより、メモリ効率を向上させることができます。モジュールをグループ化すればするほど、メモリ使用効率が上がります。レイヤーのシーケンシャルモジュールをチェックポイントする場合、`smp.set_activation_checkpointing` 関数の `strategy` の引数は、チェックポイントのためにレイヤーをグループ化します。例えば、チェックポイントのために複数のレイヤーをグループ化すると、一度に 1 つのレイヤーにチェックポイントするよりもメモリ効率が高くなります。これにより、余分な計算時間と引き換えにメモリ使用量を削減できます。 詳細については、「[アクティベーションチェックポイント](model-parallel-extended-features-pytorch-activation-checkpointing.md)」を参照してください。 **テンソル並列性** + テンソル並列性の度合いは、2 の累乗 (2、4、8、...、2n) である必要があり、最大度はノードあたりの GPU 数に等しくなければなりません。例えば、8 つの GPU を持つノードを使用する場合、テンソル並列性の度数は 2、4、8 です。テンソル並列性の度合いには、任意の数値 (3、5、6、7 など) を推奨しません。複数のノードを使用する場合、テンソルの並列性の度合いを誤って設定すると、ノード間でテンソル並列処理が実行されることがあります。これにより、ノード間のアクティベーションの通信によりオーバーヘッドが大きく増え、計算コストが高くなる可能性があります。 詳細については、「[テンソル並列処理](model-parallel-extended-features-pytorch-tensor-parallelism.md)」を参照してください。 **ノード間でのパイプラインの並列処理** + パイプラインの並列処理は、1 つのノード内と複数のノードの両方で実行できます。パイプライン並列処理をテンソル並列処理と組み合わせて使用する場合は、複数のノードにわたってパイプライン並列処理を実行し、個々のノード内でテンソル並列処理を維持することをお勧めします。 + パイプラインの並列処理には、`microbatches`、`active_microbatches`、および`prescaled_batch` の 3 つのノブがあります。 + テンソル並列処理をパイプラインの並列処理で使用する場合は、モデル並列グループごとのバッチサイズを増やし効率的なパイプライン処理を行えるよう、`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 ドキュメントの「[Parameters for `smdistributed`](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#parameters-for-smdistributed)」を参照してください。 **アクティベーションを CPU へオフロードする** + アクティベーションのチェックポイントやパイプラインの並列処理と組み合わせて使用します。オフロードとプリロードがバックグラウンドで確実に行われるよう、microbatches パラメータに 1 より大きい値を指定します。 + アクティベーションをオフロードする場合、`active_microbatches` を増やすことができ、マイクロバッチの総数と一致することもあります。これは、チェックポイントされているモジュールと、モデルのパーティション化方法によって異なります。 詳細については、「[アクティベーションオフロード](model-parallel-extended-features-pytorch-activation-offloading.md)」を参照してください。 ### リファレンス設定 SageMaker モデル並列処理トレーニングチームは、GPT-2 モデル、シーケンス長 512、語彙サイズ 50,000 での実験に基き、以下の基準点を提供します。 | モデルパラメータの数 | インスタンスタイプ | パイプラインの並列性 | テンソル並列性 | オプティマイザの状態シャーディング | アクティベーションチェックポイント | 事前スケールされたバッチ | バッチサイズ | | --- | --- | --- | --- | --- | --- | --- | --- | | 10 億 | 16 ml.p4d.24xlarge | 1 | 4 | 正 | 各トランスフォーマーレイヤー | 正 | batch\$1size=40 | | 30 億 | 16 ml.p4d.24xlarge | 1 | 8 | 正 | 各トランスフォーマーレイヤー | 正 | batch\$1size=32 | | 60 億 | 32 ml.p4d.24xlarge | 2 | 8 | 正 | 各トランスフォーマーレイヤー | 正 | batch\$1size=56, microbatches=4, active\$1microbatches=2 | 前の構成から外挿して、モデル設定の GPU メモリ使用量を推定することができます。例えば、10 億パラメータモデルのシーケンス長を増やしたり、モデルのサイズを 20 億に増やす場合は、最初にバッチサイズを小さくすることができます。それでもモデルが適合しない場合は、テンソル並列性の度合いを上ることができます。 ## トレーニングスクリプトの変更 + 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. [**トレーニングジョブ**] を選択します。 1. メインペインで、詳細を表示するトレーニングジョブ名を選択します。 1. メインペインの **[モニタリング]** セクションを参照し、自動ビジュアライゼーションを確認します。 1. トレーニングジョブログを表示するには、**[モニタリング]** セクションの **[ログの表示]** を選択します。CloudWatch でトレーニングジョブの分散トレーニングジョブログにアクセスできます。マルチノード分散トレーニングを開始した場合は、**[algo-n-1234567890]** 形式のタグを含む複数のログストリームが表示されます。**algo-1** ログストリームは、メイン (0 番目) ノードからのトレーニングログを追跡します。 詳細については、「[トレーニングジョブのモニタリングと分析を行うための Amazon CloudWatch メトリクス](training-metrics.md)」を参照してください。 ## アクセス許可 モデル並列処理または「[SageMaker distributed training example notebooks](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)」を参照してください。 ## バッチサイズとマイクロバッチ数 + このライブラリは、バッチサイズを増やすと非常に効率的になります。モデルが 1 つのデバイス内に収まるが、小さなバッチサイズでしかトレーニングできないユースケースの場合、ライブラリの統合後にバッチサイズを増やすことができ、また、そのようにしてください。モデル並列処理は大規模モデルに対するメモリを節約するので、以前はメモリに収まらなかったバッチサイズを使ってトレーニングを実行できるようになります。 + 過少または過大なマイクロバッチ数を選択すると、パフォーマンスが低下する可能性があります。ライブラリは各デバイスで各マイクロバッチを順次実行するため、マイクロバッチサイズ (バッチサイズをマイクロバッチ数で割った値) は、各 GPU を十分に利用できる大きさである必要があります。同時に、パイプラインの効率はマイクロバッチの数とともに向上するため、適切なバランスを取ることが重要です。通常は、まず、2 つまたは 4 つのマイクロバッチを試し、バッチサイズをメモリ制限まで増やしてから、より大きなバッチサイズとマイクロバッチ数を試してみることをお勧めします。マイクロバッチ数を増やすと、インターリーブパイプラインを使う場合に、より大きなバッチサイズが可能になる可能性があります。 + バッチサイズは、常にマイクロバッチ数で割り切れなければなりません。データセットのサイズによって、すべてのエポックの最後のバッチが残りのバッチよりも小さなサイズになることがあり、このより小さなバッチもマイクロバッチ数で割り切れる必要があることに注意してください。そうでない場合は、`tf.Dataset.batch()` の呼び出しで `drop_remainder=True` を設定するか (TensorFlow の場合)、`DataLoader` で `drop_last=True` を設定して (PyTorch の場合)、この最後の小さなバッチが使われないようにします。データパイプラインに異なる API を使っている場合、最後のバッチがマイクロバッチ数で割り切れないときは常に、手動でそのバッチをスキップする必要があります。 ## 手動パーティショニング + 手動パーティショニングを行う場合、トランスフォーマーアーキテクチャの埋め込みテーブルなど、モデル内の複数のオペレーションやモジュールによって使用されるパラメータに注意してください。同じパラメータを共有するモジュールは、正確性確保のために同じデバイスに配置する必要があります。自動パーティショニングが使われる場合、ライブラリはこの制約を自動的に適用します。 ## データ準備 + モデルが複数の入力を取り込む場合は、`smp.dp_rank()` を使ってデータパイプラインにランダムなオペレーション (シャッフルなど) をシードするようにしてください。データセットがデータ並列デバイス間で決定的にシャードされている場合、シャードが `smp.dp_rank()` でインデックス付けされるようにしてください。これは、モデルパーティションを形成するすべてのランクで見られるデータの順序が必ず一致するようにするためです。 ## `smp.DistributedModel` から返されるテンソル + `smp.DistributedModel.call` (TensorFlow の場合) または `smp.DistributedModel.forward` (PyTorch の場合) の関数から返されるテンソルはすべて、その特定のテンソルを計算したランクから、他のすべてのランクにブロードキャストされます。その結果、call メソッドと forward メソッドの範囲外では必要とされないテンソル (中間アクティベーションなど) はどれも、不要な通信やメモリのオーバーヘッドを発生させ、パフォーマンスを低下させるため、返さないようにしてください。 ## `@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)」を参照してください。 ## パラメータの初期化の遅延 パラメータが 1,000 億個を超える非常に大規模なモデルでは、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()` に置き換える必要がある場合があります。以下の等式は常に当てはまります: + `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 モデル並列ライブラリでの 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 モデル並列ライブラリでの SageMaker デバッガーの使用に関する考慮事項 SageMaker デバッガーは 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 を使った収束 TensorFlow とモデル並列ライブラリで SageMaker AI マルチノードトレーニングを使うと、トレーニング入力ファイルの順序がノードごとに異なる可能性があるため、損失が期待どおりに収束しない可能性があります。これにより、同じモデル並列グループ内の異なるランクが異なる入力ファイルを処理することになり、不整合を生む可能性があります。これを防ぐには、TensorFlow データセットに変換する前に、すべてのランクで入力ファイルの順序が同じになるようにします。これを実現する 1 つの方法は、トレーニングスクリプトで入力ファイル名を並べ替えることです。 ## 分散トレーニングジョブの停止またはクラッシュ トレーニングジョブに停止、クラッシュ、または応答しないなどの問題がある場合は、以下のトラブルシューティング項目を読んで、問題の原因を特定してください。さらにサポートが必要な場合は、[AWS サポート](https://aws.amazon.com/premiumsupport) を通して、SageMaker の分散トレーニングチームに連絡してください。 + **NCCL の初期化ステップで分散トレーニングジョブが停止する**場合は、次の点を考慮してください。 + カスタム VPC とそのサブネットで EFA 対応インスタンス (`ml.p4d` または `ml.p3dn` インスタンス) のいずれかを使用している場合は、使用するセキュリティグループに、同じ 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` が有効になっている場合、部分オプティマイザをチェックポイントする時にも同様の停止が発生する可能性があります。 モデル並列処理によるモデルのチェックポイント設定の詳細については、「[General Instruction for Saving and Loading](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 モデルのチェックポイント機能 (v1.6.0 から v1.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 の 1 つに障害が発生していることを意味します。テクニカルサポートが必要な場合は、ジョブ ARN を AWS チームと共有し、可能であればチェックポイントからトレーニングジョブを再開してください。 + まれに、以前に機能していたジョブ設定が 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.forward()` を呼び出す代わりに `super()._orig_forward()` を呼び出す必要があります。