Accelerate ND‑Parallel : Guide pour un entraînement multi‑GPU efficace

TL;DR

• Accelerate propose des outils ND‑Parallel pour composer plusieurs stratégies de parallélisme (DP, FSDP, TP, CP) dans un même script d’entraînement.
• L’intégration avec Axolotl fournit des configs d’exemple et des champs de configuration pour activer des combinaisons sans réécrire le code central.
• Principaux réglages : ParallelismConfig avec dp_replicate_size, dp_shard_size, tp_size (tensor_parallel_size dans Axolotl) et sharding de contexte pour des séquences très longues.

Contexte et historique

L’entraînement de très grands modèles sur plusieurs GPU impose de choisir et de composer des stratégies de parallélisme afin d’équilibrer mémoire, calcul et communication. Le parallélisme de données (DP) réplique le modèle et l’état de l’optimiseur sur chaque appareil et répartit le mini‑batch entre répliques ; les gradients sont synchronisés avant la mise à jour. DP augmente le débit mais exige que le modèle tienne sur un seul appareil. Quand le modèle est trop grand pour une seule GPU, le Fully Sharded Data Parallel (FSDP) fragmente (shard) les poids, gradients et états d’optimiseur entre les GPU, ne rassemblant les paramètres d’une couche que lorsqu’ils sont nécessaires. Le Tensor Parallelism (TP) répartit de manière permanente les paramètres de grandes couches linéaires entre appareils. Le Context Parallelism (CP) fragmente la séquence d’entrée sur la dimension temporelle pour gérer des contextes extrêmement longs. Accelerate, en collaboration avec Axolotl, offre une approche intégrée et configurable pour composer ces stratégies, permettant de passer d’un entraînement mono‑appareil à des déploiements multi‑node tout en contrôlant les compromis mémoire/communication. Le dépôt Accelerate propose un script d’entraînement end‑to‑end et Axolotl fournit des configs testées pour démarrer rapidement.

Ce qui est nouveau

• Accelerate expose un moyen simple d’utiliser n’importe quelle combinaison de stratégies de parallélisme depuis un seul script, via ParallelismConfig et des plugins.
• Axolotl intègre ces techniques ND‑Parallel afin de faciliter le fine‑tuning à grande échelle et la composition des stratégies avec diverses techniques de fine‑tuning ; il suffit d’ajouter des champs ND‑Parallel dans le fichier de config Axolotl.
• Des configs d’exemple et des scripts end‑to‑end montrent comment configurer dataloaders, optimiseurs et boucles d’entraînement tout en composant des stratégies.

Pourquoi cela compte (impact pour développeurs/entreprises)

• Itération accélérée : les développeurs peuvent expérimenter des combinaisons de DP, FSDP, TP et CP sans réécrire l’infrastructure d’entraînement.
• Monter en charge : la composition du sharding et du model‑parallel permet d’entraîner des modèles qui ne tiennent pas sur une seule GPU ou un seul nœud.
• Contrôle des coûts et des ressources : en ajustant le nombre de répliques, la taille des shards et la taille des groupes TP, les équipes peuvent échanger mémoire contre coûts de communication et aligner la configuration sur la topologie du cluster (NVLink intra‑nœud vs Infiniband inter‑nœud).
• Réutilisabilité : les champs de config Axolotl et ParallelismConfig d’Accelerate facilitent la réutilisation d’un même script pour des environnements variés.

Détails techniques ou implémentation

Principales stratégies de parallélisme

| Stratégie | Ce qu’elle fait | Quand l’utiliser | Principaux compromis |---|---|---|---| | Parallélisme de données (DP) | Réplique le modèle et l’état de l’optimiseur ; répartit les lots et synchronise les gradients | Quand le modèle tient sur un appareil et qu’on veut augmenter le débit | Faible communication par pas (all‑reduce), mais nécessite que le modèle tienne sur l’appareil ; dp_replicate_size règle le nombre de répliques |Fully Sharded Data Parallel (FSDP) | Sharde poids, gradients et état optimiseur ; rassemble les paramètres par couche | Quand le modèle est trop grand pour une seule GPU | Échange mémoire maximale contre communication (gather par couche) ; la granularité de gather influe sur mémoire vs communication ; dp_shard_size règle le degré |Tensor Parallelism (TP) | Sharde de manière permanente de grandes couches linéaires entre appareils | Quand des couches individuelles sont trop grandes pour le gather en FSDP | Réduction de mémoire proportionnelle à la taille du groupe TP, nécessite synchronisation fréquente des activations ; efficace surtout dans un nœud ; tp_size (tensor_parallel_size dans Axolotl) définit la taille |Context Parallelism (CP) | Sharde l’entrée sur la dimension séquence pour traiter des contextes très longs | Quand on entraîne sur des séquences très longues (centaines de milliers à millions de tokens) | Réduit la mémoire d’attention par appareil en divisant la matrice d’attention ; chaque query requiert encore K et V complets pour le softmax |

Composer les stratégies et considérations cluster

• DP est une stratégie de haut niveau : dp_replicate_size=2 produit deux répliques complètes, chacune pouvant être ensuite sharded par FSDP ou TP.
• FSDP peut être étendu sur plusieurs nœuds en traitant l’ensemble des appareils comme un seul domaine de sharding (ex. 4 nœuds × 8 GPUs = 32 appareils). La communication devient plus coûteuse entre nœuds, ainsi on évite souvent FSDP au‑delà d’un nœud complet.
• TP est généralement efficace au sein d’un même nœud en raison des besoins en bande passante faible latence (NVLink). TP inter‑nœuds peut être inefficace.
• Pour des contextes très longs, l’attention croît quadratiquement avec la longueur ; l’article illustre que des séquences extrêmes (p. ex. 128k tokens) peuvent rendre les activations d’attention énormes, justifiant l’usage du CP.

Réglages dans Accelerate et Axolotl

• ParallelismConfig dans Accelerate expose dp_replicate_size, dp_shard_size et tp_size pour configurer DP, FSDP et TP et leur composition.
• Dans Axolotl, des champs similaires existent (tensor_parallel_size, dp_shard_size) ; ajouter des techniques ND‑Parallel à une config Axolotl revient souvent à ajouter des champs.
• Le dépôt Accelerate contient un FullyShardedDataParallelPlugin prêt à l’emploi et des exemples end‑to‑end.

Points clés

• Accelerate + Axolotl permettent de composer DP, FSDP, TP et CP dans un seul flux d’entraînement.
• Utiliser DP si le modèle tient sur un appareil ; augmenter dp_replicate_size pour plus de débit.
• Utiliser FSDP pour sharder l’état lorsque le modèle est trop grand ; ajuster la granularité de gather.
• Utiliser TP pour sharder de grandes couches à l’intérieur d’un nœud lorsque le gather FSDP est trop coûteux.
• Utiliser CP pour entraîner sur des séquences très longues.

Références

• Blog post: https://huggingface.co/blog/accelerate-nd-parallel