TensorFlow トレーニングスクリプトで SMDDP ライブラリを使用する (非推奨)

重要

SMDDP ライブラリは v2.11.0 以降の DLC TensorFlow のサポートを終了し、DLC では使用できなくなりました。 TensorFlow SMDDP TensorFlow ライブラリがインストールされた以前の DLC を検索するには、を参照してください。サポートされるフレームワーク

次の手順は、 SageMakerの分散データparallel TensorFlow ライブラリを利用するようにトレーニングスクリプトを変更する方法を示しています。 

ライブラリの API は Horovod の API と同様になるように設計されています。ライブラリが提供する各 API の詳細については TensorFlow、SageMaker 分散データparallel TensorFlow API のドキュメントを参照してください。

注記

SageMaker 分散データparallel は、 TensorFlow tftf.kerasモジュール以外のコアモジュールで構成されるトレーニングスクリプトに適応できます。 SageMaker 分散データparallel は Keras TensorFlow 実装ではサポートされていません。

注記

SageMaker 分散データ並列処理ライブラリは、最初から自動混合精度 (AMP) をサポートしています。AMP を有効にするには、トレーニングスクリプトにフレームワークレベルの変更を加える以外、特別なアクションは必要ありません。勾配が FP16 の場合、 SageMaker データ並列処理ライブラリは FP16 で動作を実行します。AllReduceトレーニングスクリプトへの AMP API の実装に関する詳細は、次のリソースを参照してください。

• フレームワーク- TensorFlow NVIDIA ディープラーニングパフォーマンスドキュメント内

• 深層学習の自動混合精度 (NVIDIA デベロッパードキュメント)

• TensorFlow ドキュメンテーション内の混合精度 API TensorFlow

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

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

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

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

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

   learning_rate = learning_rate * sdp.size()

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

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

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

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

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

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

以下は、 TensorFlow ライブラリを使った分散トレーニング用のトレーニングスクリプトの例です。

import tensorflow as tf

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

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

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

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

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

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

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

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

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

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

    return loss_value

...

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

トレーニングスクリプトの調整が完了したら、「ステップ 2: SageMaker Python SDK を使用して分散型トレーニングジョブを起動する」に進みます。