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

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

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

重要

SMDDP 程式庫已停止對 TensorFlow 的支援，且不再於 v2.11.0 之後的 TensorFlow DLC 中提供。若要尋找已安裝 SMDDP 程式庫的舊版 TensorFlow DLC，請參閱支援的架構。

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

程式庫 API 的設計旨在近似 Horovod API。請參閱 SageMaker AI 分散式資料平行化 TensorFlow API 文件，了解程式庫為 TensorFlow 提供之每個 API 的其他詳細資訊。

注意

SageMaker AI 分散式資料平行化可配合包含 tf 核心模組的 TensorFlow 訓練指令碼 (tf.keras 模組除外) 調整。SageMaker AI 分散式資料平行化不支援包含 Keras 實作的 TensorFlow。

注意

SageMaker AI 分散式資料平行化程式庫支援開箱即用的自動混合精確度 (AMP)。除了對訓練指令碼進行架構層級修改之外，無需執行任何額外操作即可啟用 AMP。如果梯度是 FP16，SageMaker AI 資料平行化程式庫會以 FP16 執行其 AllReduce 操作。有關將 AMP API 實施到訓練腳本的詳細資訊，請參閱以下資源：

  1. 匯入程式庫的 TensorFlow 用戶端並將其初始化。

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

  2. 使用 local_rank 將每個 GPU 固定到單一 smdistributed.dataparallel 程序-這是指特定節點內程序的相對等級。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_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 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)

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