View a markdown version of this page

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

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 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 à v2.11.0. TensorFlow Para encontrar TensorFlow DLCs anteriores com a biblioteca SMDDP instalada, consulte. Frameworks compatíveis

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

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

nota

SageMaker Os dados distribuídos paralelos de IA são adaptáveis a scripts de TensorFlow treinamento compostos por módulos tf principais, exceto tf.keras módulos. SageMaker Os dados distribuídos paralelos da IA não são compatíveis TensorFlow com a implementação do Keras.

nota

A SageMaker biblioteca de paralelismo de dados distribuídos por IA oferece suporte imediato à precisão mista automática (AMP). 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 dados do SageMaker AI 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 seguintes recursos:

  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 API sdp.tensorflow.local_rank() fornece a classificação local do dispositivo para você. 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 à SageMaker IA distribuída, 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 operadores. A API sdp.tensorflow.size() fornece o número de operadores 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 AI 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 operadores. 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 AI data parallel: Import the library TF API
import smdistributed.dataparallel.tensorflow as sdp

# SageMaker AI 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 AI 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 AI 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 AI 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 AI 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 AI data parallel: Save checkpoints only from master node.
if sdp.rank() == 0:
    checkpoint.save(checkpoint_dir)

Depois de concluir a adaptação do seu script de treinamento, prossiga para Lançamento de trabalhos de treinamento distribuídos com SMDDP usando o SDK do Python SageMaker.