翻訳は機械翻訳により提供されています。提供された翻訳内容と英語版の間で齟齬、不一致または矛盾がある場合、英語版が優先します。
ステップ 1: SageMaker の分散モデル並列ライブラリを使用して独自のトレーニングスクリプトを変更する
このセクションでは、Amazon SageMaker モデル並列処理ライブラリのコア機能を使用するようにトレーニングスクリプトをカスタマイズする方法を学んでください。ライブラリ固有の API 関数とパラメータを使用するには、このドキュメントを SageMaker Python SDK SageMaker ドキュメントのモデルparallel ライブラリ API
これらのセクションで提供されるトレーニングスクリプトの例は簡略化されており、ライブラリを使うために必要な変更を強調するように考えられています。 TensorFlow PyTorch SageMaker モデル並列処理ライブラリでまたはトレーニングスクリプトを使用する方法を示す実行可能なノートブックの例については end-to-end、を参照してください。Amazon SageMaker モデル並列処理ライブラリ v2 の例
トピック
モデル並列処理ライブラリを使用してトレーニングスクリプトのモデルを分割します。 SageMaker
トレーニングスクリプトを変更してモデル分割を設定するには、自動分割と手動分割の 2 つの方法があります。
自動モデル分割
SageMakerのモデル並列処理ライブラリを使用すると、自動モデル分割 (自動モデル分割とも呼ばれる) を利用できます。このライブラリは、メモリのバランスを取り、デバイス間の通信を最小限に抑え、パフォーマンスを最適化するパーティショニングアルゴリズムを使用します。自動パーティショニングアルゴリズムを設定して、速度やメモリを最適化できます。
また、手動でモデルを分割することもできます。モデルアーキテクチャに精通していて、モデルを効率的に分割する方法について十分に理解している場合を除き、自動モデル分割をお勧めします。
仕組み
自動パーティショニングは、最初のトレーニングステップで、smp.step
で修飾された関数が最初に呼び出されたときに行われます。この呼び出し中、ライブラリは最初に CPU RAM 上にモデルのバージョンを構築し (GPU メモリの制限を回避するため)、次にモデルグラフを分析してパーティショニングを決定します。この決定に基づいて、各モデルパーティションが GPU にロードされた後、最初のステップが実行されます。これらの分析とパーティショニングのステップのため、最初のトレーニングステップには時間がかかる場合があります。
いずれのフレームワークでも、ライブラリはインフラストラクチャに最適化された独自のバックエンドを通じてデバイス間の通信を管理します。 AWS
自動パーティション設計はフレームワークの特性に適応し、ライブラリは各フレームワークでより自然な粒度レベルでパーティショニングを行います。たとえば TensorFlow、では特定の操作をそれぞれ異なるデバイスに割り当てることができますが、では PyTorch、割り当てはモジュールレベルで行われます。この場合、各モジュールは複数の操作で構成されます。次のセクションでは、各フレームワークにおける設計の詳細を確認します。
最初のトレーニングステップでは、モデル並列処理ライブラリは、モデルグラフを構築し、テンソルとパラメータの形状を決定するためのトレースステップを内部的に実行します。このトレースステップの後、ライブラリは、モデル内のネストされた nn.Module
オブジェクトと、保存されている nn.Parameters
の量や各 nn.Module
の実行時間などのトレースから収集された追加データからなるツリーを構築します。
次に、ライブラリは、ルートからこのツリーをトラバースし、各 nn.Module
をデバイスに割り当てるパーティショニングアルゴリズムを実行し、計算負荷 (モジュール実行時間で測定) とメモリ使用量 (保存されている nn.Parameter
サイズとアクティベーションの合計で測定) のバランスを取ります。複数の nn.Modules
が同じ nn.Parameter
を共有する場合、同じパラメータの複数のバージョンを維持しないように、これらのモジュールは同じデバイス上に配置されます。パーティショニングが決定されると、割り当てられたモジュールと重みがデバイスにロードされます。
smp.step
PyTorch デコレータをトレーニングスクリプトに登録する方法については、を参照してください。による自動分割 PyTorch
モデル並列処理ライブラリは、トレーニング可能な変数のサイズとグラフ構造を分析し、内部でグラフパーティショニングアルゴリズムを使用します。このアルゴリズムは、デバイス間で必要な通信量を最小化することを目的に、次の 2 つの制約条件の下での、各オペレーションに対するデバイスの割り当てを導き出します。
-
各デバイスに保存されている変数の数のバランスを取る
-
各デバイスで実行されるオペレーションの数のバランスを取る
optimize
に speed
を指定した場合 (Python SDK のモデル並列処理パラメータで)、ライブラリは各デバイスにおけるオペレーションと tf.Variable
オブジェクトの数のバランスを取ろうとします。それ以外の場合は、tf.Variables
の合計サイズのバランスを取ろうとします。
パーティショニングの決定が行われると、ライブラリは各デバイスの実行に必要なサブグラフのシリアル化された表現を作成し、各デバイスにインポートします。パーティショニング中、ライブラリは同じ tf.Variable
を使用するオペレーションと、同じ Keras レイヤーの一部であるオペレーションを同じデバイス上に配置します。また、によって課せられるコロケーションの制約も考慮されます。 TensorFlowこれは、例えば、tf.Variable
を共有する 2 つの Keras レイヤーがある場合、これらのレイヤーに含まれるすべてのオペレーションは、1 つのデバイスに配置されることを意味します。
smp.step
PyTorch デコレータをトレーニングスクリプトに登録する方法については、を参照してください。による自動分割 TensorFlow
フレームワーク間での自動モデル分割の比較
では TensorFlow、計算の基本単位は a でtf.Operation
、モデルを tf.Operation
s の有向非巡回グラフ (DAG) TensorFlow として表します。そのため、モデル並列処理ライブラリはこの DAG を分割し、各ノードが 1 つのデバイスに割り当てられるようにします。非常に重要な点は、tf.Operation
オブジェクトにはカスタマイズ可能な属性が十分豊富にあり、すべてのモデルがそのようなオブジェクトのグラフで構成されることが保証されているという意味で、それらのオブジェクトが普遍的であるということです。
PyTorch 一方、には、十分に豊富で汎用的な同等の操作概念はありません。これらの特性を持つ最も近い計算単位はnn.Module
、はるかに高い粒度レベルの計算単位です。 PyTorch これが、ライブラリがこのレベルでパーティショニングを行う理由です。 PyTorch
手動モデル分割
デバイス間でモデルのパーティショニング方法を手動で指定する場合は、smp.partition
コンテキストマネージャーを使用します。手動パーティショニング用にコンテキストマネージャーを設定する方法については、以下のページを参照してください。
変更を加えた後にこのオプションを使用するには、ステップ 2 で SageMaker Python SDK auto_partition
False
のフレームワーク推定クラスをに設定して定義する必要があります。default_partition
smp.partition
コンテキストマネージャーを介してパーティションに明示的に配置されていないオペレーションは、default_partition
で実行されます。この場合、自動分割ロジックはバイパスされ、各オペレーションはユーザーの指定に基づいて配置されます。結果のグラフ構造に基づいて、モデル並列処理ライブラリがパイプライン化された実行スケジュールを自動的に作成します。