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

As traduções são geradas por tradução automática. Em caso de conflito entre o conteúdo da tradução e da versão original em inglês, a versão em inglês prevalecerá.

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

Importante

A SMDDP biblioteca interrompeu o suporte TensorFlow e não está mais disponível para versões TensorFlow posteriores DLCs à v2.11.0. Para saber mais sobre TensorFlow DLCs a SMDDP biblioteca instalada, consulteEstruturas compatíveis.

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

A biblioteca APIs foi projetada para ser semelhante à HorovodAPIs. Para obter detalhes adicionais sobre cada API oferta oferecida pela biblioteca TensorFlow, consulte a TensorFlow APIdocumentação paralela 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 extra é necessária para habilitar AMP além das modificações em nível de estrutura em seu script de treinamento. Se houver gradientesFP16, a biblioteca de paralelismo de SageMaker dados executará sua operação em. AllReduce FP16 Para obter mais informações sobre AMP APIs a implementação do seu script de treinamento, consulte os seguintes recursos:

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

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

  2. Fixe cada GPU um em um único smdistributed.dataparallel processo com local_rank —isso se refere à classificação relativa do processo em um determinado nó. Isso 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. O sdp.tensorflow.size() API fornece o número de trabalhadores 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 o sdp.tensorflow.broadcast_variables API depois que as variáveis do modelo e do otimizador forem inicializadas. Isso é invocado no seguinte bloco de código comosdp.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 Lançamento de trabalhos de treinamento distribuídos com o SMDDP uso do SageMaker Python SDK.