Uso de la biblioteca de SMDDP en un script de entrenamiento de TensorFlow (obsoleto) - Amazon SageMaker AI

Uso de la biblioteca de SMDDP en un script de entrenamiento de TensorFlow (obsoleto)

importante

Se ha dejado de admitir la biblioteca de SMDDP para TensorFlow y ya no está disponible en los DLC de TensorFlow posteriores a v2.11.0. Para buscar DLC de TensorFlow anteriores con la biblioteca de SMDDP instalada, consulte Marcos admitidos.

En los siguientes pasos se muestra cómo modificar un script de entrenamiento de TensorFlow para utilizar la biblioteca de paralelismo de datos distribuidos de SageMaker AI. 

Las API de biblioteca están diseñadas para ser similares a las API de Horovod. Para obtener detalles adicionales sobre cada API que ofrece la biblioteca para TensorFlow, consulte la documentación de la API de TensorFlow en paralelismo de datos distribuidos de SageMaker AI.

nota

El paralelismo de datos distribuidos de SageMaker AI son adaptables a los scripts de entrenamiento de TensorFlow compuestos de módulos principales tf excepto para los módulos tf.keras. El paralelismo de datos distribuidos de SageMaker AI no admiten TensorFlow con la implementación de Keras.

nota

La biblioteca de paralelismo de datos distribuidos de SageMaker AI admite la precisión mixta automática (AMP) desde el primer momento. No se necesita ninguna acción adicional para habilitar AMP que no sean las modificaciones a nivel de marco de su script de entrenamiento. Si los gradientes están en FP16, la biblioteca de paralelismo de datos de SageMaker AI ejecuta su operación AllReduce en FP16. Para obtener más información sobre la implementación de API de AMP en su script de entrenamiento, consulte los siguientes recursos:

  1. Importe el cliente TensorFlow de la biblioteca e inicialícelo.

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

  2. Fijar cada GPU a un único proceso de smdistributed.dataparallel con local_rank: esto se refiere a la clasificación relativa del proceso dentro de un nodo determinado. La API sdp.tensorflow.local_rank() proporciona la clasificación local del dispositivo. El nodo principal es el rango 0 y los nodos de trabajo son de rango 1, 2, 3, etc. Esto se invoca en el siguiente bloque de código como sdp.local_rank(). set_memory_growth no está directamente relacionado con SageMaker AI distribuido, sino que debe configurarse para la entrenamiento distribuido con 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 la tasa de aprendizaje según el número de trabajadores. La API sdp.tensorflow.size() le proporciona el número de trabajadores del clúster. Esto se invoca en el siguiente código como sdp.size(). 

    learning_rate = learning_rate * sdp.size()

  4. Utilice el DistributedGradientTape de la biblioteca para optimizar las operaciones AllReducedurante el entrenamiento. Esto envuelve 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. Transmite las variables de modelo iniciales desde el nodo principal (rango 0) a todos los nodos de trabajo (rangos 1 a n). Esto es necesario para garantizar una inicialización coherente en todos los rangos de trabajo. Utilice la API sdp.tensorflow.broadcast_variables después de inicializar el modelo y las variables del optimizador. Esto se invoca en el siguiente bloque de código como sdp.broadcast_variables(). 

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

  6. Por último, modifique su script para guardar los puntos de control solo en el nodo principal. El nodo principal tiene un modelo sincronizado. Esto también evita que los nodos de trabajo sobrescriban los puntos de control y, posiblemente, corrompa los puntos de control. 

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

A continuación se muestra un ejemplo de script de entrenamiento de TensorFlow para entrenamiento distribuido con la 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)

Cuando haya terminado de adaptar su script de entrenamiento, pase a Inicialización de trabajos de entrenamiento distribuido con SMDDP mediante el SageMaker Python SDK.