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

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

중요

SMDDP 라이브러리는 TensorFlow에 대한 지원을 중단했으며 v2.11.0 이후 TensorFlow용 DLCs에서 더 이상 사용할 수 없습니다. SMDDP 라이브러리가 설치된 이전 TensorFlow DLC를 찾으려면 지원되는 프레임워크 섹션을 참조하세요.

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

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

참고

SageMaker 분산 데이터 병렬은 tf.keras 모듈을 제외한 tf 핵심 모듈로 구성된 TensorFlow 훈련 스크립트에 맞게 조정할 수 있습니다. SageMaker 분산 데이터 병렬은 Keras가 구현된 TensorFlow를 지원하지 않습니다.

참고

SageMaker의 분산 데이터 병렬 처리 라이브러리는 기본적으로 자동 혼합 정밀도(AMP)를 지원합니다. 훈련 스크립트를 프레임워크 수준에서 수정하는 것 외에는 AMP를 활성화하기 위한 추가 조치가 필요하지 않습니다. FP16에 그라디언트가 있는 경우, SageMaker 데이터 병렬 처리 라이브러리는 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 분산과 직접적인 관련이 없지만 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)

훈련 스크립트 조정을 완료한 후 SageMaker Python SDK를 사용하여 SMDDP로 분산 훈련 작업 시작 섹션으로 넘어갑니다.