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

重要

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

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

程式庫 API 的設計旨在近似 Horovod API。有關庫提供的每個 API 的其他詳細信息 TensorFlow，請參閱SageMaker 分佈式數據 parallel TensorFlow API 文檔。

注意

SageMaker 分佈式數據 parallel 適應於除模塊以外tf.keras的tf核心模塊組成的 TensorFlow 培訓腳本。 SageMaker 分散式資料 parallel 不支援 TensorFlow Keras 實作。

注意

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

• 架構- TensorFlow 在 NVIDIA 深度學習效能說明文件中

• NVIDIA 開發人員文件中的 Automatic Mixed Precision for Deep Learning

• TensorFlow TensorFlow文檔中的混合精度 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 分散式沒有直接相關，但必須設定為使用的分散式訓練 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 啟動分散式訓練工作。