Distribuer l’entraînement

Entraînement multi-GPU avec TRL

Les formateurs de TRL sont lancés avec torchrun, le lanceur distribué de PyTorch. Il démarre un processus par GPU :

torchrun --nproc_per_node 8 train.py <SCRIPT_ARGS>

Le CLI trl fait la même chose en interne, donc trl sft ... s’exécute sur tous les GPU de la machine ; passez --nproc_per_node pour en utiliser moins.

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 machine unique ne dispose pas de suffisamment de GPU, torchrun permet d’étendre l’entraînement sur plusieurs machines (nœuds). Chaque nœud exécute la même commande avec le nombre total de nœuds, son propre rang et un point de rendez-vous sur le nœud principal.

Option 1 : Lancement manuel (non-HPC)

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

# Node 0 (main node, IP 10.0.0.1)
torchrun --nnodes 2 --nproc_per_node 8 --node_rank 0 --rdzv_backend c10d --rdzv_endpoint 10.0.0.1:29500 train.py

# Node 1
torchrun --nnodes 2 --nproc_per_node 8 --node_rank 1 --rdzv_backend c10d --rdzv_endpoint 10.0.0.1:29500 train.py

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

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

MAIN_NODE=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
srun torchrun --nnodes 2 --nproc_per_node 8 --node_rank $SLURM_NODEID --rdzv_backend c10d --rdzv_endpoint $MAIN_NODE:29500 train.py

Puis soumettez le travail :

sbatch slurm_job.sh

srun démarre un torchrun par nœud, et chacun d’eux démarre un processus par GPU.

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 deepspeed dans la configuration d’entraînement. Consultez le guide d’intégration DeepSpeed.

Lecture complémentaire

• torchrun (Lancement élastique)
• Documentation du gestionnaire de charge SLURM - Pour l'ordonnancement des tâches sur cluster