TensorFlow 교육 스크립트에서 SMDDP 라이브러리 사용 (더 이상 사용되지 않음)

중요

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

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

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

참고

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

참고

SageMaker 분산 데이터 병렬화 라이브러리는 기본적으로 자동 혼합 정밀도 (AMP) 를 지원합니다. 훈련 스크립트를 프레임워크 수준에서 수정하는 것 외에는 AMP를 활성화하기 위한 추가 조치가 필요하지 않습니다. 그래디언트가 FP16에 있는 경우 SageMaker 데이터 병렬화 라이브러리는 FP16에서 해당 작업을 실행합니다. AllReduce 훈련 스크립트에 AMP API를 구현하는 방법에 대한 자세한 정보는 다음 리소스를 참조하세요.

• 프레임워크 - NVIDIA 딥 러닝 성능 문서에서 TensorFlow

• NVIDIA 개발자 문서의 딥 러닝을 위한 자동 혼합 정밀도

• TensorFlow 설명서 내 혼합 정밀도 API TensorFlow

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를 사용하여 분산 교육 작업 시작 섹션으로 넘어갑니다.