Distribuer l’entraînement

Entraînement multi-GPU avec TRL

Les formateurs de TRL utilisent 🤗 Accelerate pour permettre l'entraînement distribué sur plusieurs GPU ou nœuds. Pour ce faire, créez d'abord un fichier de configuration 🤗 Accelerate en exécutant

accelerate config

et de répondre aux questions en fonction de votre configuration multi-GPU / multi-nœuds. Vous pouvez ensuite lancer l’entraînement distribué en exécutant :

accelerate launch train.py

Nous fournissons également des fichiers de configuration dans le dossier examples qui peuvent être utilisés comme modèles. Pour utiliser ces modèles, il suffit de passer le chemin du fichier de configuration lors du lancement d'un job, par exemple :

accelerate launch --config_file examples/accelerate_configs/multi_gpu.yaml train.py <SCRIPT_ARGS>

Cela distribue automatiquement la charge de travail sur tous les GPU disponibles.

Sous le capot, 🤗 Accelerate crée un modèle par GPU. Chaque processus :

• Traite son propre lot de données
• Calcule la perte et les gradients pour ce lot
• Partage les mises à jour des gradients entre tous les GPU

La taille de lot effective est calculée comme suit : $$\text{Batch Size} = \text{per\_device\_train\_batch\_size} \times \text{num\_devices} \times \text{gradient\_accumulation\_steps}$$

Pour maintenir une taille de lot constante lors du passage à plusieurs GPU, pensez à mettre à jour per_device_train_batch_size et gradient_accumulation_steps en conséquence.

Par exemple, ces configurations sont équivalentes et doivent produire les mêmes résultats :

Avoir un modèle par GPU peut entraître une forte consommation de mémoire, ce qui peut ne pas être réalisable pour les grands modèles ou les GPU à faible mémoire. Dans de tels cas, vous pouvez tirer parti de DeepSpeed, qui fournit des optimisations telles que le fractionnement du modèle, le Zero Redundancy Optimizer, l'entraînement en précision mixte et le déchargement vers le CPU ou le NVMe. Consultez notre guide DeepSpeed Integration pour plus de détails.

L’entraînement sur des séquences très longues dispose de son propre guide : Training Beyond 1M Tokens.

Entraînement multi-nœuds

Lorsqu’une seule machine ne dispose pas de suffisamment de GPU, TRL peut étendre l’entraînement sur plusieurs machines (nœuds) à l’aide de 🤗 Accelerate.

Configuration d'Accelerate

Créez un fichier de configuration accelerate (par exemple, multi_node.yaml) pour l'entraînement multi-nœuds. Champs clés :

compute_environment: LOCAL_MACHINE
distributed_type: MULTI_GPU
num_machines: 2
machine_rank: 0  # 0 for main node, 1 for second node
main_process_ip: 10.0.0.1  # IP of rank 0 node
main_process_port: 29500
num_processes: 16  # total processes across nodes
mixed_precision: bf16
use_cpu: false
same_network: true

Ajustez num_processes pour qu’il corresponde au nombre total de GPU sur tous les nœuds.

Remplacez 10.0.0.1 par l'adresse IP réelle du nœud de rang 0 (principal).

Lancement

Option 1 : Lancement manuel (non-HPC)

Exécutez ce qui suit manuellement sur chaque nœud :

# Node 0 (main node)
accelerate launch --config_file multi_node.yaml --machine_rank 0 train.py

# Node 1
accelerate launch --config_file multi_node.yaml --machine_rank 1 train.py

Option 2 : Lancement SLURM (clusters HPC)

Pour les clusters utilisant le planificateur de tâches SLURM, créez un script de tâche (par exemple, slurm_job.sh) :

#!/bin/bash
#SBATCH --nodes=2
#SBATCH --gpus-per-node=8
#SBATCH --job-name=trl_multi

srun accelerate launch --config_file multi_node.yaml train.py

Puis soumettez le travail :

sbatch slurm_job.sh

SLURM distribue automatiquement l’entraînement sur tous les nœuds et GPU demandés, et srun configure les variables d’environnement nécessaires à la communication multi-nœuds.

Directives SLURM principales :

• --nodes=2 : Demande 2 nœuds de calcul
• --gpus-per-node=8 : alloue 8 GPU par nœud (16 au total)
• --job-name : Étiquette pour le suivi dans la file d’attente des tâches

Vous pouvez combiner le multi-nœud avec DeepSpeed en définissant distributed_type: DEEPSPEED et en ajoutant un bloc deepspeed_config. Consultez le guide d’intégration DeepSpeed.

Lecture complémentaire

• Accelerate : Scripts de lancement
• Accelerate : Zoo d'exemples
• Documentation du gestionnaire de charge SLURM - Pour l'ordonnancement des tâches sur cluster