Usa la libreria SMDDP nello script di formazione (obsoleto) TensorFlow

Importante

La libreria SMDDP ha interrotto il supporto TensorFlow e non è più disponibile nei DLC successivi alla versione 2.11.0. TensorFlow Per trovare i TensorFlow DLC precedenti con la libreria SMDDP installata, consulta. Framework supportati

I passaggi seguenti mostrano come modificare uno script di TensorFlow addestramento per utilizzare la libreria parallela SageMaker di dati distribuiti. 

Le API della libreria sono progettate per essere simili alle API Horovod. Per ulteriori dettagli su ciascuna API offerta dalla libreria TensorFlow, consulta la documentazione dell' TensorFlow API parallel data SageMaker distributed.

Nota

SageMaker distributed data parallel è adattabile agli script di TensorFlow addestramento composti da moduli tf principali ad eccezione tf.keras dei moduli. SageMaker distributed data parallel non supporta TensorFlow l'implementazione Keras.

Nota

La libreria di parallelismo dei dati SageMaker distribuiti supporta automaticamente Automatic Mixed Precision (AMP). 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 di parallelismo SageMaker dei dati esegue il suo funzionamento in FP16. AllReduce Per ulteriori informazioni sull'implementazione delle API AMP nello script di addestramento, consulta le seguenti risorse:

• Frameworks: nella documentazione di NVIDIA Deep Learning Performance TensorFlow

• Automatic mixed precision for deep learning nella documentazione per sviluppatori NVIDIA

• TensorFlow API a precisione mista nella documentazione TensorFlow

1. Importa il TensorFlow client 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'sdp.tensorflow.local_rank()API fornisce il rango 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 richiamato nel seguente blocco di codice comesdp.local_rank(). set_memory_growthnon è direttamente correlato alla SageMaker distribuzione, ma deve essere impostata 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 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 TensorFlow formazione per la formazione distribuita con la libreria.

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)

Dopo aver completato l'adattamento dello script di addestramento, passa a Passaggio 2: Avvia un processo di formazione distribuito utilizzando SageMaker Python SDK.