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

importante

La biblioteca SMDDP dejó de ofrecer soporte a los DLC TensorFlow y dejó de estar disponible en ellos a partir de la versión 2.11.0 TensorFlow . Para encontrar los TensorFlow DLC anteriores con la biblioteca SMDDP 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. 

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 TensorFlow, consulte la documentación de la TensorFlow API paralela de datos SageMaker distribuidos.

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 la precisión mixta automática (AMP) lista para usar. 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 el FP16, la biblioteca de paralelismo de SageMaker datos ejecuta su operación en el FP16. AllReduce Para obtener más información sobre la implementación de API de AMP en su script de entrenamiento, consulte los siguientes recursos:

• Marcos: en la documentación de rendimiento del aprendizaje TensorFlow profundo de NVIDIA

• Precisión mixta automática para aprendizaje profundo en los Documentos para desarrolladores de NVIDIA

• TensorFlow API de precisión mixta en la TensorFlowdocumentación

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

   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 sdp.tensorflow.local_rank() API le 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. 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 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)

El siguiente es un ejemplo de guion de TensorFlow formación para una formación distribuida 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 Paso 2: Inicie un trabajo de formación distribuido con el SDK de SageMaker Python.