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

기계 번역으로 제공되는 번역입니다. 제공된 번역과 원본 영어의 내용이 상충하는 경우에는 영어 버전이 우선합니다.

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

중요

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

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

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

참고

SageMaker 분산 데이터 병렬은 tf.keras 모듈을 제외한 tf 코어 모듈로 구성된 TensorFlow 훈련 스크립트에 적응할 수 있습니다. SageMaker 분산 데이터 병렬은 Keras 구현 TensorFlow 을 지원하지 않습니다.

참고

SageMaker 분산 데이터 병렬 처리 라이브러리는 즉시 사용 가능한 자동 혼합 정밀도(AMP)를 지원합니다. 훈련 스크립트에 대한 프레임워크 수준 수정 AMP 외에는 추가 작업이 필요하지 않습니다. 그라데이션이 에 있는 경우 SageMaker 데이터 병렬 처리 라이브러리FP16는 에서 AllReduce 작업을 실행합니다FP16. 훈련 스크립트 AMP APIs 구현에 대한 자세한 내용은 다음 리소스를 참조하세요.

  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 분산과 직접 관련이 없지만 를 사용한 분산 훈련을 위해 설정해야 합니다 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_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 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을 SMDDP 사용하여 를 사용하여 분산 훈련 작업 시작 SDK 섹션으로 넘어갑니다.