Entraînement distribué TensorFlow

L'entraînement distribué TensorFlow est une technique qui utilise plusieurs machines ou plusieurs périphériques de calcul (tels que des GPU/TPU) pour travailler en collaboration afin d'accomplir des tâches d'entraînement de modèle. Grâce à l'entraînement distribué, nous pouvons :

  1. Accélérer le processus d'entraînement du modèle
  2. Traiter des jeux de données à très grande échelle
  3. Entraîner des modèles complexes avec un nombre énorme de paramètres

Concepts fondamentaux

1. Stratégie de distribution

TensorFlow propose plusieurs stratégies de distribution :

Exemple

# Stratégies de distribution courantes
strategy = tf.distribute.MirroredStrategy()  # Multi-GPU sur une seule machine
strategy = tf.distribute.MultiWorkerMirroredStrategy()  # Multi-machine multi-GPU
strategy = tf.distribute.TPUStrategy()  # Cluster TPU
strategy = tf.distribute.ParameterServerStrategy()  # Architecture de serveur de paramètres

2. Parallélisme de données vs Parallélisme de modèle

Type Parallélisme de données Parallélisme de modèle
Principe Chaque appareil traite différents lots de données Le modèle est réparti sur différents appareils
Avantages Simple à mettre en œuvre, adapté à la plupart des scénarios Adapté aux modèles très volumineux
Inconvénients Nécessite une synchronisation des gradients Complexe à implémenter

3. Mise à jour synchrone vs Mise à jour asynchrone

  • Mise à jour synchrone :Tous les appareils mettent à jour le modèle collectivement après avoir terminé leurs calculs
  • Mise à jour asynchrone :Les appareils calculent et mettent à jour indépendamment sans attendre

Étapes d'implémentation

1. Configurer l’environnement distribué

Exemple

import tensorflow as tf

# Initialiser la stratégie de distribution
strategy = tf.distribute.MirroredStrategy()

# Vérifier le nombre de périphériques disponibles
print(f"Nombre de périphériques : {strategy.num_replicas_in_sync}")

2. Construire le modèle dans la portée de la stratégie

Exemple

with strategy.scope():
    # Toutes les variables définies dans cette portée seront répliquées sur tous les périphériques
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dense(10)
    ])
   
    model.compile(
        optimizer='adam',
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=['accuracy']
    )

3. Préparer le jeu de données distribué

Exemple

# Charger le jeu de données
dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))

# Regrouper par lots et fragmenter
batch_size = 64 * strategy.num_replicas_in_sync  # Ajuster la taille du lot en fonction du nombre de périphériques
dataset = dataset.shuffle(buffer_size=10000).batch(batch_size)

4. Entraîner le modèle

Exemple

# Méthode d'entraînement conventionnelle
model.fit(dataset, epochs=10)

Configuration avancée

1. Configuration multi-machine

Exemple

# Définir la variable d'environnement TF_CONFIG sur chaque nœud worker
import json
import os

os.environ['TF_CONFIG'] = json.dumps({
    'cluster': {
        'worker': ["worker1.example.com:12345", "worker2.example.com:23456"]
    },
    'task': {'type': 'worker', 'index': 0}  # Chaque worker a un index différent
})

2. Boucle d'entraînement personnalisée

Exemple

@tf.function
def train_step(inputs):
    x, y = inputs
   
    with tf.GradientTape() as tape:
        predictions = model(x, training=True)
        loss = loss_object(y, predictions)
   
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

# Étapes d'entraînement distribué
@tf.function
def distributed_train_step(dataset_inputs):
    per_replica_losses = strategy.run(train_step, args=(dataset_inputs,))
    return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)

Conseils d'optimisation des performances

  1. Ajustement de la taille du lot :Taille du lot totale = taille du lot par appareil × nombre d'appareils
  2. Prétraitement des données :Utilisezdataset.prefetch()etdataset.cache()pour améliorer l'efficacité du chargement des données
  3. Compression du gradient :Pour les communications entre appareils, envisagez d'utiliser la compression du gradient pour réduire les besoins en bande passante
  4. Entraînement en précision mixte :Combineztf.keras.mixed_precisionpour améliorer la vitesse d'entraînement

Exemple

# Exemple de précision mixte
policy = tf.keras.mixed_precision.Policy('mixed_float16')
tf.keras.mixed_precision.set_global_policy(policy)

Dépannage

1. Mémoire insuffisante

  • Réduire la taille de lot par appareil
  • Utiliser des techniques d'accumulation de gradient
  • Activer l'option de croissance de la mémoire

Exemple

gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

2. Goulot d'étranglement de communication entre appareils

  • UtiliserNCCLen tant qu'implémentation de communication entre appareils
  • Envisagez de réduire la fréquence de synchronisation (en augmentant de manière appropriée la taille du pas de mise à jour)

Exemple

# Configurer l'implémentation de communication
os.environ['TF_GPU_ALLOCATOR'] = 'cuda_malloc_async'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

Exercices pratiques

Exercice 1 : Entraînement multi-GPU sur une seule machine

  1. Préparer un modèle CNN simple
  2. UtiliserMirroredStrategypour entraîner le jeu de données CIFAR-10 sur des configurations multi-GPU locales
  3. Comparer la différence de vitesse d'entraînement entre un seul GPU et plusieurs GPU

Exercice 2 : Simulation de configuration multi-machine

  1. UtiliserMultiWorkerMirroredStrategy
  2. pour simuler un environnement multi-worker sur la même machine (via différents ports)
  3. Observez les journaux pour comprendre le processus de coordination entre les workers
Autres extensions