Use a biblioteca SMDDP em seu script de TensorFlow treinamento (obsoleto)

Importante

A biblioteca SMDDP interrompeu o suporte TensorFlow e não está mais disponível em DLCs posteriores à versão 2.11.0. TensorFlow Para encontrar TensorFlow DLCs anteriores com a biblioteca SMDDP instalada, consulte. Estruturas compatíveis

As etapas a seguir mostram como modificar um script de TensorFlow treinamento para utilizar a biblioteca paralela SageMaker de dados distribuídos. 

As APIs da biblioteca foram projetadas para serem semelhantes às APIs do Horovod. Para obter mais detalhes sobre cada API que a biblioteca oferece TensorFlow, consulte a documentação da TensorFlow API parallel de dados SageMaker distribuídos.

nota

SageMaker distributed data parallel é adaptável a scripts de TensorFlow treinamento compostos por módulos tf principais, exceto tf.keras módulos. SageMaker distributed data parallel não é compatível TensorFlow com a implementação do Keras.

nota

A biblioteca de paralelismo de dados SageMaker distribuídos oferece suporte à Precisão Mista Automática (AMP) pronta para uso. Nenhuma ação adicional é necessária para habilitar o AMP, além das modificações no nível do framework no seu script de treinamento. Se os gradientes estiverem no FP16, a biblioteca de paralelismo de SageMaker dados executará sua operação no FP16. AllReduce Para obter mais informações sobre como implementar as APIs de AMP no seu script de treinamento, consulte os recursos a seguir:

• Frameworks - TensorFlow na documentação de desempenho do NVIDIA Deep Learning

• Precisão mista automática para aprendizado profundo nos documentos de desenvolvedores da NVIDIA

• TensorFlow APIs de precisão mista na documentação TensorFlow

1. Importe o TensorFlow cliente da biblioteca e inicialize-o.

   import smdistributed.dataparallel.tensorflow as sdp 
   sdp.init()

2. Fixe cada GPU em um único smdistributed.dataparallel processo com local_rank — isso se refere à classificação relativa do processo em um determinado nó. A sdp.tensorflow.local_rank() API fornece a classificação local do dispositivo. A classificação do nó líder é 0 e as classificações dos nós de processamento são 1, 2, 3 e assim por diante. Isso é invocado no seguinte bloco de código comosdp.local_rank(). set_memory_growthnão está diretamente relacionado ao treinamento SageMaker distribuído, mas deve ser configurado para treinamento distribuído com 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. Escale a taxa de aprendizado pelo número de trabalhadores. A API sdp.tensorflow.size() fornece o número de workers no cluster. Isso é invocado no bloco de código a seguir como sdp.size().

   learning_rate = learning_rate * sdp.size()

4. Use a biblioteca DistributedGradientTape para otimizar as operações AllReduce durante o treinamento. Isso envolve 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. Transmita as variáveis iniciais do modelo do nó líder (classificação 0) para todos os nós de processamento (classificações de 1 a n). Isso é necessário para garantir uma inicialização consistente em todas as categorias de trabalhadores. Use a API sdp.tensorflow.broadcast_variables depois que as variáveis do modelo e do otimizador forem inicializadas. Isso é invocado no bloco de código a seguir como sdp.broadcast_variables().

   sdp.broadcast_variables(model.variables, root_rank=0)
   sdp.broadcast_variables(opt.variables(), root_rank=0)

6. Por fim, modifique seu script para salvar pontos de verificação somente no nó líder. O nó líder tem um modelo sincronizado. Isso também evita que os nós de processamento sobrescrevam os pontos de verificação e possivelmente os corrompam.

   if sdp.rank() == 0:
       checkpoint.save(checkpoint_dir)

Veja a seguir um exemplo de script de TensorFlow treinamento para treinamento distribuído com a biblioteca.

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)

Depois de concluir a adaptação do seu roteiro de treinamento, vá para Etapa 2: iniciar um trabalho de treinamento distribuído usando o SDK do SageMaker Python.