このページでは、AI Platform Training 上での分散トレーニング構造をカスタム コンテナでサポートする方法を説明します。
カスタム コンテナでは、分散処理をサポートしている ML フレームワークを使用して分散トレーニングを行うことができます。ここで使用する用語は TensorFlow の分散モデルに基づいていますが、同様の分散構造を持つ ML フレームワークであれば、どの ML フレームワークでも使用できます。たとえば、MXNet での分散トレーニングではスケジューラ、ワーカー、サーバーを使用します。この構造は AI Platform Training カスタム コンテナの分散トレーニング構造と一致します。AI Platform Training ではマスター、ワーカー、パラメータ サーバーを使用します。
トレーニング クラスタの構造
AI Platform Training で分散トレーニング ジョブを実行する場合、トレーニング クラスタで複数のマシン（ノード）を指定します。トレーニング サービスにより、指定したマシンタイプに応じてリソースが割り当てられます。特定のノード上で実行中のジョブは「レプリカ」と呼ばれます。分散 TensorFlow モデルに従って、トレーニング クラスタ内の各レプリカには分散トレーニングにおける役割またはタスクが 1 つ割り当てられます。
マスター ワーカー: 1 つのレプリカのみがマスター ワーカー（チーフワーカーとも呼ばれます）として指定されます。このタスクは他のレプリカを管理し、ジョブ全体のステータスを報告します。
ワーカー: 1 つ以上のレプリカを「ワーカー」として指定できます。これらのレプリカに担当させる作業は、ジョブ構成で指定します。
パラメータ サーバー: 1 つ以上のレプリカを「パラメータ サーバー」として指定できます。これらのレプリカは、モデル パラメータを格納し、ワーカー間で共有されるモデルの状態を調整します。
評価者: 1 つ以上のレプリカを「評価者」として指定できます。これらのレプリカを使用して、モデルを評価できます。なお、TensorFlow を使用している場合、TensorFlow では通常、複数の評価者を使用することを想定します。
API マッピング
トレーニング クラスタ内のマシンに割り当てることができるロールは、4 種類あります。これらのロールは、トレーニング ジョブの入力パラメータを表す TrainingInput で指定できる以下の 4 つのフィールドに対応します。
masterConfig.imageUriは、マスター ワーカー上で実行されるコンテナイ メージの URI を表します。workerConfig.imageUri、parameterServerConfig.imageUri、evaluatorConfig.imageUriはそれぞれ、ワーカー上、パラメータ サーバー上、評価者上で実行されるコンテナ イメージの URI を表します。これらのフィールドに値が設定されていない場合、AI Platform Training はmasterConfig.imageUriの値を使用します。
これらのフィールドの値は、それぞれに対応するフラグを使用して gcloud ai-platform jobs submit training で設定することもできます。
- マスター ワーカー構成には、
--master-image-uriを使用します。 - ワーカー構成には、
--worker-image-uriを使用します。 - パラメータ サーバー構成には、
--parameter-server-image-uriを使用します。 - 現在のところ、評価者のコンテナ イメージ URI を指定するフラグはありません。
evaluatorConfig.imageUriは config.yaml 構成ファイルで指定できます。
カスタム コンテナを使用する分散トレーニング ジョブを送信する方法の例をご覧ください。
CLUSTER_SPEC について
AI Platform Training はすべてのレプリカで、環境変数 CLUSTER_SPEC に値を代入します。これは、クラスタ全体がどのように設定されているかを示す環境変数です。TensorFlow の TF_CONFIG と同様に、CLUSTER_SPEC にはクラスタ内の各レプリカの詳細、たとえばインデックスとロール（マスター、ワーカー、パラメータ サーバー、または評価者）に関する情報が記述されます。
TensorFlow で分散トレーニングを実行すると、TF_CONFIG が解析されて tf.train.ClusterSpec が作成されます。同様に、他の ML フレームワークで分散トレーニングを実行するときは、CLUSTER_SPEC を解析して、フレームワークに必要なすべての環境変数や設定に値を代入する必要があります。
CLUSTER_SPEC の形式
CLUSTER_SPEC 環境変数は、次の形式の JSON 文字列です。
| キー | 説明 | |
|---|---|---|
"cluster"
|
カスタム コンテナのクラスタの説明。「TF_CONFIG」と同様に、このオブジェクトは TensorFlow クラスタ仕様として形式設定され、tf.train.ClusterSpec のコンストラクタに渡すことができます。 | |
"task"
|
コードが実行されている特定のノードのタスクを記述します。この情報を使用して、分散ジョブ内の特定のワーカー用のコードを記述できます。このエントリは、次のキーを持つ辞書です。 | |
"type"
|
このノードで実行されるタスクのタイプ。有効な値は master、worker、ps、evaluatorです。 |
|
"index"
|
ゼロから開始するタスクのインデックス。ほとんどの分散トレーニング ジョブには単一のマスタータスク、1 つ以上のパラメータ サーバー、1 つ以上のワーカーがあります。 | |
"trial"
|
現在実行中のハイパーパラメータ調整トライアルの ID。ジョブのハイパーパラメータ調整を構成する際に、トレーニングするトライアル数を構成します。この値により、コードで実行されているトライアルを識別できます。ID は、トライアル番号を含む文字列値で、1 から始まります。 | |
"job"
|
ジョブを開始したときに使用したジョブ パラメータ。ほとんどの場合、このエントリはコマンドライン引数を通じてアプリケーションに渡されるデータの複製であるため、無視できます。 | |
TF_CONFIG との比較
なお、AI Platform Training では、すべてのトレーニング ジョブの各レプリカに TF_CONFIG 環境変数も設定されます。AI Platform Training は、カスタム コンテナ トレーニング ジョブのレプリカに CLUSTER_SPEC 環境変数のみを設定します。2 つの環境変数は一部の値を共有していますが、形式が異なります。
カスタム コンテナをトレーニングする場合、マスター レプリカはデフォルトで master タスク名の付いた TF_CONFIG 環境変数でラベル付けされます。トレーニング ジョブを作成するときに trainingInput.useChiefInTfConfig フィールドを true に構成するか、トレーニング ジョブで 1 つ以上の管理者レプリカを使用することで、chief タスク名でラベル付けされるように構成できます。これは、カスタム コンテナで TensorFlow 2 を使用する場合に特に便利です。
この構成オプション以外に、TensorFlow を使用した分散トレーニングは、カスタム コンテナを使用する場合や、AI Platform Training ランタイムを使用する場合と同じように動作します。AI Platform Training 上の分散トレーニングに TF_CONFIG を使用する方法の詳細と例をご覧ください。