Utilizzare la libreria SMDDP in uno script di addestramento TensorFlow (obsoleto) - Amazon SageMaker AI

Le traduzioni sono generate tramite traduzione automatica. In caso di conflitto tra il contenuto di una traduzione e la versione originale in Inglese, quest'ultima prevarrà.

Utilizzare la libreria SMDDP in uno script di addestramento TensorFlow (obsoleto)

Importante

La libreria SMDDP ha interrotto il supporto per TensorFlow e non è più disponibile nei Container DL per le versioni di TensorFlow successive alla 2.11.0. Per trovare i precedenti Container DL TensorFlow con la libreria SMDDP installata, consulta Framework supportati.

Le fasi seguenti mostrano come modificare uno script di addestramento TensorFlow per utilizzare la libreria SageMaker AI Distributed Data Parallel. 

Le API della libreria sono progettate per essere simili alle API Horovod. Per ulteriori dettagli su ciascuna API offerta dalla libreria per TensorFlow, consulta la documentazione API TensorFlow per SageMaker AI Distributed Data Parallel.

Nota

SageMaker AI Distributed Data Parallelism è adattabile agli script di addestramento TensorFlow composti da moduli core tf, ad eccezione dei moduli tf.keras. SageMaker AI Distributed Data Parallelism non supporta l’implementazione di TensorFlow con Keras.

Nota

SageMaker AI Distributed Data Parallelism supporta Automatic Mixed Precision (AMP) per impostazione predefinita. Non è necessaria alcuna azione aggiuntiva per abilitare AMP oltre alle modifiche a livello di framework allo script di addestramento. Se i gradienti sono in FP16, la libreria SageMaker AI Distributed Data Parallelism esegue l’operazione AllReduce in FP16. Per ulteriori informazioni sull'implementazione delle API AMP nello script di addestramento, consulta le seguenti risorse:

  1. Importa il client TensorFlow della libreria e inizializzalo.

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

  2. Associa ogni GPU a un singolo processo smdistributed.dataparallel con local_rank —questo si riferisce al grado relativo del processo all'interno di un determinato nodo. L’API sdp.tensorflow.local_rank() fornisce la classificazione locale del dispositivo. Il nodo principale è di classificazione 0 e i nodi di lavoro sono di grado 1, 2, 3 e così via. Questo viene invocato nel seguente blocco di codice come sdp.local_rank(). set_memory_growth non è direttamente correlato a SageMaker AI distribuito, ma deve essere impostato per l’addestramento distribuito 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. Scala il tasso di apprendimento in base al numero di worker. L'API sdp.tensorflow.size() fornisce il numero di worker nel cluster. Questo viene richiamato nel seguente blocco di codice come sdp.size(). 

    learning_rate = learning_rate * sdp.size()

  4. Utilizza le librerie DistributedGradientTape per ottimizzare le operazioni AllReduce durante l'addestramento. Questo esegue il wrapping di 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. Trasmetti le variabili iniziali del modello dal nodo leader (rango 0) a tutti i nodi di lavoro (dai ranghi da 1 a n). Ciò è necessario per garantire un'inizializzazione coerente tra tutte le classificazioni dei worker. Utilizza l'API sdp.tensorflow.broadcast_variables dopo l'inizializzazione delle variabili del modello e dell'ottimizzatore. Questo viene richiamato nel seguente blocco di codice come sdp.broadcast_variables(). 

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

  6. Infine, modifica dello script per salvare i checkpoint solo sul nodo principale. Il nodo principale ha un modello sincronizzato. Ciò evita inoltre che i nodi di lavoro sovrascrivano i checkpoint e possano danneggiarli. 

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

Di seguito è riportato un esempio di script di addestramento TensorFlow per l’addestramento distribuito con la libreria.

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)

Dopo aver completato l'adattamento dello script di addestramento, passa a Avvio di job di addestramento distribuito con SMDDP mediante SageMaker Python SDK.