Die SMDDP-Bibliothek in Ihrem TensorFlow-Trainingsskript verwenden (veraltet) - Amazon SageMaker KI

Die vorliegende Übersetzung wurde maschinell erstellt. Im Falle eines Konflikts oder eines Widerspruchs zwischen dieser übersetzten Fassung und der englischen Fassung (einschließlich infolge von Verzögerungen bei der Übersetzung) ist die englische Fassung maßgeblich.

Die SMDDP-Bibliothek in Ihrem TensorFlow-Trainingsskript verwenden (veraltet)

Wichtig

Die SMDDP-Bibliothek hat die Unterstützung für TensorFlow eingestellt und ist in DLCs für TensorFlow ab Version 2.11.0 nicht mehr verfügbar. Frühere TensorFlow-DLCs mit installierter SMDDP-Bibliothek finden Sie unter Unterstützte Frameworks.

In den folgenden Schritten wird gezeigt, wie Sie ein TensorFlow-Trainingsskript so ändern, dass es die parallel verteilte Datenbibliothek von SageMaker AI verwendet. 

Die Bibliotheks-APIs sind so konzipiert, dass sie den Horovod-APIs ähneln. Weitere Informationen zu den einzelnen APIs, die die Bibliothek für TensorFlow anbietet, finden Sie in der Dokumentation SageMaker AI parallel verteilte Daten TensorFlow-API.

Anmerkung

SageMaker AI parallel verteilte Daten ist an TensorFlow-Trainingsskripte anpassbar, die aus tf-Kernmodulen mit Ausnahme von tf.keras-Modulen bestehen. SageMaker AI parallel verteilte Daten unterstützt TensorFlow mit Keras-Implementierung nicht.

Anmerkung

Die Datenparallelitätsbibliothek mit verteilten Daten von SageMaker AI unterstützt Automatic Mixed Precision (AMP) standardmäßig. Um AMP zu aktivieren, sind außer den Änderungen auf Framework-Ebene an Ihrem Trainingsskript keine weiteren Maßnahmen erforderlich. Wenn Gradienten in FP16 sind, führt die SageMaker-AI-Datenparallelitätsbibliothek ihren Betrieb in FP16 aus. AllReduce Weitere Informationen zum Implementieren der AMP-APIs in Ihrem Trainingsskript finden Sie in den folgenden Ressourcen:

  1. Importieren Sie den TensorFlow-Client der Bibliothek und initialisieren Sie ihn.

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

  2. Ordnen Sie jede GPU einem einzelnen smdistributed.dataparallel Prozess zu mit local_rank–das bezieht sich auf den relativen Rang des Prozesses innerhalb eines bestimmten Knotens. Die sdp.tensorflow.local_rank()-API stellt Ihnen den lokalen Rang des Geräts zur Verfügung. Der Führungsnoten hat Rang 0, und die Worker-Knoten haben Rang 1, 2, 3, usw. Dies wird im folgenden Codeblock als sdp.local_rank() aufgerufen. set_memory_growth steht nicht in direktem Zusammenhang mit SageMaker AI verteilt, muss aber für verteiltes Training mit TensorFlow eingerichtet werden. 

    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. Skalieren Sie die Lernrate nach der Anzahl der Auftragnehmer. Die sdp.tensorflow.size() API stellt Ihnen die Anzahl der Auftragnehmer im Cluster zur Verfügung. Dies wird im folgenden Codeblock als sdp.size() aufgerufen. 

    learning_rate = learning_rate * sdp.size()

  4. Verwenden Sie die DistributedGradientTape der Bibliothek, um den AllReduce Betrieb während des Trainings zu optimieren. tf.GradientTape ist damit abgeschlossen.  

    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. Senden Sie die anfänglichen Modellvariablen vom Führungsknoten (Rang 0) an alle Worker-Knoten (Ränge 1 bis n). Dies ist erforderlich, um eine konsistente Initialisierung in allen Auftragnehmer-Rängen sicherzustellen. Verwenden Sie die sdp.tensorflow.broadcast_variables API, nachdem die Modell- und Optimizer-Variablen initialisiert wurden. Dies wird im folgenden Codeblock als sdp.broadcast_variables() aufgerufen. 

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

  6. Ändern Sie abschließend Ihr Skript so, dass es Checkpoints nur auf dem Führungsknoten speichert. Der Führungsknoten hat ein synchronisiertes Modell. Dadurch wird auch vermieden, dass Worker-Knoten die Checkpoints überschreiben und die Checkpoints möglicherweise beschädigen. 

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

Im Folgenden finden Sie ein Beispiel für ein TensorFlow-Trainingsskript für verteiltes Training mit der Bibliothek.

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)

Nachdem Sie die Anpassung Ihres Trainingsskripts abgeschlossen haben, fahren Sie mit Starten von verteilten Trainingsjobs mit dem SageMaker Python SDK fort.