在 TensorFlow 訓練指令碼中使用SMDDP程式庫 （已棄用） - Amazon SageMaker AI

本文為英文版的機器翻譯版本，如內容有任何歧義或不一致之處，概以英文版為準。

在 TensorFlow 訓練指令碼中使用SMDDP程式庫 （已棄用）

重要

SMDDP 程式庫已停止支援 TensorFlow ，且不再於 中提供 2.11.0 TensorFlow 版之後DLCs的 。若要尋找已安裝程式SMDDP庫的先前 TensorFlow DLCs ，請參閱 支援的架構。

下列步驟說明如何修改 TensorFlow 訓練指令碼，以利用 SageMaker AI 的分散式資料平行程式庫。 

程式庫APIs的設計類似於 Horovod APIs。如需API程式庫提供的每個項目的其他詳細資訊 TensorFlow，請參閱 SageMaker AI 分散式資料平行 TensorFlow API文件。

注意

SageMaker AI 分散式資料平行處理適用於tf核心模組以外的 TensorFlow 訓練指令碼tf.keras。 SageMaker AI 分散式資料平行處理不支援 TensorFlow Keras 實作。

注意

SageMaker AI 分散式資料平行處理程式庫支援開箱即用的自動混合精確度 (AMP)。除了訓練指令碼的架構層級修改AMP之外，不需要額外的動作即可啟用 。如果漸層位於 中FP16， SageMaker AI 資料平行處理程式庫會在 中執行其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 AI 不直接相關，但必須使用 設定分散式訓練 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 AI 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 AI data parallel: Import the library TF API
import smdistributed.dataparallel.tensorflow as sdp

# SageMaker AI 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 AI 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 AI 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 AI 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 AI 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 AI data parallel: Save checkpoints only from master node.
if sdp.rank() == 0:
    checkpoint.save(checkpoint_dir)

完成訓練指令碼的調整後，請繼續前往 SMDDP 使用 SageMaker Python 使用 啟動分散式訓練任務 SDK。