TensorFlow トレーニングスクリプトでSMDDPライブラリを使用する (廃止） - Amazon SageMaker

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

TensorFlow トレーニングスクリプトでSMDDPライブラリを使用する (廃止）

重要

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

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

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

注記

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

注記

SageMaker 分散データ並列処理ライブラリは、自動混合精度 (AMP) をすぐにサポートします。トレーニングスクリプトのフレームワークレベルの変更AMP以外に、 を有効にするための追加のアクションは必要ありません。グラデーションが にある場合FP16、 SageMaker データ並列処理ライブラリは でそのAllReduceオペレーションを実行しますFP16。トレーニングスクリプトAMPAPIsへの の実装の詳細については、次のリソースを参照してください。

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

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

  2. 各 GPU を で単一の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_variablesAPI後に を使用します。これは、次のコードブロックで として呼び出されます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)

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