Usa la libreria SMDDP nello script di formazione (obsoleto) TensorFlow - 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à.

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

Importante

La libreria SMDDP ha interrotto il supporto TensorFlow e non è più disponibile nelle TensorFlow versioni successive alla DLCs v2.11.0. Per trovare le versioni precedenti TensorFlow DLCs con la libreria SMDDP installata, vedere. Framework supportati

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

La libreria APIs è progettata per essere simile a APIs Horovod. Per ulteriori dettagli su ciascuna API offerta dalla libreria TensorFlow, consulta la documentazione dell' TensorFlow API SageMaker AI distributed data parallel.

Nota

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

Nota

La libreria di parallelismo dei dati distribuiti SageMaker AI supporta immediatamente 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 presenti FP16, la libreria di parallelismo dei dati SageMaker AI esegue il suo funzionamento in. AllReduce FP16 Per ulteriori informazioni sull'implementazione di AMP nello script APIs di formazione, consulta le seguenti risorse:

  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 all' SageMaker intelligenza artificiale distribuita, 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 TensorFlow formazione per la formazione distribuita 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 lavori di formazione distribuiti con SMDDP utilizzando Python SDK SageMaker .