TensorFlow 훈련 스크립트에서 SMDDP 라이브러리 사용(사용되지 않음)

다음 단계에서는 SageMaker AI의 분산 데이터 병렬 라이브러리를 활용하도록 TensorFlow 훈련 스크립트를 수정하는 방법을 보여줍니다. 

라이브러리 API는 Horovod API와 유사하도록 설계되었습니다. 라이브러리가 TensorFlow에 제공하는 각 API에 대한 자세한 내용은 SageMaker AI 분산 데이터 병렬 TensorFlow API 설명서를 참조하세요.

다음은 라이브러리를 사용한 분산 훈련을 위한 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로 분산 훈련 작업 시작 섹션으로 넘어갑니다.