Use a biblioteca SMDDP em seu script de treinamento do TensorFlow (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 treinamento do TensorFlow (obsoleto)

Importante

A biblioteca SMDDP interrompeu o compatibilidade com o TensorFlow e não está mais disponível em DLCs para TensorFlow posteriores à v2.11.0. Para encontrar DLCs anteriores do TensorFlow com a biblioteca SMDDP instalada, consulte Frameworks compatíveis.

Os passos a seguir mostram como você pode modificar seu script de treinamento do TensorFlow para adaptar a biblioteca de paralelismo de dados distribuídos do SageMaker AI. 

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

nota

O paralelismo de dados distribuídos do SageMaker AI é adaptável aos scripts de treinamento do TensorFlow compostos de módulos principais tf, exceto módulos tf.keras. O paralelismo de dados distribuídos do SageMaker AI não é compatível com a implementação do TensorFlow com Keras.

nota

A biblioteca de paralelismo de dados distribuídos do SageMaker AI permite usar prontamente a técnica de 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 em FP16, a biblioteca de paralelismo de dados do SageMaker AI executará a operação AllReduce em FP16. Para obter mais informações sobre como implementar as APIs de AMP no seu script de treinamento, consulte os seguintes recursos:

  1. Importe o cliente TensorFlow 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 bloco de código a seguir como sdp.local_rank(). set_memory_growth não está diretamente relacionado ao SageMaker AI distribuído, mas deve ser configurado para treinamento distribuído com o 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 treinamento do TensorFlow 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 com o uso do SageMaker Python SDK.