Utilice la SMDDP biblioteca en su script TensorFlow de entrenamiento (obsoleto) - Amazon SageMaker

Las traducciones son generadas a través de traducción automática. En caso de conflicto entre la traducción y la version original de inglés, prevalecerá la version en inglés.

Utilice la SMDDP biblioteca en su script TensorFlow de entrenamiento (obsoleto)

importante

La SMDDP biblioteca dejó de ofrecer soporte para la versión TensorFlow 2.11.0 y ya no estará disponible en DLCs una versión TensorFlow posterior a la versión 2.11.0. Para buscar versiones anteriores TensorFlow DLCs con la SMDDP biblioteca instalada, consulte. Marcos admitidos

Los siguientes pasos muestran cómo modificar un script de TensorFlow entrenamiento para utilizar SageMaker la biblioteca paralela de datos distribuidos. 

La biblioteca APIs está diseñada para ser similar a HorovodAPIs. Para obtener detalles adicionales sobre cada uno de los API que ofrece la biblioteca TensorFlow, consulte la TensorFlow APIdocumentación de SageMaker distributed data parallel.

nota

SageMaker distributed data parallel se adapta a los guiones de TensorFlow entrenamiento compuestos por módulos tf principales, excepto tf.keras los módulos. SageMaker distributed data parallel no es compatible TensorFlow con la implementación de Keras.

nota

La biblioteca de paralelismo de datos SageMaker distribuidos admite Automatic Mixed Precision (AMP) de forma inmediata. No es necesario realizar ninguna acción adicional para permitir que se realicen AMP modificaciones en el guion de entrenamiento a nivel de marco. Si los gradientes están dentroFP16, la biblioteca de paralelismo de SageMaker datos ejecuta su operación en ella. AllReduce FP16 Para obtener más información sobre la implementación AMP APIs en su script de entrenamiento, consulte los siguientes recursos:

  1. Importa el TensorFlow cliente de la biblioteca e inicialízalo.

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

  2. Fije cada uno de ellos GPU a un solo smdistributed.dataparallel proceso con una teclalocal_rank: se refiere al rango relativo del proceso dentro de un nodo determinado. Le sdp.tensorflow.local_rank() API 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 comosdp.local_rank(). set_memory_growthno está directamente relacionado con lo SageMaker distribuido, pero debe configurarse para el 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. Le sdp.tensorflow.size() API 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 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. Úselo sdp.tensorflow.broadcast_variables API después de inicializar las variables del modelo y 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)

El siguiente es un ejemplo de guion de TensorFlow entrenamiento para un entrenamiento distribuido con la 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)

Cuando haya terminado de adaptar su script de entrenamiento, pase a Lanzamiento de trabajos de formación distribuidos SMDDP con SageMaker Python SDK.