Utilisation de la bibliothèque SMDDP dans votre script d’entraînement TensorFlow (obsolète) - Amazon SageMaker AI

Utilisation de la bibliothèque SMDDP dans votre script d’entraînement TensorFlow (obsolète)

Important

La bibliothèque SMDDP ne prend plus en charge TensorFlow et n’est plus disponible dans les DLC pour TensorFlow ultérieurs à la version 2.11.0. Pour rechercher les DLC TensorFlow précédents sur lesquels la bibliothèque SMDDP est installée, consultez Cadres de travail pris en charge.

Les étapes suivantes vous expliquent comment modifier un script d’entraînement TensorFlow pour utiliser la bibliothèque SageMaker AI de parallélisme distribué des données. 

La conception des API de la bibliothèque est similaire à celle des API Horovod. Pour plus de détails sur chaque API que la bibliothèque propose pour TensorFlow, consultez la documentation de l’API TensorFlow de la bibliothèque SageMaker AI de parallélisme distribué des données.

Note

La bibliothèque SageMaker AI de parallélisme distribué des données est adaptable aux scripts d’entraînement TensorFlow composés des modules principaux tf, à l’exception des modules tf.keras. La bibliothèque SageMaker AI de parallélisme distribué des données ne prend pas en charge l’implémentation de TensorFlow avec Keras.

Note

La bibliothèque SageMaker AI de parallélisme distribué des données prend en charge la précision mixte automatique (AMP). Pour activer l'AMP, il vous suffit de modifier le cadre de votre script d'entraînement. Si les gradients sont en FP16, la bibliothèque SageMaker AI de parallélisme des données exécute son opération AllReduce en FP16. Pour plus d'informations sur la mise en œuvre des API AMP dans votre script d'entraînement, consultez les ressources suivantes :

  1. Importez le client TensorFlow de la bibliothèque et initialisez-le.

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

  2. Épinglez chaque GPU à un processus smdistributed.dataparallel unique avec local_rank : cela fait référence au rang relatif du processus au sein d'un nœud donné. L’API sdp.tensorflow.local_rank() vous indique le rang local du dispositif. Le nœud principal est le rang 0, et les nœuds des employés sont les rangs 1, 2, 3, etc. Cela est invoqué comme sdp.local_rank() dans le bloc de code suivant. set_memory_growth n’est pas directement lié à l’entraînement distribué SageMaker, mais doit être défini pour l’entraînement distribué avec 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. Mettez à l'échelle le taux d'apprentissage en fonction du nombre d'employés. L'API sdp.tensorflow.size() vous indique le nombre d'employés dans le cluster. Cela est appelé sous sdp.size() dans le bloc de code suivant. 

    learning_rate = learning_rate * sdp.size()

  4. Utilisez le DistributedGradientTape de la bibliothèque pour optimiser les opérations AllReduce pendant l'entraînement. Cela recouvre 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. Diffusez les variables initiales du modèle, du nœud principal (rang 0) vers tous les nœuds d'employés (rangs 1 à n). Cela est indispensable pour garantir une initialisation cohérente dans tous les rangs des employés. Utilisez l'API sdp.tensorflow.broadcast_variables après l'initialisation des variables du modèle et de l'optimiseur. Ceci est invoqué dans le bloc de code suivant comme sdp.broadcast_variables(). 

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

  6. Enfin, modifiez votre script de sorte à enregistrer les points de contrôle sur le nœud principal uniquement. Le nœud principal a un modèle synchronisé. Cela évite également que les nœuds d'employés écrasent les points de contrôle et les endommagent éventuellement. 

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

Voici un exemple de script d'entraînement TensorFlow pour l'entraînement distribué avec la bibliothèque.

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)

Une fois que vous avez terminé d'adapter votre scénario d'entraînement, passez à Lancement de tâches d’entraînement distribué avec SMDDP à l’aide du kit SageMaker Python SDK.